#[path = "../lib.rs"]
mod lib;

#[derive(Clone)]
enum Cell {
    Square,
    Round,
    Empty,
}

fn parse_input() -> Vec<Vec<Cell>> {
    return lib::iter_input()
        .map(|line| {
            line.bytes()
                .map(|b| match b {
                    b'#' => Cell::Square,
                    b'O' => Cell::Round,
                    b'.' => Cell::Empty,
                    _ => unreachable!(),
                })
                .collect()
        })
        .collect();
}

fn tilt<F: Fn(usize, usize) -> (usize, usize)>(
    rows: &mut Vec<Vec<Cell>>,
    width: usize,
    height: usize,
    coords: F,
) {
    for a in 0..width {
        let mut next = 0;
        for b in 0..height {
            let (x, y) = coords(a, b);
            match rows[y][x] {
                Cell::Empty => {}
                Cell::Square => next = b + 1,
                Cell::Round => {
                    rows[y][x] = Cell::Empty;
                    let (nx, ny) = coords(a, next);
                    rows[ny][nx] = Cell::Round;
                    next += 1;
                }
            }
        }
    }
}

fn cycle(rows: &mut Vec<Vec<Cell>>) {
    let width = rows[0].len();
    let height = rows.len();
    tilt(rows, width, height, |a, b| (a, b));
    tilt(rows, height, width, |a, b| (b, a));
    tilt(rows, width, height, |a, b| (a, height - 1 - b));
    tilt(rows, height, width, |a, b| (width - 1 - b, a));
}

fn hash(rows: &Vec<Vec<Cell>>) -> Vec<(usize, usize)> {
    let mut hash = vec![];
    for y in 0..rows.len() {
        for x in 0..rows[0].len() {
            match rows[y][x] {
                Cell::Round => hash.push((x, y)),
                _ => {}
            }
        }
    }
    return hash;
}

fn weigh(rows: &Vec<Vec<Cell>>) -> usize {
    let mut sum = 0;
    for y in 0..rows.len() {
        for cell in rows[y].iter() {
            match cell {
                Cell::Round => sum += rows.len() - y,
                _ => {}
            }
        }
    }
    return sum;
}

fn part1(rows: &mut Vec<Vec<Cell>>) -> usize {
    let width = rows[0].len();
    let height = rows.len();
    tilt(rows, width, height, |a, b| (a, b));
    return weigh(rows);
}

fn part2(rows: &mut Vec<Vec<Cell>>) -> usize {
    let period;
    let mut history = vec![];
    history.push(hash(&rows));
    loop {
        cycle(rows);
        let hash = hash(&rows);
        if let Some(i) = history.iter().position(|h| *h == hash) {
            period = history.len() - i;
            break;
        } else {
            history.push(hash);
        }
    }
    for _ in 0..((1_000_000_000 - history.len()) % period) {
        cycle(rows);
    }
    return weigh(rows);
}

fn main() {
    let mut rows1 = parse_input();
    let mut rows2 = rows1.clone();

    println!("part1: {}", part1(&mut rows1));
    println!("part2: {}", part2(&mut rows2));
}