day6.hs

module Day6 where

import Data.List
import Data.Map.Strict ((!?))
import Data.Map.Strict qualified as Map
import Data.Set qualified as Set
import Debug.Trace

main :: IO ()
main = do
  input <- readFile "day6.txt"
  let map = Map.fromList [((j, i), c) | (line, i) <- zip (lines input) [0..], (c, j) <- zip line [0..]]
  let step (coord, direction) =
        let newCoord = coord `add` directionToVector direction
        in case map !? newCoord of
             Just char | char == '#' -> step (coord, turn direction)
             Just char -> (coord, direction) : step (newCoord, direction)
             Nothing -> [(coord, direction)]
  let Just (start, startDirection) = find (\(_, char) -> char `elem` "^>v<") $ Map.toList map
  let findObstructions ((coord, direction) : xs) =
        let obstruction = coord `add` directionToVector direction
         in case map !? obstruction of
              Just char | char /= '#' ->
                if isLoop obstruction (start, startDirection) Set.empty
                  then Set.singleton obstruction <> findObstructions xs
                  else findObstructions xs
              Nothing -> Set.empty
        where
          isLoop obstruction (coord, direction) seen | (coord, direction) `Set.member` seen = True
          isLoop obstruction (coord, direction) seen =
            let newCoord = coord `add` directionToVector direction
                seen' = Set.singleton (coord, direction) <> seen
            in case map !? newCoord of
                 Just char ->
                   if char == '#' || obstruction == newCoord
                     then isLoop obstruction (coord, turn direction) seen'
                     else isLoop obstruction (newCoord, direction) seen'
                 Nothing -> False
      findObstructions _ = Set.empty
  let origPath = step (start, startDirection)
  print $ length $ nubBy (\(a1, _) (a2, _) -> a1 == a2) origPath
  print $ Set.size $ findObstructions origPath
  where
    add (a1, b1) (a2, b2) = (a1 + a2, b1 + b2)
    directionToVector '^' = (0, -1)
    directionToVector '>' = (1, 0)
    directionToVector 'v' = (0, 1)
    directionToVector '<' = (-1, 0)
    directionToVector _ = error "impossible"
    turn '^' = '>'
    turn '>' = 'v'
    turn 'v' = '<'
    turn '<' = '^'
    turn _ = error "impossible"