{-# LANGUAGE NamedFieldPuns #-}

module Day16 where

import AoC

import Data.List
import Data.Maybe
import qualified Data.Set as S
import Data.Set (Set)
import Text.ParserCombinators.ReadP
import Text.Read.Lex

data Field = Field { name :: String
                   , valids :: Set Int
                   } deriving (Show)

type Ticket = [Int]

data Input = Input { fields :: [Field]
                   , myTicket :: Ticket
                   , otherTickets :: [Ticket]
                   } deriving (Show)

line :: ReadP a -> ReadP a
line p = p <* char '\n'

parseField :: ReadP Field
parseField = do
    name <- many1 $ satisfy (/=':')
    string ": "
    ranges <- parseRange `sepBy` string " or "
    let valids = S.fromList . concat . map (uncurry enumFromTo) $ ranges
    return $ Field { name, valids }
    where parseRange = do
            n1 <- readDecP
            char '-'
            n2 <- readDecP
            return (n1, n2)

parseTicket :: ReadP Ticket
parseTicket = readDecP `sepBy` char ','

parseInput :: ReadP Input
parseInput = do
    fields <- many1 $ line parseField
    line $ string ""
    line $ string "your ticket:"
    myTicket <- line parseTicket
    line $ string ""
    line $ string "nearby tickets:"
    otherTickets <- many1 $ line parseTicket
    return $ Input {fields, myTicket, otherTickets }

matchesAnyField :: [Field] -> Int -> Bool
matchesAnyField fields n = any ((n `S.member`) . valids) fields

iterateUntilNothing :: (a -> Maybe a) -> a -> a
iterateUntilNothing f x = maybe x (iterateUntilNothing f) $ f x

data Resolve a = Resolved a | Choice (Set a) deriving (Show)

resolve :: Ord a => [[a]] -> Maybe [a]
resolve = sequence . map fromResolved . iterateUntilNothing resolve' . map (Choice . S.fromList)
    where fromResolved (Resolved x) = Just x
          fromResolved _ = Nothing

findSingletonChoice :: [Resolve a] -> Maybe (a, [Resolve a])
findSingletonChoice [] = Nothing
findSingletonChoice (x:xs) =
    case x of
      Choice s | S.size s == 1 -> let [r] = S.elems s
                                  in  Just (r, Resolved r : xs)
      _ -> fmap (\(r, rest) -> (r, x:rest)) $ findSingletonChoice xs

resolve' :: Ord a => [Resolve a] -> Maybe [Resolve a]
resolve' xs = do
    (r, xs') <- findSingletonChoice xs
    return $ map (dropChoice r) xs'
    where dropChoice x (Choice s) = Choice (S.delete x s)
          dropChoice _ y = y

part1 :: Input -> Int
part1 Input { fields, otherTickets } = sum . filter (not . matchesAnyField fields) . concat $ otherTickets

part2 :: Input -> Int
part2 Input { fields, myTicket, otherTickets } = product departureValues
    where validTickets = filter (all $ matchesAnyField fields) otherTickets
          columns = transpose validTickets
          fieldNames = fromJust . resolve . map findPossibleFields $ columns
          departureIndices = findIndices ("departure" `isPrefixOf`) fieldNames
          departureValues = map (myTicket !!) departureIndices
          findPossibleFields col = map name . filter (\field -> all (`S.member` valids field) col) $ fields

main = runAoC (fromJust . oneCompleteResult parseInput) part1 part2