module Day7 where

import Util
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Debug.Trace

-- New beam positions and number of splits
splitBeams :: Grid -> (Set Point2, Int) -> (Set Point2, Int)
splitBeams g (beams,splits) =
  foldl (\(beams',splits') (x,y) ->
          if Map.lookup (x,y+1) g == Just '^' then
            (Set.insert (x-1,y+1) $ Set.insert (x+1,y+1) beams', splits' + 1)
          else
            (Set.insert (x,y+1) beams', splits'))
        (Set.empty, splits)
        beams

part1 :: String -> String
part1 input =
  let grid = toGrid input
      starts = map fst $ filter ((== 'S') . snd) (Map.toList grid)
      rows = maximum $ map snd $ Map.keys grid
      (_, splits) = foldl
                          (\(beams', splits') _ -> splitBeams grid (beams', splits'))
                          (Set.fromList starts, 0)
                          [0..rows]
   in show splits

splitTime :: Grid -> Map Point2 Int -> Map Point2 Int
splitTime g beams =
  Map.foldlWithKey (\beams' (x,y) times ->
                   if Map.lookup (x,y+1) g == Just '^' then
                     Map.insertWith (+) (x-1,y+1) times $
		     Map.insertWith (+) (x+1,y+1) times $
		     beams'
                   else
		     Map.insertWith (+) (x,y+1) times beams')
	Map.empty
	beams
		  
part2 :: String -> String
part2 input =
  let grid = toGrid input
      starts = map fst $ filter ((== 'S') . snd) (Map.toList grid)
      rows = maximum $ map snd $ Map.keys grid
      beams = foldl
                 (\beams' _ -> splitTime grid beams')
                 (Map.fromList [(s,1) | s <- starts])
                 [0..rows]
   in
     show $
     sum $
     map snd $
     Map.toList beams