-- | Compute the pushforward maps @pi_*@ and @delta_*@ between the
-- @GL2@-equivariant cohomology rings
--
-- Recall that:
--
-- * @Delta_nu : Q^d -> Q^n@
--
-- * @pi : Q^n -> P^n@
--
-- and @Q^n = P^1 x P^1 x ... x P^1@.
--  

{-# LANGUAGE 
      BangPatterns, TypeSynonymInstances, FlexibleInstances, FlexibleContexts,
      ScopedTypeVariables, TypeFamilies 
  #-}

module Math.RootLoci.CSM.Equivariant.PushForward 
  ( -- * The function tau
    tau , tauEta
    -- * pushforward along the diagonal map @Delta_{nu} : Q^d -> Q^n@
  , delta_star_ , delta_star , delta_star' 
    -- * pushforward along the order-forgetting map @pi : Q^n -> P^n@
  , pi_star_table
  , compute_pi_star
  , pi_star
    -- * Fibonacci-type recursion formula for @pi_*@
  , piStarTableAff 
  , piStarTableProj
  )
  where

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

import Math.Combinat.Numbers
import Math.Combinat.Sign
import Math.Combinat.Partitions.Integer
import Math.Combinat.Partitions.Set
import Math.Combinat.Sets
import Math.Combinat.Tuples

import Data.Array (Array)
import Data.Array.IArray

import Math.RootLoci.Algebra
import Math.RootLoci.Geometry
import Math.RootLoci.Misc

import qualified Math.Algebra.Polynomial.FreeModule as ZMod

--------------------------------------------------------------------------------
-- * The function tau

-- | @tau_k := ( a^(k+1) - b^(k+1) ) / ( a - b )@
tau :: ChernBase base => Int -> ZMod base
tau k = select1 ( tauAB k , tauChern k ) 

-- | In chern classes, the coefficients of tau are (signed) binomial coefficients; cf. A011973
tauChern :: Int -> ZMod Chern
tauChern k 
  | k <  -1    = error "tau: negative index is not implemented"
  | k == -1    = ZMod.zero 
  | otherwise  = ZMod.fromList [ ( Chern (k - 2*j) j , negateIfOdd j $ binomial (k-j) j )  | j<-[0..div k 2] ]

tauChernUnsafe :: Int -> ZMod Chern
tauChernUnsafe = icache $ \k -> abToChern (tauAB k)

tauAB :: Int -> ZMod AB
tauAB k 
  | k <  -1   = error "tau: negative index is not implemented"
  | k == -1   = ZMod.zero
  | otherwise = ZMod.fromList [ (AB j (k-j) , 1) | j <- [0..k] ]

tauEta :: ChernBase base => Int -> ZMod (Eta base)
tauEta k = injectZMod (tau k)

--------------------------------------------------------------------------------
-- * @Delta_{\nu} : Q^d -> Q^n@

-- | Input: diagonal eta indices, and whether we are pushing forward 1 or the generator u/xi
delta_star_single :: ChernBase base => [Int] -> Bool -> ZMod (Eta base)
delta_star_single ks xi = 
  if xi
    then bbb 
    else aaa 
  where
    n = length ks

    aaa = ZMod.sum [ sigma (n-1-i) * (tauEta i) | i<-[0..n-1] ]
    bbb = full - rest

    ab   = ZMod.generator $ Eta [] $ select0 (alphaBeta, c2)
    full = ZMod.generator (Eta ks mempty)      -- == sigma_n(eta)
    rest = ZMod.sum [ sigma (n-i) * tauEta (i-2) * ab | i<-[2..n] ]

    sigma k = symPoly k [ Eta [k] mempty | k<-ks ]
  
-- | a group generator on the left is a subset (=product) of U-s, which
-- we map to a linear combinaton of H-s
delta_star_1 :: ChernBase base => Partition -> Omega base -> ZMod (Eta base)
delta_star_1 part = delta_star_1' (linearIndices part)

-- | a group generator on the left is a subset (=product) of U-s, which
-- we map to a linear combinaton of H-s
delta_star_1' :: forall base. ChernBase base => [[Int]] -> Omega base -> ZMod (Eta base)
delta_star_1' idxtable (Omega us ab) = final where
  
  final = mulInjMonom ab $ ZMod.product $ go 1 idxtable
          
  go :: Int -> [[Int]] -> [ZMod (Eta base)]
  go _ []       = []
  go k (is:iss) = this : go (k+1) iss where
    this = delta_star_single is (k `elem` us)

delta_star_ :: ChernBase base => Partition -> ZMod (Omega base) -> ZMod (Eta base)
delta_star_ part = ZMod.flatMap (delta_star_1 part)

delta_star :: ChernBase base => SetPartition -> ZMod (Omega base) -> ZMod (Eta base)
delta_star setp = ZMod.flatMap (delta_star_1' (fromSetPartition setp))

-- | We can give an explicit indexing scheme (set partition), instead of the linear indexing
-- used above. This will be useful when computing the \"open\" part
delta_star' :: ChernBase base => [[Int]] -> ZMod (Omega base) -> ZMod (Eta base)
delta_star' indices = ZMod.flatMap (delta_star_1' indices)

--------------------------------------------------------------------------------
-- * @pi : Q^n -> P^n@

-- | This is upside the class where @[0:1]@ is a root with multiplicity @k@ and @[1:0]@ is a root with multiplicity l
up_root_xy :: Int -> (Int,Int) -> ZMod (Eta AB)
up_root_xy n (k,l) = as * bs where

  as = ZMod.product [ abh      i  1 0 | i<-[1..k] ]
  bs = ZMod.product [ abh (n+1-j) 0 1 | j<-[1..l] ]

  -- (eta_i + na*alpha + nb*beta)
  abh i na nb = ZMod.fromList 
    [ (Eta [i] (AB 0 0) , 1 )  
    , (Eta []  (AB 1 0) , na)  
    , (Eta []  (AB 0 1) , nb) 
    ]

-- | This is downside the class where @[0:1]@ is a root with multiplicity @k@ and @[1:0]@ is a root with multiplicity l.
-- It should be true that @pi_* up_root_xy = down_root_xy@
down_root_xy :: Int -> (Int,Int) -> ZMod (Gam AB)
down_root_xy n (k,l) = as * bs where

  as = ZMod.product [ abg (n-i) (i) | i<-[0..k-1] ]
  bs = ZMod.product [ abg (j) (n-j) | j<-[0..l-1] ]

  -- (na*alpha + nb*beta + gamma)
  abg na nb = ZMod.fromList 
    [ (Gam 1 (AB 0 0) , 1 ) 
    , (Gam 0 (AB 1 0) , fromIntegral na) 
    , (Gam 0 (AB 0 1) , fromIntegral nb) 
    ]

pi_star_0 :: Int -> Int -> ZMod (Gam AB)
pi_star_0 n k = ZMod.sum
  [ ZMod.scale 
      (negateIfOdd i $ binomial k i * factorial (n-k+i)) 
      (mulAB (AB i 0) $ down_root_xy n (k-i,0)) 
  | i<-[0..k] ]

-- | Table of @pi_*( eta_1*eta_2*...*eta_k )@, computed by breaking the symmetry.
pi_star_table :: Int -> Array Int (ZMod (Gam AB))
pi_star_table = monoCache calc where
  calc n = listArray (0,n) [ pi_star_0 n k | k<-[0..n] ]

-- | Slow implementation of @pi_star@, using @pi_star_table@
compute_pi_star 
  :: Int               -- ^ the number of points @m@ (recall the pi : @Q^m -> P^m@)
  -> ZMod (Eta AB) 
  -> ZMod (Gam AB)
compute_pi_star m = ZMod.flatMap f where 
  table = pi_star_table m
  f (Eta hs ab) = mulAB ab (table ! length hs)

--------------------------------------------------------------------------------
-- * Fibonacci-type recursion formula for @pi_*@

-- | However it should faster to just use the recursion for the @P_j(m)@ polynomials,
-- which this function does.
pi_star 
  :: forall base. (ChernBase base) 
  => Int                      -- ^ the number of points @m@ (recall the pi : @Q^m -> P^m@)
  -> ZMod (Eta base) 
  -> ZMod (Gam base)
pi_star m = ZMod.flatMap f where 
  table = piStarTableProj m :: Array Int (ZMod (Gam base))
  f (Eta hs ab) = mulInjMonom ab (table ! length hs)

piStarTableAff :: ChernBase base => Int -> Array Int (ZMod base)
piStarTableAff = polyCache2 calc where
  calc n = select2 ( aff_fibPiStar_AB n , aff_fibPiStar_Chern n )

piStarTableProj :: ChernBase base => Int -> Array Int (ZMod (Gam base))
piStarTableProj = polyCache3 calc where
  calc n = select3 ( proj_fibPiStar_AB n , proj_fibPiStar_Chern n )

{-

class ChernBase (PiStarBase tgtmonom) => PiStar tgtmonom where
  type PiStarBase tgtmonom :: *
  piStarTable :: Int -> Array Int (ZMod tgtmonom)

instance PiStar (Gam Chern) where { piStarTable = proj_fibPiStar_Chern ; type PiStarBase (Gam Chern) = Chern }
instance PiStar (Gam AB   ) where { piStarTable = proj_fibPiStar_AB    ; type PiStarBase (Gam AB   ) = AB    }
instance PiStar      Chern  where { piStarTable = aff_fibPiStar_Chern  ; type PiStarBase (Chern    ) = Chern }
instance PiStar      AB     where { piStarTable = aff_fibPiStar_AB     ; type PiStarBase (AB       ) = AB    }

-- instance PiStar (Gam Schur) where { piStarTable = proj_fibPiStar_Schur ; type PiStarBase = Gam Schur }
-- instance PiStar      Schur  where { piStarTable = aff_fibPiStar_Schur  ; type PiStarBase = Schur     }
-}

proj_fibPiStar_Chern :: Int -> Array Int (ZMod (Gam Chern))
proj_fibPiStar_Chern m = listArray (0,m) $ take (m+1) fib where

  fib :: [ZMod (Gam Chern)]
  fib = ZMod.konst                    (factorial  m   )
      : ZMod.singleton (Gam 1 mempty) (factorial (m-1))
      : zipWith3 g [1..] (tail fib) fib 

  g :: Integer -> ZMod (Gam Chern) -> ZMod (Gam Chern) -> ZMod (Gam Chern)
  g k prev1 prev2 
    = ZMod.divideByConst (mm-k)
    $ mulGam prev1 + ZMod.scale k (mulInjMonom c1 prev1) 
                   + ZMod.scale k (mulInjMonom c2 prev2) 

  mm = fromIntegral m :: Integer

--  c1 = Chern 1 0
--  c2 = Chern 0 1

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

aff_fibPiStar_Chern :: Int -> Array Int (ZMod Chern)
aff_fibPiStar_Chern m = listArray (0,m) $ take (m+1) fib where

  fib :: [ZMod Chern]
  fib = ZMod.konst (factorial m)
      : ZMod.zero
      : zipWith3 g [1..] (tail fib) fib 

  g :: Integer -> ZMod Chern -> ZMod Chern -> ZMod Chern
  g k prev1 prev2 
    = ZMod.divideByConst (mm-k)
    $ ZMod.scale         (   k)
    $ (ZMod.mulByMonom c1 prev1 + ZMod.mulByMonom c2 prev2) 

  mm = fromIntegral m :: Integer

--  c1 = Chern 1 0
--  c2 = Chern 0 1

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

proj_fibPiStar_AB :: Int -> Array Int (ZMod (Gam AB))
proj_fibPiStar_AB m = fmap (convertGam chernToAB) (proj_fibPiStar_Chern m)

proj_fibPiStar_Schur :: Int -> Array Int (ZMod (Gam Schur))
proj_fibPiStar_Schur m = fmap (convertGam chernToSchur) (proj_fibPiStar_Chern m)

aff_fibPiStar_AB :: Int -> Array Int (ZMod AB)
aff_fibPiStar_AB m =  fmap chernToAB (aff_fibPiStar_Chern m)

aff_fibPiStar_Schur :: Int -> Array Int (ZMod Schur)
aff_fibPiStar_Schur m =  fmap chernToSchur (aff_fibPiStar_Chern m)

--------------------------------------------------------------------------------
-- * helpers

-- | Multiplies by an injected monomial 
mulInjMonom :: (Functor f, Monoid ab, Ord (f ab)) => ab -> ZMod (f ab) -> ZMod (f ab)
mulInjMonom monom = ZMod.mapBase f where
  f = fmap (mappend monom)

-- | Multiplies by @(alpha^i * beta^j)@
mulAB :: (Functor f, Ord (f AB)) => AB -> ZMod (f AB) -> ZMod (f AB)
mulAB = mulInjMonom

-- | Multiplies with @gamma@
mulGam :: Ord ab => ZMod (Gam ab) -> ZMod (Gam ab)
mulGam = ZMod.mapBase f where 
  f (Gam k x) = Gam (k+1) x

{-
-- | Multiplies by alpha^i beta^j
omegaMulAB :: AB -> ZMod (Omega AB) -> ZMod (Omega AB)
omegaMulAB (AB i j) = Map.mapKeys f where
  f (Omega us (AB a b)) = Omega us (AB (a+i) (b+j))

-- | Multiplies by alpha^i beta^j
etaMulAB :: AB -> ZMod (Eta AB)-> ZMod (Eta AB)
etaMulAB (AB i j) = Map.mapKeys f where
  f (Eta hs (AB a b)) = Eta hs (AB (a+i) (b+j))

-- | Multiplies by alpha^i beta^j
gamMulAB :: AB -> ZMod (Gam AB) -> ZMod (Gam AB)
gamMulAB (AB i j) = Map.mapKeys f where
  f (Gam g (AB a b)) = Gam g (AB (a+i) (b+j))
-}

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