(* ===== EQUIVARIANT MOTIVIC CHERN CLASSES OF COINCIDENT ROOT LOCI ==== *)

<<Combinatorica`

(* === NORMALIZATION in K(P^n) === *)

(* Weights of Sym^n C^2 *)
Clear[L,X,Y]
WT[n_,i_]:=X^(n-i)*Y^i

(* the relation in K(P^n). Convention: 
L is the tautological line bundle, 
   c_1(X) = -alpha, c_1(Y) = -beta *)
KREL[n_]:=Product[1-L*WT[n,i],{i,0,n}]

(* normal form of L^(n+1) *)
Clear[Lpow$nplus1]
Lpow$nplus1[n_] := Lpow$nplus1[n]= Expand[L^(n+1)-(-1)^(n+1)KREL[n]/(X*Y)^Binomial[n+1,2]]

(* (unefficiently) normalize a polynomial in L *) 
Clear[KnormalizeSlow]
KnormalizeSlow[n_, Z0_]:=Module[{Z=Expand[Z0],m},
m=Exponent[Z,L];
If[m≤n,Z,KnormalizeSlow[n,Z/.{L^m->L^(m-n-1)*Lpow$nplus1[n]}]]
]

(* Table of normal forms of L^p *)
Clear[LPowXY]
LPowXY[n_,k_]:=LPowXY[n,k] = Expand[KnormalizeSlow[n,L^k]]

(* MoreE eficient normalization in K^(P^n) 
usage: KnormalizeVarXY[dim,L,expr] *)
Clear[KnormalizeVarXY]
KnormalizeVarXY[n_,uuu_,Z0_]:=Module[
{m,Z},
Z=Expand[Z0];
m=Exponent[Z,uuu];
Expand[Sum[Coefficient[Z,uuu,k]*(LPowXY[n,k]/.{L->uuu}),{k,0,m}]]
]

(* === PUSHFORWARD ALONG PSI (THE MULTIPLICATION MAP) === *)

(* formulas for (psi_{n,1})_! (L1^p*L2^q) *)
Clear[PSI$N1f]
PSI$N1f[n_,0,1]:=L*(X^(n+1)-Y^(n+1))/(X-Y)
PSI$N1f[n_,p_,0]:=L^p*(X^(p+1)-Y^(p+1))/(X-Y)+L^(p+1) (X^(n+1)*Y^(p+1)-Y^(n+1) X^(p+1))/(X-Y)
PSI$N1f[n_,p_,q_]:=If[p≥q,
L^q*PSI$N1f[n,p-q,0],
L^p*PSI$N1f[n,0,q-p]
]

Clear[PSI$N1]
PSI$N1[n_,p_,q_]:=PSI$N1[n,p,q]=Expand[Factor[PSI$N1f[n,p,q]]]

Clear[newPsiBang2]
newPsiBang2[{n_,1},{var1_,var2_},outvar_,ZZ_]:=Module[
{Z},
Z=Expand[ZZ];
Expand[Sum[Coefficient[Coefficient[Z,var1,i],var2,j]*(PSI$N1[n,i,j]/.{L->outvar}),{i,0,n},{j,0,1}]]
]

(* find elements in K(P^m x P^1) such that their pushforward to K(P^(m+1)) is {1,L,L^2,...L^(m+1)} *)
findBasis$M1[m_]:=findBasis$M1[m]=
Module[{vars,A,B,list,i,j,k,eqs,sols,sol},
vars=Table[Subscript[a,i],{i,0,m+1}];
A=Sum[Subscript[a,i]*L1^i,{i,0,m}]+
Subscript[a,m+1]*L1^m*L2;
B=newPsiBang2[{m,1},{L1,L2},L,A];
list=Table[Null,{i,0,m+1}];
For[k=0,k≤m+1,k++,
eqs=Table[Coefficient[B,L,i]==If[k==i,1,0],{i,0,m+1}];
sols =Solve[eqs,vars];
sol = sols[[1]];
list[[k+1]] =A/.sol;
];
list
]

Clear[PSI$NM,S,T]
PSI$NM[n_,0,p_,0]:=L^p
PSI$NM[0,m_,0,q_]:=L^q
PSI$NM[n_,1,p_,q_]:=PSI$NM[n,1,p,q]=PSI$N1[n,p,q]
PSI$NM[n_,mplus1_,p_,q_]:=PSI$NM[n,mplus1,p,q]=Module[
{m=mplus1-1,basis,A,B,C},
basis=findBasis$M1[m];
A=basis[[q+1]]/.{L1->L2,L2->T};
B=newPsiBang2[{n,m},{L1,L2},S,L1^p*A];
C=newPsiBang2[{n+m,1},{S,T},L,B];
Expand[Factor[C]]
]

newPsiBang2[{n_,m_},{var1_,var2_},outvar_,ZZ_]:=Module[
{Z,deg1,deg2},
Z=Expand[ZZ];
deg1=Exponent[Z,var1];
deg2=Exponent[Z,var2];
Expand[Sum[Coefficient[Coefficient[Z,var1,i],var2,j]*(PSI$NM[n,m,i,j]/.{L->outvar}),{i,0,deg1},{j,0,deg2}]]
]

newPsiBangMany[{n_},{uuu_},www_,Z_]:=Z/.{uuu->www}
newPsiBangMany[ns_,uuus_,www_,Z_]:=Module[{ttt,vvvs,ms},
If[Length[uuus]==1,Z/.{uuus[[1]]->www},
If[Length[uuus]==2,newPsiBang2[ns,uuus,www,Z],
ms=Join[{ns[[1]]+ns[[2]]},Drop[ns,2]];
vvvs=Join[{ttt},Drop[uuus,2]];
newPsiBangMany[ms,vvvs,www, newPsiBang2[Take[ns,2],Take[uuus,2],ttt,Z]]]]]

(* === PUSHFORwARD ALONG THE DIAGONAL MAP === *)

notk[n_,k_]:=Select[Range[0,n],#≠k&]

(* class of the diagonal in K(P^n x P^n) *)
Clear[DELTA$CLASS$2]
DELTA$CLASS$2[n_]:=DELTA$CLASS$2[n]=
Factor[Sum[ Product[ (1-L1*WT[n,i])(1-L2*WT[n,i])/(1-WT[n,i]/WT[n,k])  ,{i,notk[n,k]}] ,{k,0,n}]]

newDeltaBang2[n_,invar_,{var1_,var2_},ZZ_]:=Module[
{Z=Expand[ZZ],m,A},
m=Exponent[Z,invar];
A=Expand[(DELTA$CLASS$2[n]/.{L1->var1,L2->var2})*Sum[var1^p*Coefficient[Z,invar,p],{p,0,m}]];
KnormalizeVarXY[n,var1,A]
]

(* usage: deltaBangMany[dim,L,{L1,L2,L3},X] *)
newDeltaBangMany[n_,uuu_,vvvs_,Z_]:=Module[{ttt},
If[Length[vvvs]==1,Z/.{uuu->vvvs[[1]]},
If[Length[vvvs]==2,newDeltaBang2[n,uuu,vvvs,Z],
newDeltaBangMany[n,ttt,Drop[vvvs,1], newDeltaBang2[n,uuu,{vvvs[[1]],ttt},Z]]]]]

(* PUSHFORWARD ALONG THE POWER MAP *)

Clear[OMEGA$ND]
OMEGA$ND[n_,d_,p_]:=OMEGA$ND[n,d,p]=Module[{vars,ns,A,B},
vars=Table[Subscript[TMP$X,i],{i,1,d}];
ns= Table[n,{i,1,d}];
A=newDeltaBangMany[n,L,vars,L^p];
B=newPsiBangMany[ns,vars,L,Expand[A]];
Expand[B]
]

newOmegaBang[n_,0,var1_,var2_,ZZ_]:=Coefficient[ZZ,var1,0]
newOmegaBang[n_,1,var1_,var2_,ZZ_]:=ZZ/.{var1->var2}
newOmegaBang[n_,d_,var1_,var2_,ZZ_]:=Module[
{Z=Expand[ZZ],m,A},
m=Exponent[Z,var1];
A=Sum[(OMEGA$ND[n,d,p]/.{L->var2})*Coefficient[Z,var1,p],{p,0,m}];
Expand[A]
]

(* 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 *) 
newOmegaBangLam[ns_,ds_,uuus_,www_,Z0_]:=Module[
{Z=Expand[Z0],
m = Length[ns],
vars,W,nds,ttt,i},
vars=Table[Subscript[ttt,i],{i,1,m}];
nds=Table[ ns[[i]]*ds[[i]], {i,1,m}];
W=Z;
For[i=1,i≤m,i++,W=newOmegaBang[ns[[i]],ds[[i]],uuus[[i]],vars[[i]],W]];
newPsiBangMany[nds,vars,www,W]
]

(*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]}]]

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]

(* motivic Chern class of P^n *)
Clear[mcPn];
mcPn[n_]:=mcPn[n]=
Module[{rel,tmp},rel=Product[1-L*X^(n-i)*Y^i,{i,0,n}];
tmp=Product[1+t*L*X^(n-i)*Y^i,{i,0,n}];
Factor[(tmp-(-t)^(n+1) rel)/(1+t)]]

Clear[L,S]

LL[i_]:=Subscript[L,i];
SS[i_]:=Subscript[S,i];
LLs[n_]:=Table[LL[i],{i,1,n}]
SSs[n_]:=Table[SS[i],{i,1,n}]

Clear[mcXLam,mcDisj1,mcDisj,mcDisjSorted]

mcXLam[{}]:=mcXLam[{}]=1
mcXLam[{1}]:=mcXLam[{1}]=mcPn[1]

mcDisj1[0]:=mcDisj1[0]=1
mcDisj1[1]:=mcDisj1[1]=mcPn[1]

mcDisj[{}]:=mcDisj[{}]=1

mcDisjSorted[{}]:=mcDisjSorted[{}]=1

(* equivariant motivic chern class of D(n)=X(1^n) *)
mcDisj1[n_]:=mcDisj1[n]=Module[
{parts=Select[Combinatorica`Partitions[n],Length[#]<n&]},
Expand[mcPn[n]-Sum[mcXLam[p],{p,parts}]]
]

