-- | Wrapper around Data.Graph with support for edge labels
{-# LANGUAGE ScopedTypeVariables #-}
module Distribution.Solver.Modular.LabeledGraph (
    -- * Graphs
    Graph
  , Vertex
    -- ** Building graphs
  , graphFromEdges
  , graphFromEdges'
  , buildG
  , transposeG
    -- ** Graph properties
  , vertices
  , edges
    -- ** Operations on the underlying unlabeled graph
  , forgetLabels
  , topSort
  ) where

import Data.Array
import Data.Graph (Vertex, Bounds)
import Data.List (sortBy)
import Data.Maybe (mapMaybe)
import qualified Data.Graph as G

{-------------------------------------------------------------------------------
  Types
-------------------------------------------------------------------------------}

type Graph e = Array Vertex [(e, Vertex)]
type Edge  e = (Vertex, e, Vertex)

{-------------------------------------------------------------------------------
  Building graphs
-------------------------------------------------------------------------------}

-- | Construct an edge-labeled graph
--
-- This is a simple adaptation of the definition in Data.Graph
graphFromEdges :: forall key node edge. Ord key
               => [ (node, key, [(edge, key)]) ]
               -> ( Graph edge
                  , Vertex -> (node, key, [(edge, key)])
                  , key -> Maybe Vertex
                  )
graphFromEdges edges0 =
    (graph, \v -> vertex_map ! v, key_vertex)
  where
    max_v        = length edges0 - 1
    bounds0      = (0, max_v) :: (Vertex, Vertex)
    sorted_edges = sortBy lt edges0
    edges1       = zip [0..] sorted_edges

    graph        = array bounds0 [(v, (mapMaybe mk_edge ks))
                                 | (v, (_, _, ks)) <- edges1]
    key_map      = array bounds0 [(v, k                    )
                                 | (v, (_, k, _ )) <- edges1]
    vertex_map   = array bounds0 edges1

    (_,k1,_) `lt` (_,k2,_) = k1 `compare` k2

    mk_edge :: (edge, key) -> Maybe (edge, Vertex)
    mk_edge (edge, key) = do v <- key_vertex key ; return (edge, v)

    --  returns Nothing for non-interesting vertices
    key_vertex :: key -> Maybe Vertex
    key_vertex k = findVertex 0 max_v
      where
        findVertex a b
          | a > b     = Nothing
          | otherwise = case compare k (key_map ! mid) of
              LT -> findVertex a (mid-1)
              EQ -> Just mid
              GT -> findVertex (mid+1) b
          where
            mid = a + (b - a) `div` 2

graphFromEdges' :: Ord key
                => [ (node, key, [(edge, key)]) ]
                -> ( Graph edge
                   , Vertex -> (node, key, [(edge, key)])
                   )
graphFromEdges' x = (a,b)
  where
    (a,b,_) = graphFromEdges x

transposeG :: Graph e -> Graph e
transposeG g = buildG (bounds g) (reverseE g)

buildG :: Bounds -> [Edge e] -> Graph e
buildG bounds0 edges0 = accumArray (flip (:)) [] bounds0 (map reassoc edges0)
  where
    reassoc (v, e, w) = (v, (e, w))

reverseE :: Graph e -> [Edge e]
reverseE g = [ (w, e, v) | (v, e, w) <- edges g ]

{-------------------------------------------------------------------------------
  Graph properties
-------------------------------------------------------------------------------}

vertices :: Graph e -> [Vertex]
vertices = indices

edges :: Graph e -> [Edge e]
edges g = [ (v, e, w) | v <- vertices g, (e, w) <- g!v ]

{-------------------------------------------------------------------------------
  Operations on the underlying unlabelled graph
-------------------------------------------------------------------------------}

forgetLabels :: Graph e -> G.Graph
forgetLabels = fmap (map snd)

topSort :: Graph e -> [Vertex]
topSort = G.topSort . forgetLabels