use std::cmp::{max, min};
use std::collections::{BTreeMap, BTreeSet};

use aoc_2022::*;

type Position = Vector2D<isize>;

type Input = BTreeSet<Position>;
type Output = isize;

#[derive(Debug, Clone, Copy)]
enum Direction {
    North,
    South,
    West,
    East,
}

impl Direction {
    fn adjacent(self) -> Vec<Position> {
        match self {
            Direction::North => vec![
                Position::new(0, -1),
                Position::new(1, -1),
                Position::new(-1, -1),
            ],
            Direction::South => vec![
                Position::new(0, 1),
                Position::new(1, 1),
                Position::new(-1, 1),
            ],
            Direction::West => vec![
                Position::new(-1, 0),
                Position::new(-1, 1),
                Position::new(-1, -1),
            ],
            Direction::East => vec![
                Position::new(1, 0),
                Position::new(1, 1),
                Position::new(1, -1),
            ],
        }
    }
}

lazy_static! {
    static ref DIRECTIONS: Vec<Direction> = vec![
        Direction::North,
        Direction::South,
        Direction::West,
        Direction::East
    ];
    static ref NEIGHBOURHOOD: BTreeSet<Position> =
        DIRECTIONS.iter().flat_map(|d| d.adjacent()).collect();
}

fn propose_moves(i: usize, elves: &Input) -> BTreeMap<Position, Position> {
    let mut already_proposed: BTreeSet<Position> = BTreeSet::new();

    let directions = (0..DIRECTIONS.len())
        .map(|j| DIRECTIONS[(i + j) % DIRECTIONS.len()].adjacent())
        .collect_vec();

    let mut moves = BTreeMap::new();
    for elf in elves {
        if NEIGHBOURHOOD
            .iter()
            .map(|d| *elf + *d)
            .all(|pos| !elves.contains(&pos))
        {
            // debug!("{:?} not moving.", elf);
            continue;
        }

        for adjacent in &directions {
            if adjacent
                .iter()
                .map(|dpos| *dpos + *elf)
                .all(|pos| !elves.contains(&pos))
            {
                let new_position = adjacent[0] + *elf;

                if already_proposed.contains(&new_position) {
                    // cannot move more than one elf to the position
                    moves.remove(&new_position);

                    break;
                }

                moves.insert(new_position, *elf);
                already_proposed.insert(new_position);
                break;
            }
        }
    }

    moves
}

fn execute_moves(elves: &mut Input, moves: &BTreeMap<Position, Position>) {
    for (to, from) in moves.iter() {
        elves.remove(from);
        elves.insert(*to);
    }
}

fn round(i: usize, elves: &mut Input) -> bool {
    let moves = propose_moves(i, elves);
    execute_moves(elves, &moves);

    moves.is_empty()
}

fn _show_ground_with_dimensions(
    positions: &Input,
    min_x: isize,
    max_x: isize,
    min_y: isize,
    max_y: isize,
) {
    for y in min_y..=max_y {
        let mut row = String::new();

        for x in min_x..=max_x {
            row += if positions.contains(&Position::new(x, y)) {
                "#"
            } else {
                "."
            }
        }
        debug!("{}", row);
    }
}

fn get_bounds(positions: &Input) -> (Vector2D<isize>, Vector2D<isize>) {
    let f = |init, cmp: &dyn Fn(isize, isize) -> isize| {
        positions
            .iter()
            .fold(Vector2D::new(init, init), |acc, elf| {
                Vector2D::new(cmp(acc.x(), elf.x()), cmp(acc.y(), elf.y()))
            })
    };

    (f(isize::MAX, &min::<isize>), f(isize::MIN, &max::<isize>))
}

fn _show_ground(positions: &Input) {
    let (min_pos, max_pos) = get_bounds(positions);

    _show_ground_with_dimensions(
        positions,
        min_pos.x(),
        max_pos.x(),
        min_pos.y(),
        max_pos.y(),
    )
}

struct Day23;
impl Solution<Input, Output> for Day23 {
    fn parse_input<P: AsRef<Path>>(pathname: P) -> Input {
        let map: Vec<Vec<char>> = file_to_string(pathname)
            .lines()
            .map(|l| l.chars().collect_vec())
            .collect_vec();

        map.iter()
            .enumerate()
            .flat_map(|(y, row)| {
                row.iter()
                    .enumerate()
                    .filter(|&(_, c)| c != &'.')
                    .map(move |(x, _)| Vector2D::new(x as isize, y as isize))
            })
            .collect()
    }

    fn part_1(input: &Input) -> Output {
        let mut positions = input.clone();
        // debug!("== Initial State ==");
        // show_ground_with_dimensions(&positions, -3, 10, -2, 9);

        for i in 0..10 {
            round(i, &mut positions);
            // debug!("== End of Round {} ==", i + 1);
            // show_ground_with_dimensions(&positions, -3, 10, -2, 9);
        }

        let (min_pos, max_pos) = get_bounds(&positions);

        (max_pos.x() - min_pos.x() + 1) * (max_pos.y() - min_pos.y() + 1)
            - (positions.len() as isize)
    }

    fn part_2(input: &Input) -> Output {
        let mut positions = input.clone();

        let mut rounds = 0;
        while !round(rounds, &mut positions) {
            rounds += 1;
        }

        (1 + rounds).try_into().unwrap()
    }
}

fn main() -> Result<()> {
    // Day23::run("sample")
    Day23::main()
}

test_sample!(day_23, Day23, 110, 20);