{-

Benchmark 1: ./out
  Time (mean ± σ):      19.7 ms ±   1.7 ms    [User: 15.9 ms, System: 3.6 ms]
  Range (min … max):    17.7 ms …  22.9 ms    10 runs

-}

module Main where

import Data.List (singleton, elemIndex, nub)
import Data.Maybe (fromJust)

part1 :: String -> Int
part1 str = go (drop 1 input) 0 beams
  where
  input = lines str

  beams = [ fromJust $ elemIndex 'S' $ head input ]

  go :: [String] -> Int -> [Int] -> Int
  go []           acc _     = acc
  go (line:lines) acc beams = go lines (acc + countSplits beams) $ nub $ concatMap splitBeam beams
    where
    countSplits :: [Int] -> Int
    countSplits = length . filter (\beam -> (line !! beam) == '^')
    
    splitBeam beam | (line !! beam) == '^' = [beam-1, beam+1]
                   | otherwise             = [beam]


type Beam = (Int, Int)
part2 :: String -> Int
part2 str = go (drop 1 input) beams
  where
  input = lines str

  beams = [ (fromJust $ elemIndex 'S' $ head input, 1) ]

  go :: [String] -> [Beam] -> Int
  go []           = sum . map snd
  go (line:lines) = go lines . mergeNub . concatMap splitBeam
    where
    splitBeam (index, count) | (line !! index) == '^' = [(index-1, count), (index+1,count)]
                             | otherwise              = [(index, count)]

    mergeNub :: [Beam] -> [Beam]
    mergeNub []         =  []
    mergeNub ((i,n):xs) =  (i,n + merged) : mergeNub (filter (\(i',n') -> i /= i') xs)
      where
      merged = sum (map snd $ filter (\(i', n') -> i == i') xs)

main :: IO ()
main = do
  str <- getContents

  print $ part1 str
  print $ part2 str