(* ::Package:: *)

<<Combinatorica`
(* select the rationally equivalent "preferred" polynomial *)
NormalizePoly[poly_]:=If[IntegerQ[First[Factor[poly]]],Expand[Rest[Factor[poly]]],poly]
(* produce the symmetric basis polynomial associated to an integer partition *)
MakePoly[exps_]:=Plus@@(Times@@Table[Symbol["x"<>ToString[i]]^#[[i]],{i,1,Length[#]}]&/@Permutations[exps])

(* routines to apply the multiplicative n-coboundary map *)
(* clip a power series to a certain accuracy *)
ClipPoly[p_,deg_]:=With[{temp=Select[p//Expand/.Plus->List,Plus@@Exponent[#,Variables[#]]<=deg&]},If[Head[temp]===List,Plus@@temp,temp]]
(* invert a power series of the form 1 + higher order terms *)
InversePoly[p_,accuracy_]:=ClipPoly[With[{deg=Min[(Plus@@Exponent[#,Variables[#]])&/@Rest[p/.Plus->List]],x=-(p-1)},1+Sum[Select[x,(Plus@@Exponent[#,Variables[#]])+(i-1) deg<=accuracy&]^i,{i,Floor[accuracy/deg]}]],accuracy]
(* builds the substitution rules for the n-cocycle condition in H^*(Ga; Gm) *)
MCocycleSubstitutions[n_]:=Join[{{}},Table[Join[Table[Symbol["x"<>ToString[j]]->Symbol["x"<>ToString[j-1]],{j,1,i-1}],{Symbol["x"<>ToString[i]]->Symbol["x"<>ToString[i]]+Symbol["x"<>ToString[i-1]]}],{i,1,n}],{Table[Symbol["x"<>ToString[i]]->Symbol["x"<>ToString[i-1]],{i,1,n}]}]
(* produces the total degree of a monomial *)
ExponentTotals[poly_]:=Union[(Plus@@Exponent[#,Variables[#]])&/@(poly/.Plus->List)]
(* routines to efficiently multiply two truncated power series up to a particular accuracy *)
PolyExpPartitions[poly_,exps_]:=(With[{set=#},Select[poly,Plus@@Exponent[#,Variables[#]]==set&]]&/@exps)
PolyMultiplication[poly1_,poly2_,acc_]:=With[{exps1=ExponentTotals[poly1],exps2=ExponentTotals[poly2]},With[{sets1=PolyExpPartitions[poly1,exps1],sets2=PolyExpPartitions[poly2,exps2]},Expand[Plus@@((Times@@{sets1[[#[[1]]]],sets2[[#[[2]]]]})&/@({#[[1,2]],#[[2,2]]}&/@Select[Flatten[Table[{Thread[{exps1,Range[Length[exps1]]}][[i]],Thread[{exps2,Range[Length[exps2]]}][[j]]},{i,Length[exps1]},{j,Length[exps2]}],1],#[[1,1]]+#[[2,1]]<=acc&]))]]]
(* apply the mod-p n-coboundary map to a truncated power series in *)
MCocycleImage[poly_,n_,accuracy_,p_]:=PolynomialMod[Fold[PolynomialMod[PolyMultiplication[#1,#2,accuracy],p]&,#[[1]],#[[2;;4]]]&@With[{subs=Expand[poly/.#&/@MCocycleSubstitutions[n]]},Table[If[Divisible[i-1,2],PolynomialMod[subs[[i]],p],PolynomialMod[InversePoly[PolynomialMod[subs[[i]],p],accuracy],p]],{i,1,n+2}]],p]

(* routines for the computations in [HLP08] *)
(* perform one step of "splitting" operations to a power-of-p multiindex *)
Splits[index_,char_]:=Sort[#,Greater]&/@Select[Union[Table[If[index[[i]]!=1,Join[Drop[index,{i}],Table[index[[i]]/char,{j,char}]],{}],{i,1,Length[index]}]],#!={}&]
(* perform one step of splitting to a list of power-of-p multiindices *)
StepSplittingTree[indices_,char_]:=Union[indices,With[{splits=Splits[#,char]&/@indices},If[Length[indices]==1,splits[[1]],Join@@splits]]]
(* perform one step of gathering to a multiindex *)
GetGathers[index_,count_]:=With[{gathers=Sort[#,Greater]&/@Union[Join[Delete[index,{{#[[1]]},{#[[2]]}}],{Plus@@index[[#]]}]&/@Subsets[Range[Length[index]],{2}]]},Which[count==1,gathers,count==0,{index},True,Union[Fold[Join,{},GetGathers[#,count-1]&/@gathers]]]]
(* build the minimal power-of-p partition of a number *)
PPartition[n_,p_]:=If[n==0,{},With[{pull=Log[p,n]//Floor},Join[{p^pull},PPartition[n-p^pull,p]]]]
(* find the p-adic digital sum of a number *)
DigitalSum[n_,base_]:=Plus@@Most[DigitCount[n,base]*Range[1,base]]
(* build the complete set of grouped symmetric 2-cocycle exponent lists of a particular homogenous degree and characteristic, ranging over all numbers of variables *)
CocycleIndices[n_,p_]:=Join[With[{pops=FixedPoint[StepSplittingTree[#,p]&,{PPartition[n,p]}]},Union@@Table[GetGathers[#,i]&/@pops,{i,0,p-2}]],Table[GetGathers[PPartition[n,p],i],{i,1,DigitalSum[n,p]-1}]]
(* compute the ando-hopkins-strickland 2-cocycle, called zeta^deg_dim in [HLP08] *)
GuaranteedCocycle[deg_,dim_]:=Fold[#1+(-1)^(dim-Length[#2])*(Plus@@#2)^deg&,0,Subsets[Table[Symbol["x"<>ToString[i]],{i,1,dim}]]]//Expand
(* use the main theorem of [HLP08] to produce a basis symmetric 2-cocycle from an exponent list generated by CocycleIndices *)
CocycleIndexGroupToCocycle[indexgroup_,p_]:=With[{gc=NormalizePoly[GuaranteedCocycle[Plus@@indexgroup[[1]],Length[indexgroup[[1]]]]]},PolynomialMod[Plus@@(MakePoly[#]*Coefficient[gc,Times@@Table[Symbol["x"<>ToString[i]]^#[[i]],{i,1,Length[#]}]]&/@indexgroup)//Expand,p]]

(* automated exploratory routines for multiplicative cocycles, using the spectral sequence in [HLP09] and the additive calculation in [HLP08] *)
(* inverse functions to MakePoly *)
MakeExponentList[poly_]:=If[IntegerQ[First[#]],{{First[#]},Exponent[Rest[#],Variables[#]]},{{1},Exponent[#,Variables[#]]}]&/@#&/@If[Head[poly]===List,(Expand[poly]/.Plus->List),Expand[{poly}]/.Plus->List]
MakeExponentListSorted[poly_]:=Sort[#,Plus@@#1[[2]]<Plus@@#2[[2]]&]&/@Union/@Map[Sort[#,Greater]&,MakeExponentList[poly],3]
(* a haskell borrow, could probably be reimplemented with Union and less manual computation *)
Nub[list_]:=Fold[If[MemberQ[#1,#2],#1,Append[#1,#2]]&,{},list]
(* given an batch of MCocycleImage output, find the possible candidate extensions to delete the bottom 2-coboundary classes *)
CandidateExtensions[work_,p_]:=Nub[{Sort[#[[1]],Greater],#[[2]]}&/@Fold[Join[#1,With[{firstcarrysafe=(CarryCount[{#2[[1]],#2[[2]]},p]==0),secondcarrysafe=(CarryCount[{#2[[2]],#2[[3]]},p]==0)},Which[firstcarrysafe&&Not[secondcarrysafe],{{Join[{#2[[1]]+#2[[2]]},Drop[#2,2]],r}},secondcarrysafe&&Not[firstcarrysafe],{{Join[{#2[[1]],#2[[2]]+#2[[3]]},Drop[#2,3]],r}},True,{}]]]&,{},With[{mintotal=Min[Map[Plus@@#[[2]]&,work]]},#[[2]]&/@Select[work,Plus@@#[[2]]==mintotal&]]]]
(* compute the p-adic carry count of a multiindex, as in [HLP08] *)
CarryCount[index_,p_,c_:0]:=If[Divisible[Multinomial@@index,p^(c+1)],CarryCount[index,p,c+1],c]
(* an alternate form of CocycleIndexGroupToCocycle, attaches zeta_dim^deg coefficients to an additive 2-cocycle multiindex list *)
GetCoefficients[indexgroup_,p_]:=With[{gc=NormalizePoly[GuaranteedCocycle[Plus@@indexgroup[[1]],Length[indexgroup[[1]]]]]},{{Mod[Coefficient[gc,Times@@Table[Symbol["x"<>ToString[i]]^#[[i]],{i,1,Length[#]}]],p]},#}&/@indexgroup]
(* applies MCocycleImage to a truncated multiplicative extension and returns the bottom "noise" to be cancelled *)
GetBottomMCocycleImage[group_,n_,p_]:=With[{mel=MakeExponentList[MCocycleImage[1+Plus@@(#[[1,1]]*MakePoly[#[[2]]]&/@group),n,2*(Max@@(Plus@@#[[2]]&/@group)),p]-1][[1]]},With[{skinny=Min@@(Plus@@#[[2]]&/@mel)},Select[mel,Plus@@#[[2]]==skinny&]]]
(* FSM to pick the next cancellation action to be taken by analyzing GetBottomMCocycleImage results *)
GetNextStep[group_,n_,p_]:=Join[group,Union[{#[[1]],Sort[#[[2]],Greater]}&/@Select[With[{firstcarrysafe=CarryCount[Take[#[[2]],{1,2}],p]==0,secondcarrysafe=CarryCount[Take[#[[2]],{2,3}],p]==0},Which[firstcarrysafe&&!secondcarrysafe,{{PolynomialMod[#[[1,1]]*Binomial[#[[2,1]]+#[[2,2]],#1[[2,1]]]^-1,p]},Join[{#[[2,1]]+#[[2,2]]},Drop[#[[2]],{1,2}]]},!firstcarrysafe&&secondcarrysafe,{{PolynomialMod[#[[1,1]]*-1*Binomial[#[[2,2]]+#[[2,3]],#1[[2,2]]]^-1,p]},Join[{#[[2,2]]+#[[2,3]]},Drop[#[[2]],{2,3}]]},True,{}]]&/@GetBottomMCocycleImage[group,n,p],#!={}&]]]
(* if we don't have any required extensions, we pick one.  TODO: this may need backtracking! *)
GetNextSuggested[group_,n_,p_]:=Join[group,{With[{gbmci=GetBottomMCocycleImage[group,n,p]},With[{first=Sort[Join[{gbmci[[1,2,1]]+gbmci[[1,2,2]]},Drop[gbmci[[1,2]],2]],Greater],second=Sort[Join[{gbmci[[1,2,2]]+gbmci[[1,2,3]]},Drop[gbmci[[1,2]],{2,3}]],Greater]},If[Or@@(#[[2]]==first&/@group),{{PolynomialMod[-1*gbmci[[1,1,1]]Binomial[gbmci[[1,2,2]]+gbmci[[1,2,3]],gbmci[[1,2,2]]]^-1,p]},second},{{PolynomialMod[gbmci[[1,1,1]] Binomial[gbmci[[1,2,1]]+gbmci[[1,2,2]],gbmci[[1,2,2]]]^-1,p]},first}]]]}]

(* takes an additive 2-cocycle's exponent groups and attempts to build a multiplicative extension, stopping when the ghost degree exceeds ceil *)
FindCocycle[group_,n_,p_,ceil_]:=With[{gns=GetNextStep[group,n,p]},Which[group==gns,FindCocycle[GetNextSuggested[gns,n,p],n,p,ceil],(Max@@(#[[2]]&/@(Tally[#[[2]]&/@gns])))>1,{group,"Rerequired!"},Max@@(Plus@@#[[2]]&/@gns)>ceil,{group,"Hit ceiling!"},True,FindCocycle[gns,n,p,ceil]]]

(* controlling routines for the GBMCI mess above, proof of correctness is in the spectral sequence of [HLP09] *)
(* builds a table of scalar coefficients for for use in BuildLinearCombinations *)
Coefficients[len_,p_]:=Select[Tuples[Range[0,p-1],len],With[{incoming=#},Or@@(Take[incoming,#]==Join[Table[0,{#-1}],{1}]&/@Range[1,Length[incoming]])]&]
(* builds all possible linear combinations of additive 2-cocycles without scalar repeats (so of the form (0, ..., 0, 1, n1, ..., nk)).  useful because the multiplicative map is not an additive map *)
BuildLinearCombinations[deg_,dim_,p_]:=Map[Join@@#&,With[{cocycleindices=GetCoefficients[#,p]&/@Select[Union[CocycleIndices[deg,p]],#1!={}&&Length[#1[[1]]]==dim&]},With[{coeff=#[[1]],indexlist=#[[2]]},{{PolynomialMod[#[[1,1]]*coeff,p]},#[[2]]}&/@indexlist]&/@Transpose[{#,cocycleindices}]&/@Coefficients[Length[cocycleindices],p]],{1}]
(* given a degree and dimension, attempt to find all (obstructions to) multiplicative extensions of all symmetric additive 2-cocycles of homogenous degree deg, number of variables dim, characteristic p, and with maximal attempted degree ceil *)
FindCocycleSlot[deg_,dim_,ceil_,p_]:=Fold[(Print[FindCocycle[#2,2,p,ceil]];{})&,{},BuildLinearCombinations[deg,dim,p]]
(* runs FindCocycleSlot over a range of dimensions *)
FindCocycleRange[deg_,lowdim_,highdim_,ceil_,p_]:=Fold[(FindCocycleSlot[deg,#2,ceil,p];{})&,{},Range[lowdim,highdim]]
Print["loaded"]
