module Commons where

import GHC.IO.Handle (isEOF)
import Data.Set (Set, insert, member)
import Data.Array (Array, listArray, bounds, (!), (//), Ix (inRange))


data Direction = North | East | South | West deriving (Ord, Eq, Show)
data Tile = Empty | Mirror | AntiMirror | HSplitter | VSplitter deriving (Eq, Show)
type Cavern = Array (Int, Int) Tile


parseLine :: String -> [Tile]
parseLine = map (\case
                  '/'  -> Mirror
                  '\\' -> AntiMirror
                  '-'  -> HSplitter
                  '|'  -> VSplitter
                  '.'  -> Empty)

parseCavern :: IO [[Tile]]
parseCavern = do done <- isEOF
                 if done
                    then return []
                    else do line <- getLine
                            let cavernLine = parseLine line
                            cavern <- parseCavern
                            return (cavernLine: cavern)

parse :: IO Cavern
parse = do cavern <- parseCavern
           return (listArray ((1, 1), (length cavern, length $ head cavern)) $ concat cavern)


lightBeamEast :: (Int, Int) -> Direction -> Set (Int, Int, Direction) -> Cavern -> Set (Int, Int, Direction)
lightBeamEast (y, x) d en = lightBeam (y, x + 1) d (insert (y, x, East) en)

lightBeamSouth :: (Int, Int) -> Direction -> Set (Int, Int, Direction) -> Cavern -> Set (Int, Int, Direction)
lightBeamSouth (y, x) d en = lightBeam (y + 1, x) d (insert (y, x, South) en)

lightBeamWest :: (Int, Int) -> Direction -> Set (Int, Int, Direction) -> Cavern -> Set (Int, Int, Direction)
lightBeamWest (y, x) d en = lightBeam (y, x - 1) d (insert (y, x, West) en)

lightBeamNorth :: (Int, Int) -> Direction -> Set (Int, Int, Direction) -> Cavern -> Set (Int, Int, Direction)
lightBeamNorth (y, x) d en = lightBeam (y - 1, x) d (insert (y, x, North) en)

lightBeam :: (Int, Int) -> Direction -> Set (Int, Int, Direction) -> Cavern -> Set (Int, Int, Direction)
lightBeam (y, x) East en c
  | member (y, x, East) en              = en
  | not (inRange (bounds c) (y, x + 1)) = insert (y, x, East) en
  | otherwise                           =
      case c ! (y, x + 1) of
           Mirror     -> lightBeamEast (y, x) North en c
           AntiMirror -> lightBeamEast (y, x) South en c
           VSplitter  -> lightBeamEast (y, x) South (lightBeamEast (y, x) North en c) c
           _          -> lightBeamEast (y, x) East en c
lightBeam (y, x) West en c
  | member (y, x, West) en              = en
  | not (inRange (bounds c) (y, x - 1)) = insert (y, x, West) en
  | otherwise                           =
      case c ! (y, x - 1) of
           Mirror     -> lightBeamWest (y, x) South en c
           AntiMirror -> lightBeamWest (y, x) North en c
           VSplitter  -> lightBeamWest (y, x) South (lightBeamWest (y, x) North en c) c
           _          -> lightBeamWest (y, x) West en c
lightBeam (y, x) South en c
  | member (y, x, South) en             = en
  | not (inRange (bounds c) (y + 1, x)) = insert (y, x, South) en
  | otherwise                           =
      case c ! (y + 1, x) of
           Mirror     -> lightBeamSouth (y, x) West en c
           AntiMirror -> lightBeamSouth (y, x) East en c
           HSplitter  -> lightBeamSouth (y, x) East (lightBeamSouth (y, x) West en c) c
           _          -> lightBeamSouth (y, x) South en c
lightBeam (y, x) North en c
  | member (y, x, North) en             = en
  | not (inRange (bounds c) (y - 1, x)) = insert (y, x, North) en
  | otherwise                           =
      case c ! (y - 1, x) of
           Mirror     -> lightBeamNorth (y, x) East en c
           AntiMirror -> lightBeamNorth (y, x) West en c
           HSplitter  -> lightBeamNorth (y, x) East (lightBeamNorth (y, x) West en c) c
           _          -> lightBeamNorth (y, x) North en c