use nom::{
    branch::alt,
    character::complete::{char, line_ending},
    combinator::{map, value},
    multi::{many1, separated_list1},
    IResult,
};
use std::{collections::BTreeSet, fs, mem};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_21.txt")?;
    let garden = WalledGarden::parser(&input).unwrap().1;

    dbg!(garden.count_closed_walls_walkable_after_steps(garden.center(), 64));
    dbg!(garden.count_open_walls_walkable_after_steps(26501365));

    Ok(())
}

#[derive(Debug, Clone)]
struct WalledGarden {
    rocks: Vec<Vec<bool>>,
    walkable_to: Vec<Vec<bool>>,
    walkable_to_back: Vec<Vec<bool>>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum WalledGardenState {
    Empty,
    Walkable,
    Rock,
}

#[derive(Debug)]
struct MaxWalkable {
    odd_steps_max: usize,
    even_steps_max: usize,
    min_steps: usize,
}

#[derive(Debug)]
struct EntryPoint {
    entry: (usize, usize),
    max: MaxWalkable,
}

impl WalledGarden {
    fn new(tiles: Vec<Vec<WalledGardenState>>) -> WalledGarden {
        let rocks: Vec<Vec<bool>> = tiles
            .iter()
            .map(|row| {
                row.iter()
                    .map(|t| *t == WalledGardenState::Rock)
                    .collect::<Vec<bool>>()
            })
            .collect();
        let walkable_to: Vec<Vec<bool>> = vec![vec![false; rocks[0].len()]; rocks.len()];

        WalledGarden {
            rocks,
            walkable_to_back: walkable_to.clone(),
            walkable_to: walkable_to.clone(),
        }
    }

    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            separated_list1(line_ending, many1(WalledGardenState::parser)),
            WalledGarden::new,
        )(input)
    }

    fn count_closed_walls_walkable_after_steps(
        &self,
        start: (usize, usize),
        steps: usize,
    ) -> usize {
        let mut garden = self.clone();
        garden.walkable_to[start.1][start.0] = true;
        for _ in 0..steps {
            garden.closed_walls_next_mut();
        }
        garden.count_walkable()
    }

    fn closed_walls_find_max_walkable(&self, start: (usize, usize)) -> MaxWalkable {
        let mut odd_steps_max = 0;
        let mut even_steps_max = 0;

        let mut garden = self.clone();
        garden.walkable_to[start.1][start.0] = true;

        let mut unchanged_count = 0;
        for i in 1.. {
            garden.closed_walls_next_mut();

            if i % 2 == 0 {
                let new_even_max_countable = garden.count_walkable();
                if even_steps_max == new_even_max_countable {
                    unchanged_count += 1;
                } else {
                    even_steps_max = new_even_max_countable;
                }
            } else {
                let new_odd_max_countable = garden.count_walkable();
                if odd_steps_max == new_odd_max_countable {
                    unchanged_count += 1;
                } else {
                    odd_steps_max = new_odd_max_countable;
                }
            }

            if unchanged_count > 1 {
                return MaxWalkable {
                    odd_steps_max,
                    even_steps_max,
                    min_steps: i,
                };
            }
        }
        unreachable!()
    }

    fn closed_walls_next_mut(&mut self) {
        for (y, row) in self.walkable_to_back.iter_mut().enumerate() {
            for (x, tile) in row.iter_mut().enumerate() {
                if !self.rocks[y][x] {
                    *tile = (y > 0 && self.walkable_to[y - 1][x])
                        || (y < self.walkable_to.len() - 1 && self.walkable_to[y + 1][x])
                        || (x > 0 && self.walkable_to[y][x - 1])
                        || (x < self.walkable_to[y].len() - 1 && self.walkable_to[y][x + 1]);
                }
            }
        }

        mem::swap(&mut self.walkable_to, &mut self.walkable_to_back);
    }

    fn count_walkable(&self) -> usize {
        self.walkable_to
            .iter()
            .flat_map(|row| row.iter())
            .filter(|s| **s)
            .count()
    }

    fn count_open_walls_walkable_after_steps(&self, steps: usize) -> usize {
        // assumptions:
        // - this field is square
        // - there is a direct path from the starting point up, down, left, and right
        // - there is a direct path around the edges
        // - the start is in the center

        let (size_x, size_y) = self.size();
        let (center_x, center_y) = self.center();

        let max_chunks_deviance = if (steps - center_y) % size_y > 0 {
            1 + (steps - center_y) / size_y
        } else {
            (steps - center_y) / size_y
        };

        let mut total_walkable = 0;

        for quadrant in [
            EntryPoint::new((0, 0), &self),
            EntryPoint::new((size_x - 1, 0), &self),
            EntryPoint::new((0, size_y - 1), &self),
            EntryPoint::new((size_x - 1, size_y - 1), &self),
        ] {
            let steps_to_quadrant_alignment = center_x + center_y + 2;

            let mut distance_from_edge = 0;
            while max_chunks_deviance > distance_from_edge {
                let steps_from_alignment_to_target_chunk =
                    (max_chunks_deviance - distance_from_edge - 1) * size_y;
                if steps_to_quadrant_alignment + steps_from_alignment_to_target_chunk > steps {
                    distance_from_edge += 1;
                    continue;
                }
                let steps_in_chunk =
                    steps - steps_to_quadrant_alignment - steps_from_alignment_to_target_chunk;
                if steps_in_chunk >= quadrant.max.min_steps {
                    break;
                }

                let walkable_per_chunk =
                    self.count_closed_walls_walkable_after_steps(quadrant.entry, steps_in_chunk);
                total_walkable += walkable_per_chunk * (max_chunks_deviance - distance_from_edge);
                distance_from_edge += 1;
            }

            let remaining_diagonals = max_chunks_deviance - distance_from_edge;
            let even_length_diagonals = remaining_diagonals / 2;
            let odd_length_diagonals = even_length_diagonals + remaining_diagonals % 2;

            let even_length_diagonal_chunks = even_length_diagonals * (even_length_diagonals + 1);
            let odd_length_diagonal_chunks = odd_length_diagonals.pow(2);

            let odd_diagonal_has_even_steps_left = (steps - steps_to_quadrant_alignment) % 2 == 0;
            total_walkable += if odd_diagonal_has_even_steps_left {
                odd_length_diagonal_chunks * quadrant.max.even_steps_max
                    + even_length_diagonal_chunks * quadrant.max.odd_steps_max
            } else {
                even_length_diagonal_chunks * quadrant.max.even_steps_max
                    + odd_length_diagonal_chunks * quadrant.max.odd_steps_max
            };
        }

        for cardinal in [
            EntryPoint::new((0, center_y), &self),
            EntryPoint::new((center_x, 0), &self),
            EntryPoint::new((size_x - 1, center_y), &self),
            EntryPoint::new((center_x, size_y - 1), &self),
        ] {
            let steps_to_cardinal_alignment = center_y + 1;

            let mut distance_from_edge = 0;
            while max_chunks_deviance > distance_from_edge {
                let steps_from_alignment_to_target_chunk =
                    (max_chunks_deviance - distance_from_edge - 1) * size_y;
                let steps_in_chunk =
                    steps - steps_to_cardinal_alignment - steps_from_alignment_to_target_chunk;
                if steps_in_chunk >= cardinal.max.min_steps {
                    break;
                }

                let walkable_per_chunk =
                    self.count_closed_walls_walkable_after_steps(cardinal.entry, steps_in_chunk);
                total_walkable += walkable_per_chunk;
                distance_from_edge += 1;
            }

            let remaining_chunks = max_chunks_deviance - distance_from_edge;
            let even_index_chunks = remaining_chunks / 2;
            let odd_index_chunks = even_index_chunks + remaining_chunks % 2;

            let odd_chunk_has_even_steps_left = (steps - steps_to_cardinal_alignment) % 2 == 0;
            total_walkable += if odd_chunk_has_even_steps_left {
                odd_index_chunks * cardinal.max.even_steps_max
                    + even_index_chunks * cardinal.max.odd_steps_max
            } else {
                even_index_chunks * cardinal.max.even_steps_max
                    + odd_index_chunks * cardinal.max.odd_steps_max
            };
        }

        for center in [EntryPoint::new((center_x, center_y), &self)] {
            total_walkable += if steps >= center.max.min_steps {
                if steps % 2 == 0 {
                    center.max.even_steps_max
                } else {
                    center.max.odd_steps_max
                }
            } else {
                self.count_closed_walls_walkable_after_steps(center.entry, steps)
            };
        }

        total_walkable
    }

    fn size(&self) -> (usize, usize) {
        (self.rocks[0].len(), self.rocks.len())
    }

    fn center(&self) -> (usize, usize) {
        let (size_x, size_y) = self.size();
        (size_x / 2, size_y / 2)
    }
}

