(******************************************************************************
 *	Recurrence.m - a tool to solve recurrence relations		      *
 ******************************************************************************)


BeginPackage["Recurrence`"];


(* define public functions *)

(* define recurrence relations *)
DefineRecurrence::usage = "DefineRecurrence[F, eql ,vars ,opts]
F: symbol for the following recurrence relations
eql: list of recurrence relations
vars: list of index names
opts: options, can be Symmetric, Priority, Start, Directory and Masters"

UndefineRecurrence::usage = "UndefineRecurrence[F]"

(* find solutions to recurrence relations *)
SolveRecurrence::usage = "SolveRecurrence[F, max]
max: index range to solve"

(* load and save symbol F *)
LoadRecurrence::usage = "LoadRecurrence[F, file]"
SaveRecurrence::usage = "SaveRecurrence[F, file]"

(* is true if a symbol is a recurrence relation *)
RecurrenceQ::usage = "RecurrenceQ[F]"

(* retuns a list of indices occurring in an equation *)
Indizes::Usage = "Indizes[F, eq]"

(* options to DefineRecurrence *)
Symmetric::usage = ""
Priority::usage = ""
Start::usage = ""
Directory::usage = ""
Masters::usage = ""

(* get list of master integrals, not yet implemented *)
getMasters::Usage = "getMasters[F]"


Begin["`Private`"];


(* error messages *)
General::invarg = "invalid argument: `1` `2`"
General::invarg1 = "invalid argument: `1`"
General::invarg0 = "invalid argument."

RecurrenceQ[_] = False;

(* Indizes[x] gives a list of indices occuring in an equation x *)
SetAttributes[Indizes,Listable];
Indizes[F_,x_]:=Cases[x,F[_],{0,\[Infinity]}]/.F[i__]:>i;

(* simplify index operators *)
SetAttributes[{IndexVal,IndexAdd},HoldFirst];
IndexAdd/:Times[IndexAdd[F_,inc1_],IndexAdd[F_,inc2_]]=
			IndexAdd[F,inc1+inc2];
IndexAdd/:Power[IndexAdd[F_,inc_],p_Integer]=IndexAdd[F,inc*p];

SetAttributes[DefineRecurrence,HoldFirst];
Options[DefineRecurrence]:={Symmetric->Null,
			    Priority->Null,
			    Start->Null,
			    Directory->Null,
			    Masters->Null};

DefineRecurrence[F_,eql_,vars_,opts___] := Module[
		{sym,prio,start,dir,master,optsOK},

	{optsOK,sym,prio,start,dir,master}=checkOptions[F,eql,vars,opts];
	If[!optsOK,Return[Null]];

	initEquations[F,eql,vars];
	initData[F,dir];
	initSymmetries[F,sym];
	initPriority[F,prio,start];
	initMasters[F,master];

	RecurrenceQ[F]^=True;
	F::usage=ToString[F]<>"[{i,j,k,...}]";

	Protect[F];
	Return[F];
];

