extern crate advent_of_code_2018;
use advent_of_code_2018::*;

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

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

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

    let (round_hp_product, _) = calculate_round_hp_product(&input, 3);
    debug!(round_hp_product);

    for power in 3.. {
        let (round_hp_product_without_losses, elf_losses) = calculate_round_hp_product(&input, power);
        if (elf_losses == 0) {
            debug!(round_hp_product_without_losses);
            break;
        }
    }
    
    Ok(())
}

fn calculate_round_hp_product(input: &[String], elf_power: i32) -> (i32, usize) {
    let goblin_power = 3;
    let initial_hp = 200;

    let mut map = Vec::new();
    let mut units = Vec::new();

    for (y, line) in input.iter().enumerate() {
        let mut map_row = Vec::new();
        for (x, c) in line.chars().enumerate() {
            map_row.push(c == '#');

            let maybe_unit = match c {
                'G' => Some(Unit {
                    team: false,
                    x: x as i32,
                    y: y as i32,
                    hp: initial_hp
                }),
                'E' => Some(Unit {
                    team: true,
                    x: x as i32,
                    y: y as i32,
                    hp: initial_hp
                }),
                _ => None
            };
            if let Some(unit) = maybe_unit {
                units.push(unit);
            }
        }

        map.push(map_row);
    }
    let start_elves = units.iter().filter(|u| u.team).count();

    for round in 0.. {
        // debug!(round);
        // debug!(units);
        // debug_print(&units, &map);
        units.sort_unstable_by(|a, b| a.y.cmp(&b.y).then(a.x.cmp(&b.x)));

        for i in 0..units.len() {
            if units[i].hp > 0 {
                let battle_over = !units.iter().any(|other| other.hp > 0 && other.team != units[i].team);
                if battle_over {
                    println!("Battle over");
                    let hp_sum = units.iter().filter(|u| u.hp > 0).map(|u| u.hp).sum::<i32>();
                    debug!(round);
                    debug!(hp_sum);

                    let elf_losses = start_elves - units.iter().filter(|u| u.team).count();
                    debug!(elf_losses);
                    return (round * hp_sum, elf_losses);
                }

                let is_next_to_opponent = unit_at_point(!units[i].team, units[i].x-1, units[i].y, &units)
                    .or_else(|| unit_at_point(!units[i].team, units[i].x+1, units[i].y, &units))
                    .or_else(|| unit_at_point(!units[i].team, units[i].x, units[i].y-1, &units))
                    .or_else(|| unit_at_point(!units[i].team, units[i].x, units[i].y+1, &units)).is_some();
                if !is_next_to_opponent {
                    if let Some((next_x, next_y)) = find_next_step(i, &units, &map) {
                        units[i].x = next_x;
                        units[i].y = next_y;
                    }
                }

                let mut attack_candidates = adjacent_units(!units[i].team, units[i].x, units[i].y, &units);
                attack_candidates.sort_by(|&a, &b| units[a].hp.cmp(&units[b].hp));
                if let Some(&to_attack) = attack_candidates.first() {
                    let attack_power = if units[i].team { elf_power } else { goblin_power };
                    units[to_attack].hp -= attack_power;
                }
            }
        }

        units.retain(|u| u.hp > 0);
    }

    panic!("Match ended unexpectedly");
}

fn find_next_step(i: usize, units: &[Unit], map: &Vec<Vec<bool>>) -> Option<(i32, i32)> {
    let mut explored = vec!(Explored {
        x: units[i].x,
        y: units[i].y,
        distance: 0,
        parent: 0
    });

    let mut explored_index = 0;
    let mut distance_found = None;
    while explored_index < explored.len() && distance_found.map_or(true, |distance| explored[explored_index].distance < distance) {
        let distance = explored[explored_index].distance;
        for (new_x, new_y) in adjacent_empty_spaces(explored[explored_index].x, explored[explored_index].y, &map, &units, &explored) {
            explored.push(Explored {
                x: new_x,
                y: new_y,
                distance: distance + 1,
                parent: explored_index
            });

            if adjacent_units(!units[i].team, new_x, new_y, &units).len() > 0 {
                distance_found = Some(distance + 1);
            }
        }
        explored_index += 1;
    }

    if let Some(distance) = distance_found {
        let mut candidates = explored.iter().enumerate()
            .filter(|(_,e)| e.distance == distance)
            .filter(|(_,e)| adjacent_units(!units[i].team, e.x, e.y, &units).len() > 0)
            .collect::<Vec<_>>();
        candidates.sort_by(|(_,a), (_,b)| a.y.cmp(&b.y).then(a.x.cmp(&b.x)));
        let (best_index, _) = candidates[0];
        
        let mut result_index = best_index;
        let mut next_index = best_index;
        while next_index != 0 {
            result_index = next_index;
            next_index = explored[result_index].parent;
        }
        return Some((explored[result_index].x, explored[result_index].y));
    }
    None
}

