use itertools::Itertools;
use std::fmt::Debug;

pub fn run() {
    let input = include_str!("../input/day11.txt");
    // let input = r#"...#......
    // .......#..
    // #.........
    // ..........
    // ......#...
    // .#........
    // .........#
    // ..........
    // .......#..
    // #...#....."#;
    let image: Image = input.into();
    let expansions = image.expand();

    dbg!(calculate_distances(&image, &expansions, 2));
    dbg!(calculate_distances(&image, &expansions, 1_000_000));
}

fn calculate_distances(image: &Image, expansions: &Expansions, factor: usize) -> usize {
    image
        .iter_galaxies()
        .combinations(2)
        .map(|v| expansions.distance(factor, v[0], v[1]))
        .sum::<usize>()
}

#[derive(Clone)]
struct Image {
    chars: Vec<Vec<char>>,
    width: usize,
    height: usize,
}

impl From<&str> for Image {
    fn from(value: &str) -> Self {
        let chars = value
            .trim()
            .lines()
            .map(|l| l.trim().chars().collect::<Vec<_>>())
            .collect::<Vec<_>>();
        let width = chars[0].len();
        let height = chars.len();
        Self {
            chars,
            width,
            height,
        }
    }
}

impl Debug for Image {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let s = self
            .chars
            .iter()
            .map(|row| row.iter().collect::<String>())
            .collect::<Vec<_>>()
            .join("\n");
        write!(f, "Image(\n{})", s)
    }
}

impl Image {
    fn at(&self, col: usize, row: usize) -> char {
        if col < self.chars[0].len() && row < self.chars.len() {
            self.chars[row][col]
        } else {
            '.'
        }
    }

    fn expand(&self) -> Expansions {
        let rows = (0..self.height)
            .rev()
            .filter(|&row| (0..self.width).all(|col| self.at(col, row) == '.'))
            .collect::<Vec<_>>();

        let cols = (0..self.width)
            .rev()
            .filter(|&col| (0..self.height).all(|row| self.at(col, row) == '.'))
            .collect::<Vec<_>>();

        Expansions { rows, cols }
    }

    fn iter_galaxies<'a>(&'a self) -> impl Iterator<Item = (usize, usize)> + 'a {
        (0..self.height)
            .into_iter()
            .flat_map(|row| (0..self.width).into_iter().map(move |col| (col, row)))
            .filter(|(col, row)| self.at(*col, *row) == '#')
    }
}

#[derive(Debug, Clone)]
struct Expansions {
    rows: Vec<usize>,
    cols: Vec<usize>,
}

impl Expansions {
    fn distance(&self, factor: usize, (x1, y1): (usize, usize), (x2, y2): (usize, usize)) -> usize {
        let min_x = x1.min(x2);
        let max_x = x1.max(x2);
        let min_y = y1.min(y2);
        let max_y = y1.max(y2);
        let expanded_rows = self
            .rows
            .iter()
            .filter(|row| (min_y..=max_y).contains(row))
            .count();
        let expanded_cols = self
            .cols
            .iter()
            .filter(|col| (min_x..=max_x).contains(col))
            .count();
        (max_x - min_x) + (max_y - min_y) + (expanded_rows + expanded_cols) * (factor - 1)
    }
}