module DayFifteen (parseGridV2, executeMove, partOne, sampleInput, parseGrid, parseActions, sampleActions, neighborsInDirection, Action(..), buildGraph, Item(..), Feature(..), allReachable, isBlocked) where

import qualified Data.Graph as G
import Data.List (group, sort)
import Data.Foldable (find)
import Data.Ix
import Data.Maybe (mapMaybe, catMaybes, fromMaybe)
import qualified Data.Set as S

data Item = Wall | Box | DoubleBox | Robot  deriving (Ord, Eq, Show)
type Location = (Int, Int)
data Feature = Feature Location Item deriving (Ord, Eq, Show)
type World = (Location, S.Set Feature)
type WorldGraph = (G.Graph, G.Vertex -> Feature, Feature -> Maybe G.Vertex)
data Action = PushUp | PushDown | PushLeft | PushRight deriving (Ord, Eq, Show)

parseActions = map parseAction
  where parseAction '<' = PushLeft
        parseAction '>' = PushRight
        parseAction 'v' = PushDown
        parseAction _ = PushUp

parseGrid :: [String] -> World
parseGrid inputLines = (robotLocation, S.fromList features)
        where parseLine :: (String, Int) -> [Feature]
              parseLine (contents, lineNumber) = mapMaybe parseChar $ zip contents [0..]
                where parseChar ('O', col) = Just $ Feature (lineNumber, col) Box
                      parseChar ('#', col) = Just $ Feature (lineNumber, col) Wall
                      parseChar ('@', col) = Just $ Feature (lineNumber, col) Robot -- Keep???
                      parseChar _ = Nothing
              features = concatMap parseLine $ zip inputLines [0..]
              robotLocation = head $ [location | Feature location item <- features, item == Robot]

parseGridV2 :: [String] -> World
parseGridV2 inputLines = (robotLocation, S.fromList features)
        where parseLine :: (String, Int) -> [Feature]
              parseLine (contents, lineNumber) = concatMap parseChar $ zip contents [0..]
                where parseChar ('O', col) = [Feature (lineNumber, col * 2) DoubleBox]
                      parseChar ('#', col) = [Feature (lineNumber, col * 2) Wall, Feature (lineNumber, col * 2 + 1) Wall]
                      parseChar ('@', col) = [Feature (lineNumber, col * 2) Robot]
                      parseChar _ = []
              features = concatMap parseLine $ zip inputLines [0..]
              robotLocation = head $ [location | Feature location item <- features, item == Robot]


moveTowards (r,c) PushRight = (r, c+1)
moveTowards (r,c) PushLeft = (r, c-1)
moveTowards (r,c) PushUp = (r - 1, c)
moveTowards (r,c) PushDown = (r+1, c)

pushFeature f@(Feature _ Wall) _ = f
pushFeature (Feature loc item) action = Feature (moveTowards loc action) item


neighborsInDirection :: World -> Action -> Feature -> [Feature]
neighborsInDirection (_, features) action from = S.toList $ S.filter isHit features
        where isHit f = (footprint f) `intersects` (footprint (pushFeature from action))
              intersects a b = not $ a `S.disjoint` b

footprint (Feature loc DoubleBox) = S.fromList [loc, moveTowards loc PushRight]
footprint (Feature loc _) = S.singleton loc



buildGraph :: World -> Action -> WorldGraph
buildGraph world action = (thegraph, lookup,  vertexFromKey)
  where (thegraph, nodeFromVertex, vertexFromKey) = G.graphFromEdges edges
        edges = [(feature, feature, neighborsInDirection world action feature) | feature <- featuresList]
        featuresList = S.toList $ snd world
        lookup = snd3 . nodeFromVertex
        snd3 (a,b,c) = b

getLoc (Feature loc _) = loc

allReachable :: WorldGraph -> Feature -> Action -> [Feature]
allReachable (g, getFeature, getVertex) from direction = fromMaybe [] result
  where result = do
                   vertex <- getVertex from
                   let reachableVertices = G.reachable g vertex
                   return $ map getFeature reachableVertices

isBlocked :: WorldGraph -> Location -> Action -> Bool
isBlocked g startLocation direction = any isWall impacted
  where isWall (Feature _ Wall) = True
        isWall _ = False
        impacted = allReachable g (Feature startLocation Robot) direction


executeMove :: World -> Action -> World
executeMove world@(startLocation, features) action
  | stuck = world
  | otherwise = (getLoc (pushFeature robot action), S.map updateIndividual features)
  where stuck = isBlocked g startLocation action
        g = buildGraph world action
        affected = S.fromList $ allReachable g (Feature startLocation Robot) action
        robot = Feature startLocation Robot
        updateIndividual feature
          | feature `S.member` affected = pushFeature feature action
          | otherwise = feature



fullSimulate :: World -> [Action] -> World
fullSimulate world actions = foldl executeMove world actions

partOne :: World -> [Action] -> Int
partOne world actions = scoreWorld $ fullSimulate world actions
  where scoreWorld world@(_, features) = sum $ map scoreFeature $ S.toList features
        scoreFeature (Feature (r,c) Box) = 100 * r + c
        scoreFeature (Feature (r,c) DoubleBox) = 100 * r + c
        scoreFeature _ = 0


sampleActions = "<vv>^<v^>v>^vv^v>v<>v^v<v<^vv<<<^><<><>>v<vvv<>^v^>^<<<><<v<<<v^vv^v>^vvv<<^>^v^^><<>>><>^<<><^vv^^<>vvv<>><^^v>^>vv<>v<<<<v<^v>^<^^>>>^<v<v><>vv>v^v^<>><>>>><^^>vv>v<^^^>>v^v^<^^>v^^>v^<^v>v<>>v^v^<v>v^^<^^vv<<<v<^>>^^^^>>>v^<>vvv^><v<<<>^^^vv^<vvv>^>v<^^^^v<>^>vvvv><>>v^<<^^^^^^><^><>>><>^^<<^^v>>><^<v>^<vv>>v>>>^v><>^v><<<<v>>v<v<v>vvv>^<><<>^><^>><>^v<><^vvv<^^<><v<<<<<><^v<<<><<<^^<v<^^^><^>>^<v^><<<^>>^v<v^v<v^>^>>^v>vv>^<<^v<>><<><<v<<v><>v<^vv<<<>^^v^>^^>>><<^v>>v^v><^^>>^<>vv^<><^^>^^^<><vvvvv^v<v<<>^v<v>v<<^><<><<><<<^^<<<^<<>><<><^^^>^^<>^>v<>^^>vv<^v^v<vv>^<><v<^v>^^^>>>^^vvv^>vvv<>>>^<^>>>>>^<<^v>^vvv<>^<><<v>v^^>>><<^^<>>^v^<v^vv<>v^<<>^<^v^v><^<<<><<^<v><v<>vv>>v><v^<vv<>v^<<^"

sampleInput = ["##########",
        "#..O..O.O#",
        "#......O.#",
        "#.OO..O.O#",
        "#[email protected].#",
        "#O#..O...#",
        "#O..O..O.#",
        "#.OO.O.OO#",
        "#....O...#",
        "##########"]