-- | The equivariant Segre-Schwartz-MacPherson classes
--
-- We can recover the Segre-SM classes by dividing the CSM class
-- by the total Chern class of the tangent bundle of the (smooth)
-- ambient variety.
--
-- The Segre-SM class is useful because it behaves well wrt. pullback.
--

{-# LANGUAGE ScopedTypeVariables, BangPatterns #-}
module Math.RootLoci.Segre.Equivariant where

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

import Math.Combinat.Classes
import Math.Combinat.Numbers
import Math.Combinat.Sign
import Math.Combinat.Compositions
import Math.Combinat.Partitions.Integer
import Math.Combinat.Numbers.Series

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

import Math.RootLoci.CSM.Equivariant.Umbral       -- this is the fastest one

--------------------------------------------------------------------------------
-- * The total Chern class

-- | Total Chern class of the representation @Sym^m C^2@ 
--
-- > c(Sym^m C^2) = \prod_{i=0}^m (1 + i*a + (m-i)*b)
--
affTotalChernClass :: ChernBase base => Int -> ZMod base
affTotalChernClass m = select1 (total , abToChern total) where
  total = product [ 1 + w | w <- affineWeights m ]

-- | Parts of the total Chern class, separated by degree
affTotalChernClassByDegree :: ChernBase base => Int -> [ZMod base]
affTotalChernClassByDegree = elems . separateGradedParts . affTotalChernClass

--------------------------------------------------------------------------------
-- * Inverse of the total Chern class

-- | Infinite power series expansion (by degree) of the multiplicative
-- inverse of the total Chern class of the representation @Sym^m C^2@
--
-- This is just the sum of all complete symmetric polynomials of the sums.
--
recipTotalChernClass :: forall base. ChernBase base => Int -> [ZMod base]
recipTotalChernClass m = pseries' coeffs where

  coeffs      = zip (map ZMod.neg prodWeights) [1..]
  prodWeights = tail (affTotalChernClassByDegree m)

-- | Another implementation of 'recipTotalChernClass'
recipTotalChernClass2 :: forall base. ChernBase base => Int -> [ZMod base]
recipTotalChernClass2 m = integralReciprocalSeries (affTotalChernClassByDegree m) where

-- | A third, very slow implementation of 'recipTotalChernClass'
recipTotalChernClassSlow :: forall base. ChernBase base => Int -> [ZMod base]
recipTotalChernClassSlow m = select2 (list , map abToChern list) where

  weights = affineWeights m
  list    = [ grade d | d <- [0..] ]

  grade :: Int -> ZMod AB
  grade d = negateIfOdd d 
          $ ZMod.sum (map mkProduct $ compositions (m+1) d)

  mkProduct es = ZMod.product [ (weightPowers!i) !! e | (i,e) <- zip [0..m] es ]

  -- much faster to cache to powers of the weights!
  weightPowers :: Array Int [ZMod AB]
  weightPowers = listArray (0,m) [ wtPowList (weights !! i) | i <- [0..m] ] 

  wtPowList :: ZMod AB -> [ZMod AB]
  wtPowList w = go 1 where { go !x = x : go (x*w) }


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

-- | Divides a polynomial with the total chern class. As the result is an
-- infinite power series, we return it's homogeneous parts as an infinite list.
--
-- Equivalent (but should be faster than) to:
--
-- > separeteGradedParts what `mulSeries` (recipTotalChernClass m)
-- 
divideByTotalChernClass :: ChernBase base => Int -> ZMod base -> [ZMod base]
divideByTotalChernClass m what = convolveWithPSeries' coeffs numerList where

  numerArr  = separateGradedParts what
  numerList = elems numerArr

  coeffs      = zip (map ZMod.neg prodWeights) [1..]
  prodWeights = tail (affTotalChernClassByDegree m)

-- | Another, very slow implementation of 'divideByTotalChernClass'
divideByTotalChernClassSlow :: ChernBase base => Int -> ZMod base -> [ZMod base]
divideByTotalChernClassSlow m what = final where
  (0,n)     = bounds numerArr
  numerArr  = separateGradedParts what
  denomList = recipTotalChernClassSlow m
  final     = [ part d | d <- [0..] ]
  part deg  = ZMod.sum 
    [ (numerArr ! i) * (denomList !! j)   
    | j <- [ max 0 (deg-n) .. deg ] 
    , let i = deg - j 
    ]

--------------------------------------------------------------------------------
-- * Affine Segre-SM classes

-- | Affine equivariant Segre-SM class of the open strata
affineOpenSegreSM :: ChernBase base => Partition -> [ZMod base]
affineOpenSegreSM part = divideByTotalChernClass m (umbralAffOpenCSM part) where
  m = weight part

-- | Affine equivariant Segre-SM class of the zero orbit
affineZeroSegreSM :: ChernBase base => Int -> [ZMod base]
affineZeroSegreSM m = divideByTotalChernClass m (affineZeroCSM m)

-- | Affine equivariant Segre-SM class of the closure of the strata (including the zero orbit!)
affineClosedSegreSM :: ChernBase base => Partition -> [ZMod base]
affineClosedSegreSM part = divideByTotalChernClass m (umbralAffClosedCSM part) where
  m = weight part

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