extern crate advent_of_code_2018;
use advent_of_code_2018::*;

use std::error::Error;
use std::path::PathBuf;
use std::str::FromStr;

use std::collections::HashMap;

// cargo watch -cs "cargo run --release --bin day_24"

#[derive(Debug, Clone)]
struct Army {
    units: u32,
    hp: u32,
    weaknesses: Vec<String>,
    immunities: Vec<String>,
    damage_type: String,
    damage: u32,
    initiative: u32
}

impl Army {
    fn effective_power(&self) -> u32 {
        self.units * self.damage
    }

    fn actual_damage(&self, other: &Army) -> u32 {
        let modifier = if other.weaknesses.contains(&self.damage_type) {
            2
        } else if other.immunities.contains(&self.damage_type) {
            0
        } else {
            1
        };
        self.effective_power() * modifier
    }
}

impl FromStr for Army {
    type Err = Box<Error>;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        // 2991 units each with 8084 hit points (weak to fire) with an attack that does 19 radiation damage at initiative 11
        let mut words = s.split_whitespace();
        let units = words.next().unwrap().parse()?;
        let hp = words.nth(3).unwrap().parse()?;
        let _ = words.nth(1);
        
        let mut weaknesses = Vec::new();
        let mut immunities = Vec::new();
        
        let mut bracket_handled = false;
        while !bracket_handled {
            let maybe_weak_or_immune = words.next().unwrap().trim_matches(|c: char| !c.is_alphabetic());
            if maybe_weak_or_immune == "weak" {
                let mut weakness = words.nth(1).unwrap();
                weaknesses.push(weakness.trim_matches(|c: char| !c.is_alphabetic()).to_string());
                while weakness.ends_with(',') {
                    weakness = words.next().unwrap();
                    weaknesses.push(weakness.trim_matches(|c: char| !c.is_alphabetic()).to_string());
                }
            }
            else if maybe_weak_or_immune == "immune" {
                let mut immunity = words.nth(1).unwrap();
                immunities.push(immunity.trim_matches(|c: char| !c.is_alphabetic()).to_string());
                while immunity.ends_with(',') {
                    immunity = words.next().unwrap();
                    immunities.push(immunity.trim_matches(|c: char| !c.is_alphabetic()).to_string());
                }
            }
            else {
                bracket_handled = true;
            }
        }
        let damage = words.nth(4).unwrap().parse()?;
        let damage_type = words.next().unwrap().to_string();
        let initiative = words.nth(3).unwrap().parse()?;
        
        Ok(Army {
            units, hp, weaknesses, immunities, damage_type, damage, initiative
        })
    }
}



fn main() -> Result<(), Box<Error>> {
    let input = read_file(&PathBuf::from("inputs/24.txt"))?;

    let mut input_iter = input.iter();
    let _ = input_iter.next();
    let immune_army: Vec<Army> = input_iter.by_ref()
        .take_while(|&line| line != "Infection:")
        .map(|line| line.parse().unwrap())
        .collect();
    let infection_army: Vec<Army> = input_iter
        .map(|line| line.parse().unwrap())
        .collect();

    let (unboosted_immune_army_size, unboosted_infection_army_size) = simulate(&immune_army, &infection_army, 0);
    debug!(unboosted_immune_army_size);
    debug!(unboosted_infection_army_size);

    for boost in 0.. {
        debug!(boost);
        let (immune_army_size, infection_army_size) = simulate(&immune_army, &infection_army, boost);
        debug!(immune_army_size);
        debug!(infection_army_size);
        if immune_army_size > 0 {
            break;
        }
    }

    Ok(())
}


