From a01b1f6ac7ad62094449a5096ea9472f77d7c7b4 Mon Sep 17 00:00:00 2001
From: Alex Jercan <jercan_alex27@yahoo.com>
Date: Sat, 16 Dec 2023 10:34:03 +0200
Subject: [PATCH] day 16

---
 aoc.sh            |  3 ++
 aoc2023.cabal     |  1 +
 src/Day16.hs      | 94 ++++++++++++++++++++++++++++++++++++++++++++++-
 test/Main.hs      |  2 +
 test/TestDay16.hs | 22 +++++++++++
 5 files changed, 120 insertions(+), 2 deletions(-)
 create mode 100644 test/TestDay16.hs

diff --git a/aoc.sh b/aoc.sh
index a18bb18..6ae36fa 100755
--- a/aoc.sh
+++ b/aoc.sh
@@ -72,3 +72,6 @@ stack exec day14 < input/day14.input
 
 echo -e "${IGreen}--- Day 15: Lens Library ---${Color_Off}"
 stack exec day15 < input/day15.input
+
+echo -e "${IYellow}--- Day 16: The Floor Will Be Lava ---${Color_Off}"
+stack exec day16 < input/day16.input
diff --git a/aoc2023.cabal b/aoc2023.cabal
index 65cd193..98732a8 100644
--- a/aoc2023.cabal
+++ b/aoc2023.cabal
@@ -395,6 +395,7 @@ test-suite spec
       TestDay13
       TestDay14
       TestDay15
+      TestDay16
       Paths_aoc2023
   autogen-modules:
       Paths_aoc2023
diff --git a/src/Day16.hs b/src/Day16.hs
index 964d3dd..0f76b36 100644
--- a/src/Day16.hs
+++ b/src/Day16.hs
@@ -1,10 +1,100 @@
 module Day16 (main, part1, part2) where
 
+import qualified Data.Map as M
+import qualified Data.Set as S
+
+data Cell = Empty | MirrorS | MirrorB | SplitterH | SplitterV deriving (Show, Eq)
+data Dir = N | E | S | W deriving (Show, Eq, Ord)
+
+type Grid = (M.Map (Int, Int) Cell, Int, Int)
+type Pos = (Int, Int)
+
+parseCell :: Char -> Cell
+parseCell '.' = Empty
+parseCell '/' = MirrorS
+parseCell '\\' = MirrorB
+parseCell '-' = SplitterH
+parseCell '|' = SplitterV
+parseCell _ = error "Invalid cell"
+
+parse :: String -> Grid
+parse input = (grid, w, h)
+  where
+    grid =
+        M.fromList $
+            concat $
+                zipWith (\row cells -> zipWith (\col cell -> ((row, col), parseCell cell)) [0 ..] cells) [0 ..] $
+                    lines input
+    w = length $ head $ lines input
+    h = length $ lines input
+
+move :: Pos -> Dir -> (Pos, Dir)
+move (row, col) N = ((row - 1, col), N)
+move (row, col) E = ((row, col + 1), E)
+move (row, col) S = ((row + 1, col), S)
+move (row, col) W = ((row, col - 1), W)
+
+mirrorS :: Pos -> Dir -> (Pos, Dir)
+mirrorS (row, col) N = ((row, col + 1), E)
+mirrorS (row, col) E = ((row - 1, col), N)
+mirrorS (row, col) S = ((row, col - 1), W)
+mirrorS (row, col) W = ((row + 1, col), S)
+
+mirrorB :: Pos -> Dir -> (Pos, Dir)
+mirrorB (row, col) N = ((row, col - 1), W)
+mirrorB (row, col) E = ((row + 1, col), S)
+mirrorB (row, col) S = ((row, col + 1), E)
+mirrorB (row, col) W = ((row - 1, col), N)
+
+splitterH :: Pos -> Dir -> [(Pos, Dir)]
+splitterH (row, col) N = [((row, col - 1), W), ((row, col + 1), E)]
+splitterH (row, col) E = [((row, col + 1), E)]
+splitterH (row, col) S = [((row, col - 1), W), ((row, col + 1), E)]
+splitterH (row, col) W = [((row, col - 1), W)]
+
+splitterV :: Pos -> Dir -> [(Pos, Dir)]
+splitterV (row, col) N = [((row - 1, col), N)]
+splitterV (row, col) E = [((row - 1, col), N), ((row + 1, col), S)]
+splitterV (row, col) S = [((row + 1, col), S)]
+splitterV (row, col) W = [((row - 1, col), N), ((row + 1, col), S)]
+
+valid :: Pos -> Grid -> Bool
+valid (row, col) (_, w, h) = row >= 0 && row < h && col >= 0 && col < w
+
+step :: Grid -> (Pos, Dir) -> [(Pos, Dir)]
+step grid@(m, _, _) (p, dir) = filter ((`valid` grid) . fst) $ case m M.! p of
+    Empty -> [move p dir]
+    MirrorS -> [mirrorS p dir]
+    MirrorB -> [mirrorB p dir]
+    SplitterH -> splitterH p dir
+    SplitterV -> splitterV p dir
+
+simulate :: (Pos, Dir) -> Grid -> Int
+simulate start grid = go (S.singleton start) [start]
+  where
+    go :: S.Set (Pos, Dir) -> [(Pos, Dir)] -> Int
+    go visited [] = length $ S.map fst visited
+    go visited lights =
+        let
+            lights' = concatMap (step grid) lights
+            visited' = S.union visited $ S.fromList lights'
+            lights'' = filter (`S.notMember` visited) lights'
+         in
+            go visited' lights''
+
 part1 :: String -> String
-part1 = id
+part1 = show . simulate ((0, 0), E) . parse
 
 part2 :: String -> String
-part2 = const ""
+part2 input = show $ maximum ss
+  where
+    grid = parse input
+    (_, w, h) = grid
+    ss =
+        [simulate ((row, 0), E) grid | row <- [0 .. h - 1]]
+            ++ [simulate ((row, w - 1), W) grid | row <- [0 .. h - 1]]
+            ++ [simulate ((0, col), S) grid | col <- [0 .. w - 1]]
+            ++ [simulate ((h - 1, col), N) grid | col <- [0 .. w - 1]]
 
 solve :: String -> String
 solve input = "Part 1: " ++ part1 input ++ "\nPart 2: " ++ part2 input ++ "\n"
diff --git a/test/Main.hs b/test/Main.hs
index 496c153..078b43b 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -15,6 +15,7 @@ import TestDay12
 import TestDay13
 import TestDay14
 import TestDay15
+import TestDay16
 
 main :: IO ()
 main = do
@@ -33,3 +34,4 @@ main = do
     TestDay13.test
     TestDay14.test
     TestDay15.test
+    TestDay16.test
diff --git a/test/TestDay16.hs b/test/TestDay16.hs
new file mode 100644
index 0000000..515f2aa
--- /dev/null
+++ b/test/TestDay16.hs
@@ -0,0 +1,22 @@
+module TestDay16 (test) where
+
+import Day16
+import Test.Hspec
+
+input :: String
+input = ".|...\\....\n|.-.\\.....\n.....|-...\n........|.\n..........\n.........\\\n..../.\\\\..\n.-.-/..|..\n.|....-|.\\\n..//.|...."
+
+test1 :: Expectation
+test1 = part1 input `shouldBe` "46"
+
+test2 :: Expectation
+test2 = part2 input `shouldBe` "51"
+
+test :: IO ()
+test = hspec $ do
+    describe "day16" $ do
+        describe "part1" $ do
+            it "should work for the examples" test1
+
+        describe "part2" $ do
+            it "should work for the examples" test2