
   0 | module Days.Day8
   1 | 
   2 | import Data.String
   3 | import Data.Vect
   4 | import Data.List.Lazy
   5 | import Control.Monad.State
   6 | import System
   7 | 
   8 | import Util
   9 | 
  10 | data Cell : Type where
  11 |   Empty : Cell
  12 |   Antenna : (freq : Char) -> Cell
  13 | 
  14 | charToCell : Char -> Cell
  15 | charToCell '.' = Empty
  16 | charToCell c = Antenna c
  17 | 
  18 | cellToChar : Cell -> Char
  19 | cellToChar Empty = '.'
  20 | cellToChar (Antenna freq) = freq
  21 | 
  22 | cellToAntenna : Cell -> Maybe Char
  23 | cellToAntenna Empty = Nothing
  24 | cellToAntenna (Antenna freq) = Just freq
  25 | 
  26 | record Input where
  27 |   constructor MkInput
  28 |   rows, cols : Nat
  29 |   grid : Grid rows cols Cell
  30 |   antinodes : Grid rows cols Bool
  31 | %name Input input
  32 | 
  33 | isAntenna : Coord rows cols -> Grid rows cols Cell -> Bool
  34 | isAntenna x xs =
  35 |   case x `index` xs of
  36 |     Empty => False
  37 |     Antenna _ => True
  38 | 
  39 | antinodePositions : (input : Input) -> LazyList (Coord input.rows input.cols)
  40 | antinodePositions (MkInput rows cols grid antinodes) =
  41 |   map fst . filter snd . flat $ antinodes
  42 | 
  43 | markAntinode : (input : Input) -> Coord input.rows input.cols -> Input
  44 | markAntinode (MkInput rows cols grid antinodes) x =
  45 |   let antinodes = replaceAt x True antinodes
  46 |   in MkInput rows cols grid antinodes
  47 | 
  48 | Show Input where
  49 |   show input =
  50 |     let char_grid = mapGrid cellToChar input.grid
  51 |         (char_grid, _) = runState char_grid (markAntinodes (antinodePositions input))
  52 |     in gridToString char_grid
  53 |     where
  54 |       markAntinodes : LazyList (Coord input.rows input.cols)
  55 |                     -> State (Grid input.rows input.cols Char) ()
  56 |       markAntinodes [] = pure ()
  57 |       markAntinodes (x :: xs) = do
  58 |         unless (isAntenna x input.grid) (modify $ replaceAt x '#')
  59 |         markAntinodes xs
  60 | 
  61 | uniqueAntennas : Input -> List Char
  62 | uniqueAntennas input =
  63 |   nub . toList . mapMaybe cellToAntenna . map snd . flat $ input.grid
  64 | 
  65 | ||| Locate antennas of the given frequency
  66 | locateAntennas : (freq : Char) -> (input : Input) -> LazyList (Coord input.rows input.cols)
  67 | locateAntennas freq input =
  68 |   map fst . filter ((== Just freq) . cellToAntenna . snd) . flat $ input.grid
  69 | 
  70 | coordToIntegers : Coord rows cols -> (Integer, Integer)
  71 | coordToIntegers (x, y) = (finToInteger x, finToInteger y)
  72 | 
  73 | integersToCoords : {rows, cols : Nat} -> (Integer, Integer) -> Maybe (Coord rows cols)
  74 | integersToCoords (x, y) = do
  75 |   x <- integerToFin x (S rows)
  76 |   y <- integerToFin y (S cols)
  77 |   Just (x, y)
  78 | 
  79 | addVector : (Integer, Integer) -> (Integer, Integer) -> (Integer, Integer)
  80 | addVector (x1, y1) (x2, y2) = (x1 + x2, y1 + y2)
  81 | 
  82 | subVector : (Integer, Integer) -> (Integer, Integer) -> (Integer, Integer)
  83 | subVector (x1, y1) (x2, y2) = (x1 - x2, y1 - y2)
  84 | 
  85 | negVector : (Integer, Integer) -> (Integer, Integer)
  86 | negVector (x, y) = (-1 * x, -1 * y)
  87 | 
  88 | ||| Mark a coordinate as an antinode, if it exists
  89 | markLocation : (Integer, Integer) -> State Input ()
  90 | markLocation x = do
  91 |   input <- get
  92 |   case integersToCoords x of
  93 |     Nothing => pure ()
  94 |     Just x => put $ markAntinode input x
  95 | 
  96 | ||| Marks a location and all others inline with it according to the given offset
  97 | markLocation' : (loc, offset : (Integer, Integer)) -> State Input ()
  98 | markLocation' loc offset = do
  99 |   input <- get
 100 |   case integersToCoords loc of
 101 |     Nothing => pure ()
 102 |     Just x => do
 103 |       put $ markAntinode input x
 104 |       markLocation' (loc `addVector` offset) offset
 105 | 
 106 | ||| Mark the antinodes for a pair of antennas of the same frequency
 107 | ||| Works in integer space to make the type mangling easier
 108 | markPair : (a1, a2 : (Integer, Integer)) -> State Input ()
 109 | markPair a1 a2 =
 110 |   let diff = a1 `subVector` a2
 111 |       node1 = a1 `addVector` diff
 112 |       node2 = a2 `subVector` diff
 113 |   in do
 114 |     markLocation node1
 115 |     markLocation node2
 116 | 
 117 | markPair' : (a1, a2 : (Integer, Integer)) -> State Input ()
 118 | markPair' a1 a2 =
 119 |   let offset = a1 `subVector` a2
 120 |   in do
 121 |     markLocation' a1 offset
 122 |     markLocation' a1 (negVector offset)
 123 | 
 124 | ||| Mark the antinodes of all pairs of antennas of a given frequency
 125 | markFrequency : (freq : Char) -> State Input ()
 126 | markFrequency freq = do
 127 |   input <- get
 128 |   let antennas = map coordToIntegers $ locateAntennas freq input
 129 |   let pairs = allPairs antennas
 130 |   markPairs pairs
 131 |   where
 132 |     markPairs : LazyList ((Integer, Integer), (Integer, Integer)) -> State Input ()
 133 |     markPairs [] = pure ()
 134 |     markPairs ((x, y) :: xs) = do
 135 |       markPair x y
 136 |       markPairs xs
 137 | 
 138 | markFrequency' : (freq : Char) -> State Input ()
 139 | markFrequency' freq = do
 140 |   input <- get
 141 |   let antennas = map coordToIntegers $ locateAntennas freq input
 142 |   let pairs = allPairs antennas
 143 |   markPairs pairs
 144 |   where
 145 |     markPairs : LazyList ((Integer, Integer), (Integer, Integer)) -> State Input ()
 146 |     markPairs [] = pure ()
 147 |     markPairs ((x, y) :: xs) = do
 148 |       markPair' x y
 149 |       markPairs xs
 150 | 
 151 | ||| Mark all the antiodes in a given input
 152 | markAll : State Input ()
 153 | markAll = do
 154 |   input <- get
 155 |   let frequencies = uniqueAntennas input
 156 |   traverse_ markFrequency frequencies
 157 | 
 158 | markAll' : State Input ()
 159 | markAll' = do
 160 |   input <- get
 161 |   let frequencies = uniqueAntennas input
 162 |   traverse_ markFrequency' frequencies
 163 | 
 164 | parseInput : String -> Maybe Input
 165 | parseInput str = do
 166 |   (rows ** cols ** grid) <- stringTo2D str
 167 |   let grid = mapGrid charToCell grid
 168 |   Just $ MkInput rows cols grid (const False)
 169 | 
 170 | export
 171 | part1 : String -> IO (String, Input)
 172 | part1 str = do
 173 |   Just input <- pure $ parseInput str
 174 |     | _ => do
 175 |       putStrLn "Failed parsing input"
 176 |       exitFailure
 177 |   let (stepped, _) = runState input markAll
 178 |   pure $ (show . count (const True) $ (antinodePositions stepped), input)
 179 | 
 180 | export
 181 | part2 : Input -> String -> IO String
 182 | part2 input _ = do
 183 |   let (stepped, _) = runState input markAll'
 184 |   pure . show . count (const True) $ antinodePositions stepped
 185 |