use aoc_runner_derive::{aoc, aoc_generator};
use grid::{AsCoord2d, Grid};
use itertools::Itertools;
use std::str::FromStr;

#[aoc_generator(day7)]
fn parse(input: &str) -> Grid<u8> {
    Grid::from_str(input).unwrap()
}

fn emit_beams(grid: &mut Grid<u8>, pos: impl AsCoord2d) -> u64 {
    match grid.set(&pos, b'|') {
        None | Some(b'|') => 0,
        Some(b'.') | Some(b'S') => emit_beams(grid, (pos.x(), pos.y() + 1)),
        Some(b'^') => {
            1 + emit_beams(grid, (pos.x() - 1, pos.y())) + emit_beams(grid, (pos.x() + 1, pos.y()))
        }
        _ => panic!(),
    }
}

#[aoc(day7, part1)]
fn part1(input: &Grid<u8>) -> u64 {
    let mut grid = input.clone();
    emit_beams(&mut grid, input.find(&b'S').unwrap())
}

#[aoc(day7, part2)]
fn part2(input: &Grid<u8>) -> u64 {
    let mut grid = Grid::same_shape(input, 0u64);
    grid.set(&input.find(&b'S').unwrap(), 1);

    // Start row is already set up
    for y in 1..grid.height() {
        let row_above = grid.row_iter(y as i64 - 1).unwrap().cloned().collect_vec();
        // for each column with non-zero counts in the row above get its column index and count
        for (x, n) in row_above.into_iter().enumerate().filter(|(_i, n)| *n > 0) {
            match input.get(&(x, y)) {
                None => panic!("How did we end up outside the grid?"),
                // if our current position is a caret, add that count to our neighbours
                Some(b'^') => {
                    grid.set(&(x - 1, y), grid.get(&(x - 1, y)).unwrap() + n);
                    grid.set(&(x + 1, y), grid.get(&(x + 1, y)).unwrap() + n);
                }
                // if our current position is a . add that count here
                Some(b'.') => {
                    grid.set(&(x, y), grid.get(&(x, y)).unwrap() + n);
                }
                _ => panic!(),
            }
        }
    }

    grid.row(grid.height() as i64 - 1)
        .unwrap()
        .iter()
        .sum::<u64>()
}

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

    const TRIVIAL_EXAMPLE: &str = "...S...
.......
...^...
..^.^.^
...^.^.
.......
.......";

    const EXAMPLE: &str = ".......S.......
...............
.......^.......
...............
......^.^......
...............
.....^.^.^.....
...............
....^.^...^....
...............
...^.^...^.^...
...............
..^...^.....^..
...............
.^.^.^.^.^...^.
...............";

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse(EXAMPLE)), 21);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse(EXAMPLE)), 40);
    }

    #[test]
    fn part2_trivial_example() {
        assert_eq!(part2(&parse(TRIVIAL_EXAMPLE)), 7);
    }
}