impl WalledGardenState {
    fn parser(input: &str) -> IResult<&str, Self> {
        alt((
            value(WalledGardenState::Empty, char('.')),
            value(WalledGardenState::Walkable, char('S')),
            value(WalledGardenState::Rock, char('#')),
        ))(input)
    }
}

impl EntryPoint {
    fn new(entry: (usize, usize), garden: &WalledGarden) -> EntryPoint {
        EntryPoint {
            max: garden.closed_walls_find_max_walkable(entry),
            entry,
        }
    }
}

#[derive(Debug, Clone)]
struct OpenGarden {
    rockmap_size: (isize, isize),
    rocks: BTreeSet<(isize, isize)>,
    walkable: BTreeSet<(isize, isize)>,
}

impl OpenGarden {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            separated_list1(line_ending, many1(WalledGardenState::parser)),
            |walled_garden_map| OpenGarden {
                rockmap_size: (
                    walled_garden_map.len() as isize,
                    walled_garden_map[0].len() as isize,
                ),
                rocks: walled_garden_map
                    .iter()
                    .enumerate()
                    .flat_map(|(y, row)| {
                        row.iter().enumerate().filter_map(move |(x, s)| {
                            (*s == WalledGardenState::Rock).then(|| (y as isize, x as isize))
                        })
                    })
                    .collect(),
                walkable: walled_garden_map
                    .iter()
                    .enumerate()
                    .flat_map(|(y, row)| {
                        row.iter().enumerate().filter_map(move |(x, s)| {
                            (*s == WalledGardenState::Walkable).then(|| (y as isize, x as isize))
                        })
                    })
                    .collect(),
            },
        )(input)
    }

    fn count_open_walls_walkable_after_steps(&self, steps: usize) -> usize {
        let mut garden = self.clone();
        for _ in 0..steps {
            garden.next_mut();
        }
        garden.count_walkable()
    }

    fn next_mut(&mut self) {
        let walkable = mem::take(&mut self.walkable);
        self.walkable = walkable
            .iter()
            .flat_map(|(y, x)| [(y - 1, *x), (y + 1, *x), (*y, x - 1), (*y, x + 1)])
            .filter(|(y, x)| !self.is_rock(*y, *x))
            .collect();
    }

    fn is_rock(&self, y: isize, x: isize) -> bool {
        let y = y.rem_euclid(self.rockmap_size.0);
        let x = x.rem_euclid(self.rockmap_size.1);
        self.rocks.contains(&(y, x))
    }

    fn count_walkable(&self) -> usize {
        self.walkable.len()
    }
}

