BeginPackage["jTool`",{"loopTool`"}];

generateVarList::usage="";
generateVarList1::usage="";
generateVarList2::usage="";
generateVarList3::usage="";
generateVarList4::usage="";
setVariableProperty::usage="";
setIpProperty::usage="";

convertToJ::usage="";
convertFromJ::usage="";
getVariable::usage"";

priorityOrder::usage="";
priorityOrder2::usage="";
equationPriority::usage=""
jQuality::usage="";
sortVariable::usage="";
sortEquation::usage="";

pickJle::usage="";
pickJge::usage="";
pickJl::usage="";
pickJg::usage="";
pickJne::usage="";
pickJe::usage="";

pickJles::usage="";
pickJges::usage="";
pickJes::usage="";

pickJle2::usage="";
pickJge2::usage="";
pickJl2::usage="";
pickJg2::usage="";
pickJne2::usage="";
pickJe2::usage="";

pickJmost::usage="";
pickJrest::usage="";
pickJfirst::usage="";
pickJlast::usage="";

pickJmost2::usage="";
pickJrest2::usage="";
pickJfirst2::usage="";
pickJlast2::usage="";

jleQ::usage="";
jgeQ::usage="";
jlQ::usage="";
jgQ::usage="";
jneQ::usage="";
jeQ::usage="";

jlesQ::usage="";
jgesQ::usage="";
jesQ::usage="";

jle2Q::usage="";
jge2Q::usage="";

allVarBase::usage="";
allVarBase1::usage="";
allVarBase2::usage="";
allVarBase3::usage="";
allVarBase4::usage="";

Begin["`private`"];

Unprotect["jTool`*"];
Unprotect["jTool`private`*"];

invRule[var_,ip_]:=
Module[{tmp},
tmp=Cases[var,ip[__],{0,∞}]//Union;
tmp=Solve[Thread[Array["V",Length[var]]==var],tmp];
tmp[[1]]
];

