use std::collections::HashSet;

use color_eyre::eyre;
use eyre::eyre;

use crate::lib::Vec2;


struct Octopi {
    grid: Vec2<u8>,
    flashed: HashSet<(usize, usize)>,
}

impl Octopi {
    fn step(&mut self) -> usize {
        for row in 0..self.grid.row_count() {
            for col in 0..self.grid.col_count() {
                self.inc(row, col);
            }
        }
        let count = self.flashed.len();
        self.flashed.clear();
        count
    }

    fn inc(&mut self, row: usize, col: usize) {
        if self.flashed.contains(&(row, col)) {
            return;
        }

        let energy = self.grid.get_mut(row, col);
        *energy += 1;
        if *energy > 9 {
            self.flashed.insert((row, col));
            *energy = 0;
            for (n_row, n_col) in self.grid.neighbor_coords(row, col) {
                self.inc(n_row, n_col);
            }
        }
    }
}


fn load(data: &str) -> eyre::Result<Octopi> {
    let mut grid = Vec2::new(10);
    for line in data.lines() {
        let mut row = Vec::with_capacity(10); // kinda inefficient but whatever, it only happens once
        for c in line.chars() {
            match c.to_digit(10) {
                Some(d) => row.push(d as u8),
                None => return Err(eyre!("Invalid digit: {}", c)),
            }
        }
        grid.push(row);
    }
    Ok(Octopi {grid, flashed: HashSet::new()})
}


pub fn run(data: &str) -> eyre::Result<(usize, usize)> {
    let mut octopi = load(data)?;

    let mut num_flashes = 0;
    let mut two = None;
    for i in 0.. {
        let n = octopi.step();
        if i < 100 {
            num_flashes += n;
        }
        if n == 100 {
            two = Some(i);
            break;
        }
    }
    Ok((num_flashes, two.unwrap() + 1))
}