//! Day 11

use std::collections::{HashMap, HashSet};

use grid::Grid;

use crate::util::grid::Pos;

pub const DAY: u8 = 11;
pub const TITLE: &str = "Cosmic Expansion";

struct Universe {
    g: Grid<bool>,
    empty_rows: HashSet<usize>,
    empty_cols: HashSet<usize>,
}

impl Universe {
    fn parse(input: Vec<String>) -> Self {
        let mut g = Grid::new(input.len(), input[0].len());

        for (y, line) in input.iter().enumerate() {
            for (x, c) in line.chars().enumerate() {
                if c == '#' {
                    g[(y, x)] = true;
                }
            }
        }

        let empty_rows = g
            .iter_rows()
            .enumerate()
            .filter_map(|(y, mut row)| Some(y).filter(|_| row.all(|f| !f)))
            .collect();
        let empty_cols = g
            .iter_cols()
            .enumerate()
            .filter_map(|(x, mut col)| Some(x).filter(|_| col.all(|f| !f)))
            .collect();

        Self {
            g,
            empty_rows,
            empty_cols,
        }
    }

    fn print(&self) {
        for mut row in self.g.iter_rows() {
            row.for_each(|f| {
                if !f {
                    print!(".");
                } else {
                    print!("#");
                }
            });
            println!();
        }
    }

    /// Get the positions of all galaxies
    fn index(&self) -> HashMap<u32, Pos> {
        let mut index = HashMap::new();
        let mut n = 0;
        for (y, mut row) in self.g.iter_rows().enumerate() {
            for (x, f) in row.enumerate() {
                if *f {
                    index.insert(n, Pos::usize(x, y));
                    n += 1;
                }
            }
        }
        index
    }

    /// Calculate distance between 2 galaxies.
    fn distance(&self, p1: Pos, p2: Pos, expansion: i64) -> i64 {
        let (m1, n1) = (p1.y.min(p2.y), p1.x.min(p2.x));
        let (m2, n2) = (p1.y.max(p2.y), p1.x.max(p2.x));

        // Empty rows between g1 and g2
        let er = (m1..m2)
            .filter(|i| self.empty_rows.contains(&(*i as usize)))
            .count() as i64;

        // Empty cols between g1 and g2
        let ec = (n1..n2)
            .filter(|j| self.empty_cols.contains(&(*j as usize)))
            .count() as i64;

        let nr = (m2 - m1) - er;
        let nc = (n2 - n1) - ec;

        nr + nc + (expansion * (er + ec))
    }
}

fn solve(input: Vec<String>, expansion: i64) -> String {
    let mut gal = Universe::parse(input);
    let ix = gal.index();

    ix.iter()
        .flat_map(|(i, g1)| {
            ix.iter()
                .filter(move |(j, _)| j > &i)
                .map(|(_, g2)| gal.distance(*g1, *g2, expansion))
        })
        .sum::<i64>()
        .to_string()
}

pub fn part1(input: Vec<String>) -> String {
    solve(input, 2)
}

pub fn part2(input: Vec<String>) -> String {
    solve(input, 1_000_000)
}

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

    #[test]
    fn t_example1() {
        // possibly wrong example: answer 374
        assert_eq!(part1(crate::read_example(DAY, 1)), "538");
    }

    #[test]
    fn t_part1() {
        assert_eq!(part1(crate::read_input(DAY)), "9509330");
    }

    #[test]
    fn t_example2() {
        assert_eq!(part2(crate::read_example(DAY, 2)), "164000210");
    }

    #[test]
    fn t_part2() {
        assert_eq!(part2(crate::read_input(DAY)), "635832237682");
    }
}