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

{-| Module  : FiniteCategories
Description : A 'ColimitCategory' is the colimit of a diagram in 'FinCat'.
Copyright   : Guillaume Sabbagh 2022
License     : GPL-3
Maintainer  : guillaumesabbagh@protonmail.com
Stability   : experimental
Portability : portable

A 'ColimitCategory' is the colimit of a diagram in 'FinCat'. To compute colimits in a usual category, see Math.CocompleteCategory. To compute colimits in a custom 'FiniteCategory', see 'colimits' in Math.ConeCategory.

Note that computing a 'ColimitCategory' is MUCH slower than computing a 'LimitCategory' as coequalizers of categories are trickier than equalizers of categories. We can only compute colimits of 'CompositionGraph's as coequalizing might create new formal morphisms when gluing two objects. Therefore 'colimit' transforms all given categories into 'CompositionGraph's. If you already have a 'CompositionGraph', consider using 'colimitOfCompositionGraphs' instead of 'colimit'.
-}

module Math.FiniteCategories.ColimitCategory
(
    glueObject,
    glueObjects,
    glueMorphism,
    glueMorphisms,
    colimitOfCompositionGraphs,
    coprojectBaseCompositionGraphs
)
where
    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.List               (intercalate)
    import              Data.Simplifiable
    
    import              Math.Category
    import              Math.FiniteCategory
    import              Math.FiniteCategoryError
    import              Math.CocompleteCategory
    import              Math.Categories.FunctorCategory
    import              Math.FiniteCategories.DiscreteCategory
    import              Math.Categories.ConeCategory
    import              Math.Categories.FinCat
    import              Math.FiniteCategories.CompositionGraph
    import              Math.FiniteCategories.Parallel
    import              Math.Categories.FinGrph
    import              Math.IO.PrettyPrint
    
    import              GHC.Generics
    
    -- | Glue two objects of a 'CompositionGraph' and return an insertion diagram into the new CompositionGraph.
    glueObject :: (Eq a, Eq b) => CompositionGraph a b -> a -> a -> FinFunctor (CompositionGraph a b) (CGMorphism a b) a
    glueObject cg a b = completeDiagram Diagram{src=cg,tgt=newCG, omap = omapD, mmap = mmapD}
        where
            glueObj obj = if obj == b then a else obj
            omapD = memorizeFunction glueObj (nodes.support $ cg)
            glueArr arr = Arrow{sourceArrow = glueObj.sourceArrow $ arr, targetArrow = glueObj.targetArrow $ arr, labelArrow = labelArrow arr}
            mmapD = memorizeFunction ((unsafeArrowToCGMorphism cg).glueArr.unsafeCGMorphismToArrow) (genArrowsWithoutId cg)
            newCG = unsafeCompositionGraph (unsafeGraph (Map.values omapD) (glueArr <$> (edges.support $ cg))) (law cg)
    
    -- | Glue objects of a 'CompositionGraph' and return an insertion diagram into the new CompositionGraph.
    glueObjects :: (Eq a, Eq b) => CompositionGraph a b -> Set a -> FinFunctor (CompositionGraph a b) (CGMorphism a b) a
    glueObjects cg objects
        | Set.null objects = error "glue 0 object"
        | otherwise = Set.foldr (\obj funct -> (glueObject (tgt funct) representant obj) @ funct) (identity FinCat cg) objects
        where
            representant = anElement objects

    -- | Glue two morphisms with same source and target of a 'CompositionGraph' and return an insertion diagram into the new CompositionGraph.
    -- The first argument should not be composite if the second is a generator.
    glueMorphism :: (Eq a, Eq b) => CompositionGraph a b -> CGMorphism a b -> CGMorphism a b -> FinFunctor (CompositionGraph a b) (CGMorphism a b) a
    glueMorphism cg m1 m2
        | isComposite_ m1 && isGenerator_ m2 = error "Glue composite to generator."
        | m1 == m2 = identity FinCat cg
        | isGenerator_ m1 && isGenerator_ m2 = completeDiagram Diagram{src=cg,tgt=newCGGG,omap=omapGG,mmap=mmapGG}
        | otherwise = completeDiagram Diagram{src=cg,tgt=newCG,omap=omapGG,mmap=mmapId}
        where
            isGenerator_ m = length (snd.path $ m) == 1
            isComposite_ = not.isGenerator_
            glueMorph m = if m == m2 then m1 else m
            glueArr a = if a == (unsafeCGMorphismToArrow m2) then unsafeCGMorphismToArrow m1 else a
            omapGG = memorizeFunction id (nodes.support $ cg)
            mmapGG = memorizeFunction glueMorph (genArrowsWithoutId cg)
            newLawGG =  weakMapFromSet $ (\(x,y) -> (glueArr <$> x, glueArr <$> y)) <$> ((Map.mapToSet).law $ cg)
            newCGGG = unsafeCompositionGraph (unsafeGraph (nodes.support $ cg) ((edges.support $ cg) `Set.difference` (set [unsafeCGMorphismToArrow m2]))) newLawGG
            newLaw = Map.insert (snd.path $ m2) (snd.path $ m1) (law cg)
            mmapId = memorizeFunction (\morph -> CGMorphism{path=path morph, compositionLaw = newLaw}) (genArrowsWithoutId cg)
            newCG = unsafeCompositionGraph (support cg) newLaw
            
    -- | Glue morphisms with same source and target of a 'CompositionGraph' and return an insertion diagram into the new CompositionGraph.
    glueMorphisms :: (Eq a, Eq b) => CompositionGraph a b -> Set (CGMorphism a b) -> FinFunctor (CompositionGraph a b) (CGMorphism a b) a
    glueMorphisms cg morphisms
        | Set.null morphisms = error "glue 0 morphism"
        | Set.null generators = Set.foldr (\morph funct -> (glueMorphism (tgt funct) (anElement morphisms) morph) @ funct) (identity FinCat cg) morphisms
        | otherwise = Set.foldr (\morph funct -> (glueMorphism (tgt funct) (anElement generators) morph) @ funct) (identity FinCat cg) morphisms
        where
            generators = Set.filter (isGenerator cg) morphisms


    instance (Eq e, Eq n, Eq oIndex) => 
        HasCoproducts (FinCat (CompositionGraph n e) (CGMorphism n e) n) (FinFunctor (CompositionGraph n e) (CGMorphism n e) n) (CompositionGraph n e) (FinCat (CompositionGraph (Colimit oIndex n) (Colimit oIndex e)) (CGMorphism (Colimit oIndex n) (Colimit oIndex e)) (Colimit oIndex n)) (FinFunctor (CompositionGraph (Colimit oIndex n) (Colimit oIndex e)) (CGMorphism (Colimit oIndex n) (Colimit oIndex e)) (Colimit oIndex n)) (CompositionGraph (Colimit oIndex n) (Colimit oIndex e)) oIndex where
        coproduct diagOfCG = unsafeCocone coproductCG natCoproduct
            where
                indexingCategory = src diagOfCG
                coproductCG = unsafeCompositionGraph (nadir coproductGraphCocone) (Map.unions [weakMap $ (\(rp1,rp2) -> (transformArrowIntoArrowCoproduct i <$> rp1, transformArrowIntoArrowCoproduct i <$> rp2)) <$> (Map.mapToList (law (diagOfCG ->$ i))) | i <- (setToList (ob indexingCategory))])
                transformObjectIntoObjectCoproduct i x = CoproductElement i x
                transformArrowIntoArrowCoproduct i a = Arrow{sourceArrow = CoproductElement i (sourceArrow a), targetArrow = CoproductElement i (targetArrow a), labelArrow = CoproductElement i (labelArrow a)}
                coproductGraphCocone = coproduct (completeDiagram $ Diagram{src = indexingCategory, tgt = FinGrph, omap = Map.weakMapFromSet [(i, support (diagOfCG ->$ i))| i <- ob indexingCategory], mmap = weakMap []})
                newDiagOfCG = completeDiagram Diagram{src = indexingCategory, tgt = FinCat, omap = Map.weakMapFromSet [(i,coprojectCG cg) | (i,cg) <- Map.mapToSet (omap diagOfCG)], mmap = weakMap []}
                natCoproduct = unsafeNaturalTransformation newDiagOfCG (constantDiagram indexingCategory FinCat coproductCG) (weakMapFromSet [(i, insertionFunctor i)| i <- ob indexingCategory])
                coprojectArrow a = Arrow{sourceArrow = Coprojection $ sourceArrow a, targetArrow = Coprojection $ targetArrow a, labelArrow = Coprojection $ labelArrow a}
                coprojectGraph g = unsafeGraph ((Coprojection) <$> nodes g) ((coprojectArrow) <$> edges g)
                coprojectLaw cl = Map.weakMapFromSet [(coprojectArrow <$> rp1, coprojectArrow <$> rp2) | (rp1,rp2) <- Map.mapToSet cl]
                coprojectCG cg = unsafeCompositionGraph (coprojectGraph $ support cg) (coprojectLaw $ law cg)
                coprojectCGMorphism CGMorphism{path = (a, rp),compositionLaw = cl} = CGMorphism{path = (Coprojection a, coprojectArrow <$> rp), compositionLaw = coprojectLaw cl}
                insertionFunctor i = completeDiagram Diagram{src = coprojectCG (diagOfCG ->$ i), tgt = coproductCG, omap = weakMapFromSet [(o, (nodeMap insertionGH) |!| o) | o <- ob (newDiagOfCG ->$ i)], mmap = weakMapFromSet [(coprojectCGMorphism m, (unsafeArrowToCGMorphism coproductCG (transformArrowIntoArrowCoproduct i (unsafeCGMorphismToArrow m)))) | m <- (genArrowsWithoutId (diagOfCG ->$ i))]}
                    where
                        insertionGH = (legsCocone coproductGraphCocone) =>$ i
        

    instance (Eq n, Eq e) => HasCoequalizers (FinCat (CompositionGraph n e) (CGMorphism n e) n) (FinFunctor (CompositionGraph n e) (CGMorphism n e) n) (CompositionGraph n e) where
        coequalize parallelDiag = unsafeCocone coeq nat
            where
                selectingObjects = ob $ parallelDiag ->$ ParallelA
                idFunct = Diagram{src = (parallelDiag ->$ ParallelB), tgt = (parallelDiag ->$ ParallelB), omap = memorizeFunction id (ob (parallelDiag ->$ ParallelB)), mmap = memorizeFunction id (genArrows (parallelDiag ->$ ParallelB))}
                glueObjectsFunctor = Set.foldr (\obj funct -> (glueObjects (tgt funct) (set [(funct @ (parallelDiag ->£ ParallelF)) ->$ obj,(funct @ (parallelDiag ->£ ParallelG)) ->$ obj])) @ funct) idFunct selectingObjects
                selectingArrows = genArrowsWithoutId $ parallelDiag ->$ ParallelA
                glueArrowsFunctor = Set.foldr (\morph funct -> (glueMorphisms (tgt funct) (set [(funct @ (parallelDiag ->£ ParallelF)) ->£ morph,(funct @ (parallelDiag ->£ ParallelG)) ->£ morph])) @ funct) glueObjectsFunctor selectingArrows
                coeq = tgt glueArrowsFunctor
                nat = unsafeNaturalTransformation parallelDiag (constantDiagram Parallel FinCat coeq) (weakMap [(ParallelB, glueArrowsFunctor), (ParallelA, glueArrowsFunctor @ (parallelDiag ->£ ParallelF))])
        
    -- | Note that computing a 'ColimitCategory' is MUCH slower than computing a 'LimitCategory' as coequalizers of categories are trickier than equalizers of categories. We can only compute colimits of 'CompositionGraph's as coequalizing might create new formal morphisms when gluing two objects. Therefore 'colimit' transforms all given categories into 'CompositionGraph's. If you already have a 'CompositionGraph', consider using 'colimitOfCompositionGraphs' instead of 'colimit'.
    instance (FiniteCategory c m o, Morphism m o, Eq c, Eq m, Eq o,
              FiniteCategory cIndex mIndex oIndex, Morphism mIndex oIndex, Eq oIndex, Eq mIndex) => 
        CocompleteCategory (FinCat c m o) (FinFunctor c m o) c (FinCat (CompositionGraph (Colimit oIndex o) (Colimit oIndex m)) (CGMorphism (Colimit oIndex o) (Colimit oIndex m)) (Colimit oIndex o)) (FinFunctor (CompositionGraph (Colimit oIndex o) (Colimit oIndex m)) (CGMorphism (Colimit oIndex o) (Colimit oIndex m)) (Colimit oIndex o)) (CompositionGraph (Colimit oIndex o) (Colimit oIndex m)) cIndex mIndex oIndex where
        colimit diag = colimitFromCoproductsAndCoequalizers coprojectDiag newDiag
            where
                newDiag = Diagram{src = src diag, tgt = FinCat, omap = newOmap, mmap = newMmap}
                newOmap = Map.weakMapFromSet [(k, tgt $ finiteCategoryToCompositionGraph2 (diag ->$ k)) | k <- ob (src diag)]
                newMmap =  Map.weakMapFromSet [(a, diagramToDiagramOfCompositionGraphs2 (diag ->£ a)) | a <-  arrows (src diag)]
                coprojectArrow a = Arrow{sourceArrow = Coprojection $ sourceArrow a, targetArrow = Coprojection $ targetArrow a, labelArrow = Coprojection $ labelArrow a}
                coprojectGraph g = unsafeGraph ((Coprojection) <$> nodes g) ((coprojectArrow) <$> edges g)
                coprojectLaw cl = Map.weakMapFromSet [(coprojectArrow <$> rp1, coprojectArrow <$> rp2) | (rp1,rp2) <- Map.mapToSet cl]
                coprojectCG cg = unsafeCompositionGraph (coprojectGraph $ support cg) (coprojectLaw $ law cg)
                coprojectCGMorphism CGMorphism{path = (a, rp),compositionLaw = cl} = CGMorphism{path = (Coprojection a, coprojectArrow <$> rp), compositionLaw = coprojectLaw cl}
                coprojectDiag d = Diagram{src = coprojectCG $ src d, tgt = coprojectCG $ tgt d, omap = (\(x,y) -> (Coprojection x, Coprojection y)) <|$|> omap d, mmap = (\(x,y) -> (coprojectCGMorphism x, coprojectCGMorphism y)) <|$|> mmap d}

        coprojectBase diag = Diagram{src = FinCat, tgt = FinCat, omap = Map.weakMapFromSet [(diag ->$ i,coprojectCG $ tgt $ finiteCategoryToCompositionGraph2 $  diag ->$ i) | i <- ob $ src diag], mmap = Map.weakMapFromSet [(diag ->£ f,coprojectDiag  $ diagramToDiagramOfCompositionGraphs2 $ diag ->£ f) | f <- arrows $ src diag] }
            where
                coprojectArrow a = Arrow{sourceArrow = Coprojection $ sourceArrow a, targetArrow = Coprojection $ targetArrow a, labelArrow = Coprojection $ labelArrow a}
                coprojectGraph g = unsafeGraph ((Coprojection) <$> nodes g) ((coprojectArrow) <$> edges g)
                coprojectLaw cl = Map.weakMapFromSet [(coprojectArrow <$> rp1, coprojectArrow <$> rp2) | (rp1,rp2) <- Map.mapToSet cl]
                coprojectCG cg = unsafeCompositionGraph (coprojectGraph $ support cg) (coprojectLaw $ law cg)
                coprojectCGMorphism CGMorphism{path = (a, rp),compositionLaw = cl} = CGMorphism{path = (Coprojection a, coprojectArrow <$> rp), compositionLaw = coprojectLaw cl}
                coprojectDiag d = Diagram{src = coprojectCG $ src d, tgt = coprojectCG $ tgt d, omap = (\(x,y) -> (Coprojection x, Coprojection y)) <|$|> omap d, mmap = (\(x,y) -> (coprojectCGMorphism x, coprojectCGMorphism y)) <|$|> mmap d}
                
    -- | Computes the colimit of a 'Diagram' of 'CompositionGraph's, use this if you already have a 'Diagram' of 'CompositionGraph's. 
    colimitOfCompositionGraphs :: (FiniteCategory cIndex mIndex oIndex, Morphism mIndex oIndex, Eq cIndex, Eq mIndex, Eq oIndex, Eq n, Eq e) => 
        Diagram cIndex mIndex oIndex (FinCat (CompositionGraph n e) (CGMorphism n e) n) (FinFunctor (CompositionGraph n e) (CGMorphism n e) n) (CompositionGraph n e) -> Cocone cIndex mIndex oIndex (FinCat (CompositionGraph (Colimit oIndex n) (Colimit oIndex e)) (CGMorphism (Colimit oIndex n) (Colimit oIndex e)) (Colimit oIndex n)) (FinFunctor (CompositionGraph (Colimit oIndex n) (Colimit oIndex e)) (CGMorphism (Colimit oIndex n) (Colimit oIndex e)) (Colimit oIndex n)) (CompositionGraph (Colimit oIndex n) (Colimit oIndex e))
    colimitOfCompositionGraphs = colimitFromCoproductsAndCoequalizers coprojectDiag
            where
                coprojectArrow a = Arrow{sourceArrow = Coprojection $ sourceArrow a, targetArrow = Coprojection $ targetArrow a, labelArrow = Coprojection $ labelArrow a}
                coprojectGraph g = unsafeGraph ((Coprojection) <$> nodes g) ((coprojectArrow) <$> edges g)
                coprojectLaw cl = Map.weakMapFromSet [(coprojectArrow <$> rp1, coprojectArrow <$> rp2) | (rp1,rp2) <- Map.mapToSet cl]
                coprojectCG cg = unsafeCompositionGraph (coprojectGraph $ support cg) (coprojectLaw $ law cg)
                coprojectCGMorphism CGMorphism{path = (a, rp),compositionLaw = cl} = CGMorphism{path = (Coprojection a, coprojectArrow <$> rp), compositionLaw = coprojectLaw cl}
                coprojectDiag d = completeDiagram Diagram{src = coprojectCG $ src d, tgt = coprojectCG $ tgt d, omap = (\(x,y) -> (Coprojection x, Coprojection y)) <|$|> omap d, mmap = (\(x,y) -> (coprojectCGMorphism x, coprojectCGMorphism y)) <|$|> mmap d}
                
    -- | Uncoproject a base of 'CompositionGraph's.
    coprojectBaseCompositionGraphs :: (FiniteCategory cIndex mIndex oIndex, Morphism mIndex oIndex, Eq mIndex, Eq oIndex, Eq e, Eq n) => Diagram cIndex mIndex oIndex (FinCat (CompositionGraph n e) (CGMorphism n e) n) (FinFunctor (CompositionGraph n e) (CGMorphism n e) n) (CompositionGraph n e) -> Diagram (FinCat (CompositionGraph n e) (CGMorphism n e) n) (FinFunctor (CompositionGraph n e) (CGMorphism n e) n) (CompositionGraph n e) (FinCat (CompositionGraph (Colimit oIndex n) (Colimit oIndex e)) (CGMorphism (Colimit oIndex n) (Colimit oIndex e)) (Colimit oIndex n)) (FinFunctor (CompositionGraph (Colimit oIndex n) (Colimit oIndex e)) (CGMorphism (Colimit oIndex n) (Colimit oIndex e)) (Colimit oIndex n)) (CompositionGraph (Colimit oIndex n) (Colimit oIndex e))
    coprojectBaseCompositionGraphs diag = Diagram{src = FinCat, tgt = FinCat, omap = Map.weakMapFromSet [(diag ->$ i,coprojectCG $ diag ->$ i) | i <- ob $ src diag], mmap = Map.weakMapFromSet [(diag ->£ f,coprojectDiag $ diag ->£ f) | f <- arrows $ src diag] }
        where
            coprojectArrow a = Arrow{sourceArrow = Coprojection $ sourceArrow a, targetArrow = Coprojection $ targetArrow a, labelArrow = Coprojection $ labelArrow a}
            coprojectGraph g = unsafeGraph ((Coprojection) <$> nodes g) ((coprojectArrow) <$> edges g)
            coprojectLaw cl = Map.weakMapFromSet [(coprojectArrow <$> rp1, coprojectArrow <$> rp2) | (rp1,rp2) <- Map.mapToSet cl]
            coprojectCG cg = unsafeCompositionGraph (coprojectGraph $ support cg) (coprojectLaw $ law cg)
            coprojectCGMorphism CGMorphism{path = (a, rp),compositionLaw = cl} = CGMorphism{path = (Coprojection a, coprojectArrow <$> rp), compositionLaw = coprojectLaw cl}
            coprojectDiag d = Diagram{src = coprojectCG $ src d, tgt = coprojectCG $ tgt d, omap = (\(x,y) -> (Coprojection x, Coprojection y)) <|$|> omap d, mmap = (\(x,y) -> (coprojectCGMorphism x, coprojectCGMorphism y)) <|$|> mmap d}