use anyhow::{Context, Result};
use rand::seq::SliceRandom;
use rand::thread_rng;
use std::collections::HashMap;

struct DisjointSets {
    parent: Vec<usize>,
    rank: Vec<usize>,
    size: usize,
}

impl DisjointSets {
    fn new(size: usize) -> Self {
        let mut parent = Vec::with_capacity(size);
        for i in 0..size {
            parent.push(i);
        }
        let rank = vec![0; size];
        DisjointSets { parent, rank, size }
    }

    fn reset(&mut self) {
        for i in 0..self.parent.len() {
            self.parent[i] = i;
        }
        self.rank.fill(0);
        self.size = self.parent.len();
    }

    fn find(&mut self, member: usize) -> usize {
        let parent = self.parent[member];
        if parent == member {
            member
        } else {
            self.parent[member] = self.find(parent);
            self.parent[member]
        }
    }

    fn union(&mut self, u: usize, v: usize) {
        let mut pu = self.find(u);
        let mut pv = self.find(v);

        if pu != pv {
            if self.rank[pu] < self.rank[pv] {
                (pu, pv) = (pv, pu)
            }

            self.parent[pv] = pu;
            self.size -= 1;

            if self.rank[pu] == self.rank[pv] {
                self.rank[pu] += 1;
            }
        }
    }
}

fn main() -> Result<()> {
    let filename = std::env::args()
        .nth(1)
        .context("./day25 <path to puzzle input>")?;
    let input = std::fs::read_to_string(filename)?;

    let mut edges = Vec::new();
    let mut name_ids = HashMap::new();
    for line in input.lines() {
        let (from, to) = line.split_once(": ").context("malformed input.")?;
        let to = to.split(' ').collect::<Vec<_>>();
        let l = name_ids.len();
        let from_id: usize = *name_ids.entry(from).or_insert(l);
        for to in to {
            let l = name_ids.len();
            let to_id: usize = *name_ids.entry(to).or_insert(l);
            edges.push((from_id, to_id));
        }
    }

    let mut rng = thread_rng();
    let mut sets = DisjointSets::new(name_ids.len());
    let part1 = loop {
        edges.as_mut_slice().shuffle(&mut rng);
        let mut it = edges.iter();

        while sets.size > 2 {
            let edge = it.next().unwrap();
            sets.union(edge.0, edge.1);
        }

        let cut_size = it
            .filter(|edge| sets.find(edge.0) != sets.find(edge.1))
            .count();

        if cut_size == 3 {
            (0..sets.parent.len()).for_each(|i| {
                sets.find(i);
            });
            let p1 = sets.find(sets.parent[0]);
            let partition_size = sets.parent.iter().filter(|&&v| v == p1).count();
            break partition_size * (sets.parent.len() - partition_size);
        }

        sets.reset();
    };

    println!("1) {}", part1);
    println!("2) Merry Christmas");

    Ok(())
}