{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ViewPatterns #-}
module Main where

import Data.Either (partitionEithers)
import Data.Foldable (toList)
import Data.Function (on)
import Data.List
import qualified Text.Parsec as P

import Input
import Util

data IVMap a = IVSplit (a, a) (IVMap a) (IVMap a) | IVLeaf (a, a)
  deriving (Show)

ivmFromList :: Ord a => [(a, a)] -> IVMap a
ivmFromList =
    (\case l -> splitIntoTree (length l) l)
    . foldr mergeInto []
    . map ((,) <$> fst . head <*> maximum . map snd)
    . groupBy ((==) `on` fst)
    . sort
  where
    mergeInto :: Ord a => (a, a) -> [(a, a)] -> [(a, a)]
    mergeInto (a, b) [] = [(a, b)]
    mergeInto (a, b) ((c, d) : l) | b >= c = mergeInto (a, d) l
                                  | otherwise = (a, b) : (c, d) : l

    topRange :: IVMap a -> (a, a)
    topRange (IVSplit p _ _) = p
    topRange (IVLeaf p) = p

    splitIntoTree :: Ord a => Int -> [(a, a)] -> IVMap a
    splitIntoTree _ [] = error "splitIntoTree: empty list"
    splitIntoTree _ [p] = IVLeaf p
    splitIntoTree n l =
        let (pre, post) = splitAt (n `div` 2) l
            t1 = splitIntoTree (n `div` 2) pre
            t2 = splitIntoTree ((n + 1) `div` 2) post
        in IVSplit (fst (topRange t1), snd (topRange t2)) t1 t2

ivmContains :: Ord a => IVMap a -> a -> Bool
ivmContains (IVLeaf (a, b)) x = a <= x && x <= b
ivmContains (IVSplit (a, b) t1 t2) x | a <= x && x <= b = ivmContains t1 x || ivmContains t2 x
                                     | otherwise = False


distinctElemSAT :: Eq a => [[a]] -> [a]
distinctElemSAT =
    map (\case [x] -> x ; [] -> error "No assignment" ; _ -> error "Ambiguous?")
    . fixpoint eliminateSingles
  where
    fixpoint :: Eq a => (a -> a) -> a -> a
    fixpoint f x = let y = f x in if y == x then x else fixpoint f y

    eliminateSingles :: Eq a => [[a]] -> [[a]]
    eliminateSingles l = map (\row -> if length row == 1 then row else filter (`notElem` [x | [x] <- l]) row) l


data Rule = Rule String [(Int, Int)] (IVMap Int)
  deriving (Show)

parseTicket :: String -> [Int]
parseTicket = map read . toList . splitOn (== ',')

parseRule :: String -> Rule
parseRule = (\case Right (n,ivs) -> Rule n ivs (ivmFromList ivs)
                   Left e -> error (show e))
            . P.parse pRule ""
  where
    pRule = (,) <$> P.many1 (P.satisfy (/= ':')) <*> (P.string ": " >> pRange `P.sepBy` P.string " or ")
    pRange = (,) <$> (read <$> P.many1 P.digit) <*> (P.char '-' >> read <$> P.many1 P.digit)


main :: IO ()
main = do
    input <- getInput 16
    let [map parseRule -> rules, [""],
         ["your ticket:", parseTicket -> mytick], [""],
         "nearby tickets:" : (map parseTicket -> nbticks)] =
            groupBy (\a b -> not (null a || null b)) input

        anyValidIVM = ivmFromList (concat [ps | Rule _ ps _ <- rules])

        (concat -> invalids, valids) = partitionEithers
            [case partition (ivmContains anyValidIVM) tick of
               (yeses, []) -> Right yeses
               (_, nos) -> Left nos
            | tick <- nbticks]

    print (sum invalids)

    let fieldNames = distinctElemSAT
            [[name | Rule name _ ivm <- rules, all (ivmContains ivm) column]
            | column <- transpose valids]

    print (product [value | (value, name) <- zip mytick fieldNames, take 10 name == "departure "])