-- | Motivic classes in homology

{-# LANGUAGE 
      DataKinds, KindSignatures, TypeOperators, ScopedTypeVariables, MultiParamTypeClasses,
      TypeFamilies, FlexibleContexts, TypeSynonymInstances, FlexibleInstances
  #-}
module Math.RootLoci.Motivic.Homology where

--------------------------------------------------------------------------------

import Data.Array

import Data.Proxy
import GHC.TypeLits

import Unsafe.Coerce as Unsafe

import Math.Combinat.Classes
import Math.Combinat.Numbers
import Math.Combinat.Partitions

import qualified Math.Algebra.Polynomial.FreeModule as ZMod

import qualified Math.Algebra.Polynomial.Monomial.Infinite as XInf
import qualified Math.Algebra.Polynomial.Monomial.Indexed  as XS

import Math.Algebra.Polynomial.Multivariate.Infinite as XInf
import Math.Algebra.Polynomial.Multivariate.Indexed  as XS

import Math.Algebra.Polynomial.Class
import Math.Algebra.Polynomial.Univariate
import Math.Algebra.Polynomial.Pretty

import Math.RootLoci.Geometry.Cohomology ( G(..) )

import Math.RootLoci.Misc.Common
import Math.RootLoci.Motivic.Abstract as Abstract
import Math.RootLoci.Motivic.Classes

import Math.RootLoci.CSM.Aluffi

--------------------------------------------------------------------------------

interpretSingleLam :: (Dim -> KRing Integer) -> SingleLam -> KRing Integer
interpretSingleLam symFun (SingleLam (Bindings bindings) (Single body)) = result where
  syms   = map symFun bindings
  result = psiAny $ crossKs $ map f body 
  f (DeBruijn i, e) = omegaH e (syms!!i)  

csmPn :: Dim -> KRing Integer
csmPn (Dim d) = Uni $ ZMod.fromList [ (U k , binomial (d+1) (k+1)) | k<-[0..d] ]

-- | CSM class in homology
csm_xlam_P1 :: Partition -> KRing Integer
csm_xlam_P1 part = Uni $ ZMod.flatMap f (Abstract.xlam part) where
  f x = unUni (interpretSingleLam csmPn x)

-- | CSM class in cohomology (via Poincare duality)
csm_xlam_P1_cohom :: Partition -> ZMod.ZMod G 
csm_xlam_P1_cohom part = ZMod.mapBase f $ unUni $ csm_xlam_P1 part  where
  n   = weight part
  f (U k) = G (n-k)

-- | Compares Aluffi's CSM formula to the motivic algorithm (up to partitions of size @n@)
test_motivic_csm_vs_aluffi :: Int -> Bool
test_motivic_csm_vs_aluffi n = and
  [ csm_xlam_P1_cohom part == aluffiOpenCSM part
  | k<-[1..n] , part <- partitions k 
  ]

--------------------------------------------------------------------------------

instance SingleToMulti (KRing c) (GRing c) where
  singleToMulti = embedInf

instance Ring c => Psi (GRing c) (KRing c) where
  psi = psiAny

instance Ring c => Omega (KRing c) where
  omega = omegaH

--------------------------------------------------------------------------------

type KRing c   = Univariate c "u"     -- ^ @lim_n H_*(Sym^n(P1))@
type GRing c   = XInf.Poly  c "u"     -- ^ @lim_{n1,n2,...} H_*(Sym^n1(P1) x Sym^n2(P1) x ... )@

-- fuck Haskell's type level naturals, they are completely unusable
-- type NRing c k = XS.Poly    c "u" k   -- ^ @lim_{n1,...,nk} H_*(Sym^n1(P1) x ... x Sym^nk(P1) )@

embedInf :: KRing c -> GRing c
embedInf = XInf.Poly . ZMod.unsafeMapBase f . unUni where
  f (U k) = if k > 0 then XInf [k] else XInf []

project1 :: GRing c -> KRing c
project1 = Uni . ZMod.unsafeMapBase f . XInf.unPoly where
  f (XInf ns) = U $ head ns

delta2 :: Ring c => KRing c -> GRing c
delta2 = XInf.Poly . ZMod.flatMap f . unUni where
  f (U k) = ZMod.sum [ ZMod.generator (XInf [i,k-i]) | i<-[0..k] ] 

deltaN :: Ring c => Int -> KRing c -> GRing c
deltaN n input 
  | n <= 0    = error "deltaN: n <= 0"
  | n == 1    = embedInf input
  | n == 2    = delta2   input
  | otherwise = unify1st2nd
              $ mapCoeffP delta2 
              $ separate1st (deltaN (n-1) input)
  where
    mapCoeffP f = XInf.Poly . ZMod.mapCoeff f . XInf.unPoly

psi2 :: Ring c => GRing c -> KRing c
psi2 = Uni . ZMod.mapMaybeBaseCoeff f . XInf.unPoly where
  f (XInf xs) = let [i,j] = take 2 (xs ++ [0,0])
                in  Just ( U (i+j) , fromInteger (binomial (i+j) i) )

psiNaive :: (Ring c) => Int -> GRing c -> KRing c
psiNaive n input 
  | n <= 0    = error "psiN: n <= 0"
  | n == 1    = project1 input
  | n == 2    = psi2     input
  | otherwise = psi2 $ kkToG2 $ psiNaive (n-1) $ separate1st input

psiAny :: Ring c => GRing c -> KRing c
psiAny = Uni . ZMod.mapMaybeBaseCoeff f . XInf.unPoly where
  f (XInf is) = Just (U (sum is) , fromInteger (multinomial is))

omegaNaive :: Ring c => Int -> KRing c -> KRing c
omegaNaive n = psiAny . deltaN n

omegaH :: Ring c => Int -> KRing c -> KRing c
omegaH d = Uni . ZMod.mapMaybeBaseCoeff f . unUni where
  f (U k) = Just (U k, fromIntegral d ^ k)

separate1st :: forall c n. (Ring c) => GRing c -> GRing (KRing c) 
separate1st = XInf.Poly . ZMod.mapMaybeBaseCoeff g . ZMod.mapCoeff f . XInf.unPoly where
  f c  = scalarP c :: KRing c
  g (XInf (k:ns)) = Just (XInf ns, c) where
    c = monomP (U k)

unify1st :: forall c n. (Ring c) => GRing (KRing c) -> GRing c 
unify1st = XInf.Poly . ZMod.fromList . concatMap f . ZMod.toList . XInf.unPoly where
  f (XInf xs , Uni poly) = [ (XInf (k:xs) , c)  | (U k, c) <- ZMod.toList poly ]

unify1st2nd :: forall c n. (Ring c) => GRing (GRing c) -> GRing c 
unify1st2nd = XInf.Poly . ZMod.fromList . concatMap f . ZMod.toList . XInf.unPoly where
  f (XInf xs , XInf.Poly poly) = [ (XInf (kl++xs) , c)  | (XInf kl0, c) <- ZMod.toList poly , let kl = take 2 (kl0++[0,0]) ]

crossKs :: Ring c => [KRing c] -> GRing c
crossKs = XInf.Poly . ZMod.productWith empty cross . map (ZMod.mapBase sing) . map unUni where
  sing (U k) = XInf [k]
  cross (XInf as) (XInf bs) = XInf (as++bs)
  empty = XInf []

kkToG2 :: Ring c => KRing (KRing c) -> GRing c 
kkToG2 = XInf.Poly . ZMod.fromList . concatMap f . ZMod.toList . unUni where
  f (U k , Uni poly) = [ (XInf [k,l] , c)  | (U l, c) <- ZMod.toList poly ]

unifyKK :: Ring c => KRing (KRing c) -> KRing c 
unifyKK = Uni . ZMod.fromList . concatMap f . ZMod.toList . unUni where
  f (U k , Uni poly) = [ (U (k+l) , c)  | (U l, c) <- ZMod.toList poly ]

--------------------------------------------------------------------------------

{-
deltaN :: Num c => Int -> KRing c -> GRing c
deltaN 0 = error "deltaN: 0"
deltaN 1 = embed
deltaN 2 = delta2
deltaN 
-}