use std::collections::HashMap;

use aoc_runner_derive::{aoc, aoc_generator};

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
enum Rock {
    Round,
    Square,
    None
}
trait RockConvertable {
    fn to_rock(&self) -> Rock;
}

impl RockConvertable for char {
    fn to_rock(&self) -> Rock {
        match self {
            'O' => Rock::Round,
            '#' => Rock::Square,
            '.' => Rock::None,
            _ => panic!("Invalid rock char")
        }
    }
}

#[aoc_generator(day14)]
fn parse(input: &str) -> Vec<Vec<Rock>> {
    let input: Vec<Vec<Rock>> = input.lines()
        .map(|line| {
            line.chars().map(|char| {
                char.to_rock()
            }).collect::<Vec<_>>()
        }).collect::<Vec<_>>();
    input
}

#[aoc(day14, part1)]
fn part1(input: &Vec<Vec<Rock>>) -> usize {
    let mut load: usize = 0;
    for j in 0..input[0].len() {
        let mut cur_weight = input[0].len() + 1;

        for i in 0..input.len() {
            match input[i][j] {
                Rock::Round => {
                    cur_weight -= 1;
                    load += cur_weight;
                },
                Rock::Square => {
                    cur_weight = input[0].len() - i;
                },
                Rock::None => continue
            }
        }
    }
    load
}

#[aoc(day14, part2)]
fn part2(input: &Vec<Vec<Rock>>) -> usize {
    let mut seen_at = HashMap::new();
    let mut next_y = Vec::<usize>::new();
    let mut map = input.to_vec();

    // Functions to get a Rock using a rotated coordinate space
    let cycle_parts: &[(Box<dyn Fn(&mut Vec<Vec<Rock>>, usize, usize) -> &mut Rock>, _, _); 4] = &[
        (Box::new(|map: &mut Vec<Vec<Rock>>, x: usize, y: usize| &mut map[y][x]),
         map[0].len(), map.len()),
        (Box::new(|map: &mut Vec<Vec<Rock>>, x: usize, y: usize| &mut map[x][y]),
         map.len(), map[0].len()),
        (Box::new(|map: &mut Vec<Vec<Rock>>, x: usize, y: usize| { let h = map.len(); &mut map[h - 1 - y][x] }),
         map[0].len(), map.len()),
        (Box::new(|map: &mut Vec<Vec<Rock>>, x: usize, y: usize| { let w = map[0].len(); &mut map[x][w - 1 - y] }),
         map.len(), map[0].len()),
    ];

    let mut cycle = 0;
    const END: u32 = 1000000000;
    while cycle < END {
        // Handle tilts in each direction
        for (getter, width, height) in cycle_parts {
            next_y.clear();
            next_y.resize(*width, 0);
            for y in 0..*height {
                for x in 0..*width {
                    let item = getter(&mut map, x, y);
                    match *item {
                        Rock::None => {}
                        Rock::Square => {
                            next_y[x] = y + 1;
                        }
                        Rock::Round => {
                            *item = Rock::None;
                            *getter(&mut map, x, next_y[x]) = Rock::Round;
                            next_y[x] += 1;
                        }
                    }
                }
            }
        }

        // More compact representation of the current state, for saving in hashmap
        let key = map.iter().enumerate().flat_map(|(y, line)| {
            line.iter().enumerate().filter_map(move |(x, &ref rock)| {
                if rock == &Rock::Round {
                    Some((x as u8, y as u8))
                } else {
                    None
                }
            })
        }).collect::<Vec<_>>();

        cycle += 1;
        if let Some(seen_at_cycle) = seen_at.insert(key, cycle) {
            // Current state was identical to one we'd already seen, we can skip forward
            let diff = cycle - seen_at_cycle;
            let remaining = END - cycle;
            let skipped = remaining / diff * diff;
            cycle += skipped;
        }
    }

    let height = map.len();

    map.into_iter().enumerate().flat_map(|(y, line)| line.into_iter().filter_map(move |rock| {
        if rock == Rock::Round {
            Some(height - y)
        } else {
            None
        }
    })).sum()
}


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

    const EX: &str = r"O....#....
O.OO#....#
.....##...
OO.#O....O
.O.....O#.
O.#..O.#.#
..O..#O..O
.......O..
#....###..
#OO..#....";

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

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