#[test]
fn open_matches_optimized_for_small_steps() {
    let input = fs::read_to_string("inputs/day_21.txt").unwrap();
    let walled_garden = WalledGarden::parser(&input).unwrap().1;
    let open_garden = OpenGarden::parser(&input).unwrap().1;

    let steps = 132;
    assert_eq!(
        walled_garden.count_open_walls_walkable_after_steps(steps),
        open_garden.count_open_walls_walkable_after_steps(steps)
    );
}

#[test]
fn open_matches_optimized_for_medium_steps() {
    let input = fs::read_to_string("inputs/day_21.txt").unwrap();
    let walled_garden = WalledGarden::parser(&input).unwrap().1;
    let open_garden = OpenGarden::parser(&input).unwrap().1;

    let steps = 65 + 132;
    assert_eq!(
        walled_garden.count_open_walls_walkable_after_steps(steps),
        open_garden.count_open_walls_walkable_after_steps(steps)
    );
}

#[test]
fn open_matches_optimized_for_bigger_steps() {
    let input = fs::read_to_string("inputs/day_21.txt").unwrap();
    let walled_garden = WalledGarden::parser(&input).unwrap().1;
    let open_garden = OpenGarden::parser(&input).unwrap().1;

    let steps = 270;
    assert_eq!(
        walled_garden.count_open_walls_walkable_after_steps(steps),
        open_garden.count_open_walls_walkable_after_steps(steps)
    );
}