src/advent15/advent15.hs

   1 import Data.List
   2 -- import Data.Tuple (swap)
   3 import Data.Maybe
   4 import qualified Data.Foldable
   5 import Data.Foldable (forM_)
   6 import qualified Data.Map.Strict as M
   7 import Data.Map.Strict ((!))
   8 import qualified Data.Set as S
   9 import qualified Data.PQueue.Prio.Min as P
  10 import qualified Data.Sequence as Q
  11 import Data.Sequence ((|>), (<|), Seq (Empty, (:<|), (:|>)))
  12
  13 import Debug.Trace
  14
  15 type Coord = (Int, Int) -- row, column
  16 type HitPoints = Int
  17 data Species = Elf | Goblin deriving (Show, Eq)
  18 data Agent = Agent Species HitPoints deriving (Show, Eq)
  19 type Agents = M.Map Coord Agent
  20 type Layout = S.Set Coord
  21 type OrderedLayout = Q.Seq Coord
  22
  23 type Closed = Layout
  24 data Agendum = Agendum {_current :: Coord, _trail :: OrderedLayout, _cost :: Int} deriving (Show, Eq)
  25 type Agenda = P.MinPQueue (Int, Int) Agendum
  26
  27 newGoblin = Agent Goblin 200
  28 newElf = Agent Elf 200
  29
  30 isSpecies s (Agent s' _) = s == s'
  31
  32 isElf = isSpecies Elf
  33 isGoblin = isSpecies Goblin
  34
  35 otherSpecies Elf = Goblin
  36 otherSpecies Goblin = Elf
  37
  38 elfs :: Agents -> Agents
  39 elfs agents = M.filter isElf agents
  40
  41 goblins :: Agents -> Agents
  42 goblins agents = M.filter isGoblin agents
  43
  44
  45
  46 test1 = "#######\n#E..G.#\n#...#.#\n#.G.#G#\n#######\n"
  47 test2 = "#########\n#G..G..G#\n#.......#\n#.......#\n#G..E..G#\n#.......#\n#.......#\n#G..G..G#\n#########\n"
  48 test3 = "#######\n#.G...#\n#...EG#\n#.#.#G#\n#..G#E#\n#.....#\n#######\n"
  49
  50 main :: IO ()
  51 main = do
  52     text <- readFile "data/advent15.txt"
  53     let (layout, agents) = parse test3
  54     print layout
  55     print agents
  56     putStrLn $ showWorld layout agents
  57     print $ pathsToEnemies (1, 1) layout agents
  58     -- let (h, e) = M.findMin $ elfs agents
  59     -- print (h, e)
  60     -- forM_ (M.assocs agents) $ \(ah, a) -> do
  61     --     print (ah, a)
  62     --     print $ pathsToEnemies ah layout agents
  63     --     print $ bestMove ah layout agents
  64     let a1 = doRound layout agents
  65     putStrLn $ showWorld layout a1
  66     print a1
  67     let a2 = doRound layout a1
  68     putStrLn $ showWorld layout a2
  69     print a2
  70     let a3 = doRound layout a2
  71     putStrLn $ showWorld layout a3
  72     print a3
  73     let a4 = doRound layout a3
  74     putStrLn $ showWorld layout a4
  75     print a4
  76     let a5 = doRound layout a4
  77     putStrLn $ showWorld layout a5
  78     print a5
  79     let a6 = doRound layout a5
  80     putStrLn $ showWorld layout a6
  81     print a6
  82
  83
  84 showWorld layout agents = unlines rows
  85     where rows = map (showRow layout agents) [minRow..maxRow]
  86           minRow = fst $ S.findMin layout'
  87           maxRow = fst $ S.findMax layout'
  88           layout' = S.union layout $ S.fromList $ M.keys agents
  89
  90 showRow layout agents row = map (\col -> showCell (row, col) layout agents) [minCol..maxCol]
  91     where minCol = minimum $ map snd $ S.toList layout'
  92           maxCol = maximum $ map snd $ S.toList layout'
  93           layout' = S.union layout $ S.fromList $ M.keys agents
  94
  95 showCell c layout agents =
  96     if c `M.member` agents
  97     then if isElf a then 'E' else 'G'
  98     else if c `S.member` layout then '.' else '\x2593'
  99     where a = agents!c
 100
 101
 102 -- Parsing
 103 parse :: String -> (Layout, Agents)
 104 parse text = foldl' parseRow (S.empty, M.empty) $ zip [0..] $ lines text
 105
 106 parseRow :: (Layout, Agents) -> (Int, String) -> (Layout, Agents)
 107 parseRow (layout, agents) (r, row) = foldl' parseCellWithY (layout, agents) $ zip [0..] row
 108     where parseCellWithY = parseCell r
 109
 110 parseCell :: Int -> (Layout, Agents) -> (Int, Char) -> (Layout, Agents)
 111 parseCell r (layout, agents) (c, cell) =
 112     let here = (r, c)
 113     in case cell of
 114             'G' -> (S.insert here layout, M.insert here newGoblin agents)
 115             'E' -> (S.insert here layout, M.insert here newElf agents)
 116             '.' -> (S.insert here layout, agents)
 117             _   -> (layout, agents)
 118
 119 -- Locations
 120
 121 adjacent :: Coord -> Layout -> Layout
 122 adjacent (r, c) layout = S.intersection layout
 123                          $ S.fromList [(r+1, c), (r-1, c), (r, c+1), (r, c-1)]
 124
 125 free :: Coord -> Layout -> Agents -> Bool
 126 free here layout agents = (here `S.member` layout) && (here `M.notMember` agents)
 127
 128 adjacentFree :: Coord -> Layout -> Agents -> Layout
 129 adjacentFree here layout agents = S.filter (\c -> free c layout agents) (adjacent here layout)
 130
 131 orderedAdjacentFree :: Coord -> Layout -> Agents -> OrderedLayout
 132 orderedAdjacentFree here layout agents = Q.sort $ S.foldl' (|>) Q.empty cells
 133     where cells = adjacentFree here layout agents
 134
 135
 136 -- Searching
 137
 138 initAgenda :: Coord -> Coord -> Agenda
 139 initAgenda start goal = P.singleton ((estimateCost start goal), (fst start * 100 + snd start)) Agendum {_current = start, _trail = Q.empty, _cost = 0}
 140
 141 aStar :: Coord -> Layout -> Agents -> Agenda -> Closed -> Maybe Agendum
 142 aStar goal layout agents agenda closed
 143     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 144     | P.null agenda = Nothing
 145     | otherwise =
 146         if reached == goal then Just currentAgendum
 147         else if reached `S.member` closed
 148             then aStar goal layout agents (P.deleteMin agenda) closed
 149             else aStar goal layout agents newAgenda closed'
 150         where
 151             (_, currentAgendum) = P.findMin agenda
 152             reached = _current currentAgendum
 153             closed' = S.insert reached closed
 154             tieBreakerCost a = foldl' (\t (r, c) -> t + r * 100 + c) 0 ((_current a) <| (_trail a))
 155             newAgenda = foldl' (\q a -> P.insert ((estimatedCost a), tieBreakerCost a) a q) (P.deleteMin agenda) $ candidates layout agents currentAgendum closed'
 156             estimatedCost agendum = estimateCost reached goal + _cost agendum
 157
 158 candidates :: Layout -> Agents -> Agendum -> Closed -> Q.Seq Agendum
 159 candidates layout agents agendum closed = newCandidates
 160     where
 161         candidate = _current agendum
 162         previous = _trail agendum
 163         succs = orderedAdjacentFree candidate layout agents
 164         nonloops = Q.filter (\s -> s `S.notMember` closed) succs
 165         newCandidates = fmap (\n -> makeAgendum n) nonloops
 166         makeAgendum new = Agendum { _current = new,
 167                                     _trail = candidate <| previous,
 168                                     _cost = _cost agendum + 1}
 169
 170 estimateCost :: Coord -> Coord -> Int
 171 estimateCost (r, c) (gr, gc) = abs (r - gr) + abs(c - gc)
 172
 173 -- Move selection
 174
 175 shortestDistanceTo :: Coord -> Coord -> Layout -> Agents -> Maybe (Int, OrderedLayout)
 176 shortestDistanceTo here there layout agents =
 177     if searchResult == Nothing
 178     then Nothing
 179     else Just (_cost $ fromJust searchResult, (_current $ fromJust searchResult) <| (_trail $ fromJust searchResult))
 180     where searchResult = aStar there layout agents (initAgenda here there) S.empty
 181
 182 enemyLocations :: Coord -> Agents -> Layout
 183 enemyLocations here agents = S.fromList $ M.keys $ M.filter (isSpecies enemySpecies) agents
 184     where Agent thisSpecies _ = agents!here
 185           enemySpecies = otherSpecies thisSpecies
 186
 187 agentTargets :: Coord -> Layout -> Agents -> Layout
 188 agentTargets here layout agents = S.foldl S.union S.empty enemyAdjacents
 189     where enemies = enemyLocations here agents
 190           enemyAdjacents = S.map (\l -> adjacentFree l layout agents) enemies
 191
 192 pathsToEnemies :: Coord -> Layout -> Agents -> [(Int, OrderedLayout)]
 193 pathsToEnemies here layout agents = catMaybes $ map sdt $ S.toList targets
 194     where sdt there = shortestDistanceTo here there layout agents
 195           targets = agentTargets here layout agents
 196
 197 closestEnemies :: Coord -> Layout -> Agents -> [OrderedLayout]
 198 closestEnemies here layout agents = possibles
 199     where paths = pathsToEnemies here layout agents
 200           closest = minimum $ map fst paths
 201           possibles = map snd $ filter (\p -> fst p == closest) paths
 202
 203 bestMove :: Coord -> Layout -> Agents -> Coord
 204 bestMove here layout agents =
 205     if null paths
 206     then here
 207     else head $ sort $ map pathStep paths
 208     where paths = closestEnemies here layout agents
 209           pathStep p = if Q.length p > 1 then Q.index p (Q.length p - 2) else Q.index p 1
 210
 211 makeBestMove :: Coord -> Layout -> Agents -> Agents
 212 makeBestMove here layout agents = M.insert there agent $ M.delete here agents
 213     where agent = agents!here
 214           there = bestMove here layout agents
 215
 216 -- Attacking
 217
 218 bestTarget :: Coord -> Layout -> Agents -> Coord
 219 bestTarget here layout agents = keyOfMinHP $ M.filterWithKey (\c _ -> c `S.member` enemies) agents
 220     where enemies = touchingEnemies here layout agents
 221
 222 attack :: Coord -> Agents -> Agents
 223 attack target agents = M.insert target (Agent species (hp - 3)) agents
 224     where Agent species hp = agents!target
 225
 226 keyOfMinHP :: Agents -> Coord -- Ord b => M.Map a b -> a
 227 keyOfMinHP m = fst $ M.foldrWithKey mergeKV (M.findMin m) m
 228     where mergeKV k (Agent s v) (bestK, (Agent sb bestV)) =
 229             if v < bestV then (k, (Agent s v)) else (bestK, (Agent sb bestV))
 230
 231 makeAttack :: Coord -> Layout -> Agents -> Agents
 232 makeAttack here layout agents = attack target agents
 233     where target = bestTarget here layout agents
 234
 235 -- Game loop
 236
 237 doRound :: Layout -> Agents -> Agents
 238 doRound layout agents = agents'
 239     -- where agents' = foldl' (\a h -> makeBestMove h layout a) agents $ M.keys agents
 240     where agents' = foldl' (\a h -> agentAction h layout a) agents $ M.keys agents
 241
 242 touchingEnemies :: Coord -> Layout -> Agents -> Layout
 243 touchingEnemies here layout agents = S.intersection neighbourhood enemies
 244     where neighbourhood = adjacent here layout
 245           enemies = enemyLocations here agents
 246
 247 agentAction :: Coord -> Layout -> Agents -> Agents
 248 agentAction here layout agents =
 249     if S.null targets
 250     then makeBestMove here layout agents
 251     else makeAttack here layout agents
 252     where targets = touchingEnemies here layout agents