{-# LANGUAGE NamedFieldPuns #-}

module Lib where

import Control.Monad (join)
import Flow
import qualified Data.Map.Strict as Map
import Data.Maybe (listToMaybe, maybeToList)
import qualified Data.Set as Set
import Text.ParserCombinators.ReadP

type IntType = Int

data Vec2 = Vec2 IntType IntType
  deriving (Eq, Ord, Show)
vAdd :: Vec2 -> Vec2 -> Vec2
vAdd (Vec2 aX aY) (Vec2 bX bY) = Vec2 (aX+bX) (aY+bY)

data Direction = North | East | South | West
  deriving (Eq, Ord, Show)
dOpposite :: Direction -> Direction
dOpposite North = South
dOpposite East = West
dOpposite South = North
dOpposite West = East
dVec :: Direction -> Vec2
dVec North = Vec2 0 (-1)
dVec East = Vec2 1 0
dVec South = Vec2 0 1
dVec West = Vec2 (-1) 0

data Tile = Path | Start | End | Wall
  deriving (Eq, Ord, Show)

type WorldMap = Map.Map Vec2 Tile

data Edge = Edge { eDirFrom :: Direction, eFrom :: Vec2, eDirTo :: Direction, eTo :: Vec2, ePrice :: Int }
  deriving (Show)

type WorldGraphEdges = Map.Map Vec2 [Edge]
data WorldGraph = WorldGraph { wStart :: Vec2, wEnd :: Vec2, wEdges :: WorldGraphEdges }
  deriving (Show)

task1 :: String -> IO IntType
task1 input =
  let worldMap = parseFile input
      rawGraph = mapToGraph worldMap
  in do
    worldGraph <- simplifyGraph rawGraph
    worldGraph |> wEdges |> Map.elems |> join |> showEdgesSimple |> putStrLn
    (bestPath, _) <- dijkstra worldGraph Map.empty (Map.singleton (wStart worldGraph, East) (0, []))
    return bestPath

task2 :: String -> IO IntType
task2 input = 
  let worldMap = parseFile input
      rawGraph = mapToGraph worldMap
  in do
    (bestCost, bestPaths) <- dijkstra rawGraph Map.empty (Map.singleton (wStart rawGraph, East) (0, []))
    [wEnd rawGraph] : bestPaths |> join |> Set.fromList |> length |> return

showEdgesSimple :: [Edge] -> String
showEdgesSimple [] = ""
showEdgesSimple (edge:edges) = (show $ eFrom edge) ++ " -> " ++ (show $ eTo edge) ++ ": " ++ (show $ ePrice edge) ++ "\n" ++ (showEdgesSimple edges)

type WasBestePreis = (Int, [(Vec2, Direction)])
dijkstra :: WorldGraph -> Map.Map (Vec2, Direction) WasBestePreis -> Map.Map (Vec2, Direction) WasBestePreis -> IO (Int, [[Vec2]])
dijkstra world@WorldGraph { wEdges, wEnd, wStart } done unvisited
  | null unvisited =
    let bestePreis = [North, East, South, West] |> map (\d -> done Map.!? (wEnd, d) |> maybeToList) |> join |> foldr combineBestePreis (maxBound :: IntType, [])
    in bestePreis |> \(p, os) -> (p, map (reconstructPath done) os |> join) |> return
  | otherwise =
    let ((pos, dir), (cost, _)) = findNext
        nextDone = Map.insert (pos, dir) (unvisited Map.! (pos, dir)) done
        outgoing = Map.findWithDefault [] pos wEdges
        unvisitedOutgoing = outgoing |> filter (\Edge { eTo, eDirTo } -> Map.notMember (eTo, eDirTo) done)
        unvisitedOutgoingWithPrice = unvisitedOutgoing |> map (\e@Edge { eDirFrom, ePrice } -> (e, if dir == dOpposite eDirFrom then cost + ePrice else cost + ePrice + 1000)) :: [(Edge, IntType)]
        unvisitedWithoutCurrent = Map.delete (pos, dir) unvisited
        nextUnvisited = foldr (\(Edge { eTo, eDirTo }, price) -> Map.alter (pure . maybeCombineBestePreis (price, [(pos, dir)])) (eTo, eDirTo)) unvisitedWithoutCurrent unvisitedOutgoingWithPrice
    in dijkstra world nextDone nextUnvisited
  where
    findNext :: ((Vec2, Direction), WasBestePreis)
    findNext =
      let options = unvisited |> Map.assocs
      in foldr (\a@(_,(aC,_)) b@(_,(bC,_)) -> if aC < bC then a else b) (head options) (drop 1 options)

    maybeCombineBestePreis :: WasBestePreis -> Maybe WasBestePreis -> WasBestePreis
    maybeCombineBestePreis a Nothing = a
    maybeCombineBestePreis a (Just b) = combineBestePreis a b

    combineBestePreis :: WasBestePreis -> WasBestePreis -> WasBestePreis
    combineBestePreis (aP, aOs) (bP, bOs)
      | aP < bP = (aP, aOs)
      | aP == bP = (aP, aOs ++ bOs)
      | aP > bP = (bP, bOs)

reconstructPath :: Map.Map (Vec2, Direction) WasBestePreis -> (Vec2, Direction) -> [[Vec2]]
reconstructPath done cur@(curPos,_) =
  case done Map.!? cur of
    Nothing -> [[curPos]]
    Just (_,[]) -> [[curPos]]
    Just (_,os) -> map (reconstructPath done) os |> join |> map (\subPath -> curPos : subPath)

showEntries :: Show a => Show b => [(a,b)] -> String
showEntries [] = ""
showEntries ((k,v):rest) = show k ++ ": " ++ show v ++ "\n" ++ showEntries rest

findBestPath :: WorldGraph -> [(Vec2, Direction)] -> Maybe Int
findBestPath _ [] = error "Must call findBestPath with an initial position"
findBestPath world@WorldGraph { wEdges, wEnd } ((lastPos, lastDir):path) =
  let edgeOptions = wEdges Map.!? lastPos |> expectJust ("Position is missing in map " ++ show lastPos) |> filter (\Edge{eTo=optTo} -> not $ listContains optTo $ map (\(p,_) -> p) path)
  in case edgeOptions |> map (maybeToList . exploreOption) |> join of
    [] -> Nothing
    prices -> return $ minimum prices
  where
    exploreOption :: Edge -> Maybe Int
    exploreOption edge@Edge{ePrice, eDirFrom, eTo, eDirTo}
      | eTo == wEnd = Just hopPrize
      | otherwise = findBestPath world ((eTo, eDirTo):(lastPos, lastDir):path) |> fmap (\pr -> pr + hopPrize)
      where
        hopPrize :: Int
        hopPrize = if eDirFrom == dOpposite lastDir then ePrice else ePrice + 1000

mapToGraph :: WorldMap -> WorldGraph
mapToGraph worldMap =
  let start = Map.assocs worldMap |> filter (\(_, t) -> t == Start) |> listToMaybe |> expectJust "No start found" |> entryKey
      end = Map.assocs worldMap |> filter (\(_, t) -> t == End) |> listToMaybe |> expectJust "No end found" |> entryKey
      pathMap = worldMap |> Map.filter ((/=) Wall)
      edges = pathMap |> Map.keys |> map edgesFromPos |> join
      edgesByFrom = map (\e -> (eFrom e, [e])) edges |> Map.fromListWith (++)
  in WorldGraph { wStart = start, wEnd = end, wEdges = edgesByFrom }

  where
    edgesFromPos :: Vec2 -> [Edge]
    edgesFromPos pos =
      [North, East, South, West]
      |> map (\dir -> dVec dir |> vAdd pos |> \t -> if validTarget $ worldMap Map.!? t then [Edge { eDirTo = dir, eTo = t, eDirFrom = dOpposite dir, eFrom = pos, ePrice = 1 }] else [])
      |> join

    validTarget :: Maybe Tile -> Bool
    validTarget (Just Wall) = False
    validTarget (Just _) = True
    validTarget _ = False

simplifyGraph :: WorldGraph -> IO WorldGraph
simplifyGraph world@WorldGraph { wStart, wEnd, wEdges } =
  case Map.assocs wEdges |> filter (\(p, edges) -> p /= wStart && p /= wEnd && length edges <= 2) of
    [] -> return world
    ((pos, []):_) -> simplifyGraph WorldGraph { wStart, wEnd, wEdges = Map.delete pos wEdges }
    ((_, [e]):_) -> do
      --putStrLn $ "Removing " ++ (show e)
      simplifyGraph WorldGraph { wStart, wEnd, wEdges = removeEdge e wEdges }
    ((_, [e1, e2]):_) ->
      let newEdges = wEdges |> combineEdges e1 e2
      in do
        --putStrLn $ "Combining " ++ (show e1) ++ " : " ++ (show e2)
        simplifyGraph WorldGraph { wStart, wEnd, wEdges = newEdges }
    _ -> error "Ureachable in simplifyGraph"
  where
    combineEdges :: Edge -> Edge -> WorldGraphEdges -> WorldGraphEdges
    combineEdges e1@Edge { eDirFrom = eDirFrom1, eTo = eTo1, eDirTo = eDirTo1, ePrice = ePrice1 } e2@Edge { eDirFrom = eDirFrom2, eTo = eTo2, eDirTo = eDirTo2, ePrice = ePrice2 } edges =
      let combinedPrice = ePrice1 + ePrice2 + if eDirFrom1 == dOpposite eDirFrom2 then 0 else 1000
      in edges |> removeEdge e1 |> removeEdge e2
          |> Map.adjust (\es -> Edge { eFrom = eTo1, eDirFrom = eDirTo1, eTo = eTo2, eDirTo = eDirTo2, ePrice = combinedPrice }:es) eTo1
          |> Map.adjust (\es -> Edge { eFrom = eTo2, eDirFrom = eDirTo2, eTo = eTo1, eDirTo = eDirTo1, ePrice = combinedPrice }:es) eTo2
    
    removeEdge :: Edge -> WorldGraphEdges -> WorldGraphEdges
    removeEdge Edge { eFrom, eTo } edgesMap =
      edgesMap
      |> Map.adjust (\edges -> edges |> filter (not . (isEdgeTo eTo))) eFrom
      |> Map.adjust (\edges -> edges |> filter (not . (isEdgeTo eFrom))) eTo
 
    isEdgeTo :: Vec2 -> Edge -> Bool
    isEdgeTo testTo Edge { eTo } = testTo == eTo

parseFile :: String -> WorldMap
parseFile input =
  case readP_to_S parse input of
    [] -> error "Failed to parse"
    ((v,_):_) -> v

  where
    parse :: ReadP WorldMap
    parse = do
      worldMap <- parseWorld
      _ <- eof
      return worldMap

parseWorld :: ReadP WorldMap
parseWorld = do
  ls <- many1 parseLine
  zip [0..] ls |> map (\(y, line) -> zip [0..] line |> map (\(x, c) -> (Vec2 x y, toTile c))) |> join |> Map.fromList |> return

  where
    parseLine :: ReadP String
    parseLine = do
      content <- munch1 (not . isNewLine)
      _ <- char '\n'
      return content
    
    toTile :: Char -> Tile
    toTile '.' = Path
    toTile 'S' = Start
    toTile 'E' = End
    toTile '#' = Wall
    toTile c = error $ "Unknown tile: " ++ [c]

isNewLine :: Char -> Bool
isNewLine = (==) '\n'

listContains :: Eq a => a -> [a] -> Bool
listContains _ [] = False
listContains a (e:es) = if a == e then True else listContains a es

nothingIfEmpty :: Foldable t => t a -> Maybe (t a)
nothingIfEmpty values
  | null values = Nothing
  | otherwise = Just values

entryKey :: (a,b) -> a
entryKey (k,_) = k
entryValue :: (a,b) -> b
entryValue (_,v) = v

expectJust :: String -> Maybe a -> a
expectJust message Nothing = error message
expectJust _ (Just value) = value