module Day8 where

import Util
import Debug.Trace
import Data.Ord
import Data.Tuple (swap)
import Data.List (tails, sortBy, sort)
import Data.Map (Map)
import qualified Data.Map as Map

parseLine :: [String] -> Point3
parseLine [a,b,c] = (read a, read b, read c)

parse :: String -> [Point3]
parse s =
  map (parseLine . split ',') $ lines s

dist2 :: Point3 -> Point3 -> Int
dist2 (x,y,z) (x',y',z') = (x-x') ^ 2 + (y-y') ^ 2 + (z-z') ^ 2

pairs :: [a] -> [(a,a)]
pairs xs = [(x, y) | (x:ys) <- tails xs, y <- ys]

merge :: (Point3, Point3) -> Map Point3 Int -> Map Point3 Int
merge (p, q) m =
  let pid = m Map.! p
      qid = m Map.! q
   in if pid == qid then
        m
      else
        Map.map (\c -> if c == pid then qid else c) m

part1 :: String -> String
part1 input =
  let boxes = parse input
      connectionOrder =
        dropWhile (\(p,q) -> p == q) $
        sortBy (comparing (\(p,q) -> dist2 p q)) $
        pairs boxes
      circuits =
        foldl (\ids' edge -> merge edge ids')
          (Map.fromList $ zip boxes [1..])
          (take 1000 connectionOrder)
      sizes =
        reverse $
        sort $
        map snd $
        Map.toList $
        Map.foldl (\sizes' c -> Map.insertWith (+) c 1 sizes')
          Map.empty
          circuits
   in show $ product $ take 3 sizes

part2 :: String -> String
part2 input =
  let boxes = parse input
      connectionOrder =
        dropWhile (\(p,q) -> p == q) $
        sortBy (comparing (\(p,q) -> dist2 p q)) $
        pairs boxes
      circuits =
        scanl (\ids' edge -> merge edge ids')
          (Map.fromList $ zip boxes [1..])
          connectionOrder
      merges =
        filter (\((p,q),ids) -> ids Map.! p /= ids Map.! q)
	  $ zip connectionOrder circuits
      ((p,q),_) = merges !! (length boxes - 2)
      (x,_,_) = p
      (x',_,_) = q
   in show (x * x')