use aoc::*;
use itertools::Itertools;
use std::cmp::Reverse;
use std::collections::HashSet;

const INPUT: &str = include_str!("../../input/08");

fn main() {
    println!("Part 1: {}", part1(INPUT, 1000));
    println!("Part 2: {}", part2(INPUT));
}

#[test]
fn example() {
    assert_eq!(part1(include_str!("../../input/08-test"), 10), 40);
    assert_example!(part2, "08-test", 25272);
}

fn part1(input: &str, n: usize) -> usize {
    let boxes = parse_input(input);
    pairs(&boxes)
        .take(n)
        .fold(Vec::new(), connect)
        .into_iter()
        .map(|circuit| circuit.len())
        .sorted_by_key(|&l| Reverse(l))
        .take(3)
        .product()
}

fn part2(input: &str) -> usize {
    let boxes = parse_input(input);
    let pairs = pairs(&boxes);

    let mut circuits: Vec<Circuit> = Vec::new();

    for pair in pairs {
        circuits = connect(circuits, pair);
        let done = circuits.iter().any(|c| c.len() >= boxes.len());
        if done {
            return pair.0 .0[0] * pair.1 .0[0];
        }
    }

    unreachable!()
}

fn connect(mut circuits: Vec<Circuit>, (a, b): (JBox, JBox)) -> Vec<Circuit> {
    let with_a = circuits.iter_mut().position(|c| c.contains(&a));
    let with_b = circuits.iter_mut().position(|c| c.contains(&b));
    match (with_a, with_b) {
        (Some(with_a), Some(with_b)) => {
            // Merge into with_a. with_b will be empty, but whatever.
            let with_b_circuit = circuits[with_b].drain().collect_vec();
            circuits[with_a].extend(with_b_circuit);
        }
        (Some(with_a), None) => {
            circuits[with_a].insert(b);
        }
        (None, Some(with_b)) => {
            circuits[with_b].insert(a);
        }
        (None, None) => {
            circuits.push(Circuit::from([a, b]));
        }
    }
    circuits
}

fn pairs(boxes: &[JBox]) -> impl Iterator<Item = (JBox, JBox)> + use<'_> {
    boxes
        .iter()
        .copied()
        .combinations(2)
        .map(|v| (v[0], v[1]))
        .sorted_by_key(|&(a, b)| distance(a, b))
}

fn parse_input(input: &str) -> Vec<JBox> {
    input.lines().map(JBox::parse).collect()
}

fn distance(a: JBox, b: JBox) -> usize {
    let sum_squares: usize =
        a.0.iter()
            .zip(&b.0)
            .map(|(&a, &b)| a.abs_diff(b).pow(2))
            .sum();
    sum_squares.isqrt()
}

type Circuit = HashSet<JBox>;

#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
struct JBox([usize; 3]);

impl JBox {
    fn parse(line: &str) -> Self {
        let mut nums = line.split(',').map(|n| n.parse().unwrap());
        Self([
            nums.next().unwrap(),
            nums.next().unwrap(),
            nums.next().unwrap(),
        ])
    }
}