generateVarList[loopOrder_Integer?(#≥1&),heavy0_,massless0_,ip0_,p0_/;p0[[0]]=!=List,v0_]:=
generateVarList1[loopOrder,heavy0,massless0,ip0,p0,v0][[1;;2]];

generateVarList[loopOrder_Integer?(#≥1&),heavy0_,massless0_,ip0_,p0_List,v0_]/;Length[p0]==loopOrder+1:=
Module[{t},generateVarList[loopOrder,heavy0,massless0,ip0,t,v0]/.t->(p0[[#+1]]&)];

generateVarList1[loopOrder_Integer?(#≥1&),heavy0_,massless0_,ip0_,p0_/;p0[[0]]=!=List,v0_,qsq:(True|False):False]:=
Module[{heavy,massless,ip,p,v,momentum,pHeavy,pMassless,allVar,allVar2},
momentum=p/@Range[loopOrder];
pHeavy=DeleteCases[Total/@Subsets[momentum],0];
pMassless=Join[pHeavy,(p[0]+#)&/@pHeavy,If[qsq,{p[0]},{}]];
allVar=Join[heavy[#,v]&/@pHeavy,massless[#]&/@pMassless];
allVar2=allVar/.heavy->ip/.massless[a_]:>ip[a,a];
{allVar2,allVar,invRule[allVar2,ip],invRule[allVar,heavy|massless]}/.
{heavy->heavy0,massless->massless0,ip->ip0,p->p0,v->v0}
];

generateVarList1[loopOrder_Integer?(#≥1&),heavy0_,massless0_,ip0_,p0_List,v0_,qsq:(True|False):False]/;Length[p0]==loopOrder+1:=
Module[{t},generateVarList1[loopOrder,heavy0,massless0,ip0,t,v0,qsq]/.t->(p0[[#+1]]&)];

generateVarList2[loopOrder_?(#≥2&),heavy0_,massless0_,ip0_,p0_/;p0[[0]]=!=List,v0_,qsq:(True|False):False]:=
Module[{heavy,massless,ip,p,v,momentum,pHeavy,pMassless,allVar,allVar2,tmp},
momentum=p/@Range[loopOrder];
pHeavy=DeleteCases[Total/@Subsets[momentum],0];
pMassless=Join[pHeavy,(p[0]+#)&/@pHeavy,If[qsq,{p[0]},{}]];
allVar=Join[heavy[#,v]&/@pHeavy,massless[#]&/@pMassless];
allVar=allVar/.heavy[a_,v]:>heavy[a/.p[1|2]->0,v];
allVar=DeleteCases[allVar/.heavy[0,_]->0,0]//DeleteDuplicates;
allVar2=allVar/.heavy->ip/.massless[a_]:>ip[a,a];
{allVar2,allVar,invRule[allVar2,ip],invRule[allVar,heavy|massless]}/.
{heavy->heavy0,massless->massless0,ip->ip0,p->p0,v->v0}
];

generateVarList2[loopOrder_Integer?(#≥2&),heavy0_,massless0_,ip0_,p0_List,v0_,qsq:(True|False):False]/;Length[p0]==loopOrder+1:=
Module[{t},generateVarList2[loopOrder,heavy0,massless0,ip0,t,v0]/.t->(p0[[#+1]]&)];

generateVarList3[loopOrder_?(#≥1&),heavy0_,massless0_,ip0_,p0_/;p0[[0]]=!=List,v0_,qsq:True:True]:=
Module[{heavy,massless,ip,p,v,momentum,pHeavy,pMassless,allVar,allVar2},
momentum=p/@Range[loopOrder];
pHeavy=DeleteCases[Total/@Subsets[momentum],0];
pMassless=Join[pHeavy,(p[0]+#)&/@pHeavy,{p[0]}];
allVar=Join[heavy[#,v]&/@pHeavy,massless[#]&/@pMassless];
allVar=allVar/.massless[a_]:>massless[a]-((heavy[a/.p[1]->0,v])^2)+((heavy[a/.p[0]->0,v])^2);
allVar=allVar/.heavy[a_,v]:>heavy[a/.p[1]->p[0],v]/.{p[0]->p[1],p[1]->p[0]};
allVar=DeleteCases[allVar/.heavy[0,_]->0,0]//DeleteDuplicates;
allVar2=allVar/.heavy->ip/.massless[a_]:>ip[a,a];
{allVar2,allVar,invRule[allVar2,ip],invRule[allVar,heavy|massless]}/.
{heavy->heavy0,massless->massless0,ip->ip0,p->p0,v->v0}
];

generateVarList3[loopOrder_Integer?(#≥1&),heavy0_,massless0_,ip0_,p0_List,v0_,qsq:True:True]/;Length[p0]==loopOrder+1:=
Module[{t},generateVarList3[loopOrder,heavy0,massless0,ip0,t,v0]/.t->(p0[[#+1]]&)];

generateVarList4[loopOrder_?(#≥3&),heavy0_,massless0_,ip0_,p0_/;p0[[0]]=!=List,v0_,qsq:True:True]:=
Module[{heavy,massless,ip,p,v,momentum,pHeavy,pMassless,allVar,allVar2,tmp},
momentum=p/@Range[loopOrder];
pHeavy=DeleteCases[Total/@Subsets[momentum],0];
pMassless=Join[pHeavy,(p[0]+#)&/@pHeavy,{p[0]}];
allVar=Join[heavy[#,v]&/@pHeavy,massless[#]&/@pMassless];
allVar=allVar/.massless[a_]:>massless[a]-((heavy[a/.p[3]->0,v])^2)+((heavy[a/.p[0]->0,v])^2);
allVar=allVar/.heavy[a_,v]:>heavy[a/.p[3]->p[0],v]/.{p[0]->p[3],p[3]->p[0]};
allVar=allVar/.heavy[a_,v]:>heavy[a/.p[1|2]->0,v];
allVar=DeleteCases[allVar/.heavy[0,_]->0,0]//DeleteDuplicates;
allVar2=allVar/.heavy->ip/.massless[a_]:>ip[a,a];
{allVar2,allVar,invRule[allVar2,ip],invRule[allVar,heavy|massless]}/.
{heavy->heavy0,massless->massless0,ip->ip0,p->p0,v->v0}
];

generateVarList4[loopOrder_Integer?(#≥3&),heavy0_,massless0_,ip0_,p0_List,v0_,qsq:True:True]/;Length[p0]==loopOrder+1:=
Module[{t},generateVarList2[loopOrder,heavy0,massless0,ip0,t,v0]/.t->(p0[[#+1]]&)];

reduceMassless[h_[a_]]:=
Module[{tmp,tmp2},
tmp=Variables[Expand[a]];
If[tmp==={},Return[h[a]]];
tmp2=Coefficient[a,tmp];
((tmp2[[1]])^2) h[Expand[a/tmp2[[1]]]]
];

Options[setVariableProperty]={"qSquareIsMinus1"->True};
setVariableProperty[heavy_,massless_,momentum_,v_,OptionsPattern[]]:=(
ClearAll[heavy,massless];
If[TrueQ[OptionValue["qSquareIsMinus1"]],massless[momentum[0]]=-1];
massless[v]=1;
massless[a_?NumericQ]:=(a^2);
massless[a_]:=Module[{h,k},k=reduceMassless[h[a]];(k/.h->massless)/;k=!=h[a]];
heavy[x_. momentum[0]+a_.,v]:=heavy[x momentum[0]+a,v]=heavy[a,v];
heavy[a_?Positive b_.,c_]/;a=!=1:=a heavy[b,c];
heavy[c_,a_?Positive b_.]/;a=!=1:=a heavy[b,c];
heavy[a_?Negative b_.,c_]/;a=!=-1:=-a heavy[-b,c];
heavy[c_,a_?Negative  b_.]/;a=!=-1:=-a heavy[-b,c];
heavy[a_,a_]:=massless[a];
heavy[v,v]=1;
heavy[0,_]=0;
heavy[_,0]=0;
);

Options[setIpProperty]={"qSquareIsMinus1"->True};
setIpProperty[heavy_,massless_,ip_,momentum_,v_,OptionsPattern[]]:=(
ClearAll[ip];
SetAttributes[ip,Orderless];
If[TrueQ[OptionValue["qSquareIsMinus1"]],ip[momentum[0],momentum[0]]=-1];
ip[momentum[0],v]=0;
ip[v,v]=1;
ip[a_,b_?NumericQ]:=a b;
ip[b_?NumericQ,a_]:=a b;
ip[a_,b_]:=
Module[{redA,redB,varA,varB,signA,signB,coeffA,coeffB,
positiveVarA,positiveVarB,common,res,tmp},
(ip[a,b]=res//Expand)/;(
redA=Expand[a];
redB=Expand[b];
varA=Variables[redA];
varB=Variables[redB];
signA=Sign[Coefficient[redA,varA[[1]]]];
If[MemberQ[varB,varA[[1]]],
signB=Sign[Coefficient[redB,varA[[1]]]],
signB=Sign[Coefficient[redB,varB[[1]]]]
];
If[signA signB==0,Return[False]];
redA=redA signA;
redB=redB signB;
coeffA=Coefficient[redA,varA];
coeffB=Coefficient[redB,varB];
positiveVarA=Extract[varA,Position[coeffA,_Integer?(#>=1&),1]];
positiveVarB=Extract[varB,Position[coeffB,_Integer?(#>=1&),1]];
common=Intersection[positiveVarA,positiveVarB]//Total;
If[(!common===0)&&(!redA===redB),
res=signA signB (
ip[common,common]+ip[Expand[redA+redB-2common],common]+
ip[redA-common,redB-common]
);
Return[True];
];
If[signA!=1||signB!=1,res=signA signB ip[redA,redB];Return[True]];
tmp=Cases[{a},x_ y_+z_.:>{x,y,z},1]//Flatten;
If[!tmp==={},res=tmp[[1]]ip[tmp[[2]],b]+ip[tmp[[3]],b];Return[True]];
tmp=Cases[{b},x_ y_+z_.:>{x,y,z},1]//Flatten;
If[!tmp==={},res=tmp[[1]] ip[a,tmp[[2]]]+ip[a,tmp[[3]]];Return[True]];
If[(!a===b)&&!(a===v)&&!(b===v),
res=(ip[a+b,a+b]-ip[a,a]-ip[b,b])/2;
Return[True];
];
tmp=Cases[{a},x_ +v:>x,1];
If[!tmp==={},
If[!a===b,Return[False]];
res=ip[tmp[[1]],tmp[[1]]]+2ip[tmp[[1]],v]+1;Return[True];
];
If[a===v,res=heavy[b/.momentum[0]->0,v];Return[True]];
If[b===v,res=heavy[a/.momentum[0]->0,v];Return[True]];
If[a===b,res=massless[a];Return[True]];
False
)
];
);

convertFromJ[equation_,jVar_,variable_List]/;variable[[1,0]]=!=List:=
With[{l=Length[variable]},
equation/.HoldPattern[jVar[a__]]/;Length[{a}]==l:>Times@@(variable^-{a})
];

convertFromJ[equation_,jVar_,variable:{{___},{___},{___},{___}}]/;Equal@@(Length/@variable):=
Module[{l=Length[variable[[1]]],tmp},
equation/.HoldPattern[jVar[a__]]/;Length[{a}]==l:>
Times@@(variable[[2]]^-(If[TrueQ[#<0],0,#]&/@{a}))Times@@(variable[[1]]^-(If[!TrueQ[#<0],0,#]&/@{a}))
];

Options[convertToJ]={"vRepMonitor"->None};

convertToJ[0,_,_,OptionsPattern[]]=0;

convertToJ[equation_,jVar_,{var1_List,var2_List,rule1_List,rule2_List},OptionsPattern[]]/;
Length[var1]==Length[var2]==Length[rule1]==Length[rule2]:=
Module[{tmp,i,var,exp},
tmp=equation/.rule1/.rule2//Together//Expand;
OptionValue["vRepMonitor"][tmp];
var=Cases[{rule1,rule2},"V"[_],{0,∞}]//Union;
tmp=If[tmp[[0]]===Plus,List@@tmp,{tmp}];
tmp=Sum[
exp=-Exponent[tmp[[i]],#]&/@var;
(jVar@@exp)(tmp[[i]]Times@@(var^exp)),
{i,Length[tmp]}
];
tmp/;FreeQ[tmp,"V"]
];

getVariable[equation_,variablename_]:=Cases[equation,variablename[__],{0,Infinity}]//Union;

(*tgJ=
{
{0,0,0,1,0,0,-1,1,0,1,0,0,1,0,0,0,0,0,1,0,0},
{0,0,-1,1,0,0,0,1,0,1,0,0,1,0,0,0,0,0,1,0,0},
{0,0,0,1,0,0,-1,1,0,1,0,0,1,0,0,0,0,0,1,0,-1},
{1,1,0,0,0,0,0,0,0,1,1,1,0,0,0,0,0,0,0,0,1},
{1,0,0,0,0,1,0,1,1,0,0,1,0,0,0,0,0,1,0,0,0},
{1,0,0,0,0,1,0,0,1,0,1,1,0,0,0,0,0,0,0,0,1},
{0,0,0,1,0,0,-1,1,0,1,0,1,1,0,0,0,0,0,0,0,1}
};*)

(*priorityOrder2[_[variables__]]/;MemberQ[tgJ,{variables}]:=ConstantArray[10,46];*)

(*priorityOrder[_[variable__]]:=
Module[{positiveCoefficient},
positiveCoefficient=Select[{variable},#>0&];
{
Length[positiveCoefficient],
Total[positiveCoefficient],
Total[positiveCoefficient]-Total[{variable}],
Max[Abs[{variable}]]
}
];

priorityOrder2[_[variable__]]:={
priorityOrder[{variable}],
Abs[Reverse[{variable}]],
Reverse[{variable}]
}//Flatten;
*)

priorityOrder[_[arg__],posHeavy_:{},posHybrid_:{}]:=
Module[{heavy,massless,hybrid},
heavy={arg}[[posHeavy]];
hybrid={arg}[[posHybrid]];
massless={arg}[[Complement[Range[Length[{arg}]],posHeavy,posHybrid]]];
{
Count[{arg},_?Positive],
Count[{arg},Except[0]],
Total[Select[{arg},Positive]],
Total[Abs[{arg}]],
Count[hybrid,Except[0]],
Total[Select[hybrid,Positive]],
Total[Abs[hybrid]],
Count[heavy,Except[0|1]],
Count[heavy,Except[0]],
Total[Select[heavy,Positive]],
Total[Abs[heavy]],
Max[Abs[{arg}]]
}
];

(*priorityOrder[_[variable__],___]:=
Module[{positiveCoefficient},
positiveCoefficient=Select[{variable},#>0&];
{
Length[positiveCoefficient],
Total[positiveCoefficient],
Total[positiveCoefficient]-Total[{variable}],
Max[Abs[{variable}]]
}
];*)

priorityOrder2[_[arg__],posHeavy_:{},posHybrid_:{}]:={
priorityOrder[{arg},posHeavy,posHybrid],
Abs[Reverse[{arg}]],
Reverse[{arg}]
}//Flatten;

equationPriority[a_,b_]:=
Module[{tmp},
tmp=getVariable[a,b];
Join[tmp//picklast//First//priorityOrder,{Length[tmp]}]
];

jQuality[in_,posHeavy_:{},posHybrid_:{}]:=
Module[{tmp},
tmp=priorityOrder[#,posHeavy,posHybrid]&/@in;
Tally[tmp]//Sort
];

sortVariable[input_]:=SortBy[input,priorityOrder2];

sortEquation[input_]:=SortBy[input,equationPriority];

pickJle[a_,b_]:=Module[{tmp},tmp=priorityOrder[b];Select[a,OrderedQ[{priorityOrder[#],tmp}]&]//Union//sortVariable];
pickJge[a_,b_]:=Module[{tmp},tmp=priorityOrder[b];Select[a,OrderedQ[{tmp,priorityOrder[#]}]&]//Union//sortVariable];
pickJl[a_,b_]:=Module[{tmp},tmp=priorityOrder[b];Select[a,!OrderedQ[{tmp,priorityOrder[#]}]&]//Union//sortVariable];
pickJg[a_,b_]:=Module[{tmp},tmp=priorityOrder[b];Select[a,!OrderedQ[{priorityOrder[#],tmp}]&]//Union//sortVariable];
pickJne[a_,b_]:=Module[{tmp},tmp=priorityOrder[b];Select[a,(!priorityOrder[#]===tmp)&]//Union//sortVariable];
pickJe[a_,b_]:=Module[{tmp},tmp=priorityOrder[b];Select[a,(priorityOrder[#]===tmp)&]//Union//sortVariable];

pickJles[a_,b_]:=DeleteCases[pickJle[a,b],b];
pickJges[a_,b_]:=DeleteCases[pickJge[a,b],b];
pickJes[a_,b_]:=DeleteCases[pickJe[a,b],b];

pickJle2[a_,b_]:=Module[{tmp},tmp=priorityOrder2[b];Select[a,OrderedQ[{priorityOrder2[#],tmp}]&]//Union//sortVariable];
pickJge2[a_,b_]:=Module[{tmp},tmp=priorityOrder2[b];Select[a,OrderedQ[{tmp,priorityOrder2[#]}]&]//Union//sortVariable];
pickJl2[a_,b_]:=Module[{tmp},tmp=priorityOrder2[b];Select[a,!OrderedQ[{tmp,priorityOrder2[#]}]&]//Union//sortVariable];
pickJg2[2a_,b_]:=Module[{tmp},tmp=priorityOrder2[b];Select[a,!OrderedQ[{priorityOrder2[#],tmp}]&]//Union//sortVariable];
pickJne2[a_,b_]:=DeleteCases[a,b]//Union//sortVariable;
pickJe2[a_,b_]:=Cases[a,b]//Union;

pickJmost[a_]:=Module[{tmp},tmp=jQuality[a][[-1,1]];Select[a,!(priorityOrder[#]===tmp)&]//Union//sortVariable];
pickJrest[a_]:=Module[{tmp},tmp=jQuality[a][[1,1]];Select[a,!(priorityOrder[#]===tmp)&]//Union//sortVariable];
pickJfirst[a_]:=Module[{tmp},tmp=jQuality[a][[1,1]];Select[a,(priorityOrder[#]===tmp)&]//Union//sortVariable];
pickJlast[a_]:=Module[{tmp},tmp=jQuality[a][[-1,1]];Select[a,(priorityOrder[#]===tmp)&]//Union//sortVariable];

pickJmost2[a_]:=sortVariable[a]//Most;
pickJrest2[a_]:=sortVariable[a]//Rest;
pickJfirst2[a_]:=sortVariable[a]//First;
pickJlast2[a_]:=sortVariable[a]//Last;

jleQ[a_,b_]:=OrderedQ[{priorityOrder[a],priorityOrder[b]}];
jgeQ[a_,b_]:=OrderedQ[{priorityOrder[b],priorityOrder[a]}];
jlQ[a_,b_]:=!OrderedQ[{priorityOrder[b],priorityOrder[a]}];
jgQ[a_,b_]:=!OrderedQ[{priorityOrder[a],priorityOrder[b]}];
jneQ[a_,b_]:=!priorityOrder[a]===priorityOrder[b];
jeQ[a_,b_]:=priorityOrder[a]===priorityOrder[b];

jlesQ[a_,b_]:=OrderedQ[{priorityOrder[a],priorityOrder[b]}]&&(!a===b);
jgesQ[a_,b_]:=OrderedQ[{priorityOrder[b],priorityOrder[a]}]&&(!a===b);
jesQ[a_,b_]:=priorityOrder[a]===priorityOrder[b]&&(!a===b);

jle2Q[a_,b_]:=OrderedQ[{priorityOrder2[a],priorityOrder2[b]}];
jge2Q[a_,b_]:=OrderedQ[{priorityOrder2[b],priorityOrder2[a]}];

allVarBaseN[order_,f_]:=
Module[{varList,var,var2,tmp},
varList=f[order,"HV","ML","dot","p","v",True];
var=Array["p",order+1,0];
var2=Times@@@Tuples[var,2]//Union;
tmp=varList[[2]]/.(("HV"[a__])^2):>"HV2"[a]/.{
"HV"[a_,"v"]:>PadRight[Coefficient[a,var],Length[var]+2Length[var2]],
"ML"[a_]:>PadRight[PadLeft[Coefficient[(a^2),var2],Length[var]+Length[var2]],Length[var]+2Length[var2]],
"HV2"[a_,"v"]:>PadLeft[Coefficient[(a^2),var2],Length[var]+2Length[var2]]
};
tmp=DeleteCases[Transpose[tmp],{0..}]//Transpose;
{tmp,varList}
];

allVarBase[order_Integer?(#≥1&)]:=
With[{t=allVarBaseN[order,generateVarList1]},
{t[[1]],t[[2,2]]}/."ML"["p"[0]]->-1
];
allVarBase1[order_Integer?(#≥1&)]:=allVarBaseN[order,generateVarList1];
allVarBase2[order_Integer?(#≥2&)]:=allVarBaseN[order,generateVarList2];
allVarBase3[order_Integer?(#≥1&)]:=allVarBaseN[order,generateVarList3];
allVarBase4[order_Integer?(#≥3&)]:=allVarBaseN[order,generateVarList4];

Protect["jTool`*"];
Protect["jTool`private`*"];
Unprotect["jTool`$*"];
Unprotect["jTool`private`$*"];
Unprotect["jTool`private`c*"];

End[];

EndPackage[];