use std::fs;
use std::path::Path;
use std::collections::VecDeque;

const WIDTH: usize = 10;
const HEIGHT: usize = 10;

/* AOC21 Day 11: https://adventofcode.com/2021/day/11 */
fn main() {
    let input = Path::new("resources").join("input.txt");
    let content = fs::read_to_string(input).expect("Unable to read input file");
    ex1(&content, 100);
    ex2(&content);
}

fn ex1(content: &str, steps: u32) {
    let mut energy: Vec<u8> = content.lines().flat_map(|l| l.bytes().map(|b| b-b'0')).collect();
    let mut queue: VecDeque<usize> = VecDeque::with_capacity(2*WIDTH*HEIGHT);
    println!("Ex1: the result is {}", (0..steps).map(|_| step(&mut energy, &mut queue)).sum::<usize>())
}

fn ex2(content: &str) {
    let mut energy: Vec<u8> = content.lines().flat_map(|l| l.bytes().map(|b| b-b'0')).collect();
    let mut queue: VecDeque<usize> = VecDeque::with_capacity(2*WIDTH*HEIGHT);
    println!("Ex2: the result is {}", (1..).find(|_| step(&mut energy, &mut queue) == WIDTH*HEIGHT).unwrap());
}

fn step(energy: &mut [u8], queue: &mut VecDeque<usize>) -> usize {
    queue.extend(0..energy.len());
    while let Some(i) = queue.pop_front() {
        octopus(energy, queue, i);
    }
    energy.iter_mut().filter_map(|o| if *o > 9 { *o = 0; Some(o) } else { None }).count()
}

fn octopus(energy: &mut [u8], queue: &mut VecDeque<usize>, i: usize) {
    let len = energy.len();
    energy[i] += 1;
    if energy[i] == 10 {
        if i%WIDTH > 0 {
            queue.push_back(i-1);
            if i >= WIDTH+1    { queue.push_back(i-WIDTH-1) }
            if i+WIDTH-1 < len { queue.push_back(i+WIDTH-1) }
        }
        if (i+1)%WIDTH > 0 {
            queue.push_back(i+1);
            if i >= WIDTH-1    { queue.push_back(i+1-WIDTH) }
            if i+WIDTH+1 < len { queue.push_back(i+1+WIDTH) }
        }
        if i >= WIDTH    { queue.push_back(i-WIDTH) }
        if i+WIDTH < len { queue.push_back(i+WIDTH) }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const INPUT: [u8;100] = [
        5,4,8,3,1,4,3,2,2,3,
        2,7,4,5,8,5,4,7,1,1,
        5,2,6,4,5,5,6,1,7,3,
        6,1,4,1,3,3,6,1,4,6,
        6,3,5,7,3,8,5,4,7,8,
        4,1,6,7,5,2,4,6,4,5,
        2,1,7,6,8,4,1,7,2,1,
        6,8,8,2,8,8,1,1,3,4,
        4,8,4,6,8,4,8,5,5,4,
        5,2,8,3,7,5,1,5,2,6
    ];

    #[test]
    fn step10() {
        let mut input = INPUT.clone();
        let mut queue = VecDeque::with_capacity(200);
        let flash = (0..10).map(|_| step(&mut input, &mut queue)).sum::<usize>();
        assert_eq!(204, flash);
    }

    #[test]
    fn step100() {
        let mut input = INPUT.clone();
        let mut queue = VecDeque::with_capacity(200);
        let flash = (0..100).map(|_| step(&mut input, &mut queue)).sum::<usize>();
        assert_eq!(1656, flash);
    }

    #[test]
    fn synchronised() {
        let mut input = INPUT.clone();
        let mut queue: VecDeque<usize> = VecDeque::with_capacity(2*WIDTH*HEIGHT);
        assert_eq!(195, (1..).find(|_| step(&mut input, &mut queue) == WIDTH*HEIGHT).unwrap());
    }
}