
(* ==== 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)*\[Alpha] + i*\[Beta]

(* 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], {\[Alpha], \[Beta]}, {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], {\[Alpha], \[Beta]}, {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], {\[Alpha], \[Beta]}, {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)*\[Alpha] + i*\[Beta], {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, 2 m + 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 -> \[Alpha], b -> \[Beta]}

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

24 u^2 - 18 u^3 + 144 u \[Alpha] - 162 u^2 \[Alpha] + 180 \[Alpha]^2 - 
 444 u \[Alpha]^2 - 24 u^2 \[Alpha]^2 + 18 u^3 \[Alpha]^2 - 360 \[Alpha]^3 - 
 144 u \[Alpha]^3 + 162 u^2 \[Alpha]^3 - 180 \[Alpha]^4 + 444 u \[Alpha]^4 + 
 360 \[Alpha]^5 + 144 u \[Beta] - 162 u^2 \[Beta] + 504 \[Alpha] \[Beta] - 
 1056 u \[Alpha] \[Beta] + 48 u^2 \[Alpha] \[Beta] - 
 36 u^3 \[Alpha] \[Beta] - 1584 \[Alpha]^2 \[Beta] + 
 144 u \[Alpha]^2 \[Beta] - 162 u^2 \[Alpha]^2 \[Beta] - 
 144 \[Alpha]^3 \[Beta] + 168 u \[Alpha]^3 \[Beta] + 864 \[Alpha]^4 \[Beta] + 
 180 \[Beta]^2 - 444 u \[Beta]^2 - 24 u^2 \[Beta]^2 + 18 u^3 \[Beta]^2 - 
 1584 \[Alpha] \[Beta]^2 + 144 u \[Alpha] \[Beta]^2 - 
 162 u^2 \[Alpha] \[Beta]^2 + 648 \[Alpha]^2 \[Beta]^2 - 
 1224 u \[Alpha]^2 \[Beta]^2 - 1224 \[Alpha]^3 \[Beta]^2 - 360 \[Beta]^3 - 
 144 u \[Beta]^3 + 162 u^2 \[Beta]^3 - 144 \[Alpha] \[Beta]^3 + 
 168 u \[Alpha] \[Beta]^3 - 1224 \[Alpha]^2 \[Beta]^3 - 180 \[Beta]^4 + 
 444 u \[Beta]^4 + 864 \[Alpha] \[Beta]^4 + 360 \[Beta]^5

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

24 u^2 - 18 u^3 + 144 u \[Alpha] - 162 u^2 \[Alpha] + 180 \[Alpha]^2 - 
 444 u \[Alpha]^2 - 24 u^2 \[Alpha]^2 + 18 u^3 \[Alpha]^2 - 360 \[Alpha]^3 - 
 144 u \[Alpha]^3 + 162 u^2 \[Alpha]^3 - 180 \[Alpha]^4 + 444 u \[Alpha]^4 + 
 360 \[Alpha]^5 + 144 u \[Beta] - 162 u^2 \[Beta] + 504 \[Alpha] \[Beta] - 
 1056 u \[Alpha] \[Beta] + 48 u^2 \[Alpha] \[Beta] - 
 36 u^3 \[Alpha] \[Beta] - 1584 \[Alpha]^2 \[Beta] + 
 144 u \[Alpha]^2 \[Beta] - 162 u^2 \[Alpha]^2 \[Beta] - 
 144 \[Alpha]^3 \[Beta] + 168 u \[Alpha]^3 \[Beta] + 864 \[Alpha]^4 \[Beta] + 
 180 \[Beta]^2 - 444 u \[Beta]^2 - 24 u^2 \[Beta]^2 + 18 u^3 \[Beta]^2 - 
 1584 \[Alpha] \[Beta]^2 + 144 u \[Alpha] \[Beta]^2 - 
 162 u^2 \[Alpha] \[Beta]^2 + 648 \[Alpha]^2 \[Beta]^2 - 
 1224 u \[Alpha]^2 \[Beta]^2 - 1224 \[Alpha]^3 \[Beta]^2 - 360 \[Beta]^3 - 
 144 u \[Beta]^3 + 162 u^2 \[Beta]^3 - 144 \[Alpha] \[Beta]^3 + 
 168 u \[Alpha] \[Beta]^3 - 1224 \[Alpha]^2 \[Beta]^3 - 180 \[Beta]^4 + 
 444 u \[Beta]^4 + 864 \[Alpha] \[Beta]^4 + 360 \[Beta]^5

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