Day07.hs

-- |
-- Module      : AOC.Challenge.Day07
-- Copyright   : (c) Justin Le 2018
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Day 7.  See "AOC.Solver" for the types used in this module!

module AOC.Challenge.Day07 (
    day07a
  , day07b
  ) where

import           AOC.Solver         ((:~>)(..), dyno_)
import           Control.Lens
import           Control.Monad.RWS  (MonadReader(..), MonadWriter(..), MonadState(..), runRWS)
import           Data.Bifunctor     (second)
import           Data.Char          (ord, isUpper)
import           Data.Foldable      (toList, fold, find, forM_)
import           Data.List          (sortOn)
import           Data.List.NonEmpty (NonEmpty(..))
import           Data.Map           (Map)
import           Data.Semigroup     (Sum(..))
import           Data.Set           (Set)
import           Data.Set.NonEmpty  (NESet)
import           Data.Tuple         (swap)
import           Data.Witherable    (mapMaybe)
import           Numeric.Natural    (Natural)
import qualified Data.List.NonEmpty as NE
import qualified Data.Map           as M
import qualified Data.Set           as S
import qualified Data.Set.NonEmpty  as NES


parseAll :: String -> Maybe (Map Char (Set Char))
parseAll = fmap (M.fromListWith (<>) . (map . second) S.singleton)
         . traverse parseLine
         . lines

parseLine :: String -> Maybe (Char, Char)
parseLine = pack . filter isUpper
  where
    pack [_,a,b] = Just (a, b)
    pack _       = Nothing


makeDeps :: (Ord a, Ord b) => Map a (Set b) -> Map b (NESet a)
makeDeps = M.fromListWith (<>)
         . map (second NES.singleton . swap)
         . concatMap (traverse toList)
         . M.toList

findRoots :: Ord a => Map a (Set a) -> Set a
findRoots mp = cs `S.difference` targs
  where
    cs = M.keysSet mp
    targs = S.unions $ toList mp


lexicoTopo :: Ord a => Map a (Set a) -> [a]
lexicoTopo childs = go (makeDeps childs) (findRoots childs)
  where
    go deps active = flip foldMap (find (`M.notMember` deps) active) $ \c ->
      let newDeps   = mapMaybe (NES.nonEmptySet . NES.delete c) deps
          newActive = maybe id (<>) (M.lookup c childs)
                    . S.delete c
                    $ active
      in  c : go newDeps newActive

day07a :: Map Char (Set Char) :~> String
day07a = MkSol
    { sParse = parseAll
    , sShow  = id
    , sSolve = Just . lexicoTopo
    }


data Env a = Env { _envCap    :: Int
                 , _envWaiter :: a -> Natural
                 }

makeLenses ''Env

data Scheduler a = MkSched
    { _schedQueue  :: !(Set a)
    , _schedActive :: !(Map a Natural)
    }

makeClassy ''Scheduler

buildSleigh
    :: forall a m.
       ( Ord a
       , MonadState  (Scheduler a) m
       , MonadReader (Env a) m
       , MonadWriter (Sum Natural) m
       )
    => Map a (Set a)
    -> m ()
buildSleigh childs = go (findRoots childs) (makeDeps childs)
  where
    go :: Set a -> Map a (NESet a) -> m ()
    go toAdd deps = do
      popped <- stepScheduler toAdd
      forM_ (NES.nonEmptySet popped) $ \popped' ->
        let newDeps  = mapMaybe (NES.nonEmptySet . (`NES.difference` popped')) deps
            newToAdd = S.filter (`M.notMember` newDeps)
                     . foldMap (fold . (`M.lookup` childs))
                     $ popped'
        in  go newToAdd newDeps

day07b :: Map Char (Set Char) :~> Natural
day07b = MkSol
    { sParse = parseAll
    , sShow  = show
    , sSolve = \mp -> Just $
        let env = Env
              { _envCap    = dyno_ "cap" 5
              , _envWaiter = fromIntegral
                           . (+ 1)
                           . (+ dyno_ "wait" 60)
                           . subtract (ord 'A')
                           . ord
              }
        in  getSum . view _3 . runRWS (buildSleigh mp) env $ emptyScheduler
    }






-- | Scheduler Implementation


emptyScheduler :: Scheduler a
emptyScheduler = MkSched S.empty M.empty

stepScheduler
    :: ( Ord a
       , HasScheduler s a
       , MonadState  s m
       , MonadReader (Env a) m
       , MonadWriter (Sum Natural) m
       )
    => Set a
    -> m (Set a)    -- if empty, it means scheduler is exhausted
stepScheduler new = do
    cap      <- view envCap
    waiter   <- view envWaiter
    schedQueue <>= new
    numToAdd <- uses schedActive $ (cap -) . M.size
    toAdd    <- schedQueue %%= S.splitAt numToAdd
    active   <- schedActive <<>= M.fromSet waiter toAdd
    case NE.groupWith snd . sortOn snd $ M.toList active of
      [] -> pure S.empty
      toPop@((_,popTime):|_) : stillActive -> do
        schedActive .= ( M.map (subtract popTime)
                       . M.fromDistinctAscList
                       . concatMap toList
                       $ stillActive
                       )
        tell $ Sum popTime
        pure $ S.fromDistinctAscList . map fst . toList $ toPop