Abstract.hs

-- | The abstract motivic algorithm
--
-- See: B. Komuves: Motivic characteristic classes of discriminant strata
--
-- TODO: caching of results (otherwise it is very slow)

{-# LANGUAGE CPP, BangPatterns, FlexibleInstances, TypeSynonymInstances,
             MultiParamTypeClasses, FunctionalDependencies, GeneralizedNewtypeDeriving,
             TypeFamilies
  #-}
module Math.RootLoci.Motivic.Abstract where

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

import Data.Char
import Data.List
import Data.Ord
import Data.Maybe

import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import qualified Math.Algebra.Polynomial.FreeModule as ZMod 
import Math.Algebra.Polynomial.FreeModule (ZMod,QMod,FreeMod)
import Math.Algebra.Polynomial.Pretty

import Math.Combinat.Classes hiding (empty)
import Math.Combinat.Tuples
import Math.Combinat.Partitions
import Math.Combinat.Permutations hiding (permute)

-- import Debug.Trace
-- debug s x y = trace (">>> " ++ s ++ " -> " ++ show x) y

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

--------------------------------------------------------------------------------
-- * The abstract algorithm

-- | The (abstract) class of @Sym^n(X)@
symn :: Num c => Dim -> FreeMod c SingleLam
symn dim = ZMod.generator $ SingleLam (Bindings [dim]) (Single [(DeBruijn 0, 1)])

-- | The open stratum X(1,1,...,1)
open :: Dim -> ZMod SingleLam
open d@(Dim n) = symn d `ZMod.sub` rest where
  rest = ZMod.sum [ xlam p | p <- partitions n , width p < n ]

zeros :: Int -> ZMod MultiLam
zeros k = ZMod.generator 
        $ MultiLam (Bindings []) (Multi $ replicate k (Single []))

-- | The open stratum X(lambda)
xlam :: Partition -> ZMod SingleLam
xlam p = 
  if height p == 1   
    then open (Dim $ weight p)
    else normalize $ psi $ omega123 $ dvec $ dimVector p

-- | The open stratum D(n1,n2,...)
dvec :: [Dim] -> ZMod MultiLam
dvec dims0 = permute invperm sorted where
  invperm = inversePermutation perm
  perm    = sortingPermutationDesc dims0
  dims1   = permuteList perm dims0
  (dims2,dzeros) = span (>0) dims1        -- separate zero dimensions
  sorted  = cross (dvecSorted dims2) (zeros $ length dzeros)

-- | The open stratum D(n1,n2,...), assuming @n1 >= n2 >= n3 >= ...@
dvecSorted :: [Dim] -> ZMod MultiLam
dvecSorted []     = error "dvec: empty dimension vector shouldn't appear in the algorithm"
dvecSorted [n]    = singleToMulti (open n)
dvecSorted (p:ns) = normalize (ZMod.sub big rest) where
  big  = cross (singleToMulti $ open p) (dvecSorted ns)
  rest = ZMod.sum  
    [ theta (dvec (k : interleave ds es))
    | ds <- dimTuples ns
    , let es = zipWith (-) ns ds
    , let l = sum es
    , l > 0 
    , let k = p - l
    , k >= 0
    ] 

--------------------------------------------------------------------------------
-- * Data types and instances

-- | A variable, implemented as a /de Bruijn level/ (indexing starts from 0)
newtype Var 
  = DeBruijn Int      
  deriving (Eq,Ord,Show)

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

-- | We use de Bruijn levels to index the bound variables, and ecah bound variables has a dimension
newtype Bindings 
  = Bindings [Dim]
  deriving (Eq,Ord,Show)

numberOfBoundVariables :: Bindings -> Int
numberOfBoundVariables (Bindings ds) = length ds

dimensionTable :: Bindings -> Map Var Dim
dimensionTable (Bindings dims) = Map.fromList $ zip (map DeBruijn [0..]) dims

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

-- | An expression living on @Sym^n(X)@, with free variables
newtype Single
  = Single [(Var,Int)]
  deriving (Eq,Ord,Show)

unSingle :: Single -> [(Var,Int)]
unSingle (Single ves) = ves

-- | An expression living on @Sym^{n_1}(X) x ... x Sym^{n_r}(X)@, with free variables
newtype Multi
  = Multi [Single]
  deriving (Eq,Ord,Show)

-- | A lambda expression living on @Sym^n(X)@, with variables bound to @Sym^d(X)@ with different dimensions
data SingleLam
  = SingleLam !Bindings !Single
  deriving (Eq,Ord,Show)

-- | A lambda expression living on @Sym^{n_1}(X) x ... x Sym^{n_r}(X)@, with variables bound to @Sym^d(X)@ with different dimensions
data MultiLam
  = MultiLam !Bindings !Multi
  deriving (Eq,Ord,Show)

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

instance Pretty Bindings where
  pretty (Bindings dims) = "\\" ++ concat (zipWith f vars dims) where
    vars = map DeBruijn [0..]
    f v d = "(" ++ pretty v ++ ":S" ++ show d ++ ")"

instance Pretty Var where 
  pretty (DeBruijn i) = chr (97 + i) : []

instance Pretty (Var,Int) where
  pretty (v,0) = "1"
  pretty (v,1) = pretty v
  pretty (v,e) = pretty v ++ "^" ++ show e

instance Pretty Single where
  pretty (Single ves) = intercalate "*" $ map pretty ves

instance Pretty Multi where
  pretty (Multi ts) = "[" ++ (intercalate "," $ map pretty ts) ++ "]"

instance Pretty SingleLam where
  pretty (SingleLam binds body) = "{" ++ pretty binds ++ "->" ++ pretty body ++ "}"

instance Pretty MultiLam where
  pretty (MultiLam binds body) = "{" ++ pretty binds ++ "->" ++ pretty body ++ "}"

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

instance Degree SingleLam where
  type MultiDegree SingleLam = Int
  multiDegree (SingleLam binds (Single ves)) = sum [ (unDim $ (Map.!) dimTable v) * e | (v,e) <- ves ] where
    dimTable = dimensionTable binds
  totalDegree = multiDegree

instance Degree MultiLam where
  type MultiDegree MultiLam = [Int]
  multiDegree (MultiLam binds (Multi bodies)) = map totalDegree [ (SingleLam binds b) | b <- bodies ]
  totalDegree = sum . multiDegree

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

instance Empty Bindings where
  empty = Bindings []

instance Empty Single where
  empty = Single []

instance Empty Multi where
  empty = Multi []

instance Empty SingleLam where
  empty = SingleLam empty empty

instance Empty MultiLam where
  empty = MultiLam empty empty

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

-- | Shift de Bruijn levels
class Shift a where
  shift :: Int -> a -> a
  
instance Shift Var where
  shift k (DeBruijn l) = DeBruijn (k+l)

instance Shift Single where
  shift k (Single ves) = Single [ (shift k v, e) | (v,e) <- ves ]

instance Shift Multi where
  shift k (Multi terms) = Multi $ map (shift k) terms

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

-- | Rename variables
class Rename a where
  rename :: (Var -> Var) -> a -> a

instance Rename Var where
  rename f v = f v

instance Rename (Var,Int) where
  rename f (v,e) = (f v, e)

instance Rename Single where
  rename f (Single ves) = Single $ map (rename f) ves

instance Rename Multi where
  rename f (Multi ts) = Multi $ map (rename f) ts

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

-- | Extract the exponent of a given variable
exponentOf :: Var -> Single -> Int
exponentOf u (Single ves) = sum [ e | (v,e) <- ves , v==u ]

-- | Extract the exponent vector of a given variable
exponentVectorOf :: Var -> Multi -> [Int]
exponentVectorOf v (Multi ts) = map (exponentOf v) ts

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

instance Normalize Single where
  normalize (Single ves) = Single $ Map.toList $ Map.fromListWith (+) ves

instance Normalize Multi where
  normalize (Multi terms) = Multi (map normalize terms) 

normalizeWithExpo :: (Rename term, Normalize term, Ord expo) => (expo -> Bool) -> (Var -> term -> expo) -> (Bindings,term) -> (Bindings,term) 
normalizeWithExpo cond expo (binds,body) = (binds',body') where
  Bindings dims = binds
  vars   = map DeBruijn [0..]
  vby    = [ (v,(d,es)) | (v,d) <- zip vars dims , let es = expo v body , cond es ]
  sorted = sortOn snd vby
  dims'  = map (fst . snd) sorted
  binds' = Bindings dims'
  f v    = DeBruijn $ fromJust $ findIndex (\pair -> fst pair == v) sorted
  body'  = normalize $ rename f $ body

instance Normalize SingleLam where
  normalize (SingleLam binds body) = SingleLam binds' body' where
    (binds',body') = normalizeWithExpo (>0) exponentOf (binds,body)

instance Normalize MultiLam where
  normalize (MultiLam binds body) = MultiLam binds' body' where
    cond ds = any (>0) ds
    (binds',body') = normalizeWithExpo cond f (binds,body)
    f v = {- reverse . -} exponentVectorOf v

instance (Eq c, Num c) => Normalize (FreeMod c SingleLam) where
  normalize = ZMod.mapBase normalize

instance (Eq c, Num c) => Normalize (FreeMod c MultiLam) where
  normalize = ZMod.mapBase normalize

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

instance SuperNormalize Multi where
  superNormalize (Multi ts) = Multi $ reverse $ dropWhile isempty $ reverse $ map normalize $ ts where
    isempty (Single xs) = null xs

instance SuperNormalize MultiLam where
  superNormalize mlam = MultiLam binds (superNormalize body) where
    (MultiLam binds body) = normalize mlam

instance (Eq c, Num c) => SuperNormalize (FreeMod c MultiLam) where
  superNormalize = ZMod.mapBase superNormalize
  
--------------------------------------------------------------------------------

instance Cross Bindings where
  cross (Bindings ds) (Bindings es) = Bindings (ds++es)
  crossInterleave = error "Bindings/crossInterleave: undefined"
  
instance Cross Multi where
  cross (Multi xs) (Multi ys) = Multi (xs++ys)
  crossInterleave (Multi xs) (Multi ys) = Multi (interleave xs ys)

instance Cross MultiLam where
  cross (MultiLam binds1 bodies1) (MultiLam binds2 bodies2) = normalize $ MultiLam binds3 bodies3 where
    n1 = numberOfBoundVariables binds1
    binds3  = binds1  `cross` binds2
    bodies3 = bodies1 `cross` (shift n1 bodies2)

  crossInterleave (MultiLam binds1 bodies1) (MultiLam binds2 bodies2) = normalize $ MultiLam binds3 bodies3 where
    n1 = numberOfBoundVariables binds1
    binds3  = binds1  `cross` binds2
    bodies3 = bodies1 `crossInterleave` (shift n1 bodies2)

instance (Eq c, Num c) => Cross (FreeMod c MultiLam) where
  cross x y = normalize $ ZMod.mulWith cross x y
  crossInterleave x y = normalize $ ZMod.mulWith crossInterleave x y

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

instance SingleToMulti Single Multi where
  singleToMulti = Multi . (:[])

instance SingleToMulti SingleLam MultiLam where
  singleToMulti (SingleLam binds single) = MultiLam binds (Multi [single])

instance (Eq c, Num c) => SingleToMulti (FreeMod c SingleLam) (FreeMod c MultiLam) where
  singleToMulti = ZMod.mapBase singleToMulti

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

instance Omega (Var,Int) where
  omega 0 _     = omegaZeroError 
  omega k (v,d) = (v,d*k)

instance Omega Single where
  omega 0 _            = Single [] -- omegaZeroError 
  omega k (Single ves) = Single $ map (omega k) ves

instance Omega Multi where
  omega 0 _          = omegaZeroError 
  omega k (Multi ts) = Multi $ map (omega k) ts

instance Omega SingleLam where
  omega 0 _                      = omegaZeroError 
  omega k (SingleLam binds body) = SingleLam binds (omega k body)

instance Omega MultiLam where
  omega 0 _                     = omegaZeroError 
  omega k (MultiLam binds body) = MultiLam binds (omega k body)

instance (Eq c, Num c) => Omega (FreeMod c SingleLam) where
  omega 0 = omegaZeroError 
  omega k = ZMod.mapBase (omega k)

instance (Eq c, Num c) => Omega (FreeMod c MultiLam) where
  omega 0 = omegaZeroError 
  omega k = normalize . ZMod.mapBase (omega k)

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

instance Omega123 Multi where
  omega123 (Multi ts) = Multi $ zipWith omega [1..] ts

instance Omega123 MultiLam where
  omega123 (MultiLam binds body) = MultiLam binds (omega123 body)

instance (Eq c, Num c) => Omega123 (FreeMod c MultiLam) where
  omega123 = normalize . ZMod.mapBase omega123

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

instance Psi Multi Single where
  psi (Multi ts) = normalize $ Single $ concat $ map unSingle ts

instance Psi MultiLam SingleLam where
  psi (MultiLam binds body) = SingleLam binds (psi body)

instance (Eq c, Num c) => Psi (FreeMod c MultiLam) (FreeMod c SingleLam) where
  psi = normalize . ZMod.mapBase psi

instance (Eq c, Num c) => Psi [FreeMod c SingleLam] (FreeMod c SingleLam) where
  psi = normalize . psi . crossMany . map singleToMulti

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

instance PsiEvenOdd Multi where
  psiEvenOdd (Multi ts) = normalize $ Multi $ zipWith f (evens ts) (odds ts) where
    f (Single xs) (Single ys) = Single (xs++ys)

instance PsiEvenOdd MultiLam where
  psiEvenOdd (MultiLam binds body) = MultiLam binds (psiEvenOdd body)

instance PsiEvenOdd (ZMod MultiLam) where
  psiEvenOdd = normalize . ZMod.mapBase psiEvenOdd

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

instance Pontrjagin SingleLam where
  pontrjaginOne     = empty
  pontrjaginMul a b = psi $ cross (singleToMulti a) (singleToMulti b)

instance Pontrjagin MultiLam where
  pontrjaginOne     = empty
  pontrjaginMul a b = psiEvenOdd $ crossInterleave a' b' where
    (a',b') = extendToCommonSize (a,b)

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

instance ExtendToCommonSize Multi where
  extendToCommonSize (Multi xs, Multi ys) = (Multi xs', Multi ys') where
    (xs',ys') = extendToCommonSize (xs,ys) 

instance ExtendToCommonSize MultiLam where
  extendToCommonSize (MultiLam as xs, MultiLam bs ys) = (MultiLam as xs', MultiLam bs ys') where
    (xs',ys') = extendToCommonSize (xs,ys) 

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

instance Permute Multi where
  permute p (Multi ts) = Multi (permuteList p ts)

instance Permute MultiLam where
  permute p (MultiLam binds multi) = MultiLam binds (permute p multi)

instance Permute (ZMod MultiLam) where
  permute p = ZMod.mapBase (permute p)

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

instance Theta Multi where
  theta (Multi (u:us)) = Multi (a:bs) where
    a = psi $ Multi (u : odds us)
    bs = zipWith f (evens us) (odds us)
    f (Single u) (Single v) = normalize $ Single (u ++ v)

instance Theta MultiLam where
  theta (MultiLam binds body) = MultiLam binds (theta body)

instance Theta (ZMod MultiLam) where
  theta = normalize . ZMod.mapBase theta

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