(* Author: Mark Shimozono   Date: August 8, 2005 *)

(* Program for computing the generalized Kostka polynomial of Shimozono-Weyman.

   la_ is a partition, given as a list of integers,
   and R_ is a list of partitions.

Example call:

GeneralizedKostkaPolynomial[{5,4,3,3,3,2,1},{{3, 3, 3},{2, 2},{2},{2},{2},{1, 1}}]

answer:
2 q^5 + 9 q^6 + 24 q^7 + 47 q^8 + 69 q^9 + 80 q^10 +
    75 q^11 + 55 q^12 + 31 q^13 + 13 q^14 + 3 q^15
*)

GeneralizedKostkaPolynomial[la_,R_] := Module[{n=Plus@@Map[Length,R],lat},
(* Check the input data to see if they are comprised of partitions. *)
  If[!CheckPartitionQ[la],Return[0]];
  If[!(And@@Map[CheckPartitionQ,R]),Return[0]];
  lat=TrimZeroes[la];
(* If the shape la_ is too tall, return 0. *)
  If[Length[lat]>n,Return[0]];
(* If la and R don't have the same total number of cells, return 0. *)
  If[Plus@@lat!=Plus@@(Join@@R),Return[0]];
(* Call the recursive function that computes the genKostka. *)
  BP[ZeroExtend[lat,n],R]
]

(* This recursive function computes the genKostka using a generalized Morris recurrence. *)
BP[la_,R_] := BP[la,R] =
  Expand[
    If[Length[R]==0,If[ZeroWtQ[la],1,0],
      If[Length[R]==1,If[la==First[R],1,0],
        Plus @@ Map[First[#]*q^(Plus@@(#[[2]]-First[R]))*BLRPoly[#[[2]],#[[3]],R]&,
                  BShiftedOrbit[la,First[R]]]
      ]
    ]
  ]

BLRPoly[al_,beta_,R_] := Module[{SToK},
  SToK[la_] := If[Length[la]>Length[beta],0,
    BP[ZeroExtend[la,Length[beta]],Rest[R]]
  ];
  If[Last[beta]<0,
    BLRPoly[al,beta-Last[beta],Prepend[Rest[R]-Last[beta],First[R]]],
    TransitionSkewExpansion[
      Join[beta+First[al],al],
      Join[RepeatLetter[First[al],Length[beta]],First[R]]]/.S->SToK
  ]
]

BShiftedOrbit[a_,firstshape_] :=
  Module[{m=Length[firstshape]},
    Select[
      Map[Prepend[SingleChop[ShiftedInversePermAction[#,a],m],Signature[#]]&,
        YoungCosetReps[Length[a],m]],
      (IsPartitionQ[#[[2]]]&&
       ContainsPartitionQ[#[[2]],firstshape])&
    ]
  ]

YoungCosetReps[n_,r_] := YoungCosetReps[n,r] =
  If[r==0||r==n,{Range[n]},
    Join[Map[Append[#,n]&,YoungCosetReps[n-1,r]],
      Map[Insert[#,n,r]&,
        YoungCosetReps[n-1,r-1]]]
  ]

(* Compute the left action of the inverse of w_ on a_. *)
InversePermAction[w_,a_] := Array[a[[w[[#]]]]&,Length[w]]

ShiftedInversePermAction[w_,u_] :=
  (InversePermAction[w,u+#]-#)&@Range[Length[w]-1,0,-1]

(************************************************************************)
(*  Littlewood-Richardson stuff *)
(************************************************************************)

(* Computes the Schur function expansion of the skew Schur function
of shape lambda/mu, in the form of a list of items {nu,m}, meaning
that the nu-th Schur function occurs m times. *)
TransitionSkewExpansion[la_,mu_] :=
TransitionTree[SkewShapeToPerm[la,mu]]

(* Given a list of partitions, return a pair {la,mu} whose
associated skew shape gives the partitions in R corner to corner. *)
ShapeListToSkew[R_] :=
  If[Length[R]==0,{{},{}},
    If[Length[R]==1,{First[R],{}},
      {Join[First[R]+First[#[[1]]],#[[1]]],
      Join[RepeatLetter[First[#[[1]]],Length[First[R]]],#[[2]]]}&@ShapeListToSkew[Rest[R]]
    ]
  ]

IncreasingQ[a_] := (Length[a]==0 || LastDescent[a]==0)

(* If Length[w] > 0 then the default last descent is 0 *)
LastDescent[w_] := Module[{i},
  For[i=Length[w]-1,i>0&&w[[i+1]]>w[[i]],i--];
  i
]

(* Return the position of the first descent in the nonempty sequence
w. If there is no descent, return Length[w]. *) FirstDescent[w_] :=
Module[{i},
  For[i=1,i<Length[w]&&w[[i]]<w[[i+1]],i++];
  i
]

TrimLeadingZeroes[a_] := Module[{i},
  For[i=1,i<=Length[a]&&a[[i]]==0,i++];
  Drop[a,i-1]
]

(* If w is grassmannian with descent r, return its shape *)
GrassShape[w_,r_] := Reverse[TrimLeadingZeroes[Take[w,r]-Range[r]]]

(* If w is a dominant permutation, return its shape *)
DominantPermShape[w_] :=
  If[Length[w]==0||First[w]==1,{},
    Join[RepeatLetter[First[w]-1,#],DominantPermShape[Drop[w,#]]]&@
      FirstDescent[w]
  ]

RepeatLetter[x_,n_] := Array[x&,n]

LastInversion[w_,r_] := Module[{i,a=Array[1&,w[[r]]-1]},
  For[i=r-1,i>0,i--,
    If[w[[i]]<w[[r]],a[[w[[i]]]]=0]];
  For[i=Length[a],i>0&&a[[i]]==0,i--];
  i
]

GetTransitions[v_,r_,ws_] := Module[{i,m=0,trx={}},
  For[i=r-1,i>0,i--,
    If[v[[i]]>m&&v[[i]]<ws,
      trx=Prepend[trx,i];
      m=v[[i]];
    ];
  ];
  trx
]

RightTrans[w_,r_,s_] :=
ReplacePart[w,{w[[s]],w[[r]]},{{r},{s}},{{1},{2}}]

(* Return a linear combination of S[list] representing the Schur
function expansion of the ww-th Stanley symmetric function. *)
TransitionTree[ww_] := TransitionTree[ww] = Module[{w,r,ws,v},
  If[Length[ww]==0,Return[S[{}]]];
  w=If[First[ww]==1,ww,Prepend[ww+1,1]];
  r = LastDescent[w];
  If[r==0,Return[S[{}]]];
  If[IncreasingQ[Take[w,r]],Return[S[GrassShape[w,r]]]];
  ws = LastInversion[w,r];
  v = RightTrans[w,r,r + Position[Drop[w,r],ws][[1,1]]];
  Plus@@Map[TransitionTree[RightTrans[v,#,r]]&,GetTransitions[v,r,ws]]
]

XF[mu_] := If[Length[mu]==0,0,First[mu]]

SafeRest[mu_] := If[Length[mu]==0,{},Rest[mu]]

SkewShapeToPerm[la_,mu_] := Module[{f,ov},
  If[Length[la]==0,Return[{}]];
  f=First[la]-XF[mu];
  If[Length[la]==1,Return[PrependToPerm[{},f+1]]];
  ov=Max[0,la[[2]]-XF[mu]];
  PrependToPerm[
    PrependFixedPoints[SkewShapeToPerm[Rest[la],SafeRest[mu]],f-ov],
    f+1]
]

PrependFixedPoints[w_,m_] := Join[Range[m],w+m]

PrependToPerm[p_,x_] :=
  If[x>Length[p],
    Join[{x},p,Range[Length[p]+1,x-1]],
    Prepend[Map[If[#>=x,#+1,#]&,p],x]]




(************************************************************************)
(* General purpose stuff *)
(************************************************************************)

(* Is a_ a partition? *)
IsPartitionQ[a_] := Module[{i},
  If[Length[a]==0,Return[True]];
  If[First[a]<0,Return[False]];
  For[i=Length[a],i>1&&a[[i]]<=a[[i-1]],i--];
  i==1
]

(* check if la_ is a partition; if not return false and print an error message. *)
CheckPartitionQ[la_] :=
  If[IsPartitionQ[la],True,Print["Error: ", la, " is not a partition."];False]

(* Is a_ before or equal to b_ in reverse lex order on partitions of the same
size? *)
RLexQ[a_,b_] := Module[{d},
  If[Length[a]==0,Return[True]];
  If[Length[b]==0,Return[False]];
  d=First[a]-First[b];
  If[d==0,RLexQ[Rest[a],Rest[b]],
    d>0
  ]
]

(* Does the partition a_ contain the partition b_?
It is used only when a_ and b_ have the same length. *)
ContainsPartitionQ[a_,b_] := Module[{i},
  For[i=Length[a],i>0&&a[[i]]>=b[[i]],i--];
  i==0
]

(* If the sequence a_ has length less than n, pad it with zeroes until it has
length n. *)
ZeroExtend[a_,n_] := If[Length[a]<n,Join[a,RepeatLetter[0,n-Length[a]]],a]

(* If a sequence has some trailing zeroes, get rid of them. *)
TrimZeroes[a_] := Module[{i},
  For[i=Length[a],i>0&&a[[i]]==0,i--];
  Take[a,i]
]

(* Make a list with n copies of x. *)
RepeatLetter[x_,n_] := Array[x&,n]

(* Is a_ a sequence of zeroes? *)
ZeroWtQ[a_] :=
  Module[{i},
    For[i=Length[a],i>0&&a[[i]]==0,i--];
    i==0
  ]

(* Chop the sequence a_ into two subsequences just after position q_. *)
SingleChop[a_,q_] :=
    {Take[a,q],Drop[a,q]}