{-# LANGUAGE DeriveGeneric #-}

module DaySixteen (scorePath, realWorld, partTwo, distance, wander, parseGrid, sampleWorld, rotations, Direction(..)) where

import Debug.Trace
import GHC.Generics (Generic)
import Data.Graph.AStar (aStar)
import Data.Tree (unfoldTree, flatten)
import Data.Maybe (mapMaybe, fromMaybe)
import Data.Hashable
import qualified Data.Set as S
import Data.HashSet

newtype Wall = Wall (Int, Int) deriving (Eq, Ord, Show)

type Location = (Int, Int)

type SearchState = (Location, Direction)
data Direction = North | East  | South | West deriving (Eq, Ord, Show, Generic, Enum)
instance Hashable Direction

type BudgetedSearchState = (SearchState, Int)

data World = World (S.Set Wall) Location Location deriving (Eq, Ord, Show)


close North East = True
close North West = True
close South East = True
close South West = True
close East North = True
close East South = True
close West North = True
close West South = True
close _ _ = False

rotations :: SearchState -> [ SearchState]
rotations searchState@(currentLocation, currentDirection) = [(currentLocation, dir) | dir <- enumFrom North, close currentDirection dir]

inc :: Int -> HashSet Int
inc x = singleton $ x + 1

moveTowards (r,c) East = (r, c+1)
moveTowards (r,c) West = (r, c-1)
moveTowards (r,c) North = (r - 1, c)
moveTowards (r,c) South = (r+1, c)

parseGrid :: [String] -> World
parseGrid inputLines = World (S.fromList features) startLocation endLocation
        where parseLine :: (String, Int) -> [Wall]
              parseLine (contents, lineNumber) = mapMaybe parseChar $ zip contents [0..]
                where parseChar ('#', col) = Just $ Wall (lineNumber, col)
                      parseChar _ = Nothing
              features = concatMap parseLine $ zip inputLines [0..]
              maxColIdx = length (head inputLines) - 1
              maxRowIdx = length inputLines - 1
              startLocation = ((maxRowIdx - 1), 1)
              endLocation = (1, (maxColIdx - 1))


optimizedNeighbors :: World -> SearchState -> HashSet SearchState
optimizedNeighbors world@(World walls startLoc endLoc) state@(location, direction) = fromList $ [r  | r <- rotations state, not (rotationIsObviouslyBad r)]  ++ steps
  where naiveStep = moveTowards location direction
        rotationIsObviouslyBad rotated = Wall ((uncurry moveTowards) rotated) `S.member` walls
        steps
          | Wall naiveStep `S.member` walls = []
          | otherwise = [(naiveStep, direction)]


distance :: SearchState -> SearchState -> Int
distance left@(_, ldir) right@(_, rdir)
         | ldir == rdir = 1
         | otherwise = 1000

wanderFrom :: World -> SearchState -> Maybe [SearchState]
wanderFrom world startState = fmap addStart $ aStar (optimizedNeighbors world) distance manhattan isGoal startState
  where targetCoords (World _ _ tc) = tc
        manhattan (loc, _) = manhattanDistance loc (targetCoords world)
        manhattanDistance (a,b) (c,d) = (abs (a-b)) + (abs (c-d))
        isGoal (loc, _) = loc == (targetCoords world)
        addStart x = startState : x

wander :: World -> Maybe [SearchState]
wander world = wanderFrom world startState
  where startCoords (World _ sc _) = sc
        startState = (startCoords world, East)
        addStart x = startState : x

canReachInBudget :: World -> SearchState -> Int -> Bool
canReachInBudget world currentState maxBudget = (maxBudget > 0) && (bestBudget `maybeLt` (Just maxBudget)) == Just True
  where bestBudget = do
                       path <- wanderFrom world currentState
                       return $ scorePath path
        maybeLt = liftA2 (<=)

scorePath path = sum $ fmap (uncurry distance) $  zip path (tail path)



unfolder :: World -> BudgetedSearchState -> (BudgetedSearchState, [BudgetedSearchState])
unfolder world cb@(currentState@(loc, direction), budget) = (cb, budgetedNeighbors)
  where nbrs = toList $ optimizedNeighbors world (traceShowId currentState)
        budgetedNeighbors
          | length nbrs == 1 = [(neighbor, (budget - (distance currentState neighbor))) | neighbor <- nbrs]
          | otherwise = [(neighbor, (budget - (distance currentState neighbor))) | neighbor <- nbrs, canReachInBudget world neighbor (budget - (distance currentState neighbor))]


partTwo world budget = S.size $ S.fromList $ fmap (fst . fst) $ flatten $ buildTree world budget
        where buildTree world budget = unfoldTree (unfolder world) (startState, budget)
                where startCoords (World _ sc tc) = sc
                      startState = (startCoords world, East)




sampleWorld = parseGrid sampleInput
realWorld = parseGrid realInput
sampleInput = ["#################",
  "#...#...#...#..E#",
  "#.#.#.#.#.#.#.#^#",
  "#.#.#.#...#...#^#",
  "#.#.#.#.###.#.#^#",
  "#>>v#.#.#.....#^#",
  "#^#v#.#.#.#####^#",
  "#^#v..#.#.#>>>>^#",
  "#^#v#####.#^###.#",
  "#^#v#..>>>>^#...#",
  "#^#v###^#####.###",
  "#^#v#>>^#.....#.#",
  "#^#v#^#####.###.#",
  "#^#v#^........#.#",
  "#^#v#^#########.#",
  "#S#>>^..........#",
  "#################"]


realInput = ["secret"]