module Part2 where

import Data.Map ((!), Map, foldl, empty, filterWithKey, insert, member, fromList)
import Commons


type Intervals = Map Char (Int, Int)

applyRule :: Rule -> Intervals -> (String, Intervals, Intervals)
applyRule rule i
  | r rule == '*'     = (destination rule, i, empty)
  | member (r rule) i = let (start, end) = i ! r rule
                        in case s rule of
                                '>' -> if start > v rule
                                       then (destination rule, i, empty)
                                       else if end <= v rule
                                            then (destination rule, empty, i)
                                            else (destination rule, insert (r rule) (v rule + 1, end) i,
                                                                    insert (r rule) (start, v rule) i)
                                '<' -> if end < v rule
                                       then (destination rule, i, empty)
                                       else if start >= v rule
                                            then (destination rule, empty, i)
                                            else (destination rule, insert (r rule) (start, v rule - 1) i,
                                                                    insert (r rule) (v rule, end) i)
  | otherwise         = (destination rule, empty, i)

applyRules :: [Rule] -> Intervals -> [(String, Intervals)]
applyRules [] _ = []
applyRules (h: t) i = let (destination, accepted, left) = applyRule h i
                      in if destination == "R" then applyRules t left
                         else (destination, accepted): applyRules t left

applyWorkflow :: Workflows -> [(String, Intervals)] -> Int
applyWorkflow _ [] = 0
applyWorkflow workflows ((wName, i): t) | wName == "A" = Data.Map.foldl (\ a (start, end) -> a * (end - start + 1)) 1 i
                                                         + applyWorkflow workflows t
                                        | otherwise    = let afterRules = applyRules (workflows ! wName) i
                                                         in applyWorkflow workflows afterRules
                                                            + applyWorkflow workflows t

getNCombinations :: Workflows -> Int
getNCombinations w = applyWorkflow w [("in", fromList [('x', (1, 4000)), ('m', (1, 4000)), ('a', (1, 4000)),
                                                       ('s', (1, 4000))])]