(* ==== COMPUTE THR EQUIVARIANT CHERN-SCHWARTZ-MACPHERSON 
CLASS OF COINCIDENT ROOT LOCI VIA THE MOTIVIC ALGORITHM ==== *)

<<Combinatorica`

(* convention: u = -c_1(L) where L is the tautological line bundle on P^n *)
uu[i_]:=Subscript[u,i]
vv[i_]:=Subscript[u,i]
ww[i_]:=Subscript[u,i]

(*only works for lists of equal length!!*)Zip[as_,bs_]:=MapThread[{#1,#2}&,{as,bs}]
Zip3[as_,bs_,cs_]:=MapThread[{#1,#2,#3}&,{as,bs,cs}]

Fst[pair_]:=pair[[1]]
Snd[pair_]:=pair[[2]]

extendListWithZeros[L_,n_]:=Join[L,Table[0,{i,1,n-Length[L]}]]

SumList[L_]:=Sum[x,{x,L}]

(* === EQUIVARIANT COHOMOLOGY RING OF P^n === *)

(* Weights of Sym^n C^2 *)
wt[n_,i_]:=(n-i)*α+i*β

(* the relation in H^*(P^n). Convention: 
u = -c1(L) where L is the tautological line bundle *)
rel[n_]:=Product[u+wt[n,i],{i,0,n}]

Clear[upow$nplus1]
upow$nplus1[n_] := upow$nplus1[n]= Expand[u^(n+1)-rel[n]]
(* (unefficiently) normalize a /polynomial/ in u *) 

Clear[normalizeSlow]
normalizeSlow[n_, X0_]:=Module[{X=Expand[X0],m},
m=Exponent[X,u];
If[m≤n,X,normalizeSlow[n,X/.{u^m->u^(m-n-1)*upow$nplus1[n]}]]
]

(* table of normalized powers of u *)
Clear[UPowAB]
UPowAB[n_,k_]:=UPowAB[n,k] = Expand[normalizeSlow[n,u^k]]

Clear[normalizeVarAB]
normalizeVarAB[n_,uuu_,X_]:=Module[
{m=Exponent[X,uuu]},
Expand[Sum[Coefficient[X,uuu,k]*(UPowAB[n,k]/.{u->uuu}),{k,0,m}]]
]

(* === pushforward along the multiplication (single monom) === *)

(* Psi_* (u^k * v^l) = ? cohomology indexing *)
Clear[psiStarAB]
psiStarAB[n_,m_,0,0]:=psiStarAB[n,m,0,0]=Binomial[n+m,n]
psiStarAB[n_,m_,k_,l_]:=psiStarAB[n,m,k,l] = Module[
{A,B,AB,F,R,IJ,sel,fun,cft},
A=Product[u+wt[n,i],{i,0,k-1}];
B=Product[v+wt[m,j],{j,0,l-1}];
AB =Expand[A*B];
sel[{i_,j_}]:=(i<k) || (j<l);
cft[{i_,j_}]:=Coefficient[Coefficient[AB,u,i],v,j];
fun[{i_,j_}]:=psiStarAB[n,m,i,j];
IJ = Flatten[Table[{i,j},{i,0,k},{j,0,l}],1];
IJ =Select[ IJ,sel];
(* Print[IJ];
Print[AB]; *)
F = Binomial[n+m-k-l,n-k]*Product[w+wt[n+m,i],{i,0,k+l-1}];
R =Sum[ cft[ij]*fun[ij],{ij,IJ}];  
Expand[F-R]
]

Clear[psiStarChern]
psiStarChern[n_,m_,k_,l_]:=psiStarChern[n,m,k,l] = 
SymmetricReduction[psiStarAB[n,m,k,l],{α,β},{c,d}][[1]]

(* === class of the diagonal === *)

notk[n_,k_]:=Select[Table[i,{i,0,n}],#≠k&]

(* class of the diagonal in P^n x P^n *)
Clear[deltaClassAB]
deltaClassAB[n_]:=deltaClassAB[n] =
Expand[Factor[Sum[ Product[(v+wt[n,i])(w+wt[n,i])/(wt[n,i]-wt[n,k]),{i,notk[n,k]}],{k,0,n}]]]

Clear[deltaClassCh]
deltaClassCh[n_]:=deltaClassCh[n] =
Expand[SymmetricReduction[deltaClassAB[n],{α,β},{c,d}][[1]]]

(* small diagonal in P^n x P^n x P^n *)
deltaClassTri[n_]:=deltaClassTri[n] =
Expand[Factor[Sum[ Product[(uu[1]+wt[n,i])(uu[2]+wt[n,i])(uu[3]+wt[n,i])/(wt[n,i]-wt[n,k]),{i,notk[n,k]}],{k,0,n}]]]

(* === pushforward along the diagonal map (single monom) === *)

(* pushforward of u^k along Delta : P^n -> P^n x P^n *)
Clear[deltaStarAB]
deltaStarAB[n_,0]  := deltaStarAB[n,0] = deltaClassAB[n]
deltaStarAB[n_,k_]:=deltaStarAB[n,k]= Module[
{prod,Y,preY,Delta,rest},
Y = Expand[Product[v+wt[n,i],{i,0,k-1}]];
preY = Y/.{v->u};
Delta = deltaClassAB[n];
prod=normalizeVarAB[n,v,Y*Delta];
rest= Sum[Coefficient[preY,u,i]*deltaStarAB[n,i],  {i,0,k-1} ];
Expand[ prod - rest]
]

Clear[deltaStarCh]
deltaStarCh[n_]:=deltaStarCh[n] =
Expand[SymmetricReduction[deltaStarAB[n],{α,β},{c,d}][[1]]]

(* === pushforwards for polynomials (not just single monoms) === *)

delta2[n_,uuu_,{vvv_,www_},X0_]:=Module[{X=Expand[X0]},
Sum[Coefficient[X,uuu,i]*(deltaStarAB[n,i]/.{v->vvv,w->www}) ,{i,0,n}]
]

deltaMany[n_,uuu_,vvvs_,X_]:=Module[{ttt},
If[Length[vvvs]==1,X/.{uuu->vvvs[[1]]},
If[Length[vvvs]==2,delta2[n,uuu,vvvs,X],
deltaMany[n,ttt,Drop[vvvs,1], delta2[n,uuu,{vvvs[[1]],ttt},X]]]]]

psi2[{n1_,n2_},{uuu_,vvv_},www_,X0_]:=Module[{X=Expand[X0]},
Sum[Coefficient[Coefficient[X,uuu,i],vvv,j]*(psiStarAB[n1,n2,i,j]/.{w->www}) ,{i,0,n1},{j,0,n2}]
]

psiMany[ns_,uuus_,www_,X_]:=Module[{ttt,vvvs,ms},
If[Length[uuus]==1,X/.{uuus[[1]]->www},
If[Length[uuus]==2,psi2[ns,uuus,www,X],
ms=Join[{ns[[1]]+ns[[2]]},Drop[ns,2]];
vvvs=Join[{ttt},Drop[uuus,2]];
psiMany[ms,vvvs,www, psi2[Take[ns,2],Take[uuus,2],ttt,X]]]]]

(* === pushforward along the power map === *)

(* compute the pushforward by composing the diagonal with the merging map *)
Clear[slowOmegaAB]
slowOmegaAB[n_,d_,k_]:=slowOmegaAB[n,d,k]=Module[
{vars=Table[Subscript[ttt,i],{i,1,d}],
dims=Table[n,{i,1,d}]},
Factor[psiMany[dims,vars,u,deltaMany[n,u,vars,u^k]]]
]

(* d-th power of u^k : P^n -> P^(n*d) *)
(* NOTE: this formula is valid for k>n too! 
you can experimentally check this 
by pre-normalizing and post-normalizing *)
Clear[omegaAB]
omegaAB[n_,d_,k_]:=
Module[
{idxs= Select[Table[i,{i,0,n*d}],Not[Divisible[#,d]]&]}
,u^k*d^(n-k)*Product[u+(n*d-i)*α+i*β,{i,idxs}]
]

(* pushforward along d-th power : P^n -> P^(n*d) *)
omega1[n_,d_,uuu_,www_,X0_]:=Module[{X=Expand[X0]},
Sum[Coefficient[X,uuu,i]*(omegaAB[n,d,i]/.{u->www}) ,{i,0,n}]
]

(* P^n1 x P^n2 x P^n3 -> P^(d1*n1) x P^(d2*n2) x P^(d3*n3) -> P^(d1*n1 + d2*n2 + d3*n3 *) 
omegaLam[ns_,ds_,uuus_,www_,X0_]:=Module[
{X=Expand[X0],
m = Length[ns],
vars,Y,nds,ttt,i},
vars=Table[Subscript[ttt,i],{i,1,m}];
nds=Table[ ns[[i]]*ds[[i]], {i,1,m}];
Y=X;
For[i=1,i≤m,i++,Y=omega1[ns[[i]],ds[[i]],uuus[[i]],vars[[i]],Y]];
psiMany[nds,vars,www,Y]
]

(* === EQUIVARIANT CSM CLASS === *)

uus[n_]:=Table[uu[i],{i,1,n}]
vvs[n_]:=Table[vv[i],{i,1,n}]

(* total chern class of P^n *)
Clear[chernPnAB]
chernPnAB[n_]:=chernPnAB[n]=normalizeVarAB[n,u,
Expand[Product[1+u+wt[n,i],{i,0,n}]]]

chernPnABVar[n_,uuu_]:=chernPnAB[n]/.{u->uuu}

EmptyPartQ[part_]:=Length[part]==0

DualPart[{}]:={}
DualPart[lam_]:=With[{m=lam[[1]]},Table[Length[Select[lam,#≥i&]],{i,1,m}]]
 
toExpoForm[{}]:={} 
toExpoForm[part_]:=Module[{k=Max[part]},Table[Length[Select[part,#==j&]],{j,1,k}]]

posVectorQ[as_]:=Map[#≥0&,as]/.{List->And};
kdeTriples[p_,ns_]:=Module[
{m = Length[ns],
posQ,oneK,A},
oneK[k_]:=Table[{k,ns-es,es},{es,Combinatorica`Compositions[p-k,m]}];
A=Table [oneK[k],{k,0,p-1}];
A=Select[Flatten[A,1],posVectorQ[Snd[#]]&];
A]

Clear[csmXLam,csmDisj1,csmDisj,csmDisjSorted]

csmXLam[{}]:=csmXLam[{}]=1
csmXLam[{1}]:=csmXLam[{1}]=chernPnAB[1]

csmDisj1[0]:=csmDisj1[0]=1
csmDisj1[1]:=csmDisj1[1]=chernPnAB[1]

csmDisj[{}]:=csmDisj[{}]=1

csmDisjSorted[{}]:=csmDisjSorted[{}]=1

(* equiv CSM of D(n) *)
csmDisj1[n_]:=csmDisj1[n]=Module[
{parts=Select[Combinatorica`Partitions[n],Length[#]<n&]},
Expand[chernPnAB[n]-Sum[csmXLam[p],{p,parts}]]
]

(* equiv CSM of X_lambda *)
csmXLam[lambda_]:=csmXLam[lambda]=Module[
{es = toExpoForm[lambda],m,m1,ns,pairs,X},
m=Length[es];
ns = Range[m];
pairs = Zip[ns,es];   (* i^e_ *)
pairs = Select[pairs,Snd[#]>0&]; (* !!! *)
m1 = Length[pairs];
ns =Map[Fst,pairs];
es =Map[Snd,pairs];
X=csmDisj[es];
(* Print["xlam1 - ",pairs];
Print["xlam2 - ",es," | ",ns," | ",uus[m1]," | ",X]; *)
Expand[omegaLam[es,ns,uus[m1],u,X]]
]

(* equiv CSM of D(d1,d2,...) *)
csmDisj[{n_}]:=csmDisj[{n}] = csmDisj1[n]/.{u->uu[1]}
csmDisj[ns0_]:=csmDisj[ns0]=Module[
{m=Length[ns0],
nis0,nis1,ns1,idxs,X,ttt},
nis0 =Zip[ns0,Range[m]];
nis1=SortBy[nis0,-Fst[#]&];
idxs=Map[Snd,nis1];
ns1 =Select[Map[Fst,nis1],#>0&];
(* Print["nis1 - ",nis0," | ",nis1];
Print["nis2 - ",ns1]; *)
X=csmDisjSorted[ns1];
X=X/.Table[uu[i]->Subscript[ttt,i],{i,1,m}];
X/.Table[Subscript[ttt,i]->uu[idxs[[i]]],{i,1,m}]
]

(* a single term corresponding to a triple (k,ds,es) *)
Clear[singleKDE]
singleKDE[{k_,ds_,es_}]:=singleKDE[{k,ds,es}]=Module[
{A,B,
m,vars,dims,
pp,qq,rr,ss,
pps,qqs,rrs,sss
},
m=Length[ds];
pp[i_]:=Subscript[pp$p,i];
qq[i_]:=Subscript[qq$q,i];
rr[i_]:=Subscript[rr$r,i];
ss[i_]:=Subscript[ss$s,i];
pps=Table[pp[i],{i,1,m}];
qqs=Table[qq[i],{i,1,m}];
rrs=Table[rr[i],{i,1,m}];
sss=Table[ss[i],{i,1,m}];
vars=Join[{zzz},pps,qqs];
dims=Join[{k},ds,es];
A=csmDisj[dims]/.Table[uu[i]->vars[[i]],{i,1,2m+1}];
B=A;
For[i=1,i≤m,i++,B=Expand[delta2[es[[i]],qq[i],{rr[i],ss[i]},B]]];
For[i=1,i≤m,i++,B=Expand[psi2[{ds[[i]],es[[i]]},{pp[i],ss[i]},uu[i],B]]];
B=psiMany[Join[{k},es],Join[{zzz},rrs],z,B];
B
]

(* equiv CSM of D(d1,d2,...), but we require d1≥d2≥d3≥...≥dn>0 *)
csmDisjSorted[{n_}]:=csmDisjSorted[{n}]=csmDisj1[n]/.{u->uu[1]}
csmDisjSorted[pns_]:=csmDisjSorted[pns]=Module[
{p=pns[[1]],
ns=Drop[pns,1],
A,B,rest,
KDE
},
KDE=kdeTriples[p,ns];
(* Print["sorted1 - ",p," | ",ns];
Print["sorted2 - ",KDE]; *)
A=(csmDisj1[p]/.{u->z})*csmDisj[ns];
rest=Sum[singleKDE[kde],{kde,KDE}];
B = Expand[A-rest];
B =B/.Table[uu[i]->uu[i+1],{i,1,Length[ns]}];
B=B/.{z->uu[1]}
]

hs2mat={g->u,a->α,b->β}

csmXLam[{2,2,1,1}]

ref2211=(180*b^2-360*b^3-180*b^4+360*b^5+504*a*b-1584*a*b^2-144*a*b^3+864*a*b^4+180*a^2-1584*a^2*b+648*a^2*b^2-1224*a^2*b^3-360*a^3-144*a^3*b-1224*a^3*b^2-180*a^4+864*a^4*b+360*a^5+144*g^1*b-444*g^1*b^2-144*g^1*b^3+444*g^1*b^4+144*g^1*a-1056*g^1*a*b+144*g^1*a*b^2+168*g^1*a*b^3-444*g^1*a^2+144*g^1*a^2*b-1224*g^1*a^2*b^2-144*g^1*a^3+168*g^1*a^3*b+444*g^1*a^4+24*g^2-162*g^2*b-24*g^2*b^2+162*g^2*b^3-162*g^2*a+48*g^2*a*b-162*g^2*a*b^2-24*g^2*a^2-162*g^2*a^2*b+162*g^2*a^3-18*g^3+18*g^3*b^2-36*g^3*a*b+18*g^3*a^2)/.hs2mat

ref2211-csmXLam[{2,2,1,1}]