use std::{cmp::Reverse, collections::BinaryHeap, fmt::Display};

use aoc_runner_derive::{aoc, aoc_generator};
use itertools::Itertools;

#[derive(PartialEq, Clone, Debug)]
struct Junction {
    pos: (i64, i64, i64),
}

fn squared_distance(a: &Junction, b: &Junction) -> u64 {
    // if a.pos == b.pos {
    //     0
    // } else {
    (a.pos.0 - b.pos.0).pow(2) as u64
        + (a.pos.1 - b.pos.1).pow(2) as u64
        + (a.pos.2 - b.pos.2).pow(2) as u64
    // }
}

impl Junction {
    fn squared_distance(&self, other: &Junction) -> u64 {
        squared_distance(self, other)
    }
}

impl Display for Junction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!(
            "({},{},{})",
            self.pos.0, self.pos.1, self.pos.2
        ))
    }
}

impl From<&str> for Junction {
    fn from(value: &str) -> Self {
        Self {
            pos: value
                .split(',')
                .map(|v| v.parse().unwrap())
                .collect_tuple()
                .unwrap(),
        }
    }
}

#[derive(Clone)]
struct Circuits {
    junctions: Vec<Junction>,
    circuits: Vec<Vec<usize>>,
}

impl Circuits {
    fn find_circuit(&self, junction: usize) -> usize {
        self.circuits
            .iter()
            .enumerate()
            .find(|(_i, c)| c.contains(&junction))
            .map(|(i, _c)| i)
            .unwrap()
    }
    fn merge_circuits(&mut self, a: usize, b: usize) {
        if a == b {
            return;
        }
        let mut items = std::mem::take(&mut self.circuits[b]);
        self.circuits[a].append(&mut items);
        self.circuits.remove(b);
    }

    fn connect(&mut self, a: usize, b: usize) {
        let a_circuit = self.find_circuit(a);
        let b_circuit = self.find_circuit(b);

        self.merge_circuits(a_circuit, b_circuit); // both are already in circuits, merge them
    }
}

#[aoc_generator(day8)]
fn parse(input: &str) -> Circuits {
    let junctions = input.lines().map(|l| l.into()).collect_vec();

    Circuits {
        circuits: Vec::from_iter((0..junctions.len()).map(|i| vec![i])),
        junctions,
    }
}

fn part1_impl(input: &Circuits, n: usize) -> u64 {
    let mut circuits = input.clone();
    for (a, b, _d) in circuits
        .junctions
        .iter()
        .enumerate()
        .tuple_combinations()
        .map(|((a_pos, a), (b_pos, b))| (a_pos, b_pos, a.squared_distance(b)))
        .sorted_unstable_by_key(|(_a_pos, _b_pos, d)| *d)
        .take(n)
    {
        circuits.connect(a, b)
    }

    circuits
        .circuits
        .iter()
        .map(|c| c.len())
        .sorted_unstable()
        .rev()
        .take(3)
        .reduce(|a, b| a * b)
        .unwrap() as u64
}

#[aoc(day8, part1, Sorted)]
fn part1(input: &Circuits) -> u64 {
    part1_impl(input, 1000)
}

#[aoc(day8, part2, Sorted)]
fn part2(input: &Circuits) -> u64 {
    let mut circuits = input.clone();

    for (a, b, _d) in circuits
        .junctions
        .iter()
        .enumerate()
        .tuple_combinations()
        .map(|((a_pos, a), (b_pos, b))| (a_pos, b_pos, a.squared_distance(b)))
        .sorted_unstable_by_key(|(_a_pos, _b_pos, d)| *d)
    {
        circuits.connect(a, b);
        if circuits.circuits.len() == 1 {
            return (circuits.junctions[a].pos.0 * circuits.junctions[b].pos.0) as u64;
        }
    }

    panic!()
}

#[derive(PartialEq, Eq, PartialOrd, Ord)]
struct JunctionPair {
    d: u64,
    a: usize,
    b: usize,
}

fn make_heap(circuits: &Circuits) -> BinaryHeap<Reverse<JunctionPair>> {
    BinaryHeap::from_iter(
        circuits
            .junctions
            .iter()
            .enumerate()
            .tuple_combinations()
            .map(|((a_pos, a), (b_pos, b))| {
                Reverse(JunctionPair {
                    a: a_pos,
                    b: b_pos,
                    d: a.squared_distance(b),
                })
            }),
    )
}

fn part1_heaped_impl(input: &Circuits, n: usize) -> u64 {
    let mut circuits = input.clone();
    let mut distances = make_heap(&circuits);

    for _ in 0..n {
        let pair = distances.pop().unwrap().0;
        circuits.connect(pair.a, pair.b);
    }

    circuits
        .circuits
        .iter()
        .map(|c| c.len())
        .sorted_unstable()
        .rev()
        .take(3)
        .reduce(|a, b| a * b)
        .unwrap() as u64
}

#[aoc(day8, part1, Heaped)]
fn part1_heaped(input: &Circuits) -> u64 {
    part1_heaped_impl(input, 1000)
}

#[aoc(day8, part2, Heaped)]
fn part2_heaped(input: &Circuits) -> u64 {
    let mut circuits = input.clone();
    let mut distances = make_heap(&circuits);

    while let Some(Reverse(jp)) = distances.pop() {
        circuits.connect(jp.a, jp.b);
        if circuits.circuits.len() == 1 {
            return (circuits.junctions[jp.a].pos.0 * circuits.junctions[jp.b].pos.0) as u64;
        }
    }

    panic!()
}

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

    const EXAMPLE: &str = "162,817,812
57,618,57
906,360,560
592,479,940
352,342,300
466,668,158
542,29,236
431,825,988
739,650,466
52,470,668
216,146,977
819,987,18
117,168,530
805,96,715
346,949,466
970,615,88
941,993,340
862,61,35
984,92,344
425,690,689";

    #[test]
    fn part1_example() {
        assert_eq!(part1_impl(&parse(EXAMPLE), 10), 40);
    }

    #[test]
    fn part1_heaped_example() {
        assert_eq!(part1_heaped_impl(&parse(EXAMPLE), 10), 40);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse(EXAMPLE)), 25272);
    }

    #[test]
    fn part2_heaped_example() {
        assert_eq!(part2_heaped(&parse(EXAMPLE)), 25272);
    }
}