module Data.GraphViz.Types.Graph
       ( DotGraph
       , GraphID(..)
       , Context(..)
         
       , toCanonical
       , unsafeFromCanonical
       , fromDotRepr
         
       , isEmpty
       , hasClusters
       , isEmptyGraph
       , graphAttributes
       , parentOf
       , clusterAttributes
       , foundInCluster
       , attributesOf
       , predecessorsOf
       , successorsOf
       , adjacentTo
       , adjacent
         
       , mkGraph
       , emptyGraph
       , (&)
       , composeList
       , addNode
       , DotNode(..)
       , addDotNode
       , addEdge
       , DotEdge(..)
       , addDotEdge
       , addCluster
       , setClusterParent
       , setClusterAttributes
         
       , decompose
       , decomposeAny
       , decomposeList
       , deleteNode
       , deleteAllEdges
       , deleteEdge
       , deleteDotEdge
       , deleteCluster
       , removeEmptyClusters
       ) where
import           Data.GraphViz.Algorithms            (CanonicaliseOptions (..),
                                                      canonicaliseOptions)
import           Data.GraphViz.Algorithms.Clustering
import           Data.GraphViz.Attributes.Complete   (Attributes)
import           Data.GraphViz.Attributes.Same
import           Data.GraphViz.Internal.Util         (groupSortBy,
                                                      groupSortCollectBy)
