0 | module Days.Day6

  2 | import Data.String

  3 | import Data.Vect

  4 | import Data.Fin

  5 | import Data.SortedSet

  6 | import Data.List.Lazy

  7 | import Control.Monad.State

  8 | import System

 10 | import Util

 11 | %hide Util.Direction

 15 | data Direction = Up | Down | Left | Right

 16 | %name Direction dir

 18 | Show Direction where

 19 |   show Up = "Up"

 20 |   show Down = "Down"

 21 |   show Left = "Left"

 22 |   show Right = "Right"

 24 | Eq Direction where

 25 |   (==) Up Up = True

 26 |   (==) Up _ = False

 27 |   (==) Down Down = True

 28 |   (==) Down _ = False

 29 |   (==) Left Left = True

 30 |   (==) Left _ = False

 31 |   (==) Right Right = True

 32 |   (==) Right _ = False

 34 | Ord Direction where

 35 |   (<) Up Up = False

 36 |   (<) Up _ = True

 37 |   (<) Down Up = False

 38 |   (<) Down Down = False

 39 |   (<) Down _ = True

 40 |   (<) Left Left = False

 41 |   (<) Left Right = True

 42 |   (<) Left _ = False

 43 |   (<) Right Right = False

 44 |   (<) Right _ = True

 46 | turnRight : Direction -> Direction

 47 | turnRight Up = Right

 48 | turnRight Right = Down

 49 | turnRight Down = Left

 50 | turnRight Left = Up

 52 | data Cell = Empty | Visited | Obstruction

 53 | %name Cell cell, cell2

 55 | visited : Cell -> Bool

 56 | visited Visited = True

 57 | visited _ = False

 59 | charToCell : Char -> Cell

 60 | charToCell '#' = Obstruction

 61 | charToCell _ = Empty

 63 | Grid : (rows : Nat) -> (cols : Nat) -> Type

 64 | Grid rows cols = Vect (S rows) (Vect (S cols) Cell)

 66 | step : {rows, cols : Nat} -> Direction -> Coord rows cols -> Maybe (Coord rows cols)

 67 | step Up (FZ, col) = Nothing

 68 | step Up ((FS x), col) = Just (weaken x, col)

 69 | step Down (row, col) =

 70 |   if row < last then Just $ (finS row, col) else Nothing

 71 | step Left (row, FZ) = Nothing

 72 | step Left (row, (FS x)) = Just (row, weaken x)

 73 | step Right (row, col) =

 74 |   if col < last then Just $ (row, finS col) else Nothing

 76 | isObstruction : Coord rows cols -> Grid rows cols -> Bool

 77 | isObstruction (row, col) xs =

 78 |   let row' = row `index`xs

 79 |   in case col `index` row' of

 80 |     Obstruction => True

 81 |     _ => False

 83 | visitCell : Coord rows cols -> Grid rows cols -> Grid rows cols

 84 | visitCell (row, col) xs =

 85 |   let row' = row `index` xs

 86 |       row'' = replaceAt col Visited row'

 87 |   in replaceAt row row'' xs

 89 | isVisit : Grid rows cols -> Coord rows cols -> Bool

 90 | isVisit xs (row, col) =

 91 |   let row' = row `index` xs

 92 |       cell = col `index` row'

 93 |   in case cell of

 94 |     Visited => True

 95 |     _ => False

 97 | obstructCell : Coord rows cols -> Grid rows cols -> Grid rows cols

 98 | obstructCell (row, col) xs =

 99 |   let row' = row `index` xs

100 |       row'' = replaceAt col Obstruction row'

101 |   in replaceAt row row'' xs

103 | record Room where

104 |   constructor MkRoom

105 |   rows, cols : Nat

106 |   grid : Grid rows cols

107 |   guard_pos : Coord rows cols

108 |   guard_dir : Direction

109 | %name Room room

112 | advanceGuard : Room -> Maybe Room

113 | advanceGuard (MkRoom rows cols grid guard_pos guard_dir) = do

114 |   -- Get the nominal next location, returning nothing if this runs out of bounds

115 |   next_pos <- step guard_dir guard_pos

116 |   -- Figure out if we are obstructed

117 |   case isObstruction next_pos grid of

118 |     -- Visit the next cell and step in our current direction

119 |     False =>

120 |       let next_grid = visitCell next_pos grid

121 |       in Just $ MkRoom rows cols next_grid next_pos guard_dir

122 |     -- Turn 90 degrees and try again