fn target_selection(targeting_army: &Vec<Army>, targeted_army: &Vec<Army>) -> HashMap<usize, usize> {
    let mut targets = HashMap::new();
    for (i, army) in targeting_army.iter().enumerate() {
        let best_target = (0..targeted_army.len())
            .filter(|potential_target| !targets.values().any(|targeted| targeted == potential_target))
            .fold(None, |best_target, next| {
                let next_actual_damage = army.actual_damage(&targeted_army[next]);
                let last_best_actual_damage = best_target.map(|best_target| army.actual_damage(&targeted_army[best_target])).unwrap_or(0);
                let next_effective_power = targeted_army[next].effective_power();
                let last_best_effective_power = best_target.map(|best_target: usize| targeted_army[best_target].effective_power()).unwrap_or(0);
                let next_initiative = targeted_army[next].initiative;
                let last_best_initiative = best_target.map(|best_target| targeted_army[best_target].initiative).unwrap_or(0);
                if next_actual_damage == 0 && last_best_actual_damage == 0 {
                    None
                } else if next_actual_damage > last_best_actual_damage {
                    Some(next)
                } else if next_actual_damage == last_best_actual_damage && next_effective_power > last_best_effective_power {
                    Some(next)
                } else if next_actual_damage == last_best_actual_damage && next_effective_power == last_best_effective_power && next_initiative > last_best_initiative {
                    Some(next)
                } else {
                    best_target
                }
            });

        if let Some(best_target) = best_target {
            targets.insert(i, best_target);
        }
    }

    targets
}

fn attack(aggressor: &Army, target: &mut Army) {
    let damage =  aggressor.actual_damage(target);
    let units_destroyed = damage / target.hp;
    target.units = target.units.saturating_sub(units_destroyed);
}


fn simulate(immune_army: &[Army], infection_army: &[Army], boost: u32) -> (u32, u32) {
    let mut infection_army = Vec::from(infection_army);
    
    let mut immune_army: Vec<Army> = immune_army.iter().map(|army| Army {
        damage: army.damage + boost,
        ..army.clone()
    }).collect();

    let mut watchdog_sum = immune_army.iter().map(|a| a.units).sum::<u32>() + infection_army.iter().map(|a| a.units).sum::<u32>();
    while immune_army.len() > 0 && infection_army.len() > 0 {
        immune_army.sort_by(|a, b| b.effective_power().cmp(&a.effective_power()).then(b.initiative.cmp(&a.initiative)));
        infection_army.sort_by(|a, b| b.effective_power().cmp(&a.effective_power()).then(b.initiative.cmp(&a.initiative)));

        let mut immune_targets = target_selection(&immune_army, &infection_army);
        let mut infection_targets = target_selection(&infection_army, &immune_army);

        while immune_targets.len() > 0 || infection_targets.len() > 0 {
            let next_immune = immune_targets.keys()
                .map(|&k| (k, immune_army[k].initiative.clone()))
                .max_by_key(|(_, initiative)| *initiative);
            let next_infection = infection_targets.keys()
                .map(|&k| (k, infection_army[k].initiative.clone()))
                .max_by_key(|(_, initiative)| *initiative);
            match (next_immune, next_infection) {
                (None, None) => {
                    panic!("No targets left")
                },
                (Some((immune_key, immune_init)), Some((_infect_key, infect_init))) if immune_init > infect_init => {
                    attack(&immune_army[immune_key], &mut infection_army[*immune_targets.get(&immune_key).unwrap()]);
                    immune_targets.remove(&immune_key);
                },
                (Some((_immune_key, _immune_init)), Some((infect_key, _infect_init))) => {
                    attack(&infection_army[infect_key], &mut immune_army[*infection_targets.get(&infect_key).unwrap()]);
                    infection_targets.remove(&infect_key);
                },
                (Some((immune_key, _immune_init)), None) => {
                    attack(&immune_army[immune_key], &mut infection_army[*immune_targets.get(&immune_key).unwrap()]);
                    immune_targets.remove(&immune_key);
                },
                (None, Some((infect_key, _infect_init))) => {
                    attack(&infection_army[infect_key], &mut immune_army[*infection_targets.get(&infect_key).unwrap()]);
                    infection_targets.remove(&infect_key);
                }
            };
        }

        immune_army.retain(|a| a.units > 0);
        infection_army.retain(|a| a.units > 0);

        let new_watchdog_sum = immune_army.iter().map(|a| a.units).sum::<u32>() + infection_army.iter().map(|a| a.units).sum::<u32>();
        if watchdog_sum == new_watchdog_sum {
            // Stalemate! Everyone loses.
            return (0, 0);
        }
        watchdog_sum = new_watchdog_sum;
    }

    (immune_army.iter().map(|a| a.units).sum(), infection_army.iter().map(|a| a.units).sum())
}