(* ::Package:: *)

BeginPackage["LHAPDFreader`"]

SetLibraryPath::usage = "SetLHAPDFPath[path], set the path for LHAPDF storage. Use SetLHAPDFPath[] to set default";
ListGrids::usage = "ListGrids[], Return the list of grids found in the current library path";
ReadInfo::usage = "ReadInfo[name] read the .info file of the set (and fill internals)";
LoadLHAPDFgrid::usage = "LoadLHAPDFgrid[name] load the information on particular LHAPDF grid and prepare interpolation functions";
LoadTMDLIBgrid::usage = "LoadLHAPDFgrid[name] load the information on particular TMDLIB grid and prepare interpolation functions";
ShowInfo::usage = "ShowInfo[param] show info associated with label param (if exists). ShowInfo[] shows all avalible information";
AlphaS::usage = "AlphaS[Q] returns the values of Alpha_s (taken from info-file)";
xPDF::usage = "xPDF[x,Q,f] returns the values xf (taken from LHAPDFgrid)";
xTMDPDF::usage = "xTMDPDF[x,qT,Q,f] returns the values xf (taken from TMDlib grid, Pavia convention)";

Begin["`Private`"]

(*--------------Internal variables---------------------------------*)

(*The path to LHAPDF directory*)
LIBpath="ERROR";
(*List of avalible PDFs*)
GRIDlist={};
(*Values of alpha-grid*)
(*Default version is from NNPDF31_nnlo_as_0118*)
AlphaGrid={{1.65`,0.33074891`},{1.7874388`,0.31762459`},{1.9429053`,0.30507081`},{2.1193749`,0.29305875`},{2.32041`,0.28156114`},{2.5502945`,0.27055221`},{2.8142025`,0.26000761`},{3.1184123`,0.24990438`},{3.4705775`,0.24022086`},{3.8800752`,0.23093662`},{4.3584516`,0.22203241`},{4.92`,0.21377883`},{5.5493622`,0.20671584`},{6.2897452`,0.19986334`},{7.1650687`,0.19321629`},{8.2052867`,0.1867697`},{9.4481248`,0.1805186`},{10.941378`,0.17445809`},{12.745972`,0.16858329`},{14.940062`,0.16288942`},{17.624572`,0.15737175`},{20.930715`,0.15202563`},{25.030298`,0.14684647`},{30.149928`,0.14182979`},{36.590777`,0.13697116`},{44.756282`,0.13226626`},{55.191298`,0.12771084`},{68.63794`,0.12330074`},{86.115921`,0.11903189`},{109.03923`,0.11490031`},{139.38725`,0.1109021`},{179.95815`,0.10703345`},{234.7482`,0.10329062`},{309.52544`,0.099669997`},{412.70732`,0.096168008`},{556.71861`,0.092781186`},{760.11795`,0.089506143`},{1050.9694`,0.086339571`},{1472.2574`,0.083278245`},{2090.6996`,0.080319017`},{3011.2909`,0.077458816`},{4401.6501`,0.074694649`},{6533.3918`,0.0720236`},{9853.5186`,0.069442824`},{15109.614`,0.066949551`},{23573.066`,0.064541082`},{37444.017`,0.062214786`},{60599.32`,0.059968105`},{100000.`,0.057798546`}};
AlphaInter=Interpolation[AlphaGrid];
AlphaS[Q_?NumericQ]:=AlphaInter[Q];
(*Correspondance between flavors*)
flavorList={};
flavorIndex[i_?IntegerQ]:=Module[{k=i},If[i==0,k=21];FirstPosition[flavorList,k][[1]]]
(*Grid limits*)
XMin=0.1;
XMax=1;
QMin=1;
QMax=5;
DeltaMin=0.;
DeltaMax=2.;
(*Info*)
RestInfo={}
ShowInfo[]:=RestInfo
ShowInfo[param_]:=Select[RestInfo,First[#]==param&,1]
(*Interpolations*)
Qranges={};
fInterpolation={};
TMDinterpolation={};

(*-------------------------------------------------------------------------------------*)
(*----Routine for setting the LHAPDF paths---*)

SetLibraryPath[path_]:=If[DirectoryQ[path],
LIBpath=path;
GRIDlist=ListSetsINTERNAL[];
path
,Print["The argument should be a valid path"]]

(*-------------------------------------------------------------------------------------*)
(*------Read the list of PDF inside datadir directory-----*)
ListGrids[]:=GRIDlist

ListSetsINTERNAL[]:=Module[{originalPath,result},
If[LIBpath=="ERROR",Print["Path to grids library is not defined"]; Return[]];
(*Save initial path; Needed to restore the original path after*)
originalPath=Directory[];
(*Go to GRIDlist*)
SetDirectory[LIBpath];
result=Select[FileNames[],DirectoryQ];
(*Restore original settings*)
SetDirectory[originalPath];
result
]

(*-------------------------------------------------------------------------------------*)
(*------Read info file of a PDF set-----*)
Options[ReadInfo]={"Silent"->False};
ReadInfo[name_,OptionsPattern[]]:=Module[{Q0,fileName,fIn,interOrder=3,
AlphaQs,AlphaVal},
If[Not[MemberQ[GRIDlist,name]],Print["The grid "<>name<>" is not present in the current library"];Return[]];
fileName=FileNameJoin[{LIBpath,name,name<>".info"}];
If[Not@FileExistsQ[fileName],Print["The .info file for "<>name<>" is not found"];Return[]];
(*Initial parsing*)
fIn=Import[fileName,"Text"];
(*If all OK, file is imported as a list of {{String},{String},...}*)
Q0=StringSplit[fIn,"\n"];
(*Result here is {String,String,...}, where each string is entry of info-file "key: value"*)
(*Parsing*)
Q0={#[[1]],StringRiffle[Rest[#],":"]}&/@(StringSplit[#,":"]&/@Q0);
(*Sometimes Key-values have spaces infront. Remove*)
Q0=FixedPoint[If[StringStartsQ[#[[1]]," "],{StringDrop[#[[1]],1],#[[2]]},#]&,#]&/@Q0;
(*-----the list QO has all information from .info in the format {{key,value},...} all in strings*)
(*Next we extract some important values and save them to separate variables*)
(* 1. Alpha values*)
(*Some set do not have alphaS values. In this case I do not update the previous interpolation*)
With[{case=Select[Q0,First[#]=="AlphaS_Qs"&]},
If[case=={},
If[Not[OptionValue["Silent"]],
Print["Info-file for "<>name<>" does not contain AlphaS values. AlphaS values are not updated."]];
,
AlphaQs=ImportString[StringReplace[Select[Q0,First[#]=="AlphaS_Qs"&][[1,2]],{"["->"","]"->"",","->"\n"}],"List"];
AlphaVal=ImportString[StringReplace[Select[Q0,First[#]=="AlphaS_Vals"&][[1,2]],{"["->"","]"->"",","->"\n"}],"List"];
(*LHA PDF grid can have repeating values of Q, due to threasholds, which is forbidden in Mathematica. To avoid problems I shift lower-energy value by 10^{-3}. 
In this case the interpolation order for grid should be =1 (linear)*)
AlphaGrid=
With[{alTab=Transpose[{AlphaQs,AlphaVal}]},
(*Some tables have empty values, they parsed as {" "," "}. Delete them*)
DeleteCases[
MapThread[If[#1[[1]]==#2[[1]],interOrder=1; {#1[[1]]-0.0001,#1[[2]]},#1]& ,{Most[alTab],Rest[alTab]}]~Join~{Last[alTab]},
{" "," "}]];
(*Interpolating alpha*)
AlphaInter=Interpolation[AlphaGrid,InterpolationOrder->interOrder];
AlphaS[Q_?NumericQ]:=AlphaInter[Q];
]];
(* 2. Flavor numbering*)
flavorList=ImportString[StringReplace[Select[Q0,First[#]=="Flavors"&][[1,2]],{"["->"","]"->"",","->"\n"}],"List"];
(* 3. Grid limits*)
XMin=First[ImportString[StringReplace[Select[Q0,First[#]=="XMin"&][[1,2]],{"["->"","]"->"",","->"\n"}],"List"]];
XMax=First[ImportString[StringReplace[Select[Q0,First[#]=="XMax"&][[1,2]],{"["->"","]"->"",","->"\n"}],"List"]];
QMin=First[ImportString[StringReplace[Select[Q0,First[#]=="QMin"&][[1,2]],{"["->"","]"->"",","->"\n"}],"List"]];
QMax=First[ImportString[StringReplace[Select[Q0,First[#]=="QMax"&][[1,2]],{"["->"","]"->"",","->"\n"}],"List"]];
(* 4. Rest I save into single info list *)
RestInfo=Rule@@#&/@Select[Q0,Not@MemberQ[{"AlphaS_Qs","AlphaS_Vals","Flavors","XMin","XMax","QMin","QMax"},#[[1]]]&];
]

(*-------------------------------------------------------------------------------------*)
(*Parser for subgrid    (between --- in .data file)*)
ParseSubGrid[g_]:=Module[{xVal,qVal,gVal},
(*Values of x*)
xVal=ImportString[g[[1]],"Table"][[1]];
(*Values of Q*)
qVal=ImportString[g[[2]],"Table"][[1]];
(*Values of f are not needed*)
(*Check the length*)
If[Length[g]-3!=Length[xVal]Length[qVal],Print["Inconsistency in the sizes of subgrids"];Return[]];
(*Parse The rest; This is the fastest method of parsing C++ number coding*)
gVal=ImportString[StringRiffle[Drop[g,3],"\n"],"Table"];
{xVal,qVal,gVal}
]

(*Generate interpolation functions from sub grids*)
GenerateInterpolationFunctions[g_,n_]:=Module[{fNum},
fNum=Length[g[[1,3]]];
Table[
Interpolation[Transpose[{g[[All,1]],g[[All,2]],g[[All,3,i]]}]]
,{i,fNum}]]


(*------Load LHAPDF grid-----*)
Options[LoadLHAPDFgrid]={"Silent"->False};
LoadLHAPDFgrid[name_,opt:OptionsPattern[LoadLHAPDFgrid]]:=LoadLHAPDFgrid[name,0,opt]
LoadLHAPDFgrid[name_,num_?IntegerQ,OptionsPattern[LoadLHAPDFgrid]]:=Module[{str,Q0,QQ0,fileName},
(*Reading of info fills all neccesary common values*)
ReadInfo[name,Silent->OptionValue["Silent"]];
(*Reading grid file*)
fileName=FileNameJoin[{LIBpath,name,name<>"_"<>StringPadLeft[ToString[num],4,"0"]<>".dat"}];
If[Not@FileExistsQ[fileName],Print["The "<>StringPadLeft[ToString[num],4,"0"]<>".dat file is not found"];Return[]];
(*
The grid is parsed as follows:
1. It is read line by line, collecting series between "---" into separate lists.
2. Each series is parsed by ParseSubGrid
3. They collected in a huge grid {x,Q,val}
4. For each grid determine Qmin,Qmax, and make interpolation function FF[flavor,range][x,Q]
*)
Clear[fInterpolation];
(* Part 1 *)
str=OpenRead[fileName];
line="";
Q0={};
QQ0={};
While[line=!=EndOfFile,
line=ReadLine[str];
If[line=="---",AppendTo[Q0,QQ0];QQ0={},AppendTo[QQ0,line]];
]
Close[str];
(* Part 2 *)
Q0=ParseSubGrid/@Rest[Q0];
(* Part 3 *)
Qranges={};
Q0=Table[
With[{xx=Q0[[n,1]],qq=Q0[[n,2]],vv=Q0[[n,3]],qLength=Length[Q0[[n,2]]]},
AppendTo[Qranges,{qq[[1]],qq[[-1]],n}];
Flatten[Table[{xx[[x]],qq[[q]],vv[[(x-1)qLength+q]]},{x,Length[xx]},{q,qLength}],1]]
,{n,Length[Q0]}];
(*Part 4*)
fIntepolation=Table[
GenerateInterpolationFunctions[Q0[[n]],n],{n,Length[Q0]}];
If[Not[OptionValue["Silent"]],
Print["Set "<>name<>"_"<>StringPadLeft[ToString[num],4,"0"]<>" loaded."]];
]

(*-------------------------------------------------------------------------------------*)
(*-------- the function the return xPDF ----------*)
xPDF[x_,q_,f_]:=Module[{ff=flavorIndex[f],qNum},
If[x<XMin || x>XMax,Return[0]];
If[q<QMin || q>QMax,Return[0]];
qNum=SelectFirst[Qranges,q<#[[2]]&][[3]];
fIntepolation[[qNum,ff]][x,q]
]


(*------Load Pavia TMD grid-----*)
Options[LoadTMDLIBgrid]={"Silent"->False,"Format"->"Pavia"};
LoadTMDLIBgrid[name_,opt:OptionsPattern[LoadTMDLIBgrid]]:=LoadTMDLIBgrid[name,0,opt]
LoadTMDLIBgrid[name_,num_?IntegerQ,OptionsPattern[LoadTMDLIBgrid]]:=Module[{str,Q0,QQ0,fileName,
Qlist,xlist,Deltalist,tmdList
},
(*Reading of info fills all neccesary common values*)
ReadInfo[name,Silent->OptionValue["Silent"]];
(*Reading grid file*)
fileName=FileNameJoin[{LIBpath,name,name<>"_"<>StringPadLeft[ToString[num],4,"0"]<>".dat"}];
If[Not@FileExistsQ[fileName],Print["The "<>StringPadLeft[ToString[num],4,"0"]<>".dat file is not found"];Return[]];
(*
The grid is parsed as follows:
1. Pavia grid has 4 lines, each is a list with given keys. The last (4th) line is a list of lists. Read lines and identify the keys. Parse
2. Determine min-max ranges for parameters.
3. The values are collected in a huge grid {Q,x,qT,val}, each is interpolated
*)
Clear[fInterpolation];
(* Part 1 *)
Q0=Import[fileName,"Text"];
Q0=StringSplit[Q0,"\n"];
(*Result here is {String,String,...}, where each string is entry of info-file "key: value"*)
(*Parsing*)
Q0={StringDelete[#[[1]]," "],StringRiffle[Rest[#],":"]}&/@(StringSplit[#,":"]&/@Q0);
(*Identify the keys*)
Qlist=First@ImportString[StringDelete[Select[Q0,First[#]=="Qg"&,1][[1,2]],{"[","]"}],"CSV"];
xlist=First@ImportString[StringDelete[Select[Q0,First[#]=="xg"&,1][[1,2]],{"[","]"}],"CSV"];
Deltalist=First@ImportString[StringDelete[Select[Q0,First[#]=="qToQg"&,1][[1,2]],{"[","]"}],"CSV"];
tmdList=With[{tr={"["->"{","]"->"}",":"->"->","e"->"*10^","E"->"*10^"}},
ToExpression@StringReplace[Select[Q0,First[#]=="TMDs"&,1][[1,2]],tr]];
(*Update flavor list*)
flavorList=First/@tmdList;
tmdList=#[[2]]&/@tmdList;

(* Part 2 *)
QMin=Min[Qlist];
QMax=Max[Qlist];
XMin=Min[xlist];
XMax=Max[xlist];
DeltaMin=Min[Deltalist];
DeltaMax=Max[Deltalist];

(*TMD1=Table[
Flatten[Table[{{Qlist[[i]],xlist[[j]],Deltalist[[k]]},tmdList[[R,i,j,k]]},{i,Length[Qlist]},{j,Length[xlist]},{k,Length[Deltalist]}],2],{R,Length[flavorList]}];*)
(* Part 3 *)
TMDinterpolation=Table[
Interpolation[Flatten[Table[{{Qlist[[i]],xlist[[j]],Deltalist[[k]]},tmdList[[R,i,j,k]]},{i,Length[Qlist]},{j,Length[xlist]},{k,Length[Deltalist]}],2],InterpolationOrder->3]
,{R,Length[flavorList]}];
If[Not[OptionValue["Silent"]],
Print["TMD grid (Pavia convention) "<>name<>"_"<>StringPadLeft[ToString[num],4,"0"]<>" loaded."]];
]

(*-------------------------------------------------------------------------------------*)
(*-------- the function the return xPDF ----------*)
xTMDPDF[x_,qT_,Q_,f_]:=Module[{ff=flavorIndex[f],qNum,delta},
delta=qT/Q;
If[x<XMin || x>XMax,Return[0]];
If[Q<QMin || Q>QMax,Return[0]];
If[delta<DeltaMin || delta>DeltaMax,Return[0]];
TMDinterpolation[[ff]][Q,x,delta]
]


(*-------------------Package initialization--------------------------------------*)

Print[Style["Use SetLibraryPath[path] to specify the path of grids library",FontColor->Red,Bold]]

Print[Style["LHAPDFreader: v.1.1.0",Bold],"  (by Alexey Vladimirov)"];
Print[Style["Avalible commands:",Bold]," SetLibraryPath, ListGrids, ReadInfo, LoadLHAPDFgrid, ShowInfo, AlphaS, xPDF,xTMDPDF"];
End[]
EndPackage[]