(* equivariant motivic chern class of X_lambda *)
mcXLam[lambda_]:=mcXLam[lambda]=Module[
{es = toExpoForm[lambda],m,m1,ns,pairs,Z},
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];
Z=mcDisj[es];
(* Print["xlam1 - ",pairs];
Print["xlam2 - ",es," | ",ns," | ",uus[m1]," | ",X]; *)
Expand[newOmegaBangLam[es,ns,LLs[m1],L,Z]]
]

(* equivariant motivic chern class of D(d1,d2,...) *)
mcDisj[{n_}]:=mcDisj[{n}] =mcDisj1[n]/.{L->LL[1]}
mcDisj[ns0_]:=mcDisj[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=mcDisjSorted[ns1];
X=X/.Table[LL[i]->Subscript[ttt,i],{i,1,m}];
X/.Table[Subscript[ttt,i]->LL[idxs[[i]]],{i,1,m}]
]

(* a single term corresponding to a triple (k,ds,es) *)
Clear[singleKDE]
singleKDE[{k_,ds_,es_}]:=ssingleKDE[{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=mcDisj[dims]/.Table[LL[i]->vars[[i]],{i,1,2m+1}];
B=A;
For[i=1,i≤m,i++,B=Expand[newDeltaBang2[es[[i]],qq[i],{rr[i],ss[i]},B]]];
For[i=1,i≤m,i++,B=Expand[newPsiBang2[{ds[[i]],es[[i]]},{pp[i],ss[i]},LL[i],B]]];
B=newPsiBangMany[Join[{k},es],Join[{zzz},rrs],z,B];
B
]

(* equiv mc of D(d1,d2,...), but we require d1≥d2≥d3≥...≥dn>0 *)
mcDisjSorted[{n_}]:=mcDisjSorted[{n}]=mcDisj1[n]/.{L->LL[1]}
mcDisjSorted[pns_]:=mcDisjSorted[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=(mcDisj1[p]/.{L->z})*mcDisj[ns];
rest=Sum[singleKDE[kde],{kde,KDE}];
B = Expand[A-rest];
B =B/.Table[LL[i]->LL[i+1],{i,1,Length[ns]}];
B=B/.{z->LL[1]}
]

(* export the classes of X(lambda) for |lambda|≤n *)

ExportMC[n_]:=Module[
{h,i,p,parts,k,m,s,j,A},
h=OpenWrite["equivariant_mc_classes.txt"];
For[i=1,i≤n,i++,
Print["\nn = ",i];
WriteString[h,"\n(* =================== *)"];
WriteString[h,"\n(* ----   n = "<>ToString[i]<>"   ---- *)\n\n"];
parts=Partitions[i];
m=Length[parts];
For[j=1,j≤m,j++,
p =parts[[j]];
Print["part = ",p];
A =Expand[mcXLam[p]];
(* Print[A]; *)
s=ToString[A,FormatType->InputForm,PageWidth->Infinity,TotalWidth->Infinity];
s=StringJoin["mc[",ToString[p],"] = ",s," ;\n\n"];
(* Print[s]; *)
WriteString[h,s];
]
]
Close[h];
]

mcXLam[{1,1,1}]

mcXLam[{1,1,1,1,1}]

mcXLam[{2,2,1,1}]