{-# LANGUAGE TypeFamilies, BangPatterns #-}
-- | This graph is based on the implementation in fgl (using
-- big-endian patricia-tries -- IntMap).
--
-- This formulation does not support parallel edges.
module Data.Graph.Haggle.PatriciaTree ( PatriciaTree ) where

import Control.DeepSeq
import Control.Monad ( guard )
import Data.Foldable ( toList )
import Data.IntMap ( IntMap )
import qualified Data.IntMap as IM
import Data.Maybe ( fromMaybe )
import Data.Monoid

import           Prelude

import qualified Data.Graph.Haggle.Classes as I
import qualified Data.Graph.Haggle.Internal.Basic as I

data Ctx nl el = Ctx !(IntMap el) I.Vertex nl !(IntMap el)

instance (NFData nl, NFData el) => NFData (Ctx nl el) where
  rnf (Ctx p v nl s) =
    p `deepseq` s `deepseq` nl `deepseq` v `seq` ()

-- | The 'PatriciaTree' is a graph implementing the 'I.InductiveGraph'
-- interface (as well as the other immutable graph interfaces).  It is
-- based on the graph type provided by fgl.
--
-- Inductive graphs support more interesting decompositions than the
-- other graph interfaces in this library, at the cost of less compact
-- representations and some additional overhead on some operations, as
-- most must go through the 'I.match' operator.
--
-- This graph type is most useful for incremental construction in pure
-- code.  It also supports node removal from pure code.
data PatriciaTree nl el = Gr { graphRepr :: IntMap (Ctx nl el) }

instance (NFData nl, NFData el) => NFData (PatriciaTree nl el) where
  rnf (Gr im) = im `deepseq` ()

instance I.Graph (PatriciaTree nl el) where
  vertices = map I.V . IM.keys . graphRepr
  isEmpty = IM.null . graphRepr
  maxVertexId (Gr g)
    | IM.null g = 0
    | otherwise = fst $ IM.findMax g
  edgesBetween (Gr g) (I.V src) (I.V dst) = toList $ do
    Ctx _ _ _ ss <- IM.lookup src g
    guard (IM.member dst ss)
    return (I.E (-1) src dst)
  edges g = concatMap (I.outEdges g) (I.vertices g)
  successors (Gr g) (I.V v) = fromMaybe [] $ do
    Ctx _ _ _ ss <- IM.lookup v g
    return $ map I.V $ IM.keys ss
  outEdges (Gr g) (I.V v) = fromMaybe [] $ do
    Ctx _ _ _ ss <- IM.lookup v g
    return $ map toEdge (IM.keys ss)
    where
      toEdge d = I.E (-1) v d

instance I.HasEdgeLabel (PatriciaTree nl el) where
  type EdgeLabel (PatriciaTree nl el) = el
  edgeLabel (Gr g) (I.E _ src dst) = do
    Ctx _ _ _ ss <- IM.lookup src g
    IM.lookup dst ss
  labeledEdges gr = map toLabEdge (I.edges gr)
    where
      toLabEdge e =
        let Just lab = I.edgeLabel gr e
        in (e, lab)
  labeledOutEdges (Gr g) (I.V s) = fromMaybe [] $ do
    Ctx _ _ _ ss <- IM.lookup s g
    return $ IM.foldrWithKey toOut [] ss
    where
      toOut d lbl acc = (I.E (-1) s d, lbl) : acc

instance I.HasVertexLabel (PatriciaTree nl el) where
  type VertexLabel (PatriciaTree nl el) = nl
  vertexLabel (Gr g) (I.V v) = do
    Ctx _ _ lbl _ <- IM.lookup v g
    return lbl
  labeledVertices gr = map toLabVert (I.vertices gr)
    where
      toLabVert v =
        let Just l = I.vertexLabel gr v
        in (v, l)

instance I.Bidirectional (PatriciaTree nl el) where
  predecessors (Gr g) (I.V v) = fromMaybe [] $ do
    Ctx pp _ _ _ <- IM.lookup v g
    return $ map I.V (IM.keys pp)
  inEdges (Gr g) (I.V v) = fromMaybe [] $ do
    Ctx pp _ _ _ <- IM.lookup v g
    return $ map toEdge (IM.keys pp)
    where
      toEdge s = I.E (-1) s v

instance I.BidirectionalEdgeLabel (PatriciaTree nl el) where
  labeledInEdges (Gr g) (I.V d) = fromMaybe [] $ do
    Ctx pp _ _ _ <- IM.lookup d g
    return $ IM.foldrWithKey toIn [] pp
    where
      toIn s lbl acc = (I.E (-1) s d, lbl) : acc

instance I.InductiveGraph (PatriciaTree nl el) where
  emptyGraph = Gr IM.empty
  insertLabeledVertex gr@(Gr g) lab =
    let vid = I.maxVertexId gr + 1
        v = I.V vid
        g' = IM.insert vid (Ctx mempty v lab mempty) g
    in (v, Gr g')
  insertLabeledEdge gr@(Gr g) v1@(I.V src) v2@(I.V dst) lab = do
    guard (IM.member src g && IM.member dst g)
    guard (not (I.edgeExists gr v1 v2))
    let e = I.E (-1) src dst
    Ctx spp sv sl sss <- IM.lookup src g
    Ctx dpp dv dl dss <- IM.lookup dst g
    let sctx' = Ctx spp sv sl (IM.insert dst lab sss)
        dctx' = Ctx (IM.insert src lab dpp) dv dl dss
        !g' = IM.insert src sctx' g
        !g'' = IM.insert dst dctx' g'
    return (e, Gr g'')
  deleteEdge g (I.E _ s d) = I.deleteEdgesBetween g (I.V s) (I.V d)
  deleteEdgesBetween gr@(Gr g) (I.V src) (I.V dst) = fromMaybe gr $ do
    Ctx spp sv sl sss <- IM.lookup src g
    Ctx dpp dv dl dss <- IM.lookup dst g
    let sctx' = Ctx spp sv sl (IM.delete dst sss)
        dctx' = Ctx (IM.delete src dpp) dv dl dss
        !g' = IM.insert src sctx' g
        !g'' = IM.insert dst dctx' g'
    return (Gr g'')
  context (Gr g) (I.V v) = do
    Ctx pp _ l ss <- IM.lookup v g
    return $ I.Context (toAdj pp) l (toAdj ss)
  match (Gr g) (I.V v) = do
    Ctx pp _ l ss <- IM.lookup v g
    let g' = foldr (IM.adjust (removeSucc v)) g (IM.keys pp)
        g'' = foldr (IM.adjust (removePred v)) g' (IM.keys ss)
        g''' = IM.delete v g''
    return $ (I.Context (toAdj pp) l (toAdj ss), Gr g''')

toAdj :: IntMap a -> [(a, I.Vertex)]
toAdj = IM.foldrWithKey f []
  where
    f dst lbl acc = (lbl, I.V dst) : acc

removeSucc :: Int -> Ctx nl el -> Ctx nl el
removeSucc v (Ctx pp vert lbl ss) =
  Ctx pp vert lbl (IM.delete v ss)

removePred :: Int -> Ctx nl el -> Ctx nl el
removePred v (Ctx pp vert lbl ss) =
  Ctx (IM.delete v pp) vert lbl ss

{- Note [Representation]

Since this graph does not support parallel edges, the edge ID does not
actually matter.  This implementation will let it always be zero.  Edge
identity can be recovered with just (src, dst).

-}