import           Data.GraphViz.Types
import qualified Data.GraphViz.Types.Canonical       as C
import qualified Data.GraphViz.Types.Generalised     as G
import           Data.GraphViz.Types.Internal.Common (partitionGlobal)
import qualified Data.GraphViz.Types.State           as St
import           Control.Applicative             (liftA2, (<$>), (<*>))
import           Control.Arrow                   ((***))
import qualified Data.Foldable                   as F
import           Data.List                       (delete, foldl', unfoldr)
import           Data.Map                        (Map)
import qualified Data.Map                        as M
import           Data.Maybe                      (fromMaybe, mapMaybe)
import qualified Data.Sequence                   as Seq
import qualified Data.Set                        as S
import           Text.ParserCombinators.ReadPrec (prec)
import           Text.Read                       (Lexeme (Ident), lexP, parens,
                                                  readPrec)
data DotGraph n = DG { strictGraph   :: !Bool
                     , directedGraph :: !Bool
                     , graphAttrs    :: !GlobAttrs
                     , graphID       :: !(Maybe GraphID)
                     , clusters      :: !(Map GraphID ClusterInfo)
                     , values        :: !(NodeMap n)
                     }
                deriving (Eq, Ord)
instance (Ord n, Show n) => Show (DotGraph n) where
  showsPrec d dg = showParen (d > 10) $
                   showString "fromCanonical " . shows (toCanonical dg)
instance (Ord n, Read n) => Read (DotGraph n) where
  readPrec = parens . prec 10
             $ do Ident "fromCanonical" <- lexP
                  cdg <- readPrec
                  return $ fromCanonical cdg
data GlobAttrs = GA { graphAs :: !SAttrs
                    , nodeAs  :: !SAttrs
                    , edgeAs  :: !SAttrs
                    }
               deriving (Eq, Ord, Show, Read)
data NodeInfo n = NI { _inCluster    :: !(Maybe GraphID)
                     , _attributes   :: !Attributes
                     , _predecessors :: !(EdgeMap n)
                     , _successors   :: !(EdgeMap n)
                     }
                deriving (Eq, Ord, Show, Read)
data ClusterInfo = CI { parentCluster :: !(Maybe GraphID)
                      , clusterAttrs  :: !GlobAttrs
                      }
                 deriving (Eq, Ord, Show, Read)
type NodeMap n = Map n (NodeInfo n)
type EdgeMap n = Map n [Attributes]
data Context n = Cntxt { node         :: !n
                         
                         
                         
                       , inCluster    :: !(Maybe GraphID)
                       , attributes   :: !Attributes
                       , predecessors :: ![(n, Attributes)]
                       , successors   :: ![(n, Attributes)]
                       }
               deriving (Eq, Ord, Show, Read)
adjacent :: Context n -> [DotEdge n]
adjacent c = mapU (`DotEdge` n) (predecessors c)
             ++ mapU (DotEdge n) (successors c)
  where
    n = node c
    mapU = map . uncurry
emptyGraph :: DotGraph n
emptyGraph = DG { strictGraph   = False
                , directedGraph = True
                , graphID       = Nothing
                , graphAttrs    = emptyGA
                , clusters      = M.empty
                , values        = M.empty
                }
emptyGA :: GlobAttrs
emptyGA = GA S.empty S.empty S.empty
(&) :: (Ord n) => Context n -> DotGraph n -> DotGraph n
(Cntxt n mc as ps ss) & dg = withValues merge dg'
  where
    ps' = toMap ps
    ps'' = M.delete n ps'
    ss' = toMap ss
    ss'' = M.delete n ss'
    dg' = addNode n mc as dg
    merge = addSucc n ps'' . addPred n ss''
            . M.adjust (\ni -> ni { _predecessors = ps', _successors = ss' }) n
infixr 5 &
composeList :: (Ord n) => [Context n] -> DotGraph n
composeList = foldr (&) emptyGraph
addSucc :: (Ord n) => n -> EdgeMap n -> NodeMap n -> NodeMap n
addSucc = addPS niSucc
addPred :: (Ord n) => n -> EdgeMap n -> NodeMap n -> NodeMap n
addPred = addPS niPred
addPS :: (Ord n) => ((EdgeMap n -> EdgeMap n) -> NodeInfo n -> NodeInfo n)
         -> n -> EdgeMap n -> NodeMap n -> NodeMap n
addPS fni t fas nm = t `seq` foldl' addSucc' nm fas'
  where
    fas' = fromMap fas
    addSucc' nm' (f,as) = f `seq` M.alter (addS as) f nm'
    addS as = Just
              . maybe (error "Node not in the graph!")
                      (fni (M.insertWith (++) t [as]))
addNode :: (Ord n)
           => n
           -> Maybe GraphID 
                            
           -> Attributes
           -> DotGraph n
           -> DotGraph n
addNode n mc as dg
  | n `M.member` ns = error "Node already exists in the graph"
  | otherwise       = addEmptyCluster mc
                      $ dg { values   = ns' }
  where
    ns = values dg
    ns' = M.insert n (NI mc as M.empty M.empty) ns
addDotNode                :: (Ord n) => DotNode n -> DotGraph n -> DotGraph n
addDotNode (DotNode n as) = addNode n Nothing as
addEdge :: (Ord n) => n -> n -> Attributes -> DotGraph n -> DotGraph n
addEdge f t as = withValues merge
  where
    
    
    merge = addPred t (M.singleton f [as]) . addSucc f (M.singleton t [as])
addDotEdge                  :: (Ord n) => DotEdge n -> DotGraph n -> DotGraph n
addDotEdge (DotEdge f t as) = addEdge f t as
addCluster :: GraphID          
              -> Maybe GraphID 
                               
                               
              -> [GlobalAttributes]
              -> DotGraph n
              -> DotGraph n
addCluster c mp gas dg
  | c `M.member` cs = error "Cluster already exists in the graph"
  | otherwise       = addEmptyCluster mp
                      $ dg { clusters = M.insert c ci cs }
  where
    cs = clusters dg
    ci = CI mp $ toGlobAttrs gas
addEmptyCluster :: Maybe GraphID -> DotGraph n -> DotGraph n
addEmptyCluster = maybe id (withClusters . (`dontReplace` defCI))
  where
    dontReplace = M.insertWith (const id)
    defCI = CI Nothing emptyGA
setClusterParent     :: GraphID -> Maybe GraphID -> DotGraph n -> DotGraph n
setClusterParent c p = withClusters (M.adjust setP c) . addCs
  where
    addCs = addEmptyCluster p . addEmptyCluster (Just c)
    setP ci = ci { parentCluster = p }
setClusterAttributes       :: GraphID -> [GlobalAttributes]
                              -> DotGraph n -> DotGraph n
setClusterAttributes c gas = withClusters (M.adjust setAs c)
                             . addEmptyCluster (Just c)
  where
    setAs ci = ci { clusterAttrs = toGlobAttrs gas }
mkGraph :: (Ord n) => [DotNode n] -> [DotEdge n] -> DotGraph n
mkGraph ns es = flip (foldl' $ flip addDotEdge) es
                $ foldl' (flip addDotNode) emptyGraph ns
toCanonical :: (Ord n) => DotGraph n -> C.DotGraph n
toCanonical dg = C.DotGraph { C.strictGraph     = strictGraph dg
                            , C.directedGraph   = directedGraph dg
                            , C.graphID         = graphID dg
                            , C.graphStatements = stmts
                            }
  where
    stmts = C.DotStmts { C.attrStmts = fromGlobAttrs $ graphAttrs dg
                       , C.subGraphs = cs
                       , C.nodeStmts = ns
                       , C.edgeStmts = getEdgeInfo False dg
                       }
    cls = clusters dg
    pM = clusterPath' dg
    clustAs = maybe [] (fromGlobAttrs . clusterAttrs) . (`M.lookup`cls)
    lns = map (\ (n,ni) -> (n,(_inCluster ni, _attributes ni)))
          . M.assocs $ values dg
    (cs,ns) = clustersToNodes pathOf (const True) id clustAs snd lns
    pathOf (n,(c,as)) = pathFrom c (n,as)
    pathFrom c ln = F.foldr C (N ln) . fromMaybe Seq.empty $ (`M.lookup`pM) =<< c
decompose :: (Ord n) => n -> DotGraph n -> Maybe (Context n, DotGraph n)
decompose n dg
  | n `M.notMember` ns = Nothing
  | otherwise          = Just (c, dg')
  where
    ns = values dg
    (Just (NI mc as ps ss), ns') = M.updateLookupWithKey (const . const Nothing) n ns
    c = Cntxt n mc as (fromMap $ n `M.delete` ps) (fromMap ss)
    dg' = dg { values = delSucc n ps . delPred n ss $ ns' }
decomposeAny :: (Ord n) => DotGraph n -> Maybe (Context n, DotGraph n)
decomposeAny dg
  | isEmpty dg = Nothing
  | otherwise  = decompose (fst . M.findMin $ values dg) dg
decomposeList :: (Ord n) => DotGraph n -> [Context n]
decomposeList = unfoldr decomposeAny
delSucc :: (Ord n) => n -> EdgeMap n -> NodeMap n -> NodeMap n
delSucc = delPS niSucc
delPred :: (Ord n) => n -> EdgeMap n -> NodeMap n -> NodeMap n
delPred = delPS niPred
delPS :: (Ord n) => ((EdgeMap n -> EdgeMap n) -> NodeInfo n -> NodeInfo n)
         -> n -> EdgeMap n -> NodeMap n -> NodeMap n
delPS fni t fm nm = foldl' delE nm $ M.keys fm
  where
    delE nm' f = M.adjust (fni $ M.delete t) f nm'
deleteNode      :: (Ord n) => n -> DotGraph n -> DotGraph n
deleteNode n dg = maybe dg snd $ decompose n dg
deleteAllEdges          :: (Ord n) => n -> n -> DotGraph n -> DotGraph n
deleteAllEdges n1 n2 = withValues (delAE n1 n2 . delAE n2 n1)
  where
    delAE f t = delSucc f t' . delPred f t'
      where
        t' = M.singleton t []
deleteEdge :: (Ord n) => n -> n -> Attributes -> DotGraph n -> DotGraph n
deleteEdge n1 n2 as dg = withValues delEs dg
  where
    delE f t = M.adjust (niSucc $ M.adjust (delete as) t) f
               . M.adjust (niPred $ M.adjust (delete as) f) t
    delEs | directedGraph dg = delE n1 n2
          | otherwise        = delE n1 n2 . delE n2 n1
deleteDotEdge :: (Ord n) => DotEdge n -> DotGraph n -> DotGraph n
deleteDotEdge (DotEdge n1 n2 as) = deleteEdge n1 n2 as
deleteCluster      :: (Ord n) => GraphID -> DotGraph n -> DotGraph n
deleteCluster c dg = withValues (M.map adjNode)
                     . withClusters (M.map adjCluster . M.delete c)
                     $ dg
  where
    p = parentCluster =<< c `M.lookup` clusters dg
    adjParent p'
      | p' == Just c = p
      | otherwise    = p'
    adjNode ni = ni { _inCluster = adjParent $ _inCluster ni }
    adjCluster ci = ci { parentCluster = adjParent $ parentCluster ci }
removeEmptyClusters :: (Ord n) => DotGraph n -> DotGraph n
removeEmptyClusters dg = dg { clusters = cM' }
  where
    cM = clusters dg
    cM' = (cM `M.difference` invCs) `M.difference` invNs
    invCs = usedClustsIn $ M.map parentCluster cM
    invNs = usedClustsIn . M.map _inCluster $ values dg
    usedClustsIn = M.fromAscList
                   . map ((,) <$> fst . head <*> map snd)
                   . groupSortBy fst
                   . mapMaybe (uncurry (fmap . flip (,)))
                   . M.assocs
isEmpty :: DotGraph n -> Bool
isEmpty = M.null . values
hasClusters :: DotGraph n -> Bool
hasClusters = M.null . clusters
isEmptyGraph :: DotGraph n -> Bool
isEmptyGraph = liftA2 (&&) isEmpty (not . hasClusters)
graphAttributes :: DotGraph n -> [GlobalAttributes]
graphAttributes = fromGlobAttrs . graphAttrs
foundInCluster :: (Ord n) => DotGraph n -> n -> Maybe GraphID
foundInCluster dg n = _inCluster $ values dg M.! n
attributesOf :: (Ord n) => DotGraph n -> n -> Attributes
attributesOf dg n = _attributes $ values dg M.! n
predecessorsOf :: (Ord n) => DotGraph n -> n -> [DotEdge n]
predecessorsOf dg t
  | directedGraph dg = emToDE (`DotEdge` t)
                       . _predecessors $ values dg M.! t
  | otherwise        = adjacentTo dg t
successorsOf :: (Ord n) => DotGraph n -> n -> [DotEdge n]
successorsOf dg f
  | directedGraph dg = emToDE (DotEdge f)
                       . _successors $ values dg M.! f
  | otherwise        = adjacentTo dg f
adjacentTo :: (Ord n) => DotGraph n -> n -> [DotEdge n]
adjacentTo dg n = sucs ++ preds
  where
    ni = values dg M.! n
    sucs = emToDE (DotEdge n) $ _successors ni
    preds = emToDE (`DotEdge` n) $ n `M.delete` _predecessors ni
emToDE :: (Ord n) => (n -> Attributes -> DotEdge n)
          -> EdgeMap n -> [DotEdge n]
emToDE f = map (uncurry f) . fromMap
parentOf :: DotGraph n -> GraphID -> Maybe GraphID
parentOf dg c = parentCluster $ clusters dg M.! c
clusterAttributes :: DotGraph n -> GraphID -> [GlobalAttributes]
clusterAttributes dg c = fromGlobAttrs . clusterAttrs $ clusters dg M.! c
instance (Ord n) => DotRepr DotGraph n where
  fromCanonical = fromDotRepr
  getID = graphID
  setID i g = g { graphID = Just i }
  graphIsDirected = directedGraph
  setIsDirected d g = g { directedGraph = d }
  graphIsStrict = strictGraph
  setStrictness s g = g { strictGraph = s }
  mapDotGraph = mapNs
  graphStructureInformation = getGraphInfo
  nodeInformation = getNodeInfo
  edgeInformation = getEdgeInfo
  unAnonymise = id 
instance (Ord n) => G.FromGeneralisedDot DotGraph n where
  fromGeneralised = fromDotRepr
instance (Ord n, PrintDot n) => PrintDotRepr DotGraph n
instance (Ord n, ParseDot n) => ParseDotRepr DotGraph n
instance (Ord n, PrintDot n, ParseDot n) => PPDotRepr DotGraph n
instance (Ord n, PrintDot n) => PrintDot (DotGraph n) where
  unqtDot = unqtDot . toCanonical
instance (Ord n, ParseDot n) => ParseDot (DotGraph n) where
  parseUnqt = fromGDot <$> parseUnqt
    where
      
      fromGDot = fromDotRepr . (`asTypeOf` (undefined :: G.DotGraph n))
  parse = parseUnqt 
cOptions :: CanonicaliseOptions
cOptions = COpts { edgesInClusters = False
                 , groupAttributes = True
                 }
fromDotRepr :: (DotRepr dg n) => dg n -> DotGraph n
fromDotRepr = unsafeFromCanonical . canonicaliseOptions cOptions . unAnonymise
unsafeFromCanonical :: (Ord n) => C.DotGraph n -> DotGraph n
unsafeFromCanonical dg = DG { strictGraph   = C.strictGraph dg
                            , directedGraph = dirGraph
                            , graphAttrs    = as
                            , graphID       = mgid
                            , clusters      = cs
                            , values        = ns
                            }
  where
    stmts = C.graphStatements dg
    mgid = C.graphID dg
    dirGraph = C.directedGraph dg
    (as, cs, ns) = fCStmt Nothing stmts
    fCStmt p stmts' = (sgAs, cs', ns')
      where
        sgAs = toGlobAttrs $ C.attrStmts stmts'
        (cs', sgNs) = (M.unions *** M.unions) . unzip
                      . map (fCSG p) $ C.subGraphs stmts'
        nNs = M.fromList . map (fDN p) $ C.nodeStmts stmts'
        ns' = sgNs `M.union` nNs
    fCSG p sg = (M.insert sgid ci cs', ns')
      where
        msgid@(Just sgid) = C.subGraphID sg
        (as', cs', ns') = fCStmt msgid $ C.subGraphStmts sg
        ci = CI p as'
    fDN p (DotNode n as') = ( n
                            , NI { _inCluster    = p
                                 , _attributes   = as'
                                 , _predecessors = eSel n tEs
                                 , _successors   = eSel n fEs
                                 }
                            )
    es = C.edgeStmts stmts
    fEs = toEdgeMap fromNode toNode es
    tEs = delLoops $ toEdgeMap toNode fromNode es
    eSel n es' = fromMaybe M.empty $ n `M.lookup` es'
    delLoops = M.mapWithKey M.delete
toEdgeMap     :: (Ord n) => (DotEdge n -> n) -> (DotEdge n -> n) -> [DotEdge n]
                 -> Map n (EdgeMap n)
toEdgeMap f t = M.map eM . M.fromList . groupSortCollectBy f t'
  where
    t' = liftA2 (,) t edgeAttributes
    eM = M.fromList . groupSortCollectBy fst snd
mapNs :: (Ord n, Ord n') => (n -> n') -> DotGraph n -> DotGraph n'
mapNs f (DG st d as mid cs vs) = DG st d as mid cs
                                 $ mapNM vs
  where
    mapNM = M.map mapNI . mpM
    mapNI (NI mc as' ps ss) = NI mc as' (mpM ps) (mpM ss)
    mpM = M.mapKeys f
getGraphInfo    :: DotGraph n -> (GlobalAttributes, ClusterLookup)
getGraphInfo dg = (gas, cl)
  where
    toGA = GraphAttrs . unSame
    (gas, cgs) = (toGA *** M.map toGA) $ globAttrMap graphAs dg
    pM = M.map pInit $ clusterPath dg
    cl = M.mapWithKey addPath $ M.mapKeysMonotonic Just cgs
    addPath c as = ( maybe [] (:[]) $ c `M.lookup` pM
                   , as
                   )
    pInit p = case Seq.viewr p of
                (p' Seq.:> _) -> p'
                _             -> Seq.empty
getNodeInfo             :: (Ord n) => Bool -> DotGraph n
                           -> NodeLookup n
getNodeInfo withGlob dg = M.map toLookup ns
  where
    (gGlob, aM) = globAttrMap nodeAs dg
    pM = clusterPath dg
    ns = values dg
    toLookup ni = (pth, as')
      where
        as = _attributes ni
        mp = _inCluster ni
        pth = fromMaybe Seq.empty $ mp `M.lookup` pM
        pAs = fromMaybe gGlob $ (`M.lookup` aM) =<< mp
        as' | withGlob  = unSame $ toSAttr as `S.union` pAs
            | otherwise = as
getEdgeInfo             :: (Ord n) => Bool -> DotGraph n -> [DotEdge n]
getEdgeInfo withGlob dg = concatMap (uncurry mkDotEdges) es
  where
    gGlob = edgeAs $ graphAttrs dg
    es = concatMap (uncurry (map . (,)))
         . M.assocs . M.map (M.assocs . _successors)
         $ values dg
    addGlob as
      | withGlob  = unSame $ toSAttr as `S.union` gGlob
      | otherwise = as
    mkDotEdges f (t, ass) = map (DotEdge f t . addGlob) ass
globAttrMap       :: (GlobAttrs -> SAttrs) -> DotGraph n
                     -> (SAttrs, Map GraphID SAttrs)
globAttrMap af dg = (gGlob, aM)
  where
    gGlob = af $ graphAttrs dg
    cs = clusters dg
    aM = M.map attrsFor cs
    attrsFor ci = as `S.union` pAs
      where
        as = af $ clusterAttrs ci
        p = parentCluster ci
        pAs = fromMaybe gGlob $ (`M.lookup` aM) =<< p
clusterPath :: DotGraph n -> Map (Maybe GraphID) St.Path
clusterPath = M.mapKeysMonotonic Just . M.map (fmap Just) . clusterPath'
clusterPath' :: DotGraph n -> Map GraphID (Seq.Seq GraphID)
clusterPath' dg = pM
  where
    cs = clusters dg
    pM = M.mapWithKey pathOf cs
    pathOf c ci = pPth Seq.|> c
      where
        mp = parentCluster ci
        pPth = fromMaybe Seq.empty $ (`M.lookup` pM) =<< mp
withValues      :: (Ord n) => (NodeMap n -> NodeMap n)
                   -> DotGraph n -> DotGraph n
withValues f dg = dg { values = f $ values dg }
withClusters      :: (Map GraphID ClusterInfo -> Map GraphID ClusterInfo)
                     -> DotGraph n -> DotGraph n
withClusters f dg = dg { clusters = f $ clusters dg }
toGlobAttrs :: [GlobalAttributes] -> GlobAttrs
toGlobAttrs = mkGA . partitionGlobal
  where
    mkGA (ga,na,ea) = GA (toSAttr ga) (toSAttr na) (toSAttr ea)
fromGlobAttrs :: GlobAttrs -> [GlobalAttributes]
fromGlobAttrs (GA ga na ea) = filter (not . null . attrs)
                              [ GraphAttrs $ unSame ga
                              , NodeAttrs  $ unSame na
                              , EdgeAttrs  $ unSame ea
                              ]
niSucc      :: (Ord n) => (EdgeMap n -> EdgeMap n) -> NodeInfo n -> NodeInfo n
niSucc f ni = ni { _successors = f $ _successors ni }
niPred      :: (Ord n) => (EdgeMap n -> EdgeMap n) -> NodeInfo n -> NodeInfo n
niPred f ni = ni { _predecessors = f $ _predecessors ni }
toMap :: (Ord n) => [(n, Attributes)] -> EdgeMap n
toMap = M.fromAscList . groupSortCollectBy fst snd
fromMap :: EdgeMap n -> [(n, Attributes)]
fromMap = concatMap (uncurry (map . (,))) . M.toList