{-# LANGUAGE ScopedTypeVariables #-}

module Graph.DijkstraSimple
  (
                            -- * How to use this library
                            -- $use
    lightestPaths
  , findPath
  , dijkstraSteps
  , EdgeTo(..)
  , Graph(..)
  , Weighter(..)
  , Path(..)
  , Paths(..)
  )
where

import qualified Data.Map.Lazy                 as M
import qualified Data.List.NonEmpty            as NE
import           Data.Maybe                     ( fromJust
                                                , isJust
                                                , isNothing
                                                )
import           Data.Ord                       ( comparing )
import qualified Data.PriorityQueue.FingerTree as P

-- | Edge to an arbitrary vertex and the associated input weight
data EdgeTo v e = EdgeTo { edgeTo :: v, edgeToWeight :: e } deriving (Eq, Show)

-- | All vertices and outgoing edges
newtype Graph v e = Graph { graphAsMap :: M.Map v [EdgeTo v e] } deriving (Eq, Show)

-- | Convert an input weight (edge-dependant) to an output weight
-- (path-dependant) for the algorithm work.
data Weighter v e a = Weighter { initialWeight :: a, weight :: EdgeTo v e -> Path v e a -> a }
-- | The lightest found path with reverse ordered list of traversed
-- vertices and output weight.
data Path v e a = Path { pathVertices :: NE.NonEmpty v, pathWeight :: a } deriving (Eq, Show)
-- | Reachable vertices and associated lightest paths
newtype Paths v e a = Paths { pathsAsMap :: M.Map v (Path v e a) } deriving (Eq, Show)

-- | Explore all the reachable edges
lightestPaths
  :: forall v e a
   . (Ord v, Ord a)
  => Graph v e
  -> v
  -> Weighter v e a
  -> Paths v e a
lightestPaths graph origin weighter =
  NE.last $ dijkstraSteps graph origin weighter

-- | Find the eventual path between two edges
findPath
  :: forall v e a
   . (Ord v, Ord a)
  => Graph v e
  -> v
  -> Weighter v e a
  -> v
  -> Maybe (Path v e a)
findPath graph origin weighter target =
  pathsAsMap (lightestPaths graph origin weighter) M.!? target

type StatePQ v e a = P.PQueue a (Path v e a, EdgeTo v e)
type State v e a
  = (M.Map v (Path v e a), ((a, (Path v e a, EdgeTo v e)), StatePQ v e a))

-- | Details each step of the Dijkstra algorithm
dijkstraSteps
  :: forall v e a
   . (Ord v, Ord a)
  => Graph v e
  -> v
  -> Weighter v e a
  -> NE.NonEmpty (Paths v e a)
dijkstraSteps graph origin weighter =
  Paths <$> maybe (M.empty NE.:| []) (NE.unfoldr nextStep) init
 where
  init :: Maybe (State v e a)
  init =
    (\p -> (M.empty, p))
      <$> (P.minViewWithKey $ findEdges
            (Path (origin NE.:| []) (initialWeight weighter))
            origin
          )
  nextStep :: State v e a -> (M.Map v (Path v e a), Maybe (State v e a))
  nextStep (paths, ((w, (path, e)), pq)) =
    let npq = if M.notMember (edgeTo e) paths
          then pq `P.union` findEdges path (edgeTo e)
          else pq
        nps = M.alter (updatePath path) (edgeTo e) paths
    in  (nps, (\q -> (nps, q)) <$> P.minViewWithKey npq)
  updatePath :: Path v e a -> Maybe (Path v e a) -> Maybe (Path v e a)
  updatePath p prev = case prev of
    Nothing -> Just p
    Just op -> Just $ if pathWeight op <= pathWeight p then op else p
  findEdges :: Path v e a -> v -> StatePQ v e a
  findEdges path vertice =
    P.fromList
      $ map (buildPath path)
      $ M.findWithDefault [] vertice
      $ graphAsMap graph
  buildPath :: Path v e a -> EdgeTo v e -> (a, (Path v e a, EdgeTo v e))
  buildPath path e = let np = addEdge e path in (pathWeight np, (np, e))
  addEdge :: EdgeTo v e -> Path v e a -> Path v e a
  addEdge e p = Path { pathVertices = edgeTo e NE.<| pathVertices p
                     , pathWeight   = weight weighter e p
                     }

-- $use
--
-- This section contains basic step-by-step usage of the library.
--
-- The first step is to build a direct graph:
--
-- > exampleGraph :: Graph Char Int
-- > exampleGraph = Graph $ M.fromList [
-- >                                     ('A', [EdgeTo 'B' 3, EdgeTo 'C' 1])
-- >                                   , ('B', [EdgeTo 'A' 3, EdgeTo 'C' 7, EdgeTo 'D' 5, EdgeTo 'E' 1])
-- >                                   , ('C', [EdgeTo 'A' 1, EdgeTo 'B' 7, EdgeTo 'D' 2])
-- >                                   , ('D', [EdgeTo 'B' 5, EdgeTo 'C' 2, EdgeTo 'E' 5])
-- >                                   , ('E', [EdgeTo 'B' 1, EdgeTo 'D' 7])
-- >                                   ]
--
-- Then pick or create a weighter (see @Graph.DijkstraSimple.Weighters@)
-- and apply it all:
--
-- > lightestPaths exampleGraph 'C' weighter
--
-- It will give all the reacheable vertices from @'C'@ and associated shortest path:
--
-- > Paths $ M.fromList [
-- >                      ('A', Path (fromList "AC") 1)
-- >                    , ('B', Path (fromList "BAC") 3)
-- >                    , ('C', Path (fromList "CAC") 1)
-- >                    , ('D', Path (fromList "DC") 2)
-- >                    , ('E', Path (fromList "EBAC") 3)
-- >                    ]