fn adjacent_empty_spaces(x: i32, y: i32, map: &Vec<Vec<bool>>, units: &[Unit], explored: &[Explored]) -> Vec<(i32, i32)> {
    [(x, y-1), (x-1, y), (x+1, y), (x, y+1)].iter()
        .filter(|(x, y)| *x >= 0 && *y >= 0)
        .filter(|(x, y)| !explored.iter().any(|e| e.x == *x && e.y == *y))
        .filter(|(x, y)| !units.iter().any(|e| e.hp > 0 && e.x == *x && e.y == *y))
        .filter(|(x, y)| !map[*y as usize][*x as usize])
        .cloned()
        .collect()
}

fn adjacent_units(team: bool, x: i32, y: i32, units: &[Unit]) -> Vec<usize> {
    [(x, y-1), (x-1, y), (x+1, y), (x, y+1)].iter()
        .filter_map(|(x, y)| unit_at_point(team, *x, *y, units))
        .collect()
}

fn unit_at_point(team: bool, x: i32, y: i32, units: &[Unit]) -> Option<usize> {
    units.iter().enumerate()
        .find(|(_, u)| u.hp > 0 && u.team == team && u.x == x && u.y == y)
        .map(|(i, _)| i)
}

#[derive(Debug, Clone)]
struct Unit {
    team: bool,
    x: i32,
    y: i32,
    hp: i32
}

#[derive(Debug)]
struct Explored {
    x: i32,
    y: i32,
    distance: u32,
    parent: usize
}


fn debug_print(units: &[Unit], map: &Vec<Vec<bool>>) {
    for (y, line) in map.iter().enumerate() {
        for (x, wall) in line.iter().enumerate() {
            let unit = units.iter().find(|u| u.hp > 0 && u.x == x as i32 && u.y == y as i32);
            print!("{}", match (wall, unit) {
                (true, _) => '#',
                (false, Some(unit)) if unit.team => 'E',
                (false, Some(_)) => 'G',
                (false, None) => '.'
            });
        }
        println!();
    }
    println!();
    println!();
}

fn part_1_test(input: &[&str]) -> i32 {
    let owned_input = input.iter()
        .map(|&line| line.to_owned())
        .collect::<Vec<String>>();
    let (product, _) = calculate_round_hp_product(&owned_input, 3);
    product
}

#[test]
fn example_1() {
    let input = vec!(
        "#######",
        "#.G...#",
        "#...EG#",
        "#.#.#G#",
        "#..G#E#",
        "#.....#",
        "#######");
    assert_eq!(part_1_test(&input), 27730);
}

#[test]
fn example_2() {
    let input = vec!(
        "#######",
        "#G..#E#",
        "#E#E.E#",      
        "#G.##.#",
        "#...#E#",
        "#...E.#",
        "#######"
    );
    assert_eq!(part_1_test(&input), 36334);
}

#[test]
fn example_3() {
    let input = vec!(
        "#######",
        "#E.G#.#",
        "#.#G..#",      
        "#G.#.G#",
        "#G..#.#",
        "#...E.#",
        "#######"
    );
    assert_eq!(part_1_test(&input), 27755);
}

#[test]
fn example_4() {
    let input = vec!(
        "#######",
        "#.E...#",
        "#.#..G#",      
        "#.###.#",
        "#E#G#G#",
        "#...#G#",
        "#######"
    );
    assert_eq!(part_1_test(&input), 28944);
}