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

use crate::{Problem, Solution};

pub struct Day11;

impl Problem for Day11 {
    const DAY: u8 = 11;

    const INPUT: &'static str = include_str!("../input/day_11.txt");
}

impl Solution for Day11 {
    type Answer1 = usize;

    type Answer2 = usize;

    fn part_1(input: &str) -> anyhow::Result<Self::Answer1> {
        solve(input, 2)
    }

    fn part_2(input: &str) -> anyhow::Result<Self::Answer2> {
        solve(input, 1000000)
    }
}

fn solve(input: &str, rate: usize) -> anyhow::Result<usize> {
    let mut graph: Graph = input.parse()?;

    graph.expand(rate);

    let mut nodes = graph.into_nodes();

    let mut edges = HashMap::new();
    while let Some(node) = nodes.pop_first() {
        for other in nodes.iter() {
            edges.insert((node, *other), node.dist(other));
        }
    }

    Ok(edges.values().sum())
}

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct Node {
    id: usize,
    row: usize,
    col: usize,
}

impl Node {
    fn dist(&self, other: &Node) -> usize {
        self.row.abs_diff(other.row) + self.col.abs_diff(other.col)
    }
}

#[derive(Debug, Default)]
struct Graph {
    inner: Vec<Vec<bool>>,
    row_values: Vec<usize>,
    col_values: Vec<usize>,
}

impl Graph {
    fn rows(&self) -> usize {
        self.inner.len()
    }

    fn cols(&self) -> usize {
        self.inner.first().map(|v| v.len()).unwrap_or_default()
    }

    fn expand(&mut self, rate: usize) {
        let mut rows = vec![false; self.rows()];
        let mut cols = vec![false; self.cols()];
        for row in 0..self.rows() {
            for col in 0..self.cols() {
                rows[row] = rows[row] || self.inner[row][col];
                cols[col] = cols[col] || self.inner[row][col];
            }
        }

        self.row_values = rows
            .iter()
            .map(|not_expanded| (!not_expanded).then_some(rate - 1).unwrap_or_default())
            .collect();

        self.col_values = cols
            .iter()
            .map(|not_expanded| (!not_expanded).then_some(rate - 1).unwrap_or_default())
            .collect();
    }

    fn into_nodes(self) -> BTreeSet<Node> {
        let mut nodes = BTreeSet::new();
        let mut id = 0;

        for (row, v) in self.inner.into_iter().enumerate() {
            for (col, value) in v.into_iter().enumerate() {
                if value {
                    id += 1;
                    nodes.insert(Node {
                        id,
                        row: row + self.row_values[0..row].iter().sum::<usize>(),
                        col: col + self.col_values[0..col].iter().sum::<usize>(),
                    });
                }
            }
        }

        nodes
    }
}

impl std::str::FromStr for Graph {
    type Err = std::convert::Infallible;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(Self {
            inner: s
                .lines()
                .map(|s| s.chars().map(|c| c == '#').collect::<Vec<_>>())
                .collect::<Vec<_>>(),
            ..Default::default()
        })
    }
}

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

    const INPUT: &str = indoc::indoc! {"
        ...#......
        .......#..
        #.........
        ..........
        ......#...
        .#........
        .........#
        ..........
        .......#..
        #...#.....
    "};

    #[test]
    fn test_part_1() -> anyhow::Result<()> {
        Ok(assert_eq!(374, Day11::part_1(INPUT)?))
    }

    #[test]
    fn test_part_2() -> anyhow::Result<()> {
        Ok(assert_eq!(8410, solve(INPUT, 100)?))
    }
}