#! /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 possibilities
-- 1 1 1 -> 4 possibilities
-- 1 1 1 1 -> 7 (!) possibilities
-- 1 1 1 1 1 -> ? possibilities
--
--------------------------------------------------------------------------------
cleanup :: [[Int]] -> [[Int]]
cleanup = filter (\v -> (any (/= 3) v) && (length v > 1))


-- That is the actual problem right here:
-- ··· -> 2^3 - 1 = 7
-- ···· -> 2^4 - 1 - 2 = 13
-- ····· -> 2^5 - 1 - 2 - 3 = 26
combinations :: [Int] -> Int
combinations v =
    (2 ^ n)           -- All the ways to take any on/off combination in n
  - (sum [1..(n-2)])  -- All the ways to take n+ consecutive in n (illegal),
                      -- *except* for n=1 and n=2 which are legal
  where
    n = (length v) - 1

estimatePermutations :: [Int] -> Int
estimatePermutations = getProduct
                     . foldMap Product
                     . map combinations
                     . 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 $ (map combinations . cleanup . group) $ distanceMap biggerExample
  print $ solvePart2 biggerExample
  print $ solvePart2 input