use aoc2024::matrix;
use itertools::Itertools;
use std::collections::{HashMap, HashSet};
use std::fmt;

type Freq = char;

enum Dot {
    Empty,
    Antenna(Freq),
}
use Dot::*;

impl fmt::Display for Dot {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let ch = match self {
            Empty => &'.',
            Antenna(c) => c,
        };
        write!(f, "{}", ch)
    }
}

impl Dot {
    fn parse(c: char) -> Dot {
        match c {
            '.' => Empty,
            c => Antenna(c),
        }
    }
}

struct M {
    matrix: matrix::Matrix<Dot>,
    antennas: HashMap<Freq, HashSet<matrix::Pos>>,
    antinodes: HashSet<matrix::Pos>,
}

impl M {
    fn new(text: &str) -> Self {
        let matrix = matrix::Matrix::new(text, Dot::parse);
        let mut antennas: HashMap<Freq, HashSet<matrix::Pos>> = HashMap::new();
        for i in 0..matrix.limit.0 {
            for j in 0..matrix.limit.1 {
                match matrix.get((i, j)) {
                    Empty => (),
                    Antenna(freq) => {
                        antennas.entry(*freq).or_default().insert((i, j));
                    }
                }
            }
        }
        M {
            matrix,
            antennas,
            antinodes: HashSet::new(),
        }
    }

    fn compute_antinodes(&mut self, version: u8) {
        for (_freq, antennas) in self.antennas.iter() {
            // println!("Freq {freq}: {antennas:?}");
            for (a1, a2) in antennas
                .into_iter()
                .permutations(2)
                .map(|x| x.into_iter().collect_tuple().unwrap())
            {
                match version {
                    0 => {
                        let an = antinode(a1, a2, &self.matrix.limit);
                        // println!("  ({a1:?}, {a2:?}) -> {an:?}");
                        if let Some((x, y)) = an {
                            self.antinodes.insert((x, y));
                        }
                    }
                    1 => {
                        for an in antinodes(a1, a2, &self.matrix.limit) {
                            self.antinodes.insert(an);
                        }
                    }
                    _ => panic!(),
                }
            }
        }
        // println!("  {:?}", self.antinodes);
    }
}

fn antinodes(a1: &matrix::Pos, a2: &matrix::Pos, limit: &matrix::Pos) -> Vec<matrix::Pos> {
    let mut result = vec![*a1, *a2];
    let mut a: matrix::Pos = result[0];
    let mut b: matrix::Pos = result[1];
    while let Some(an) = antinode(&a, &b, limit) {
        result.push(an.clone());
        a = result[result.len() - 2];
        b = result[result.len() - 1];
    }
    result
}

fn antinode(a1: &matrix::Pos, a2: &matrix::Pos, limit: &matrix::Pos) -> Option<matrix::Pos> {
    let (x1, y1) = (a1.0 as isize, a1.1 as isize);
    let (x2, y2) = (a2.0 as isize, a2.1 as isize);
    let x = x2 - (x1 - x2);
    let y = y2 - (y1 - y2);
    if y >= 0 && x >= 0 && y < limit.1 as isize && x < limit.0 as isize {
        Some((x as usize, y as usize))
    } else {
        None
    }
}

impl fmt::Display for M {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // write!(f, "{}[2J", 27 as char)?;
        write!(f, "{}", self.matrix.to_string())
    }
}

fn p1(input: &str) -> String {
    let mut m: M = M::new(input);
    m.compute_antinodes(0);

    let result = m.antinodes.len();
    result.to_string()
}

fn p2(input: &str) -> String {
    let mut m: M = M::new(input);
    m.compute_antinodes(1);

    let result = m.antinodes.len();
    result.to_string()
}

fn main() {
    aoc2024::run_day("8", p1, p2);
}