use crate::common::geometry::Point3;
use crate::solvers::Solver;
use std::collections::{BTreeMap, HashSet};
use std::io::BufRead;

pub struct Playground {
    connections: usize,
}

impl Playground {
    const CIRCUITS_TAKE: usize = 3;

    pub fn new(connections: Option<usize>) -> Self {
        Self {
            connections: connections.unwrap_or(1000),
        }
    }

    // Collects all junction box coordinates.
    fn junction_boxes<R: BufRead>(reader: R) -> Vec<Point3> {
        let mut junction_boxes = Vec::new();
        for line in reader.lines().map_while(Result::ok) {
            junction_boxes.push(Point3::from_line(&line));
        }
        junction_boxes
    }

    // Maps the closest pairs of junction boxes to their distances, up to the maximum number of
    // connections.
    fn distances<'a>(
        &self,
        junction_boxes: &'a Vec<Point3>,
    ) -> BTreeMap<u64, (&'a Point3, &'a Point3)> {
        let mut distances = BTreeMap::new();
        for i in 0..junction_boxes.len() {
            let (_, right) = junction_boxes.split_at(i + 1);
            for b in right {
                let d = junction_boxes[i].distance_sqr(b);
                distances.insert(d, (&junction_boxes[i], b));
                if distances.len() > self.connections {
                    distances.pop_last();
                }
            }
        }
        distances
    }

    // Aggregates all circuits and returns the product of the sizes of the 3 largest ones.
    fn circuit_sizes(distances: &BTreeMap<u64, (&Point3, &Point3)>) -> u64 {
        let mut circuits: Vec<HashSet<&Point3>> = Vec::new();
        for junction_boxes in distances.values() {
            let index_a = circuits.iter().position(|c| c.contains(&junction_boxes.0));
            let index_b = circuits.iter().position(|c| c.contains(&junction_boxes.1));
            if let Some(index_a) = index_a {
                if let Some(index_b) = index_b {
                    if index_a != index_b {
                        let (low, high) = if index_a < index_b {
                            (index_a, index_b)
                        } else {
                            (index_b, index_a)
                        };
                        let mut c = circuits.remove(high);
                        circuits[low].extend(c.drain());
                    }
                } else {
                    circuits[index_a].insert(junction_boxes.1);
                }
            } else {
                if let Some(index_b) = index_b {
                    circuits[index_b].insert(junction_boxes.0);
                } else {
                    let mut set = HashSet::new();
                    set.insert(junction_boxes.0);
                    set.insert(junction_boxes.1);
                    circuits.push(set);
                }
            }
        }
        circuits.sort_by(|a, b| b.len().cmp(&a.len()));
        circuits
            .iter()
            .take(Self::CIRCUITS_TAKE)
            .map(|c| c.len() as u64)
            .product()
    }
}

impl Solver for Playground {
    const PUZZLE_INDEX: u8 = 8;
    const PUZZLE_NAME: &'static str = "Playground";

    fn process_data<R: BufRead>(&self, reader: R) -> (u64, u64) {
        let junction_boxes = Self::junction_boxes(reader);
        let distances = self.distances(&junction_boxes);

        (Self::circuit_sizes(&distances), 0)
    }
}