{-# LANGUAGE UnicodeSyntax #-}
module Graph
  ( Distance, Graph, Vertex
  , fromList, findShortestPath, findShortestPaths, findShortestPaths_
  ) where

import Control.Monad (when)
import Control.Monad.Reader (Reader, ask, runReader)
import Control.Monad.State (StateT, execStateT, get, put)
import Data.Function (on)
import Data.IntMap (fromList)
import Data.List (delete, intersect, (\\), nub, sortBy)
import Data.Maybe (catMaybes, listToMaybe)
import qualified Data.IntMap as IM

import Debug.Trace (trace)

type Distance = Int
type Graph = IM.IntMap (IM.IntMap Vertex)
type Vertex = Int

findShortestPaths ∷ Graph → Vertex → IM.IntMap Distance
findShortestPaths γ v = runReader (execStateT allSteps (fromList [(v,0)])) γ

findShortestPaths_ ∷ Graph → Vertex → [(Vertex,Distance)]
findShortestPaths_ γ v = IM.toList $ findShortestPaths γ v

findShortestPath ∷ Graph → Vertex → Vertex → Maybe Distance
findShortestPath γ s t = IM.lookup t $ findShortestPaths γ s

step ∷ StateT (IM.IntMap Distance) (Reader Graph) Bool
step = do
  γ ← ask
  distance ← get
  let candidates ∷ [Vertex]
      candidates = concatMap nub . map (\v → (IM.keys $ γ IM.! v) \\ (IM.keys distance)) $ IM.keys distance ∷ [Vertex]
      nearestWay ∷ Vertex → Distance
      nearestWay c = minimum . catMaybes . map (\v → (+ distance IM.! v) `fmap` IM.lookup c $ γ IM.! v) $ IM.keys distance
      newVertex ∷ Maybe (Vertex, Distance)
      newVertex = listToMaybe . sortBy (compare `on` snd) . map (\c → (c, nearestWay c)) $ candidates
  case newVertex of
    Just (v, d) → do
      put $ IM.insert v d distance
      return True
    Nothing → return False

allSteps ∷ StateT (IM.IntMap Vertex) (Reader Graph) ()
allSteps = do
  r ← step
  when r allSteps

---------------------------

Graph.hs:34:14:
    Couldn't match type `Int' with `Maybe Int'
    When using functional dependencies to combine
      Control.Monad.State.Class.MonadState s (StateT s m),
        arising from the dependency `m -> s'
        in the instance declaration in `Control.Monad.State.Class'
      Control.Monad.State.Class.MonadState
        (IM.IntMap (Maybe Int))
        (StateT
           (IM.IntMap Int)
           (Control.Monad.Trans.Reader.ReaderT
              (IM.IntMap (IM.IntMap Int)) Data.Functor.Identity.Identity)),
        arising from a use of `get' at Graph.hs:34:14-16
    In a stmt of a 'do' block: distance <- get
    In the expression:
      do { γ <- ask;
           distance <- get;
           let candidates :: [Vertex]
               candidates = ...
               ....;
           case newVertex of {
             Just (v, d) -> do { ... }
             Nothing -> return False } }

Graph.hs:43:27:
    Couldn't match expected type `Distance'
                with actual type `Maybe Distance'
    Expected type: IM.IntMap Distance
      Actual type: IM.IntMap (Maybe Distance)
    In the third argument of `IM.insert', namely `distance'
    In the second argument of `($)', namely `IM.insert v d distance'