{-# LANGUAGE MonadComprehensions, MultiParamTypeClasses  #-}
{-| Module  : FiniteCategories
Description : Utility functions for set-valued functors.
Copyright   : Guillaume Sabbagh 2023
License     : GPL-3
Maintainer  : guillaumesabbagh@protonmail.com
Stability   : experimental
Portability : portable

Kan extensions for set-valued functors. Inspired by DBC of M. Fleming, R. Gunther, R. Rosebrugh.
-}

module Math.Functors.SetValued 
(
    -- * Right Kan extension
    LimitObject(..),
    rightKanSetValued,
    -- ** Formatting
    formatLimitObject,
    formatSetOfLimitObjects,
    formatFunctionOfLimitObjects,
    -- * Left Kan extension
    ColimitObject(..),
    leftKanSetValued,
    -- ** Formatting
    formatColimitObject,
    formatSetOfColimitObjects,
    formatFunctionOfColimitObjects,
)
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              Math.FiniteCategory
    import              Math.Categories.FinSet
    import              Math.Categories.FunctorCategory
    import              Math.Categories.CommaCategory
    import              Math.FiniteCategories.One
    import              Math.IO.PrettyPrint
    
    -- | A 'LimitObject' is a map from a comma object to an element a, the maps should be seen as elements of the cartesian products indexed by comma objects.
    type LimitObject o1 m2 a = Map (CommaObject One o1 m2) a
    
    -- | Format a 'LimitObject' to be readable.
    formatLimitObject :: (Eq o1, Eq m2, Eq a, PrettyPrint a) => LimitObject o1 m2 a -> String
    formatLimitObject mapping = "(" ++ intercalate "," [ (pprint v) | (k,v) <- Map.mapToList mapping] ++ ")"
    
    -- | Format a set of 'LimitObject's to be readable.
    formatSetOfLimitObjects :: (Eq o1, Eq m2, Eq a, PrettyPrint a) => Set (LimitObject o1 m2 a) -> String
    formatSetOfLimitObjects setOfMaps = "{" ++ intercalate "," (formatLimitObject <$> setToList setOfMaps) ++ "}"
    
    -- | Format a 'LimitObject' to be readable.
    formatFunctionOfLimitObjects :: (Eq o1, Eq m2, Eq a, PrettyPrint a) => Function (LimitObject o1 m2 a) -> String
    formatFunctionOfLimitObjects func = intercalate "," [(formatLimitObject k) ++ " -> " ++ (formatLimitObject v) | (k,v) <- (Map.mapToList (function func))]
    
    -- | Compute the image of the right Kan extension on an object of /B/. The image is a set of maps from the comma objects to elements a (you should see the maps as cartesian products indexed by comma objects).
    rightKanSetValuedObjectwise :: (FiniteCategory c1 m1 o1, Morphism m1 o1, Eq m1, Eq o1,
                                    FiniteCategory c2 m2 o2, Morphism m2 o2, Eq m2, Eq o2,
                                    Eq a) =>
                                    Diagram c1 m1 o1 c2 m2 o2 -> Diagram c1 m1 o1 (FinSet a) (Function a) (Set a) -> o2 -> Set (LimitObject o1 m2 a)
    rightKanSetValuedObjectwise f x b = filteredMaps
        where
            commaCat = CommaCategory{leftDiagram=selectObject (tgt f) b, rightDiagram=f}
            allMaps = weakMap <$> (cartesianProductOfSets.setToList $ [[(a, elemXa) | elemXa <- (x ->$ (indexTarget a))] | a <- ob commaCat])
            filters = [\mapping -> ((x ->£(indexSecondArrow m)) ||!|| (mapping |!| (source m))) == (mapping |!| (target m))  | m <- setToList $ arrows commaCat]
            filteredMaps = foldr Set.filter allMaps filters
            
    -- | Compute the image of the right Kan extension on a morphism of /B/.
    rightKanSetValuedObjectwiseMorphism :: (FiniteCategory c1 m1 o1, Morphism m1 o1, Eq m1, Eq o1,
                                            FiniteCategory c2 m2 o2, Morphism m2 o2, Eq m2, Eq o2,
                                            Eq a) =>
                                            Diagram c1 m1 o1 c2 m2 o2 -> Diagram c1 m1 o1 (FinSet a) (Function a) (Set a) -> m2 -> Function (LimitObject o1 m2 a)
    rightKanSetValuedObjectwiseMorphism f x m = result
        where
            precompose commaObj = unsafeCommaObject (indexSource commaObj) (indexTarget commaObj) ((selectedArrow commaObj) @ m)
            imageSrc = rightKanSetValuedObjectwise f x (source m)
            imageTgt = rightKanSetValuedObjectwise f x (target m)
            commaCat2 = CommaCategory{leftDiagram=selectObject (tgt f) (target m), rightDiagram=f}
            subsetIndexingObjects = ob commaCat2
            result = Function{function=weakMapFromSet [(mapping,anElement [mapping2 | mapping2 <- imageTgt, Map.isSubmapOf (Map.mapKeys precompose mapping2) mapping]) | mapping <- imageSrc],
                              codomain=imageTgt}
    
    -- | The right Kan extension of @X@ along @F@ where @X@ is a a set-valued functor.
    --
    -- We transform the @X@ functor so that its target becomes @(FinSet (Map (CommaObject One o1 m2) a))@, each object @A@ of @c1@ is mapped to @{Delta : {<One,id,A>} -> e | e <- X(A)}@.
    rightKanSetValued :: (FiniteCategory c1 m1 o1, Morphism m1 o1, Eq c1, Eq m1, Eq o1,
                          FiniteCategory c2 m2 o2, Morphism m2 o2, Eq m2, Eq o2,
                          Eq a) =>
                          Diagram c1 m1 o1 c2 m2 o2 -> Diagram c1 m1 o1 (FinSet a) (Function a) (Set a) -> (Diagram c2 m2 o2 (FinSet (LimitObject o1 m2 a)) (Function (LimitObject o1 m2 a)) (Set (LimitObject o1 m2 a)), NaturalTransformation c1 m1 o1 (FinSet (LimitObject o1 m2 a)) (Function (LimitObject o1 m2 a)) (Set (LimitObject o1 m2 a)))
    rightKanSetValued f x = (r,nat)
        where
            r = Diagram{src=tgt f,tgt=FinSet, omap=memorizeFunction (rightKanSetValuedObjectwise f x) (ob (tgt f)), mmap=memorizeFunction (rightKanSetValuedObjectwiseMorphism f x) (arrows (tgt f))}
            rof = r  <-@<- f
            transformAIntoMap k a = weakMap [(unsafeCommaObject One k (identity (tgt f) (f ->$ k)), a)]
            transformSetOfAIntoSetOfMap k v = [transformAIntoMap k e | e <- v]
            transformFunctionOfAIntoFunctionOfMaps m1 func = Function{function=weakMap [(transformAIntoMap (source m1) t,transformAIntoMap (target m1) u) | (t,u) <- Map.mapToList (function func)], codomain=transformSetOfAIntoSetOfMap (target m1) (codomain func)}
            newX = Diagram{src=src x, tgt=FinSet, omap=Map.mapWithKey transformSetOfAIntoSetOfMap (omap x), mmap=Map.mapWithKey transformFunctionOfAIntoFunctionOfMaps (mmap x)}
            Right nat = naturalTransformation rof newX (weakMapFromSet [(o,Function{function=weakMapFromSet [(mapping,transformAIntoMap o ((snd.anElement) (Set.filter (\y -> (indexTarget.fst) y == o) (Map.mapToSet mapping))) ) | mapping <- (rof ->$ o)], codomain=newX ->$ o }) | o <- ob (src f)])
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
    -- | A 'ColimitObject' is  a set of couples (comma objects, elements a), the couples should be seen as elements of a disjoint union, the set of couples as equivalence classes of elements of a disjoint union.
    type ColimitObject o1 m2 a = Set ((CommaObject o1 One m2), a)
    
    -- | Format a 'ColimitObject' to be readable.
    formatColimitObject :: (PrettyPrint a) => ColimitObject o1 m2 a -> String
    formatColimitObject = pprint.snd.anElement
    
    -- | Format a set of 'ColimitObject's to be readable.
    formatSetOfColimitObjects :: (Eq o1, Eq m2, Eq a, PrettyPrint a) => Set (ColimitObject o1 m2 a) -> String
    formatSetOfColimitObjects setOfEquivClasses = "{" ++ intercalate "," (formatColimitObject <$> setToList setOfEquivClasses) ++ "}"
    
    -- | Format a 'ColimitObject' to be readable.
    formatFunctionOfColimitObjects :: (Eq o1, Eq m2, Eq a, PrettyPrint a) => Function (ColimitObject o1 m2 a) -> String
    formatFunctionOfColimitObjects func = intercalate "," [(formatColimitObject k) ++ " -> " ++ (formatColimitObject v) | (k,v) <- (Map.mapToList (function func))]
 
    transformGraphToEquivRelation :: (Eq a) => Set (a,a) -> Set a -> Set (Set a)
    transformGraphToEquivRelation couples initialNodes = Set.foldr (\candidate clusters -> if Set.any (isIn candidate) clusters then clusters else (Set.insert (Set.insert candidate (dfs candidate [candidate] (setToList initialNodes))) clusters)) (set []) initialNodes
        where
            dfs node blacklist [] = set []
            dfs node blacklist (candidate:nodes)
                | elem candidate blacklist = (dfs node blacklist nodes)
                | (node,candidate) `isIn` couples || (candidate,node) `isIn` couples = Set.insert candidate ((dfs node blacklist nodes) ||| (dfs candidate (node:blacklist) (setToList initialNodes)))
                | otherwise = (dfs node blacklist nodes)
    
    -- | Compute the image of the left Kan extension on an object of /B/. The image is a set of set of couples (comma objects, elements a) (you should see the couples as element of a disjoint union, the set of couples as equivalence classes of elements).
    leftKanSetValuedObjectwise :: (FiniteCategory c1 m1 o1, Morphism m1 o1, Eq m1, Eq o1,
                                   FiniteCategory c2 m2 o2, Morphism m2 o2, Eq m2, Eq o2,
                                   Eq a) =>
                                   Diagram c1 m1 o1 c2 m2 o2 -> Diagram c1 m1 o1 (FinSet a) (Function a) (Set a) -> o2 -> Set (ColimitObject o1 m2 a)
    leftKanSetValuedObjectwise f x b = equivClasses
        where
            commaCat = CommaCategory{leftDiagram=f, rightDiagram=selectObject (tgt f) b}
            allCouples = ((Set.unions).setToList $ [[(a, elemXa) | elemXa <- (x ->$ (indexSource a))] | a <- ob commaCat])
            graphRelations =  Set.unions [[((source m,a),(target m,b)) | (a,b) <- (Map.mapToSet (function (x ->£(indexFirstArrow m))))] | m <- (setToList (arrows commaCat)), not (isIdentity commaCat m)]
            equivClasses = transformGraphToEquivRelation graphRelations allCouples
            
    -- | Compute the image of the left Kan extension on a morphism of /B/. 
    leftKanSetValuedObjectwiseMorphism :: (FiniteCategory c1 m1 o1, Morphism m1 o1, Eq m1, Eq o1,
                                            FiniteCategory c2 m2 o2, Morphism m2 o2, Eq m2, Eq o2,
                                            Eq a) =>
                                            Diagram c1 m1 o1 c2 m2 o2 -> Diagram c1 m1 o1 (FinSet a) (Function a) (Set a) -> m2 -> Function (ColimitObject o1 m2 a)
    leftKanSetValuedObjectwiseMorphism f x m = result
        where
            postcompose commaObj = unsafeCommaObject (indexSource commaObj) (indexTarget commaObj) (m @ (selectedArrow commaObj))
            imageSrc = leftKanSetValuedObjectwise f x (source m)
            imageTgt = leftKanSetValuedObjectwise f x (target m)
            result = Function{function=weakMapFromSet [(equivClass1, anElement [equivClass2 | equivClass2 <- imageTgt, Set.any (\(co, a) -> (postcompose co, a) `isIn` equivClass2) equivClass1 ]) | equivClass1 <- imageSrc],
                              codomain=imageTgt}
    
    -- | The left Kan extension of @X@ along @F@ where @X@ is a a set-valued functor.
    --
    -- We transform the @X@ functor so that its target becomes @(FinSet (Set ((CommaObject One o1 m2), a)))@, each object @A@ of @c1@ is mapped to @{(<One,id,A>},e) | e <- X(A)}@.
    leftKanSetValued :: (FiniteCategory c1 m1 o1, Morphism m1 o1, Eq c1, Eq m1, Eq o1,
                         FiniteCategory c2 m2 o2, Morphism m2 o2, Eq m2, Eq o2,
                         Eq a) =>
                         Diagram c1 m1 o1 c2 m2 o2 -> Diagram c1 m1 o1 (FinSet a) (Function a) (Set a) -> (Diagram c2 m2 o2 (FinSet (ColimitObject o1 m2 a)) (Function (ColimitObject o1 m2 a)) (Set (ColimitObject o1 m2 a)), NaturalTransformation c1 m1 o1 (FinSet (ColimitObject o1 m2 a)) (Function (ColimitObject o1 m2 a)) (Set (ColimitObject o1 m2 a)))
    leftKanSetValued f x = (l,nat)
        where
            l = Diagram{src=tgt f,tgt=FinSet, omap=memorizeFunction (leftKanSetValuedObjectwise f x) (ob (tgt f)), mmap=memorizeFunction (leftKanSetValuedObjectwiseMorphism f x) (arrows (tgt f))}
            lof = l  <-@<- f
            transformAIntoEquivClass k a = set [(unsafeCommaObject k One (identity (tgt f) (f ->$ k)), a)]
            transformSetOfAIntoSetOfEquivClasses k v = [transformAIntoEquivClass k e | e <- v]
            transformFunctionOfAIntoFunctionOfSetsOfEquivClasses m1 func = Function{function=weakMap [(transformAIntoEquivClass (source m1) t,transformAIntoEquivClass (target m1) u) | (t,u) <- Map.mapToList (function func)], codomain=transformSetOfAIntoSetOfEquivClasses (target m1) (codomain func)}
            newX = Diagram{src=src x, tgt=FinSet, omap=Map.mapWithKey transformSetOfAIntoSetOfEquivClasses (omap x), mmap=Map.mapWithKey transformFunctionOfAIntoFunctionOfSetsOfEquivClasses (mmap x)}
            Right nat = naturalTransformation newX lof (weakMapFromSet [(o,Function{function=weakMapFromSet [(equivClass,anElement [equivClass2 | equivClass2 <- (lof ->$ o), Set.any (\(co,a2) -> indexSource co == o && (snd.anElement $ equivClass) == a2) equivClass2]) | equivClass <- (newX ->$ o)], codomain=lof ->$ o }) | o <- ob (src f)])