{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}

module Main where

import Data.Function (on)
import Data.Ix (range)
import Data.List (foldl', groupBy, sort, sortBy)
import Data.Map (Map)
import Data.Map qualified as M
import Data.Maybe (mapMaybe)
import Debug.Trace
import System.Environment (getArgs)

main :: IO ()
main = do
  input <- readFile . head =<< getArgs
  putStrLn $ unwords ["Part 1:", show $ part1 input]
  putStrLn $ unwords ["Part 2:", show $ part2 input]

data Cell = Empty | Start | Splitter Int | Beam
  deriving (Eq, Show)

instance Read Cell where
  readsPrec _ "" = []
  readsPrec _ (c : rest) = case c of
    '.' -> [(Empty, rest)]
    'S' -> [(Start, rest)]
    '^' -> [(Splitter 0, rest)]
    '|' -> [(Beam, rest)]
    _ -> []
  readList s = case readsPrec @Cell 0 s of
    [(c, rest)] -> case readList rest of
      [] -> [([c], rest)]
      [(acc, rest')] -> [(c : acc, rest')]
    [] -> []

data Grid = Grid
  { cells :: Map (Int, Int) Cell,
    size :: (Int, Int),
    currentBeams :: [(Int, Int)],
    paths :: Int
  }
  deriving (Show)

showGrid :: Grid -> String
showGrid (Grid cs (rows, cols) _ _) =
  unlines
    . map (foldl' ((<>)) mempty . map showCell . (\r -> map (\c -> cs M.! (r, c)) [0 .. cols]))
    $ [0 .. rows]
  where
    showCell c = show $ case c of
      Empty -> '.'
      Start -> 'S'
      Splitter _ -> '^'
      Beam -> '|'

instance Read Grid where
  readsPrec _ =
    (\cs -> [(Grid (M.fromList cs) (size cs) (start (cs)) 1, "")])
      . concat
      . map parseLine
      . zip [0 ..]
      . lines
    where
      start :: [((Int, Int), Cell)] -> [(Int, Int)]
      start = map fst . filter (\(_, c) -> c == Start)

      size :: [((Int, Int), Cell)] -> (Int, Int)
      size cs = fst $ last cs

      parseLine :: (Int, String) -> [((Int, Int), Cell)]
      parseLine (r, line) = zipWith (\c cell -> ((r, c), cell)) [0 ..] . fst . head $ readList @Cell line

step :: Grid -> Grid
step g = foldl' go (g {currentBeams = []}) (filter valid $ currentBeams g)
  where
    go :: Grid -> ((Int, Int)) -> Grid
    go g (row, col) = case cells g M.! (row, col) of
      Empty ->
        g
          { cells = M.adjust (const Beam) (row, col) (cells g),
            currentBeams = (row + 1, col) : currentBeams g
          }
      Splitter 0 ->
        g
          { cells = M.adjust (const $ Splitter 1) (row, col) (cells g),
            currentBeams = (row, col + 1) : (row, col - 1) : currentBeams g,
            paths = paths g + 1
          }
      Splitter n -> g {cells = M.adjust (const $ Splitter (n + 1)) (row, col) (cells g), paths = paths g + n}
      Beam -> g {currentBeams = (row + 1, col) : currentBeams g}
      Start -> g {currentBeams = (row + 1, col) : currentBeams g}

    valid :: (Int, Int) -> Bool
    valid p = p < size g

    incrementCounter :: Cell -> Cell
    incrementCounter (Splitter n) = Splitter (n + 1)

part1 :: String -> Int
part1 = M.size . M.filter nonZero . cells . final . read
  where
    nonZero :: Cell -> Bool
    nonZero (Splitter n) = n > 0
    nonZero _ = False

    final :: Grid -> Grid
    final gr =
      let gr' = step gr
       in if null (currentBeams gr') then gr' else final gr'

part2 :: String -> Int
part2 = (\g -> go g mempty) . read
  where
    go :: Grid -> Map (Int, Int) Int -> Int
    go g acc = snd . M.findMin . foldl' (foldl' (onSplitter g)) acc $ splitters g

    splitters :: Grid -> [[(Int, Int)]]
    splitters =
      groupBy ((==) `on` fst)
        . sortBy (flip (compare `on` fst))
        . M.keys
        . M.filter
          isSplitter
        . cells

    -- onSplitter sums the number of descendants in each leg of the
    -- splitter and stores it in acc
    onSplitter :: Grid -> Map (Int, Int) Int -> (Int, Int) -> Map (Int, Int) Int
    onSplitter g acc (r, c) =
      let h c = halfSplitter c (r + 1, fst . size $ g) acc
       in M.insert (r, c) (sum . map (h . ($ c)) $ [succ, pred]) acc

    halfSplitter :: Int -> (Int, Int) -> Map (Int, Int) Int -> Int
    halfSplitter c range acc = case col c range acc of
      [] -> 1
      (n : _) -> n

    -- col returns all elements in a column c
    col :: Int -> (Int, Int) -> Map (Int, Int) a -> [a]
    col c rows g = mapMaybe ((g M.!?) . (,c)) $ range rows

    isSplitter (Splitter _) = True
    isSplitter _ = False

testInput :: String
testInput =
  unlines
    [ ".......S.......",
      "...............",
      ".......^.......",
      "...............",
      "......^.^......",
      "...............",
      ".....^.^.^.....",
      "...............",
      "....^.^...^....",
      "...............",
      "...^.^...^.^...",
      "...............",
      "..^...^.....^..",
      "...............",
      ".^.^.^.^.^...^.",
      "..............."
    ]