123 |     True =>

124 |       advanceGuard (MkRoom rows cols grid guard_pos (turnRight guard_dir))

127 | advanceAll : Room -> Room

128 | advanceAll room =

129 |   case advanceGuard room of

130 |     Nothing => room

131 |     Just next => advanceAll next

133 | Cache : Type

134 | Cache = SortedSet (Nat, Nat, Direction)

137 | andvanceAll' : Room -> (Room, Bool)

138 | andvanceAll' room =

139 |   snd . runState empty $ (advanceAll'' room)

140 |   where

141 |     advanceAll'' : Room -> State Cache (Room, Bool)

142 |     advanceAll'' room = do

143 |       cache <- get

144 |       let (row, col) = room.guard_pos

145 |       let (row, col) = (finToNat row, finToNat col)

146 |       let key = (row, col, room.guard_dir)

147 |       if contains key cache

148 |         then pure (room, True)

149 |         else do

150 |           modify (insert key)

151 |           case advanceGuard room of

152 |             Nothing => pure (room, False)

153 |             Just next => advanceAll'' next

155 | visitedPositions : Room -> Nat

156 | visitedPositions (MkRoom rows cols grid guard_pos guard_dir) =

157 |   count visited . Data.Vect.concat $ grid

159 | visitedSet : Room -> SortedSet (Nat, Nat)

160 | visitedSet (MkRoom rows cols grid guard_pos guard_dir) =

161 |   let canidates =

162 |     map (\(x, y) => (finToNat x, finToNat y)) . toList $ filter (isVisit grid) $ allCords

163 |   in fromList canidates

165 | loops : Room -> List (Nat, Nat)

166 | loops room@(MkRoom rows cols grid guard_pos guard_dir) =

167 |   -- Todo: filter out coordinates that aren't visited

168 |   let

169 |     visits = visitedSet . advanceAll $ room

170 |     canidates = filter (\x => not $ x == guard_pos) $ allCords

171 |     canidates = filter (\(x, y) => contains (finToNat x, finToNat y) visits) canidates

172 |   in

173 |     toList .

174 |     map (\(x, y) => (finToNat x, finToNat y)) .

175 |     filter

176 |       (snd .

177 |         andvanceAll' .

178 |         (\x => MkRoom rows cols x guard_pos guard_dir) .

179 |         (\x => obstructCell x grid)

180 |       ) $

181 |     canidates

184 | Show Room where

185 |   show room@(MkRoom rows cols grid guard_pos guard_dir) =

186 |     "Grid with dimensions: \{show rows}, \{show cols}\n" ++

187 |     "Guard position: \{show guard_pos}\n" ++

188 |     "Guard direction: \{show guard_dir}\n" ++

189 |     "Visited positions: \{show $ visitedPositions room}"

191 | showGrid : Room -> String

192 | showGrid (MkRoom rows cols grid guard_pos guard_dir) =

193 |   unlines . toList . map (pack . toList . map cellToChar) $ grid

194 |   where

195 |     cellToChar : Cell -> Char

196 |     cellToChar Empty = '.'

197 |     cellToChar Visited = 'X'

198 |     cellToChar Obstruction = '#'

200 | parseInput : String -> Maybe Room

201 | parseInput str = do

202 |   (rows ** cols ** char_grid) <- stringTo2D str

203 |   let (grid, guard_rows) = unzip . map parseRow $ char_grid

204 |   guard_row <- findIndex isJust guard_rows

205 |   guard_col <- join . find isJust $ guard_rows

206 |   let guard_pos = (guard_row, guard_col)

207 |   -- visit the inital position

208 |   let grid = visitCell guard_pos grid

209 |   Just $ MkRoom rows cols grid guard_pos Up

210 |   where

211 |     parseRow : Vect (S cols) Char -> (Vect (S cols) Cell, Maybe (Fin (S cols)))

212 |     parseRow xs =

213 |       let guard_col = findIndex (=='^') xs

214 |           grid_row = map charToCell xs

215 |       in (grid_row, guard_col)

218 | part1 : String -> IO (String, Room)

219 | part1 str = do

220 |   Just room <- pure $ parseInput str

221 |     | Nothing => do putStrLn "Failed to parse input"

222 |                     exitFailure

223 |   let stepped = advanceAll room

224 |   pure $ (show . visitedPositions $ stepped, room)

227 | part2 : Room -> String -> IO String

228 | part2 room _ = do

229 |   pure . show . length $ loops room