{-|
Module      : Kh
Description : Implements the Khovanov "functor".
Copyright   : Adam Saltz
License     : BSD3
Maintainer  : saltz.adam@gmail.com
Stability   : experimental

Longer description to come.
-}
module Kh
where
import Util
import Complex
import Braids
import Data.Set (Set, (\\))
import Data.List (sort)
import qualified Data.Set as S
import Data.Map (Map, (!))
import qualified Data.Map as M
import Control.Monad

-- | Return the diagram underlying a 'Generator'.
gToD :: Generator -> BDiagram
gToD g = BDiagram {resolution' = resolution g, components' = components g}
    
-- | Compute the homological grading of a 'Generator'.
homGrading :: Generator -> Int
homGrading gen = sum . resolution $ gen

-- | Compute the q-grading of a 'Generator'.  In this convention, the Khovanov differential *lowers* the q-grading by 1.
qGrading :: Generator -> Int
qGrading gen = sum . fmap signToNum . signs $ gen where

-- | Data type to represent whether a circle is trivial or not.  Used to be Bool' but I got confused about which was @True@ and which was @False@.
data Triviality = Trivial | NonTrivial deriving (Ord, Eq, Show)

{- | Determine if a component is non-trivial.
     To compute the mod 2 winding number about the braid axis, check how many of the \'top arcs\' live in a component. -}
nonTrivialCircle :: Int -> Component -> Triviality
nonTrivialCircle width arcs = if odd $ length (filter (`elem` [1..width]) (S.toList arcs))
                              then NonTrivial
                              else Trivial

-- | Compute the k-grading of a 'Generator'.
kGrading :: (Resolution, Set Component, Map Component Sign) -> Int -> Int
kGrading (_,_,ss) width = ks ! NonTrivial where
    ks = M.mapKeysWith (+) (nonTrivialCircle width) . fmap signToNum $ ss

{- |  "Apply the filtered Khovanov functor to a diagram."
  We still need the braid width as input to get the k-grading. -}
khovanov :: Int -> BDiagram -> [Generator]
khovanov width d = do
    ss <- generateSigns . components' $ d
    return Generator {resolution = resolution' d,
                      components = components' d, 
                      signs = M.fromList $ zip (S.toList $ components' d) ss, 
                      kgrade = kGrading (resolution' d, components' d, M.fromList $ zip (S.toList $ components' d) ss) width} 
    where
        generateSigns cs = Control.Monad.replicateM (length cs) [Minus,Plus] -- this is a slick and dumb way of making sequences of (-1) and 1
{-
-- | Returns the sign at mark.
markedSign :: Int -> Generator -> Int
markedSign mark gen = case findIndex (elem mark) (components gen) of
    Just x -> (!!) (signs gen) x
    Nothing -> -1 -- should thrown a exception


-- | "Apply the reduced Khovanov functor to a diagram"
reducedKhovanov :: Int -> Int -> BDiagram -> [Generator]
reducedKhovanov mark width cube  = filter (\g -> (-1) == markedSign mark g) (khovanov width cube)

-- | "Apply the quotient Khovanov functor to a diagram"
quotientKhovanov :: Int -> Int -> BDiagram -> [Generator]
quotientKhovanov mark width cube = filter (\g -> (1) == markedSign mark g) (khovanov width cube)
-}

-- | The next three functions apply the Khovanov functor to the vertices of a cube of resolutions.
khovanovComplex :: Int -> [BDiagram] -> Map Int (Set Generator)
khovanovComplex width cube = M.fromListWith S.union [(homGrading g, S.singleton g) | g <- concatMap (khovanov width) cube]
{-
reducedKhovanovComplex :: Int -> Int -> [BDiagram] -> Map Int (Set Generator)
reducedKhovanovComplex mark width cube = M.fromListWith S.union [(homGrading g, S.singleton g) | g <- concatMap (reducedKhovanov mark width) cube]

quotientKhovanovComplex :: Int -> Int -> [BDiagram] -> Map Int (Set Generator)
quotientKhovanovComplex mark width cube = M.fromListWith S.union [(homGrading g, S.singleton g) | g <- concatMap (quotientKhovanov mark width) cube]
-}
-- | An elementary morphism is determined by 'Components'.  We distinguish between 'Merge' and 'Split' 'ElMo's
data ElMo -- | elementary morphisms.  The abbreviation is borrowed from Milatz.
    = Merge (Set Component) (Set Component) -- ^ merges the first set to the second
    | Split (Set Component) (Set Component) -- ^ splits the first set into the second
    deriving (Show)

-- | Take two diagrams and returns the 'ElMo's between them, if there is one.  
whichMorphism :: BDiagram -> BDiagram -> Maybe ElMo
whichMorphism d d'     | not (succRes d d')                   = Nothing
                       | length cs2' == 2 && length cs1' == 1 = Just $ Split cs1' cs2'
                       | length cs1' == 2 && length cs2' == 1 = Just $ Merge cs1' cs2'
                       | otherwise                            = Nothing
                       where
                         cs2' = components' d' \\ components' d
                         cs1' = components' d  \\ components' d'
                         succRes :: BDiagram -> BDiagram -> Bool
                         succRes e e' = all (>=0) diff && (sum diff == 1) where 
                           diff = zipWith subtract (resolution' e)  (resolution' e')

-- | Take a morphism and two generators and returns @True@ if there should be a 'Morphism from one to the other.
morphismAction :: Maybe ElMo -> Generator -> Generator -> Bool
morphismAction (Just (Merge cs12 c3)) g g'  | (cs12 `isNotASubsetOf` gcs) || (c3 `isNotASubsetOf` g'cs)   = False
                                            | (g'cs \\ c3) /= (gcs \\ cs12)                             = False
                                            | ss `deleteKeys` cs12 /= ss' `deleteKeys` c3               = False
                                            | M.elems (ss `getSubmap` cs12) == [Plus,Plus] && M.elems (ss' `getSubmap` c3) == [Plus] = True
                                            | sort (M.elems (ss `getSubmap` cs12)) == [Plus,Minus] && M.elems (ss' `getSubmap` c3) == [Minus] = True
                                            | otherwise = False
                                            where
                                                    gcs = components g 
                                                    g'cs = components g'
                                                    ss = signs g 
                                                    ss' = signs g'

morphismAction (Just (Split c1 cs23)) g g'  | (c1 `isNotASubsetOf` gcs) || (cs23 `isNotASubsetOf` g'cs)   = False
                                            | (g'cs \\ cs23) /= (gcs \\ c1)                             = False
                                            | ss `deleteKeys` c1 /= ss' `deleteKeys` cs23               = False
                                            | M.elems (ss `getSubmap` c1) == [Plus] && sort (M.elems (ss' `getSubmap` cs23)) == [Plus, Minus] = True
                                            | M.elems (ss `getSubmap` c1) == [Minus] && M.elems (ss' `getSubmap` cs23) == [Minus, Minus] = True
                                            | otherwise = False
                                            where
                                                    gcs = components g 
                                                    g'cs = components g'
                                                    ss = signs g 
                                                    ss' = signs g'

morphismAction (Nothing) _ _  = False

-- | Compute the difference in k-grading between two 'Generator's.
kDrop :: Generator -> Generator -> Int
kDrop g g' = kgrade g - kgrade g'

-- | Compute the filtration level of an 'AlgGen'.
kgrade' :: AlgGen -> Int
kgrade' (AlgGen gs) = maximum . S.map kgrade $ gs

kDrop' :: AlgGen -> AlgGen -> Int
kDrop' gs  gs' = kgrade' gs - kgrade' gs'

-- | Like 'morphismAction', but only connects two 'Generators' if the drop in k-grading from one to the other is less than or equal to \a\.
filteredMorphismAction :: Int -> Maybe ElMo -> Generator -> Generator -> Bool
filteredMorphismAction k e g g' | kDrop g g' <= k = morphismAction e g g'
                                | otherwise = False

-- | Return 'Morphisms' from the 'Generator' into the set with 'kDrop' less than or equal to \k\.
filteredMorphismsFrom :: Int -> Generator -> Set Generator -> Morphisms
filteredMorphismsFrom k g gs = M.singleton (wrapGen g) gs' where
                                 gs'   =   S.map wrapGen
                                         . S.filter (\g' -> filteredMorphismAction k (mor g') g g') 
                                         $ gs
                                 mor g' = whichMorphism (gToD g) (gToD g')

-- | Applies 'filteredMorphismsFrom' to every 'Generator' in a list into the same list.
filteredComplexLevel :: Int -> Map Int (Set Generator) -> Int -> Morphisms
filteredComplexLevel k gs i = case M.lookup (i+1) gs of
    Nothing -> M.empty
    otherwise -> M.unionsWith S.union . S.toList . S.map (\g -> filteredMorphismsFrom k g gsi1) $ gsi where
        gsi = gs ! i
        gsi1 = gs ! (i+1)

-- | Produces the 'Generator' corresponding to the transverse invariant of a braid.
psi :: Braid -> Generator
psi b = Generator {resolution = res, components = comps, signs = ss, kgrade = k} where
    res = fmap (\x -> if x >= 0 then 0 else 1) (braidWord b)
    comps = resolutionToComponents . flip resolve res . braidToPD $ b 
    ss = M.fromList $ zip (S.toList comps) (repeat Minus)
    k = (-1)* braidWidth b
{-
-- | Produces the generator corresponding to the QUOTIENT transverse invariant of a braid
quotPsi :: Braid -> Int -> Generator
quotPsi b p = Generator {resolution = res, components = comps, signs = ss, kgrade = k} where
    res = fmap (\x -> if x >= 0 then 0 else 1) (braidWord b)
    comps = resolutionToComponents . flip resolve res . braidToPD $ b 
    Just c = findIndex (elem p) comps
    ss = uncurry (++) . second (\xs -> 1:xs) . second tail . splitAt c $ replicate (length comps)(-1)
    k = kGrading (res, comps, ss) (braidWidth b)
-}
{- |Produces the portion of the cube of resolutions of a braid which is relevant for computing kappa.
   This means only using resolutions whose weights are less than or equal to psi's.
   Note that this only uses the homological grading of psi, so we don't need a separate function for the quotient. -}
psiCube :: Braid -> [BDiagram]
psiCube b =  do
    res <- allRes (braidToPD b)
    let diagram = resolutionToComponents . resolve (braidToPD b) 
    guard (sum res <= (sum . resolution . psi $ b)) 
    return BDiagram {resolution' = res, components' = diagram res}