#! /usr/bin/env -S"ANSWER=42" nix-shell
#! nix-shell -p ghcid
#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [pretty-simple])"
#! nix-shell -i "ghcid -c 'ghci' -T main"

{-# OPTIONS_GHC -Wall -Wincomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wno-unused-top-binds -Wno-unused-imports -Wno-type-defaults #-}
{-# LANGUAGE OverloadedStrings #-}

import Data.List
import Data.Monoid
import Debug.Trace (trace)
import Text.Pretty.Simple

smallExample :: [Int]
smallExample = [16,10,15,5,1,11,7,19,6,12,4]

biggerExample :: [Int]
biggerExample = [28,33,18,42,31,14,46,20,48,47,24,23,49,45,19,38,39,11,1,32,25
                ,35,8,17,7,9,4,2,34,10,3]

solvePart1 :: [Int] -> Int
solvePart1 xs = finally
  $ span (== 1)
  $ sort
  $ map (\(x,y) -> y - x)
  $ zip (0:sorted) (sorted ++ [l+3])
  where
    l = last sorted
    sorted = sort xs
    finally (x,y) = length x * length y

--------------------------------------------------------------------------------
-- 1 -> 1 possibility
-- 1 1 -> 2 poss
-- 1 1 1 -> 4 poss
-- 1 1 1 1 -> 6 poss
-- 1 1 1 1 1 -> 8 poss
--
--------------------------------------------------------------------------------
cleanup :: [[Int]] -> [[Int]]
cleanup = filter (\v -> (any (/= 3) v) && (length v > 1))

est :: [[Int]] -> [Int]
est = map combinations
  where
    combinations v = (2 ^ l) - (2 ^ (l `div` 3))
      where
        l = (length v) - 1

estimatePermutations :: [Int] -> Int
estimatePermutations = getProduct . foldMap Product . est . cleanup . group

distanceMap :: [Int] -> [Int]
distanceMap xs = map (\(x,y) -> y - x) $ zip (0:sorted) (sorted ++ [(last sorted)+3])
  where
    sorted = sort xs

solvePart2 :: [Int] -> Int
solvePart2 = estimatePermutations . distanceMap
  where

main :: IO ()
main = do
  input' <- lines <$> readFile "day10/input"
  let input = read <$> input'
  putStrLn ":: Tests"
  print $ smallExample
  print $ biggerExample
  putStrLn ":: Day 10 - Part 1"
  print $ solvePart1 smallExample
  print $ solvePart1 biggerExample
  print $ solvePart1 input
  putStrLn ":: Tests"

  putStrLn ":: Day 10 - Part 2"
  print $ distanceMap smallExample
  print $ solvePart2 smallExample
  print $ distanceMap biggerExample
  print $ (est . cleanup . group) $ distanceMap biggerExample
  print $ solvePart2 biggerExample
  print $ solvePart2 input