use ndarray::*;


use std::fmt::Display;
use std::fs::File;
use std::io::{BufRead, BufReader, Lines};
use std::time::Instant;

// BOILERPLATE
type InputIter = Lines<BufReader<File>>;

fn get_input() -> InputIter {
    let f = File::open("input").unwrap();
    let br = BufReader::new(f);
    br.lines()
}

fn main() {
    let start = Instant::now();
    let ans1 = problem1(get_input());
    let duration = start.elapsed();
    println!("Problem 1 solution: {} [{}s]", ans1, duration.as_secs_f64());

    let start = Instant::now();
    let ans2 = problem2(get_input());
    let duration = start.elapsed();
    println!("Problem 2 solution: {} [{}s]", ans2, duration.as_secs_f64());
}

// PARSE

#[derive(Debug)]
struct Platform {
    matrix: Array2<char>,
    width: usize,
    height: usize,
}

impl<T: BufRead> From<Lines<T>> for Platform {
    fn from(lines: Lines<T>) -> Self {
        let rows: Vec<Vec<char>> = lines.map(|line| line.unwrap().chars().collect()).collect();
        let width = rows[0].len();
        let height = rows.len();
        Self {
            matrix: Array2::from_shape_vec(
                (height, width),
                rows.iter().flat_map(|row| row.iter().map(|c| *c)).collect(),
            )
            .unwrap(),
            width,
            height,
        }
    }
}

#[derive(Debug, PartialEq)]
enum Axis {
    Rows,
    Columns,
}

struct North;
struct South;
struct East;
struct West;

trait Direction {
    const REVERSED: bool;
    const AXIS: Axis;
    const ROTATIONS: usize;
}

impl Direction for North {
    const REVERSED: bool = false;
    const AXIS: Axis = Axis::Columns;
    const ROTATIONS: usize = 0;
}
impl Direction for South {
    const REVERSED: bool = true;
    const AXIS: Axis = Axis::Columns;
    const ROTATIONS: usize = 2; // 180
}
impl Direction for East {
    const REVERSED: bool = true;
    const AXIS: Axis = Axis::Rows;
    const ROTATIONS: usize = 1; // CCW
}
impl Direction for West {
    const REVERSED: bool = false;
    const AXIS: Axis = Axis::Rows;
    const ROTATIONS: usize = 3; // CW
}

impl<'a> Platform {
    fn roll<T: Direction>(&mut self) {
        // Get a view into the rotated matrix with the target direction to the north
        let mut view = self.matrix.view_mut();

        match T::ROTATIONS % 4 {
            0 => {}
            1 => {
                view.invert_axis(Axis(1));
                view.swap_axes(0, 1);
            }
            2 => {
                view.invert_axis(Axis(0));
                view.invert_axis(Axis(1));
            }
            3 => {
                view.swap_axes(0, 1);
                view.invert_axis(Axis(1));
            }
            _ => unreachable!(),
        }
        Self::roll_rocks(view);
    }
    fn roll_rocks(mut view: ArrayBase<ViewRepr<&mut char>, Dim<[usize; 2]>>) {
        let axis_len = view.len_of(Axis(1));
        for mut col in view.columns_mut() {
            for inner_idx in 1..axis_len {
                if col[inner_idx] == 'O' {
                    let lower_limit = (0..inner_idx).rev().find(|i| col[*i] == '#').unwrap_or(0);
                    if let Some(empty_pos) = (lower_limit..inner_idx).find(|i| col[*i] == '.') {
                        col.swap(inner_idx, empty_pos);
                    }
                }
            }
        }
    }
    fn score<T: Direction>(&self) -> u64 {
        match T::AXIS {
            Axis::Columns => self.score_columns::<T>(),
            Axis::Rows => self.score_rows::<T>(),
        }
    }
    fn row_or_column_score<T: Direction>(&self, idx: usize) -> usize {
        if T::REVERSED {
            idx + 1
        } else {
            match T::AXIS {
                Axis::Rows => self.width - idx,
                Axis::Columns => self.height - idx,
            }
        }
    }
    fn score_columns<T: Direction>(&self) -> u64 {
        self.matrix
            .rows()
            .into_iter()
            .enumerate()
            .map(|(idx, row)| {
                row.iter().filter(|c| **c == 'O').count() * self.row_or_column_score::<T>(idx)
            })
            .sum::<usize>() as u64
    }
    fn roll_cycle(&mut self) {
        self.roll::<North>();
        self.roll::<West>();
        self.roll::<South>();
        self.roll::<East>();
    }

    // find the first loop, return the iteration count when we first saw it and when we saw it again
    fn find_loop(&mut self) -> (usize, usize) {
        let mut first_seen = Vec::new();
        first_seen.push((self.matrix.clone(), 0));
        let mut i = 0;
        loop {
            self.roll_cycle();
            i += 1;
            if let Some((_, first_idx)) = first_seen.iter().find(|(val, _)| *val == self.matrix) {
                return (*first_idx, i);
            } else {
                first_seen.push((self.matrix.clone(), i));
            }
        }
    }
    fn score_rows<T: Direction>(&self) -> u64 {
        unimplemented!()
    }
}

impl Display for Platform {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.matrix.fmt(f)
    }
}

// PROBLEM 1 solution

fn problem1<T: BufRead>(input: Lines<T>) -> u64 {
    let mut p = Platform::from(input);
    p.roll::<North>();
    p.score::<North>()
}

// PROBLEM 2 solution
fn problem2<T: BufRead>(input: Lines<T>) -> u64 {
    const ITERATIONS: usize = 1000000000;
    let mut p = Platform::from(input);
    let first_loop = p.find_loop();
    let loop_length = first_loop.1 - first_loop.0;
    let cycles_to_skip = ((ITERATIONS - first_loop.0) / loop_length) * loop_length;
    let iterations_remaining = ITERATIONS - first_loop.0 - cycles_to_skip;

    for _ in 0..iterations_remaining {
        p.roll_cycle();
    }

    p.score::<North>()
}

#[cfg(test)]
mod tests {
    use crate::*;
    use std::io::Cursor;

    const EXAMPLE: &str = &"O....#....
O.OO#....#
.....##...
OO.#O....O
.O.....O#.
O.#..O.#.#
..O..#O..O
.......O..
#....###..
#OO..#....";

    #[test]
    fn problem1_example() {
        let c = Cursor::new(EXAMPLE);
        assert_eq!(problem1(c.lines()), 136);
    }

    #[test]
    fn problem2_example() {
        let c = Cursor::new(EXAMPLE);
        assert_eq!(problem2(c.lines()), 64);
    }
}