{-# LANGUAGE LambdaCase #-}
module Main where

import Data.Functor
import Data.List
import Data.List.Split
import Data.List.Extra hiding (split)
import Debug.Trace

data ℜ where
  None :: ℜ
  Num  :: Int -> ℜ
    deriving Eq

instance Show ℜ where
  show None = show '.'
  show (Num n) = show n

parse :: String -> [ℜ]
parse s = zip [0..] (trim s) >>= \(i,a) -> replicate (read [a]) (if even i then Num (i `div` 2) else None)

sort1 :: [ℜ] -> [ℜ]
sort1 s
  | None `notElem` s = s
  | otherwise = case last s of
                  None -> sort1 $ init s
                  n    -> sort1 $ takeWhile (/= None) s ++ [n] ++ drop 1 (dropWhile (/= None) (init s))

chunks :: [ℜ] -> [[ℜ]]
chunks =
  init . foldr (\a (x:xs) -> if a `elem` x then (a:x):xs else [a]:x:xs) [[]]

fit :: [[ℜ]] -> [ℜ] -> [ℜ]
fit s t = let (p,a) = break (elem None) s
              (n,a') = span (elem None) a in
            if null n then concat s ++ t else
  concat p ++ if length (head n) >= length t then t ++ replicate (length n) None ++ concat a' else concat n ++ fit a' t

sort2 :: [ℜ] -> [ℜ]
sort2 s = let c = chunks s -- TODO: needs more work here
              fitted = fit (init c) (last c) in
            if s == fitted then s else sort2 s

solve :: ([ℜ] -> [ℜ]) -> [ℜ] -> Int
solve f s = let fixed = (\case
                            Num n -> n
                            None  -> 0) <$> f s in
  sum $ zipWith (*) fixed [0..]

main :: IO ()
main = readFile "inputs/9.example" <&> parse >>= \i ->
       print (solve sort1 i) >>
       print (solve sort2 i)