(* ::Package:: *)

(* :Name: pcGetParam` *)
(* :Author: Hongzhe Zhou, Stockholm, 2021*)
(* :Version: 0.1 *)
(* :Summary:
    Defines getParam[] which produces various useful parameters.
*)


BeginPackage["pcGetParam`","pcReadBasic`","pcRead1D`"]


(* ::Chapter:: *)
(*Usage messages*)


pcReload::usage="Reload all pc`variables that use Set[].";

getParam::usage="getParam[sim,var] gives various parameters of sim.
Input:
  sim: String. Directory of the simulation folder
  var: String. Specifies the needed variable
Example:
  getParam[dir,\"ReN\"] returns the Reynolds number.";

LuNspec::usage="LuNspec[sim,i:1,m:All,Options] computes time-dependent Lundquist numbers from power_mag.dat.
The definition B*l^(n-1)/eta(t), where:
B is the rms magnetic energy exlucding the k=0 mode;
The length scale l is the energy-weighted value of 1/k;
n is the power in the diffusion operator del^n (normal diffusion n=2).
eta(t) is the time-dependent resistivity.
Input:
  sim: String. Directory of the simulation folder
  i: Integer. Optional. Specify to use the ith IRESISTIVITY.
  m: Optional. Take only the first m snapshots.
Options:
  \"l1/n\": By default ->False. If ->True, then compute Lu^(1/n).
Output:
  A time series of the Lundquist number."

paramMatrix::usage="paramMatrix[sim,{file1,var1},{file2,var2},...] gives in the 
MatrixForm values of {var1,var2,...} that can be retrived by readParamNml."

infoMatrix::usage="infoMatrix[sim,var1,var2,...] gives in the MatrixForm values of
{var1,var2,...} that can be retrieved by getParam[sim,\"var\"].";


Begin["`Private`"]


(* ::Chapter:: *)
(*Functions*)


(* ::Section:: *)
(*Private variables (Set[] is only allowed here!)*)


(*some frequently read variables; better to dynamically define*)
(*when data file is updated, pcReload[] is needed to clear existing definitions*)
pcReload[]:=Module[{},
  Clear[urms,urmskf];
  urms[sim_]:=urms[sim]=With[{u=readTS[sim,"urms"]},u[[-Round[Length[u]/3];;-1]]//Mean];
  (*urms of modes with k\[GreaterEqual]kf*)
  urmskf[sim_]:=urmskf[sim]=With[{file=sim<>"/data/power_kin.dat"},
    If[FileExistsQ[file],
      Sqrt[2*Total[(read1D2Scale[sim,"power_kin.dat"])//Last//Last]],
      "NaN"
    ]];
  
  Clear[hrms];
  hrms[sim_]:=hrms[sim]=With[{h=readTS[sim,"oum"]},h[[-Round[Length[h]/3];;-1]]//RootMeanSquare];
  
  Clear[kf];
  kf[sim_]:=kf[sim]=
    If[readParamNml[sim,"start.in","lforcing"] && readParamNml[sim,"run.in","FORCE"]!=0.,
      If[Length[#[[1]]]==1,#[[2,1]],#[[1,2]]]&@Import[FileNameJoin[{sim,"k.dat"}]],
      If[MemberQ[readParamNml[sim,"run.in","IFORCING_CONT"],x_/;StringStartsQ[x,"'GP_TC13"]],
        readParamNml[sim,"run.in","KMIN_GP"],
        "No forcing"
      ]
    ];
  
  Clear[ted,tedkf];
  ted[sim_]:=ted[sim]=1/(kf[sim]*urms[sim]);	
  tedkf[sim_]:=tedkf[sim]=1/(kf[sim]*urmskf[sim]);
  
  Clear[nu,ivisc];
  nu[sim_]:=nu[sim]=readParamNml[sim,"run.in","NU"];
  ivisc[sim_,i_Integer:1]:=ivisc[sim,i]=readParamNml[sim,"run.in","IVISC"][[i]];
  
  Clear[eta,etaTFM,iresis];
  eta[sim_]:=eta[sim]=
    If[readParamNml[sim,"start.in","lmagnetic"],
      readParamNml[sim,"run.in","ETA"],
      "No magnetic"
    ];
  etaTFM[sim_]:=readParamNml[sim,"run.in","ETATEST"];
  iresis[sim_,i_Integer:1]:=iresis[sim,i]=ires=readParamNml[sim,"run.in","IRESISTIVITY"][[i]];
  
  Clear[kappaTFM];
  kappaTFM[sim_]:=readParamNml[sim,"run.in","KAPPATEST"];
  
  Clear[PrM];
  PrM[sim_]:=PrM[sim]=nu[sim]/eta[sim];
  
  Clear[omega];
  omega[sim_]:=omega[sim]=Module[{os,or},
    os=readParamNml[sim,"start.in","OMEGA"];
    or=readParamNml[sim,"run.in","OMEGA"];
    Which[
      os==or, Return[os],
      os==0 && or!=0, Return[or],
      os!=0 && or==0, Return[os],
      os!=0 && or!=0 && os!=or, 
        Print["Warning: different OMEGA in start.in and run.in; returning the latter."]];Return[or]
    ]
]
pcReload[]


(* ::Section:: *)
(*getParam*)


(*directly readables*)
getParam[sim_,"folder"]:=StringJoin["...",StringTake[sim,-7]]
getParam[sim_,"id"]:=
  Module[{f},
    If[And[FileExistsQ[FileNameJoin[{sim,"sim_comment"}]],(f=Import[FileNameJoin[{sim,"sim_comment"}]])=!=""],
      First@If[Length[f]==2,Flatten/@f,StringSplit[f,"\n"]],
      "NaN"
    ]
  ]				
getParam[sim_,"comment"]:=
  Module[{f},
    If[And[FileExistsQ[FileNameJoin[{sim,"sim_comment"}]],(f=Import[FileNameJoin[{sim,"sim_comment"}]])=!=""],
      Rest@If[Length[f]==2,Flatten/@f,StringSplit[f,"\n"]],
      "No comment."
    ]
  ]

getParam[sim_,"Nx"]:=readDim[sim]["nx"]
getParam[sim_,"Ny"]:=readDim[sim]["ny"]
getParam[sim_,"Nz"]:=readDim[sim]["nz"]
getParam[sim_,"Nxyz"]:=readDim[sim]/@{"nx","ny","nz"}
getParam[sim_,"Lxyz"]:=readParamNml[sim,"start.in","LXYZ"]

getParam[sim_,"urms"]:=urms[sim]
getParam[sim_,"urmskf"]:=urmskf[sim]
getParam[sim_,"hrms"]:=hrms[sim]
getParam[sim_,"kf"]:=kf[sim]
getParam[sim_,"nu"]:=nu[sim]
getParam[sim_,"eta"]:=eta[sim]

(*dimensionless numbers*)
(*Reynolds numbers*)
getParam[sim_,"ReN"]:=urms[sim]/kf[sim]/nu[sim]
getParam[sim_,"ReN",k2_]:=urms[sim]/k2/nu[sim] (*supply kf by hand*)
getParam[sim_,"ReNkf"]:=urmskf[sim]/kf[sim]/nu[sim]

getParam[sim_,"ReM"]:=urms[sim]/kf[sim]/eta[sim]
getParam[sim_,"ReM",k2_]:=urms[sim]/k2/eta[sim] (*supply kf by hand*)
getParam[sim_,"PrM"]:=PrM[sim]
getParam[sim_,"PrMTFM"]:=nu[sim]/etaTFM[sim]

(*hyper Reynolds numbers*)
getParam[sim_,"nHyperVisc",i_Integer:1]:=Switch[ivisc[sim,i],
  "'nu-const'",          1,
  "'hyper2-simplified'", 2,
  "'hyper3-simplified'", 3,
  _,                     Print[sim,": unknown ivisc=",ivisc[sim,i]];Return[$Failed]
];
getParam[sim_,"ReNn",i_Integer:1]:=Switch[ivisc[sim,i],
  "'nu-const'",          getParam[sim,"ReN"],
  "'hyper2-simplified'", urms[sim]/kf[sim]^3/readParamNml[sim,"run.in","NU_HYPER2"],
  "'hyper3-simplified'", urms[sim]/kf[sim]^5/readParamNml[sim,"run.in","NU_HYPER3"],
  _,                     Print[sim,": unknown ivisc=",ivisc[sim,i]];Return[$Failed]
];

getParam[sim_,"nHyperEta",i_Integer:1]:=Switch[iresis[sim,i],
  "'eta-const'",   1,
  "'eta-tdep'",    1,
  "'hyper2'",      2,
  "'hyper2-tdep'", 2,
  "'hyper3'",      3,
  "'hyper3-tdep'", 3,
  _,               Print[sim,": unknown iresistivity=",iresis[sim,i]];Return[$Failed]
];

getParam[sim_,"ReMn",i_Integer:1]:=Switch[iresis[sim,i],
  "'eta-const'", getParam[sim,"ReM"],
  "'hyper2'",    urms[sim]/kf[sim]^3/readParamNml[sim,"run.in","ETA_HYPER2"],
  "'hyper3'",    urms[sim]/kf[sim]^5/readParamNml[sim,"run.in","ETA_HYPER3"],
  _,             Print[sim,": unknown iresistivity=",iresis];Return[$Failed]
]

(*rotation and shear*)
getParam[sim_,"Ro"]:=If[omega[sim]==0,"No rotation",kf[sim]*urms[sim]/2/omega[sim]]
getParam[sim_,"Ro",k2_]:=If[omega[sim]==0,"No rotation",k2*urms[sim]/2/omega[sim]]
getParam[sim_,"Rokf"]:=If[omega[sim]==0,"No rotation",kf[sim]*urmskf[sim]/2/omega[sim]]

getParam[sim_,"Co"]:=If[omega[sim]==0,0,1/getParam[sim,"Ro"]]
getParam[sim_,"Co",k2_]:=If[omega[sim]=0,0,1/getParam[sim,"Ro",k2]]
getParam[sim_,"Cokf"]:=If[omega[sim]==0,0,1/getParam[sim,"Rokf"]]

getParam[sim_,"Sh"]:=
  If[readParamNml[sim,"start.in","lshear"],
    readParamNml[sim,"run.in","SSHEAR"]/urms[sim]/kf[sim],
    0
  ]
getParam[sim_,"Shkf"]:=
  If[readParamNml[sim,"start.in","lshear"],
    readParamNml[sim,"run.in","SSHEAR"]/urmskf[sim]/kf[sim],
    0
  ]
getParam[sim_,"-Sh"]:=-getParam[sim,"Sh"]
getParam[sim_,"-Shkf"]:=-getParam[sim,"Shkf"]
getParam[sim_,"Sh",k2_]:=
  If[readParamNml[sim,"start.in","lshear"],
    readParamNml[sim,"run.in","SSHEAR"]/urms[sim]/k2,
    0
  ]
getParam[sim_,"-Sh",k2_]:=-getParam[sim,"Sh",k2]
getParam[sim_,"qshear"]:=
  Module[{s=readParamNml[sim,"run.in","SSHEAR"],o=readParamNml[sim,"start.in","OMEGA"]},
    If[s===$Failed,s=0];
    If[s==0&&o==0,"/"//Return];
    If[s!=0&&o==0,"No rotatoin"//Return];
    Rationalize[-s/o]/.Rational[x_,y_]:>ToString[x]<>"/"<>ToString[y]
]

getParam[sim_,"ScTFM"]:=nu[sim]/kappaTFM[sim]

(*secondaries*)
(*eddy turnover time*)
getParam[sim_,"ted"]:=ted[sim]
getParam[sim_,"ted",k2_]:=2\[Pi]/urms[sim]/k2
getParam[sim_,"tedkf"]:=tedkf[sim]
(*diffusive timescales*)
getParam[sim_,"teta1",k1_:1]:=1/(eta[sim]*k1^2)
(**)
getParam[sim_,"knu"]:=kf[sim]*getParam[sim,"ReN"]^(3/4)
getParam[sim_,"knu",k2_]:=k2*getParam[sim,"ReN",k2]^(3/4) (*supply kf by hand*)
getParam[sim_,"knumax"]:=(Max[readTS[sim,"urms"]]/nu[sim])^(3/4)*kf[sim]^(1/4)
getParam[sim_,"knumax",k2_]:=(Max[readTS[sim,"urms"]]/nu[sim])^(3/4)*k2^(1/4) (*supply kf by hand*)
getParam[sim_,"keta"]:=
  If[eta[sim]=="No magnetic","No magnetic",
    getParam[sim,"knu"]/If[PrM[sim]>=1,PrM[sim]^(-1/2),PrM[sim]^(-3/4)]
  ]
getParam[sim_,"keta",k2_]:=
  If[eta[sim]=="No magnetic","No magnetic",
    getParam[sim,"knu",k2]/If[PrM[sim]>=1,PrM[sim]^(-1/2),PrM[sim]^(-3/4)]
  ]
getParam[sim_,"ketamax"]:=
  If[eta[sim]=="No magnetic","No magnetic",
    getParam[sim,"knumax"]/If[PrM[sim]>=1,PrM[sim]^(-1/2),PrM[sim]^(-3/4)]
  ]
getParam[sim_,"ketamax",k2_]:=
  If[eta[sim]=="No magnetic","No magnetic",
    getParam[sim,"knumax",k2]/If[PrM[sim]>=1,PrM[sim]^(-1/2),PrM[sim]^(-3/4)]
  ]
getParam[sim_,"kRo"]:=If[omega[sim]==0,"No rotation",kf[sim]*getParam[sim,"Ro"]^(-3/2)]
getParam[sim_,"kRo",k2_]:=If[omega[sim]==0,"No rotation",k2*getParam[sim,"Ro"]^(-3/2)]


(* ::Section:: *)
(*Dimensionless parameters from spectra files*)


Options[LuNspec]={"l1/n"->False};
LuNspec[sim_,i_Integer:1,m_:All,OptionsPattern[]]:=Module[{t,spec,Eb,k,l,n,eta,ires,lnorm,toffset,t0,exp},
  (*error messages*)
  LuNspec::nofile="power_mag.dat not found from `1`.";
  LuNspec::nores="Unfamiliar iresistivity for `1`: `2`";
  
  (**)
  If[!FileExistsQ[sim<>"/data/power_mag.dat"],Message[LuNspec::nofile,sim];Return[$Failed]];
  {t,spec}=Take[#,m]&/@read1D[sim,"power_mag.dat"];
  Eb=2Total/@spec;
  spec=Rest/@spec;
  k=Range[Length[spec[[1]]]];
  l=Total[1/k*#]/Total[#]&/@spec;
  
  {eta,n}=Switch[ires=readParamNml[sim,"run.in","IRESISTIVITY"][[i]],
    "'eta-const'",  {readParamNml[sim,"run.in","ETA"],       2},
    "'eta-tdep'",   {readParamNml[sim,"run.in","ETA"],       2},
    "'hyper2'",     {readParamNml[sim,"run.in","ETA_HYPER2"],4},
    "'hyper2-tdep'",{readParamNml[sim,"run.in","ETA"],       4},
    "'hyper3'",     {readParamNml[sim,"run.in","ETA_HYPER3"],6},
    "'hyper3-tdep'",{readParamNml[sim,"run.in","ETA"],       6},
    _,               Message[LuNspec::nores,sim,ires];Return[$Failed]
  ];
  
  lnorm=readParamNml[sim,"run.in","LRESI_ETA_TDEP_T0_NORM"];
  toffset=readParamNml[sim,"run.in","ETA_TDEP_TOFFSET"];
  t0=readParamNml[sim,"run.in","ETA_TDEP_T0"];
  exp=readParamNml[sim,"run.in","ETA_TDEP_EXPONENT"];
  eta=eta*If[lnorm,Max[(#-toffset)/t0,1]^exp&,Max[#-toffset,t0]^exp&]/@t;
  
  If[OptionValue["l1/n"],
    {t,(Sqrt[Eb]*l^(n-1)/eta)^(1/n)}//Transpose,
    {t,Sqrt[Eb]*l^(n-1)/eta}//Transpose
  ]
]


(* ::Section:: *)
(*Handy functions*)


paramMatrix[sim_,vars__]:=Module[{files,prms,data},
  {files,prms}=Transpose@{vars};
  data=MapThread[readParamNml[sim,#1,#2]&,{files,prms}];
  MatrixForm[Transpose@{prms,data}]
]

infoMatrix[sim_,vars__]:=MatrixForm[Transpose[{{vars},(getParam@@Flatten[{sim,#}])&/@{vars}}]]


(* ::Chapter:: *)
(*End*)


End[]


Protect[
  pcReload,getParam,
  LuNspec,
  paramMatrix,infoMatrix
]


EndPackage[]