use ansi_term::{Color, Style};
use std::collections::{HashSet, VecDeque};
use std::fmt::{Debug, Formatter};
use std::ops::RangeInclusive;

pub fn day11() {
    // let input = r#"
    // 11111
    // 19991
    // 19191
    // 19991
    // 11111
    //     "#;

    // let grid: Grid<5> = input.into();
    // dbg!(&grid);
    // dbg!(&grid.step());
    // dbg!(&grid.step().step());

    // let input = r#"
    // 5483143223
    // 2745854711
    // 5264556173
    // 6141336146
    // 6357385478
    // 4167524645
    // 2176841721
    // 6882881134
    // 4846848554
    // 5283751526
    //     "#;
    let input = include_str!("../input/day11.txt");

    let mut grid: Grid<10> = input.into();
    let mut total_flashes = 0;

    println!("{:?}", grid);

    for step in 0..100 {
        let (new, flashed) = grid.step();
        // println!("after step: {}, new grid: {:?}", step + 1, new);
        grid = new;
        total_flashes += flashed;
    }

    println!("grid: {:?}, total flashes: {}", grid, total_flashes);

    let mut grid2: Grid<10> = input.into();
    let mut step = 0;
    loop {
        step += 1;
        let (new, flashed) = grid2.step();
        grid2 = new;
        if flashed == grid.size() {
            break;
        }
    }
    println!("grid: {:?}, after step: {}", grid2, step);
}

struct Grid<const N: usize> {
    grid: [[u32; N]; N],
}

impl<const N: usize> Debug for Grid<N> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "\n").expect("write");
        for row in 0..N {
            for col in 0..N {
                let mut style = Style::new();
                if self.grid[row][col] == 0 {
                    style = Color::White.bold();
                }
                write!(
                    f,
                    "{}{}{}",
                    style.prefix(),
                    self.grid[row][col],
                    style.suffix()
                )
                .expect("write");
            }
            write!(f, "\n").expect("write");
        }
        Ok(())
    }
}

impl<const N: usize> From<&str> for Grid<N> {
    fn from(s: &str) -> Self {
        let mut grid = [[0; N]; N];
        for (row, line) in s.trim().lines().enumerate() {
            for (col, c) in line.trim().chars().enumerate() {
                grid[row][col] = c.to_digit(10).unwrap();
            }
        }
        Self { grid }
    }
}

impl<const N: usize> Grid<N> {
    fn step(&self) -> (Grid<N>, usize) {
        let mut new_grid = [[0; N]; N];

        for row in 0..N {
            for col in 0..N {
                new_grid[row][col] = self.grid[row][col] + 1;
            }
        }

        let mut to_flash = (0..N)
            .flat_map(|row| (0..N).map(move |col| (row, col)))
            .filter(|&(row, col)| new_grid[row][col] > 9)
            .collect::<VecDeque<(usize, usize)>>();
        let mut already_flashed: HashSet<(usize, usize)> = HashSet::new();
        while let Some((row, col)) = to_flash.pop_front() {
            if already_flashed.contains(&(row, col)) {
                continue;
            }

            already_flashed.insert((row, col));
            for r in adjacent(row, N) {
                for c in adjacent(col, N) {
                    new_grid[r][c] += 1;
                    if new_grid[r][c] > 9 && !already_flashed.contains(&(r, c)) {
                        to_flash.push_back((r, c));
                    }
                }
            }
        }

        for &(row, col) in already_flashed.iter() {
            new_grid[row][col] = 0;
        }

        (Self { grid: new_grid }, already_flashed.len())
    }

    fn size(&self) -> usize {
        N * N
    }
}

fn adjacent(n: usize, max: usize) -> RangeInclusive<usize> {
    let min = if n > 0 { n - 1 } else { n };
    let max = if n < max - 1 { n + 1 } else { n };
    min..=max
}