{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MonadComprehensions #-}
{-# LANGUAGE MultiParamTypeClasses #-}


{-| Module  : FiniteCategories
Description : The __'FinGrph'__ category has finite multidigraphs as objects and multidigraph homomorphisms as morphisms.
Copyright   : Guillaume Sabbagh 2022
License     : GPL-3
Maintainer  : guillaumesabbagh@protonmail.com
Stability   : experimental
Portability : portable

The __'FinGrph'__ category has finite multidigraphs as objects and multidigraph homomorphisms as morphisms.
-}

module Math.Categories.FinGrph
(
    -- * Graph
    Arrow(..),
    Graph,
    -- ** Getters
    nodes,
    edges,
    -- ** Smart constructors
    graph,
    unsafeGraph,
    -- ** Transformation
    mapOnNodes,
    mapOnEdges,
    -- * Graph homomorphism
    GraphHomomorphism,
    -- ** Getters
    nodeMap,
    edgeMap,
    -- ** Smart constructor
    checkGraphHomomorphism,
    graphHomomorphism,
    unsafeGraphHomomorphism,
    -- * FinGrph
    FinGrph(..),
    underlyingGraph,
    underlyingGraphFormat,
)
where
    import              Math.Category
    import              Math.FiniteCategory
    import              Math.CompleteCategory
    import              Math.CocompleteCategory
    import              Math.IO.PrettyPrint
    import              Math.Categories.FunctorCategory
    import              Math.Categories.ConeCategory
    import              Math.FiniteCategories.Parallel
    
    import              Data.WeakSet        (Set)
    import qualified    Data.WeakSet    as  Set
    import              Data.WeakSet.Safe
    import              Data.WeakMap        (Map)
    import qualified    Data.WeakMap    as  Map
    import              Data.WeakMap.Safe
    import              Data.Simplifiable
    
    import              GHC.Generics
    
    -- | An 'Arrow' is composed of a source node, a target node and a label.
    data Arrow n e = Arrow{
                            sourceArrow :: n,
                            targetArrow :: n,
                            labelArrow :: e
                          }
                          deriving (Eq, Show, Generic, Simplifiable)
    
    instance (PrettyPrint n, PrettyPrint e) => PrettyPrint (Arrow n e) where
        pprint v a = (pprint v $ sourceArrow a)++"-"++(pprint v $ labelArrow a)++"->"++(pprint v $ targetArrow a)
        
        -- pprintWithIndentations cv ov indent a = indentation (ov - cv) indent ++  (pprint cv $ sourceArrow a)++"-"++(pprint cv $ labelArrow a)++"->"++(pprint cv $ targetArrow a) ++ "\n"
    
    -- | A 'Graph' is a set of nodes and a set of 'Arrow's.
    -- 
    -- 'Graph' is private, use smart constructor 'graph'.
    data Graph n e = Graph {
                        nodes :: Set n, -- ^ The set of nodes of the graph.
                        edges :: Set (Arrow n e) -- ^ The set of arrows of the graph.
                        } deriving (Eq, Generic, PrettyPrint, Simplifiable)
    
    instance (Show n, Show e) => Show (Graph n e) where
        show g = "(unsafeGraph "++(show $ nodes g)++" "++(show $ edges g)++")"
    
    -- | Smart constructor of 'Graph'. The only error possible when creating a 'Graph' is that the source or target of an arrow is not in the set of nodes of the 'Graph'.
    graph :: (Eq n) => Set n -> Set (Arrow n e) -> Maybe (Graph n e)
    graph ns es
        | (sourceArrow <$> es) `isIncludedIn` ns && (targetArrow <$> es) `isIncludedIn` ns = Just Graph{nodes=ns, edges=es}
        | otherwise = Nothing
        
    -- | Unsafe constructor of 'Graph', does not check the 'Graph' structure.
    unsafeGraph :: Set n -> Set (Arrow n e) -> Graph n e
    unsafeGraph n e = Graph{nodes=n, edges=e}
    
    -- | Map a function on nodes of a 'Graph'.
    mapOnNodes :: (n1 -> n2) -> Graph n1 e -> Graph n2 e
    mapOnNodes transformNode g = Graph {nodes = transformNode <$> nodes g, edges = transformArrow <$> edges g}
        where
            transformArrow arr = Arrow{sourceArrow = transformNode $ sourceArrow arr, targetArrow = transformNode $ targetArrow arr, labelArrow = labelArrow arr}
    
    -- | Map a function on edges of a 'Graph'.
    mapOnEdges :: (e1 -> e2) -> Graph n e1 -> Graph n e2
    mapOnEdges transformEdge g = Graph {nodes = nodes g, edges = transformArrow <$> edges g}
        where
            transformArrow arr = Arrow{sourceArrow = sourceArrow arr, targetArrow = targetArrow arr, labelArrow = transformEdge $ labelArrow arr}
    
    -- | A 'GraphHomomorphism' is composed of a map between the nodes of the graphs, a map between the edges of the graphs, and the target 'Graph' so that we can recover it from the morphism.
    --
    -- It must follow axioms such that the image of an arrow is not torn appart, that is why the constructor is private. Use the smart constructor 'graphHomomorphism' instead.
    data GraphHomomorphism n e = GraphHomomorphism {
                                    nodeMap :: Map n n, -- ^ The mapping of nodes.
                                    edgeMap :: Map (Arrow n e) (Arrow n e), -- ^ The mapping of edges.
                                    targetGraph :: Graph n e -- ^ The target graph.
                                    } deriving (Eq, Generic, Simplifiable)
    
    -- | Check wether the structure of 'GraphHomomorphism' is respected or not.
    checkGraphHomomorphism :: (Eq n, Eq e) => GraphHomomorphism n e -> Bool
    checkGraphHomomorphism gh = imageInTarget && Set.and noTear
        where
            noTear = [(nodeMap gh) |!| (sourceArrow arr) == sourceArrow ((edgeMap gh) |!| arr) && (nodeMap gh) |!| (targetArrow arr) == targetArrow ((edgeMap gh) |!| arr)| arr <- (domain.edgeMap) gh]
            imageInTarget = (image.nodeMap) gh `isIncludedIn` (nodes.targetGraph) gh && (image.edgeMap) gh `isIncludedIn` (edges.targetGraph) gh
    
    -- | The smart constructor of 'GraphHomomorphism'.
    graphHomomorphism :: (Eq n, Eq e) => Map n n -> Map (Arrow n e) (Arrow n e) -> Graph n e -> Maybe (GraphHomomorphism n e)
    graphHomomorphism nm em tg
        | checkGraphHomomorphism gh = Just gh
        | otherwise = Nothing
        where
            gh = GraphHomomorphism{nodeMap=nm, edgeMap=em, targetGraph=tg}
    
    -- | Unsafe constructor of 'GraphHomomorphism' which does not check the structure of the 'GraphHomomorphism'.
    unsafeGraphHomomorphism :: Map n n -> Map (Arrow n e) (Arrow n e) -> Graph n e -> GraphHomomorphism n e
    unsafeGraphHomomorphism nm em tg = GraphHomomorphism{nodeMap=nm, edgeMap=em, targetGraph=tg}
    
    instance (Show n, Show e) => Show (GraphHomomorphism n e) where
        show gh = "(unsafeGraphHomomorphism "++(show $ nodeMap gh)++" "++(show $ edgeMap gh)++ " " ++ (show $ targetGraph gh) ++")"
    
    instance (PrettyPrint n, PrettyPrint e, Eq n, Eq e) => PrettyPrint (GraphHomomorphism n e) where
        pprint v gh = "GH("++(pprint (v-1) $ nodeMap gh)++", "++(pprint (v-1) $ edgeMap gh)++")"
        
        -- pprintWithIndentations 0 ov indent gh = indentation ov indent ++ "...\n"
        -- pprintWithIndentations cv ov indent gh = indentation (ov - cv) indent ++  "GH\n" ++ (pprintWithIndentations (cv-1) ov indent (nodeMap gh)) ++ (pprintWithIndentations (cv-1) ov indent (edgeMap gh))
        
    instance (Eq n, Eq e) => Morphism (GraphHomomorphism n e) (Graph n e) where
        source gh = Graph {nodes = (domain.nodeMap) gh, edges = (domain.edgeMap) gh}
        target = targetGraph
        (@) gh2 gh1 =  GraphHomomorphism {nodeMap = (nodeMap gh2) |.| (nodeMap gh1), edgeMap = (edgeMap gh2) |.| (edgeMap gh1), targetGraph = target gh2}
    
    -- | The category of finite graphs.
    data FinGrph n e = FinGrph deriving (Eq, Show, Generic, PrettyPrint, Simplifiable)
        
    instance (Eq n, Eq e) => Category (FinGrph n e) (GraphHomomorphism n e) (Graph n e) where
        identity _ g = GraphHomomorphism {nodeMap = (idFromSet.nodes) g, edgeMap = (idFromSet.edges) g, targetGraph = g}
        ar _ s t = [GraphHomomorphism
            {
                nodeMap = appO, edgeMap = appF, targetGraph = t
            } | appO <- appObj, appF <- ((fmap $ (Map.unions)).cartesianProductOfSets $ [twoObjToEdgeMaps x y appO | x <- (setToList $ nodes s), y <- (setToList $ nodes s)])]
            where
                appObj = Map.enumerateMaps (nodes s) (nodes t)
                twoObjToEdgeMaps n1 n2 nMap = Map.enumerateMaps (Set.filter (\a -> sourceArrow a == n1 && targetArrow a == n2) (edges s)) (Set.filter (\a -> sourceArrow a == nMap |!| n1 && targetArrow a == nMap |!| n2) (edges t))
    
    instance (Eq n, Eq e, Eq oIndex) => HasProducts (FinGrph n e) (GraphHomomorphism n e) (Graph n e) (FinGrph (Limit oIndex n) (Limit oIndex e)) (GraphHomomorphism (Limit oIndex n) (Limit oIndex e)) (Graph (Limit oIndex n) (Limit oIndex e)) oIndex where
        product discreteDiag = unsafeCone productGraph nat
            where
                indexingCat = src discreteDiag
                productGraph = Graph{nodes = productNodes, edges = productEdges}
                productNodes = (ProductElement).weakMap <$> cartesianProductOfSets (setToList [(\x -> (i,x)) <$> (nodes (discreteDiag ->$ i)) | i <- ob indexingCat])
                productEdges = (\tupleEdges -> Arrow{sourceArrow = ProductElement (sourceArrow <$> tupleEdges), targetArrow = ProductElement (targetArrow <$> tupleEdges), labelArrow = ProductElement (labelArrow <$> tupleEdges)}) <$> weakMap <$> cartesianProductOfSets (setToList [(\x -> (i,x)) <$> (edges (discreteDiag ->$ i)) | i <- ob indexingCat])
                newDiag = completeDiagram Diagram{src = indexingCat, tgt = FinGrph, omap = projectGraph <$> omap discreteDiag, mmap = weakMap []}
                nat = unsafeNaturalTransformation (constantDiagram (src discreteDiag) FinGrph productGraph) newDiag (Map.weakMapFromSet [(i, leg i) | i <- ob indexingCat])
                projectArrow a = Arrow{sourceArrow = Projection $ sourceArrow a, targetArrow = Projection $ targetArrow a, labelArrow = Projection $ labelArrow a}
                projectGraph g = Graph{nodes = Projection <$> nodes g, edges = projectArrow <$> edges g}
                leg i = GraphHomomorphism{targetGraph = projectGraph (discreteDiag ->$ i), nodeMap = Map.weakMapFromSet [(n, Projection $ tuple |!| i) | n@(ProductElement tuple)  <- nodes productGraph], edgeMap = Map.weakMapFromSet [(e, Arrow{sourceArrow = Projection $ (extractProd.sourceArrow $ e) |!| i , targetArrow = Projection $  (extractProd.targetArrow $ e) |!| i, labelArrow = Projection $ (extractProd.labelArrow $ e) |!| i}) | e <- edges productGraph]}
                extractProd (ProductElement x) = x
    
    instance (Eq n, Eq e) => HasEqualizers (FinGrph n e) (GraphHomomorphism n e) (Graph n e) where
        equalize parallelDiag = unsafeCone equalizedGraph nat
            where
                equalizedGraph = Graph{nodes = Set.filter (\n -> (nodeMap (parallelDiag ->£ ParallelF)) |!| n == (nodeMap (parallelDiag ->£ ParallelG)) |!| n)  (nodes (parallelDiag ->$ ParallelA)), edges = Set.filter (\e -> (edgeMap (parallelDiag ->£ ParallelF)) |!| e == (edgeMap (parallelDiag ->£ ParallelG)) |!| e)  (edges (parallelDiag ->$ ParallelA))}
                mappingNode i = memorizeFunction id (nodes equalizedGraph)
                mappingEdge i = memorizeFunction id (edges equalizedGraph)
                constDiag = constantDiagram Parallel FinGrph equalizedGraph
                nat = (unsafeNaturalTransformation constDiag parallelDiag (weakMap [(ParallelA,legA), (ParallelB, (parallelDiag ->£ ParallelF) @ legA) ]))
                legA = GraphHomomorphism {nodeMap=mappingNode ParallelA, edgeMap = mappingEdge ParallelA, targetGraph = parallelDiag ->$ ParallelA}
                
    instance (Eq n, Eq e, Eq mIndex, Eq oIndex) => CompleteCategory (FinGrph n e) (GraphHomomorphism n e) (Graph n e) (FinGrph (Limit oIndex n) (Limit oIndex e)) (GraphHomomorphism (Limit oIndex n) (Limit oIndex e)) (Graph (Limit oIndex n) (Limit oIndex e)) cIndex mIndex oIndex where
        limit = limitFromProductsAndEqualizers projectGraphHomomorphism
            where
                projectArrow a = Arrow{sourceArrow = Projection $ sourceArrow a, targetArrow = Projection $ targetArrow a, labelArrow = Projection $ labelArrow a}
                projectGraph g = Graph{nodes = Projection <$> nodes g, edges = projectArrow <$> edges g}
                projectGraphHomomorphism gh = GraphHomomorphism{nodeMap = doubleProject <|$|> nodeMap gh, edgeMap = doubleProjectArrow <|$|> edgeMap gh, targetGraph = projectGraph $ targetGraph gh}
                doubleProject (x,y) = (Projection x, Projection y)
                doubleProjectArrow (x,y) = (projectArrow x, projectArrow y)
                
        projectBase diag = Diagram{src = FinGrph, tgt = FinGrph, omap = memorizeFunction projectGraph (Map.values (omap diag)), mmap = memorizeFunction projectGraphHomomorphism (Map.values (mmap diag))}
            where
                projectArrow a = Arrow{sourceArrow = Projection $ sourceArrow a, targetArrow = Projection $ targetArrow a, labelArrow = Projection $ labelArrow a}
                projectGraph g = Graph{nodes = Projection <$> nodes g, edges = projectArrow <$> edges g}
                projectGraphHomomorphism gh = GraphHomomorphism{nodeMap = doubleProject <|$|> nodeMap gh, edgeMap = doubleProjectArrow <|$|> edgeMap gh, targetGraph = projectGraph $ targetGraph gh}
                doubleProject (x,y) = (Projection x, Projection y)
                doubleProjectArrow (x,y) = (projectArrow x, projectArrow y)
                
    instance (Eq n, Eq e, Eq oIndex) => HasCoproducts (FinGrph n e) (GraphHomomorphism n e) (Graph n e) (FinGrph (Colimit oIndex n) (Colimit oIndex e)) (GraphHomomorphism (Colimit oIndex n) (Colimit oIndex e)) (Graph (Colimit oIndex n) (Colimit oIndex e)) oIndex where
        coproduct discreteDiag = result
            where
                indexingCat = src discreteDiag
                coprod = Graph{nodes = Set.unions (setToList [CoproductElement i <$> nodes (discreteDiag ->$ i) | i <- ob indexingCat]), edges = Set.unions (setToList [coproductArrow i <$> edges (discreteDiag ->$ i) | i <- ob indexingCat])}
                coproductArrow i a = Arrow{sourceArrow = CoproductElement i (sourceArrow a), targetArrow = CoproductElement i (targetArrow a), labelArrow = CoproductElement i (labelArrow a)}
                constDiag = constantDiagram indexingCat FinGrph coprod
                coprojectArrow a = Arrow{sourceArrow = Coprojection $ sourceArrow a, targetArrow = Coprojection $ targetArrow a, labelArrow = Coprojection $ labelArrow a}
                transformGraph g = Graph{nodes = Coprojection <$> nodes g, edges = coprojectArrow <$> edges g}
                transformGH GraphHomomorphism{nodeMap = nm, edgeMap = em, targetGraph = tg} = GraphHomomorphism{nodeMap = weakMapFromSet [(Coprojection k, Coprojection v) | (k,v) <- Map.mapToSet nm], edgeMap = weakMapFromSet [(coprojectArrow k, coprojectArrow v) | (k,v) <- Map.mapToSet em], targetGraph = transformGraph tg}
                newDiag = Diagram{src = indexingCat, tgt = FinGrph, omap = transformGraph <$> (omap discreteDiag), mmap = transformGH <$> (mmap discreteDiag)}
                mapping i = GraphHomomorphism{nodeMap = memorizeFunction (\(Coprojection x) -> CoproductElement i x) (nodes (newDiag ->$ i)), edgeMap = memorizeFunction (\Arrow{sourceArrow = Coprojection s, targetArrow = Coprojection t, labelArrow = Coprojection l} -> Arrow{sourceArrow = CoproductElement i s, targetArrow = CoproductElement i t, labelArrow = CoproductElement i l}) (edges (newDiag ->$ i)), targetGraph = coprod}
                result = unsafeCocone coprod $ unsafeNaturalTransformation newDiag constDiag (memorizeFunction mapping (ob indexingCat))
                
    -- | BEWARE, for the coequalizer to be correct, ALL arrow labels should be different (two arrows with different source and target might have the same source and target after the coequalization process).
    instance (Eq e, Eq n) => HasCoequalizers (FinGrph n e) (GraphHomomorphism n e) (Graph n e) where
        coequalize parallelDiag = result
            where
                glueEdge edge gh
                    | imageEdgeByF == imageEdgeByG = gh
                    | otherwise = GraphHomomorphism{nodeMap = nodeMap gh, edgeMap = Map.adjust (const $ imageEdgeByG) imageEdgeByF (edgeMap gh), targetGraph = newGraph}
                    where
                        imageEdgeByF = (edgeMap (parallelDiag ->£ ParallelF)) |!| edge
                        imageEdgeByG = (edgeMap (parallelDiag ->£ ParallelG)) |!| edge
                        newGraph = Graph{nodes = nodes (target gh), edges = Set.delete imageEdgeByF (edges (target gh))}
                glueNode node gh
                    | imageNodeByF == imageNodeByG = gh
                    | otherwise = GraphHomomorphism{nodeMap = Map.adjust (const $ imageNodeByG) imageNodeByF (nodeMap gh), edgeMap = updateArrow <$> edgeMap gh, targetGraph = newGraph}
                    where
                        imageNodeByF = (nodeMap (parallelDiag ->£ ParallelF)) |!| node
                        imageNodeByG = (nodeMap (parallelDiag ->£ ParallelG)) |!| node
                        updateNode n = if n == imageNodeByF then imageNodeByG else n
                        updateArrow a = Arrow{sourceArrow = updateNode (sourceArrow a), targetArrow = updateNode (targetArrow a), labelArrow = labelArrow a}
                        newGraph = Graph{nodes = Set.delete imageNodeByF (nodes (target gh)), edges = updateArrow <$> edges (target gh)}
                gh1 = Set.foldr glueEdge (identity FinGrph (parallelDiag ->$ ParallelB)) (edges (parallelDiag ->$ ParallelA))
                gh2 = Set.foldr glueNode gh1 (nodes (parallelDiag ->$ ParallelA))
                constDiag = constantDiagram Parallel FinGrph (target gh2)
                result = unsafeCocone (target gh2) (unsafeNaturalTransformation parallelDiag constDiag (weakMap [(ParallelA,gh2 @ (parallelDiag ->£ ParallelF)), (ParallelB, gh2) ]))
        
    instance (Eq e, Eq n, Eq mIndex, Eq oIndex) => CocompleteCategory (FinGrph n e) (GraphHomomorphism n e) (Graph n e) (FinGrph (Colimit oIndex n) (Colimit oIndex e)) (GraphHomomorphism (Colimit oIndex n) (Colimit oIndex e)) (Graph (Colimit oIndex n) (Colimit oIndex e)) cIndex mIndex oIndex where
        colimit = colimitFromCoproductsAndCoequalizers transformGHToColimGH
            where
                transformGHToColimGH gh = GraphHomomorphism{nodeMap = both Coprojection <|$|> nodeMap gh, edgeMap = both coprojectArrow <|$|> edgeMap gh, targetGraph = coprojectTargetGraph (targetGraph gh)}
                coprojectArrow a = Arrow{sourceArrow = Coprojection $ sourceArrow a, targetArrow = Coprojection $ targetArrow a, labelArrow = Coprojection $ labelArrow a}
                both f (x,y) = (f x,f y)
                coprojectTargetGraph g = Graph{nodes = Coprojection <$> nodes g, edges = coprojectArrow <$> edges g}
        
        coprojectBase diag = Diagram{src = FinGrph, tgt = FinGrph, omap = memorizeFunction coprojectGraph (Map.values (omap diag)), mmap = memorizeFunction transformGHToColimGH (Map.values (mmap diag))}
            where
                transformGHToColimGH gh = GraphHomomorphism{nodeMap = both Coprojection <|$|> nodeMap gh, edgeMap = both coprojectArrow <|$|> edgeMap gh, targetGraph = coprojectGraph (targetGraph gh)}
                coprojectArrow a = Arrow{sourceArrow = Coprojection $ sourceArrow a, targetArrow = Coprojection $ targetArrow a, labelArrow = Coprojection $ labelArrow a}
                both f (x,y) = (f x,f y)
                coprojectGraph g = Graph{nodes = Coprojection <$> nodes g, edges = coprojectArrow <$> edges g}
        
        
        
        
        
        
        
        
        
        
    -- | Return the underlying graph of a 'FiniteCategory'.
    underlyingGraph :: (FiniteCategory c m o, Morphism m o) => c -> Graph o m
    underlyingGraph c = Graph{
                                nodes = ob c,
                                edges = (\m -> Arrow{sourceArrow=source m, targetArrow=target m, labelArrow=m}) <$> arrows c
                            }
                            
    -- | Return the underlying graph of a 'FiniteCategory' and apply formatting functions on objects and arrows.
    underlyingGraphFormat :: (FiniteCategory c m o, Morphism m o) => (o -> a) -> (m -> b) -> c -> Graph a b
    underlyingGraphFormat formatObj formatAr c = Graph{
                                                        nodes = formatObj <$> ob c,
                                                        edges = (\m -> Arrow{sourceArrow=formatObj.source $ m, targetArrow=formatObj.target $ m, labelArrow=formatAr m}) <$> arrows c
                                                    }