initEquations[F_,eql_,vars_]:=Module[
		{eqs,incOne,range,minInd,maxInd,numInd},

	numInd=Length[vars];
	numIndizes[F]^=numInd;

	F/:IndexVal[F[i_],n_Integer]:=i[[n]];
	F/:IndexAdd[F[i_],inc_]:=F[i+inc];
	eqs=eql;
	Do[
		incOne=Table[Which[i==j,1,True,0],{j,numInd}];
		eqs=eqs*IndexAdd[#,Array[0&,numInd]]
			/.vars[[i]][n_Integer]->IndexAdd[#,incOne*n]
			/.vars[[i]]->IndexVal[#,i];
	,{i,numInd}];
	Equations[F]^=Expand[eqs];

	range=Indizes[F,ApplyEqs[F,vars]]/.Array[vars[[#]]->0&,numInd];
	minInd=(Min/@Transpose[#])&/@range;
	maxInd=(Max/@Transpose[#])&/@range;
	EqRange[F]^=Max/@Transpose[maxInd-minInd];
	Do[
		eqs[[j]]=eqs[[j]]
			/.IndexVal[#,i_]:>(IndexVal[#,i]-minInd[[j,i]]);
	,{j,Length[range]}];
	Equations[F]^=Simplify[Expand[eqs*Thread[IndexAdd[#,-minInd]]]];
];

ApplyEqs[F_,v_]:=Evaluate[Equations[F]]&[Unevaluated[F[v]]];

findFirst[list_,element_]:=Module[{p},
	p=Position[list,element,1,1];
	Return[If[Length[p]==0,0,p[[1,1]]]];
];

initData[F_,dir_]:=Module[{},
	indexList[F]^={};
	dependencies[F]^={};
	reverseDepend[F]^={};
	properties[F]^={};
	loadData[F];

	Directory[F]^=dir;
	If[MatchQ[dir,Null],
		initDataRAM[F];
	,
		initDataHD[F];
	];

	F[indexNr_Integer]:=getEquation[F,indexNr]/;isKnown[F,indexNr];
	F[index_?VectorQ]:=(F[indexNumber[F,SortIndex[F,index]]]
				/.F[i_Integer]:>F[getIndex[F,i]])
			/;isKnown[F,index];

	F/:IndexAdd[F[i_],inc_]:=F[indexNumber[F,SortIndex[F,i+inc]]];
];

loadData[F_]:=Module[{iList,deps,revDeps,props},

	iList=indexList[F];
	deps=dependencies[F];
	revDeps=reverseDepend[F];
	props=properties[F];
	F/:indexNumber[F,index_]:=Module[{nr},
		If[!OrderedIndexQ[F,index],Print["Fehler: ",index]];
		nr=findFirst[iList,index];
		If[nr==0,
			iList=Append[iList,index];
			deps=Append[deps,{}];
			revDeps=Append[revDeps,{}];
			props=Append[props,{False}];
			Return[Length[iList]];
		,
			Return[nr];
		];
	];
	F/:getIndex[F,i_]:=iList[[i]];
	F/:probeIndexNumber[F,index_]:=findFirst[iList,index];

	F/:setKnown[F,indexNr_]:=(props[[indexNr,1]]=True;);
	F/:isKnown[F,indexNr_Integer]:=props[[indexNr,1]];
	F/:isKnown[F,index_?VectorQ]:=Module[{indexNr},
		If[Length[index]!=numIndizes[F],Return[False];];
		indexNr=probeIndexNumber[F,SortIndex[F,index]];
		Return[(indexNr!=0)&&isKnown[F,indexNr]];
	];

	F/:setRevDep[F,indexNr_,revDep_]:=(revDeps[[indexNr]]=revDep;);
	F/:getRevDep[F,indexNr_]:=revDeps[[indexNr]];
	F/:setDependencies[F,indexNr_,depend_]:=(deps[[indexNr]]=depend;);
	F/:getDependencies[F,indexNr_]:=deps[[indexNr]];

	saveData[F]:=(
		indexList[F]^=iList;
		dependencies[F]^=deps;
		reverseDepend[F]^=revDeps;
		properties[F]^=props;
	);
];

SaveRecurrence[F_,file_]:=Module[{},
	Unprotect[F];
	saveData[F];
	Save[file,F];
];

LoadRecurrence[F_,file_]:=Module[{},
	Get[file];
	loadData[F];
	Protect[F];
];

initDataRAM[F_]:=Module[{},
	F/:storeEquation[F,indexNr_,eq_]:=
			(equation[F,indexNr]=eq;
			setKnown[F,indexNr];
			Update[F];);
	F/:getEquation[F,indexNr_]:=equation[F,indexNr];
];

initDataHD[F_]:=Module[{maxFilesPerDir},
	maxFilesPerDir=500;
	fileName[F,indexNr_]:=ToFileName[{Directory[F],
			StringJoin[ToString[F],ToString[
			Floor[indexNr/maxFilesPerDir]*maxFilesPerDir]]},
			StringJoin[ToString[F],ToString[indexNr],".m"]];
	F/:storeEquation[F,indexNr_,equation_]:=Module[{dname,fname},
			fname=fileName[F,indexNr];
			dname=DirectoryName[fname];
			If[!MatchQ[FileType[dname],Directory],
				CreateDirectory[dname];
			];
			Export[fname,equation];
			setKnown[F,indexNr];
			Update[F];
	];
	F/:getEquation[F,indexNr_]:=Import[fileName[F,indexNr]];
];

removeDoubles[list_]:=Module[{i,j,l},
	l=list;
	For[i=1,i<Length[l],i++,
		For[j=i+1,j<=Length[l],
			If[MatchQ[l[[i]],l[[j]]],
				l=Delete[l,j];
			,
				j++;
			];
		];
	];
	Return[l];
];

initMasters[F_,master_]:=Module[{i,masters},
	masters=removeDoubles[indexNumber[F,SortIndex[F,#]]&/@master];
	F/:getMasters[F]=masters;
	F/:isMaster[F,indexNr_]:=MemberQ[getMasters[F],indexNr];
];

initSymmetries[F_,sym_]:=Module[{i,j,sortF,orderF,tk,st,or},
	If[Length[sym]!=0,
		sortF={};
		orderF=True;
		j=1;
		For[i=1,i<=Length[sym],i++,
			If[sym[[i,1]]!=j,
				sortF=Append[sortF,tk[#,{j,sym[[i,1]]-1}]]
			];
			sortF=Append[sortF,st[tk[#,{sym[[i,1]],
					sym[[i,2]]}]]];
			orderF=orderF&&or[tk[#,{sym[[i,1]],sym[[i,2]]}]];
			j=sym[[i,2]]+1;
		];
		If[j<=numIndizes[F],
			sortF=Append[sortF,tk[#,{j,numIndizes[F]}]]
		];
		sort[F]^=sortF;
		order[F]^=orderF;
		F/:SortIndex[F,i_]:=Flatten[Evaluate[sort[F]]&[i]
				/. tk->Take/.st->Sort];
		F/:OrderedIndexQ[F,i_]:=Evaluate[order[F]]&[i]
				/. tk->Take/.or->OrderedQ;
	,
		F/:SortIndex[F,i_]=i;
		OrderedIndexQ[F,_]^=True;
	];
];

initPriority[F_,prio_,st_]:=Module[{i,j,k,loopInd,loopIndF,sum,start},
	start=SortIndex[F,st];
	loopIndF=Array[0&,numIndizes[F]];
	j=1;
	For[i=1,i<=Length[prio],i++,
		sum=0;
		For[k=1,k<Length[prio[[i]]],k++,
			j++;
			loopIndF[[prio[[i,k]]]]=loopInd[j];
			sum+=loopInd[j];
		];
		j++;
		loopIndF[[prio[[i,Length[prio[[i]]]]]]]=
			loopInd[j-Length[prio[[i]]]]-sum;
	];
	F/:LoopToIndex[F,multi_]:=(loopIndF/.loopInd->multi)+start;
	F/:LoopRange[F,i_,min_,max_]:=
			Sequence@@LoopPriority[prio,i,min,max];
	F/:ChooseIndex[F,l_]:=Module[{ind},
		ind=DeleteCases[l,_?(isMaster[F,#]&)];
		If[Length[ind]==0,ind=l];
		Return[ind[[MinPriority[prio,
			(getIndex[F,#]-start)&/@ind]]]];
	];
	F/:checkMax[F,max_]:=VectorQ[max,(IntegerQ[#]&&NonNegative[#])&]
			&&(Length[max]==Length[prio]);
	F/:checkMin[F,min_]:=VectorQ[min,(IntegerQ[#]&&NonPositive[#])&]
			&&(Length[min]==Length[prio]);
];

(* determine priority to loop over indices *)
LoopPriority[prio_,li_,min_,max_]:=Module[{ret,i,j,k,sum},
	ret={};
	k=1;
	For[i=1,i<=Length[prio],i++,
		If[Length[prio[[i]]]==1,
			ret=Append[ret,{li[k],min[[i]],max[[i]]}];
		,
			ret=Append[ret,{li[k],0,max[[i]]}];
		];
		sum=li[k];
		k++;
		For[j=1,j<Length[prio[[i]]],j++,
			ret=Append[ret,{li[k],0,sum}];
			sum-=li[k];
			k++;
		];
	];
	Return[ret];
];

(* determine which index in an equation has the least priority, called by
   ChooseIndex *)
MinPriority[prio_,list_]:=Module[{i,j,k,l,v,lsum,nsum},
	l=list;
	For[i=1,i<=Length[prio],i++,
		lsum=Sum[l[[1,prio[[i,j]]]],{j,Length[prio[[i]]]}];
		For[k=2,k<=Length[l],
			nsum=Sum[l[[k,prio[[i,j]]]],
				{j,Length[prio[[i]]]}];
			Which[lsum>nsum,
				l=Delete[l,k];
			,lsum==nsum,
				k++;
			,lsum<nsum,
				l=Drop[l,k-1];
				lsum=nsum;
			];
		];
		If[Length[l]==1,Return[findFirst[list,l[[1]]]];];
	];
	For[i=1,i<=Length[prio],i++,
		If[Length[prio[[i]]]==1,Continue[];];
		For[j=1,j<=Length[prio[[i]]],j++,
			v=prio[[i,j]];
			For[k=2,k<=Length[l],
				Which[l[[k-1,v]]>l[[k,v]],
					l=Delete[l,k];
				,l[[k-1,v]]==l[[k,v]],
					k++;
				,l[[k-1,v]]<l[[k,v]],
					l=Drop[l,k-1];
				];
			];
		];
		If[Length[l]==1,Return[findFirst[list,l[[1]]]];];
	];
	Return[findFirst[list,l[[1]]]];
];

SetAttributes[{checkOptions,SymbolQ},HoldFirst];
listAnd[{i__}]:=And[i];
SymbolQ[s_]:=MatchQ[Head[Unevaluated[s]],Symbol];

checkOptions[F_,eql_,vars_,opts___] := Module[
		{sym,prio,start,dir,master,error,nsym,nprio},
	error={False,0,0,0,0,0};
	{sym,prio,start,dir,master}=
			({Symmetric,Priority,Start,Directory,Masters}
			/.{opts}/.Options[DefineRecurrence]);

	If[!SymbolQ[F],
		Message[DefineRecurrence::invarg,F,"is not a symbol."];
		Return[error];
	];
	If[MemberQ[Attributes[F],Protected]&&!RecurrenceQ[F],
		Message[DefineRecurrence::invarg,F,"is protected."];
		Return[error];
	];

	If[!VectorQ[eql],
		Message[DefineRecurrence::invarg1,
			"Equations have wrong format."];
		Return[error];
	];
	If[!VectorQ[vars,SymbolQ],
		Message[DefineRecurrence::invarg1,vars];
		Return[error];
	];

	If[!MatchQ[dir,Null],
		If[!MatchQ[FileType[dir],Directory],
			Message[DefineRecurrence::invarg,dir,
				"does not exist."];
			Return[error];
		];
		dir=ToFileName[dir,ToString[F]];
		If[!MatchQ[FileType[dir],Directory],CreateDirectory[dir];];
	];

	If[start==Null,start=Array[0&,Length[vars]]];
	If[(!VectorQ[start])||(Length[start]!=Length[vars]),
		Message[DefineRecurrence::invarg,"Start ->",start];
		Return[error];
	];
	If[master==Null,master={};];
	If[!checkMaster[master,vars],
		Message[DefineRecurrence::invarg,"Masters ->",master];
		Return[error];
	];
	nsym=checkSym[sym,vars];
	If[nsym==Null,
		Message[DefineRecurrence::invarg,"Symmetric ->",sym];
		Return[error];
	];
	nprio=checkPrio[prio,vars];
	If[nprio==Null,
		Message[DefineRecurrence::invarg,"Priority ->",prio];
		Return[error];
	];

	Unprotect[F];
	ClearAll[F];
	Return[{True,nsym,nprio,start,dir,master}];
];

checkMaster[master_,vars_]:=Module[{i},
	If[!MatchQ[Head[master],List],Return[False]];
	For[i=1,i<=Length[master],i++,
		If[(!VectorQ[master[[i]]])||
			(Length[master[[i]]]!=Length[vars]),
				Return[False];
		];
	];
	Return[True];
];

checkSym[sym_,vars_]:=Module[{nsym,i,last},
	If[sym==Null,Return[{}];];
	If[!MatchQ[Head[sym],List],Return[Null];];
	If[Length[sym]==0,Return[{}];];
	nsym=sym/.Array[(vars[[#]]->#)&,Length[vars]];

	If[(Length[nsym]==2)&&(VectorQ[nsym,IntegerQ]),nsym={nsym};];
	If[(Length[nsym[[1]]]!=2)||(!MatrixQ[nsym,IntegerQ]),
		Return[Null];
	];
	nsym=Sort[Sort/@nsym];

	last=0;
	For[i=1,i<=Length[nsym],i++,
		If[nsym[[i,1]]<=last,Return[Null];];
		last=nsym[[i,2]];
	];
	If[last>Length[vars],Return[Null];];

	Return[nsym];
];

checkPrio[prio_,vars_]:=Module[{nprio,clist,i,j},
	If[prio==Null,Return[Array[{#}&,Length[vars]]];];
	If[!MatchQ[Head[prio],List],Return[Null];];
	nprio=prio/.Array[(vars[[#]]->#)&,Length[vars]];
	clist=Array[False&,Length[vars]];

	For[i=1,i<=Length[nprio],i++,
		If[IntegerQ[nprio[[i]]],nprio[[i]]={nprio[[i]]};];
		If[(!VectorQ[nprio[[i]],(IntegerQ[#]
				&&(0<#<=Length[vars]))&])
				||(Length[nprio[[i]]]==0),
			Return[Null];
		];
		For[j=1,j<=Length[nprio[[i]]],j++,
			If[clist[[nprio[[i,j]]]],Return[Null];];
			clist[[nprio[[i,j]]]]=True;
		];
	];
	If[!listAnd[clist],Return[Null];];

	Return[nprio];
];

UndefineRecurrence[F]:=Module[{},
	If[!RecurrenceQ[F],
		Message[UndefineRecurrence::invarg0];
		Return[Null];
	];
	Unprotect[F];
	ClearAll[F];
	Return[Null];
];

Masters[F_]:=getIndex[F,#]&/@getMasters[F];

SolveRecurrence[F_,min_,max_]:=Module[
		{multi,range,index},

	If[!RecurrenceQ[F],
		Message[SolveRecurrence::invarg1,F];
		Return[Null];
	];
	If[!checkMax[F,max],
		Message[SolveRecurrence::invarg1,max];
		Return[Null];
	];
	If[!checkMin[F,min],
		Message[SolveRecurrence::invarg1,min];
		Return[Null];
	];
	Unprotect[F];

	range=LoopRange[F,multi,min,max];
	Do[
		index=LoopToIndex[F,multi];

		(* respect symmetries *)
		If[!OrderedIndexQ[F,index],Continue[]];

		(* solve for index with least priority and add equation *)
		AddEquations[F,index];
	,Evaluate[range]];

	Protect[F];
	Return[F];
];

SolveRecurrence[F_,max_]:=SolveRecurrence[F,Array[0&,Length[max]],max];

changeDependencies[F_,indexNr_,new_]:=Module[{i,j,p,rd,del},
	del=getDependencies[F,indexNr];
	For[i=1,i<=Length[new],i++,
		j=new[[i]];
		p=findFirst[del,j];
		If[p==0,
			setRevDep[F,j,Append[getRevDep[F,j],indexNr]];
		,
			del=Delete[del,p];
		];
	];
	For[i=1,i<=Length[del],i++,
		j=del[[i]];
		rd=getRevDep[F,j];
		p=findFirst[rd,indexNr];
		If[p!=0,setRevDep[F,j,Delete[rd,p]];];
	];
	setDependencies[F,indexNr,new];
];

AddEquations[F_,index_]:=Module[{i,k,l,ind,eqs,eq,known,revdep},
	eqs=ApplyEqs[F,index];
	For[i=1,i<=Length[eqs],i++,
		eq=eqs[[i]];
		eq=Collect[eq,F/@Indizes[F,eq],Together];
		ind=Indizes[F,eq];
		If[Length[ind]==0,Continue[];];

		(* choose index with least priority... *)
		k=ChooseIndex[F,ind];
		(* and solve for it *)
		eq=-(eq/.F[k]->0)/Coefficient[eq,F[k]];
		ind=DeleteCases[ind,k];
		eq=Collect[eq,F/@ind,Together];

		storeEquation[F,k,eq];
		changeDependencies[F,k,ind];

		(* reverse dependencies and back substitution *)
		revdep=getRevDep[F,k];
		For[l=1,l<=Length[revdep],l++,
			k=revdep[[l]];
			eq=F[k];
			eq=Collect[eq,F/@Indizes[F,eq],Together];
			storeEquation[F,k,eq];
			ind=Indizes[F,eq];
			changeDependencies[F,k,ind];
		];
	];
];


End[(*"`Private`"*)];

EndPackage[(*Recurrence`*)];
