Day24.hs

{-# LANGUAGE OverloadedStrings #-}

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

module AOC.Challenge.Day24 (
    day24a
  , day24b
  ) where

import           AOC.Common                 ()
import           AOC.Common.Search          (exponentialFindMin)
import           AOC.Solver                 ((:~>)(..))
import           AOC.Util                   (eitherToMaybe)
import           Control.Lens               (ix, at, uses, (.~), (.=), non)
import           Control.Monad.State        (evalState)
import           Data.Char                  (isDigit, isLetter)
import           Data.Foldable              (fold)
import           Data.Function              ((&))
import           Data.Map                   (Map)
import           Data.Maybe                 (listToMaybe)
import           Data.Ord                   (Down(..))
import           Data.OrdPSQ                (OrdPSQ)
import           Data.Traversable           (forM)
import           Data.Void                  (Void)
import           Data.Witherable            (forMaybe)
import           GHC.Exts                   (sortWith)
import           Text.Megaparsec.Char.Lexer (decimal)
import qualified Data.Map                   as M
import qualified Data.OrdPSQ                as PSQ
import qualified Data.Set                   as S
import qualified Text.Megaparsec            as P
import qualified Text.Megaparsec.Char       as P

data Resist = RImmune | RWeak
  deriving (Show, Eq, Ord)

type Resistance = Map String Resist

data Team = TImm | TInf
  deriving (Show, Eq, Ord)

data Grp = G { _gHP         :: Int
             , _gResist     :: Resistance
             , _gAtk        :: Int
             , _gAtkType    :: String
             , _gInitiative :: Down Int
             , _gTeam       :: Team
             }
  deriving (Show, Eq, Ord)

type Arena = Map Grp Int

effPower :: Grp -> Int -> Int
effPower g n = _gAtk g * n

stab :: Grp -> Grp -> Int
stab g1 g2 = case M.lookup (_gAtkType g1) (_gResist g2) of
    Nothing      -> 1
    Just RImmune -> 0
    Just RWeak   -> 2

selectTargets
    :: Arena
    -> Map Grp Grp      -- ^ targets
selectTargets a = M.fromList . flip evalState candidates . forMaybe queue $ \g -> do
    targ <- uses (at (_gTeam g) . non M.empty) $ \cands -> listToMaybe
        [ h
        | (h, n) <- M.toList cands
        , let dmg = stab g h
        , dmg > 0
        , then sortWith by (Down dmg, Down (effPower h n), _gInitiative h)
        ]
    forM targ $ \t -> do
      ix (_gTeam g) . at t .= Nothing
      pure (g, t)
  where
    queue :: [Grp]
    queue = [ g
            | (g, n) <- M.toList a
            , then sortWith by (Down (effPower g n), _gInitiative g)
            ]
    candidates :: Map Team Arena
    candidates = flip M.fromSet (S.fromDistinctAscList [TImm, TInf]) $ \t ->
        M.filterWithKey (\g _ -> _gTeam g /= t) a

makeAttacks
    :: Map Grp Grp
    -> Arena
    -> Arena
makeAttacks targs a = go queue0 M.empty
  where
    go  :: OrdPSQ Grp (Down Int) Int
        -> Map    Grp            Int
        -> Arena
    go queue finished = case PSQ.minView queue of
      Nothing             -> finished
      Just (g,_,n,queue') -> case M.lookup g targs of
        Nothing   -> go queue' (M.insert g n finished)
        Just targ -> case PSQ.lookup targ queue' of
          Nothing -> case M.lookup targ finished of
            Nothing -> go queue' (M.insert g n finished)
            Just m  ->
              let totDamg   = stab g targ * n * _gAtk g
                  newM      = m - (totDamg `div` _gHP targ)
                  finished'
                    | newM > 0  = finished & ix targ .~ newM
                    | otherwise = M.delete targ finished
              in  go queue' (M.insert g n finished')
          Just (_,m) ->
              let totDamg   = stab g targ * n * _gAtk g
                  newM      = m - (totDamg `div` _gHP targ)
                  queue''
                    | newM > 0  = queue' & ix targ .~ newM
                    | otherwise = PSQ.delete targ queue'
              in  go queue'' (M.insert g n finished)
    queue0 :: OrdPSQ Grp (Down Int) Int
    queue0 = PSQ.fromList [ (g, _gInitiative g, n)
                          | (g, n) <- M.toList a
                          ]

fightBattle :: Arena -> Either Arena (Team, Map Grp Int)
fightBattle a
    | a' == a             = Left a
    | all (== TImm) teams = Right (TImm, a')
    | all (== TInf) teams = Right (TInf, a')
    | otherwise           = fightBattle a'
  where
    a' = makeAttacks (selectTargets a) a
    teams = _gTeam <$> M.keys a'

day24a :: Arena :~> Int
day24a = MkSol
    { sParse = P.parseMaybe parse24
    , sShow  = show
    , sSolve = fmap (sum . snd) . eitherToMaybe . fightBattle
    }

day24b :: Arena :~> Int
day24b = MkSol
    { sParse = P.parseMaybe parse24
    , sShow  = show
    , sSolve = \a ->
        let goodEnough i = case fightBattle (boost i a) of
              Right (TImm, b) -> Just (sum b)
              _               -> Nothing
        in  exponentialFindMin goodEnough 1   -- note: this might fail for some inputs
    }
  where
    boost :: Int -> Arena -> Arena
    boost i = M.mapKeys $ \g -> case _gTeam g of
        TImm -> g { _gAtk = _gAtk g + i }
        TInf -> g



type Parser_ = P.Parsec Void String

parse24 :: Parser_ Arena
parse24 = M.union <$> ("Immune System:" *> P.space *> teamParser TImm <* P.space)
                  <*> ("Infection:"     *> P.space *> teamParser TInf)

teamParser :: Team -> Parser_ Arena
teamParser t = M.fromList <$> (P.try (groupParser t) `P.sepEndBy1` P.newline)

groupParser :: Team -> Parser_ (Grp, Int)
groupParser _gTeam = do
    n <- decimal
    P.skipMany (P.satisfy (not . isDigit))
    _gHP <- decimal <* P.space
    "hit points" <* P.space
    _gResist <- fmap fold . P.optional . P.try $ (P.char '(' `P.between` P.char ')') resistanceParser
    P.skipMany (P.satisfy (not . isDigit))
    _gAtk <- decimal <* P.space
    _gAtkType <- P.some (P.satisfy isLetter)
    P.skipMany (P.satisfy (not . isDigit))
    _gInitiative <- Down <$> decimal
    pure (G{..}, n)

resistanceParser :: Parser_ Resistance
resistanceParser = M.unions <$> (resistSpec `P.sepBy1` (P.char ';' *> P.space))
  where
    res   = (RImmune <$ P.try "immune")
      P.<|> (RWeak   <$ P.try "weak")
    resistSpec = do
      r <- res <* P.space
      "to" <* P.space
      ts <- P.some (P.satisfy isLetter) `P.sepBy1` (P.char ',' *> P.space)
      pure . M.fromList $ (,r) <$> ts