(*
  This script is used to get a analytical formula of \psi_4,
  and save to a file.
*)

(*
	Usage: Please set the fowllowing variables from outside as
    command line arguments, like this:
		  wolframscript -file Teuk_psi4.wl m a L n in|out [filename]
    for example:
      wolframscript -file Teuk_psi4.wl 0 1e-8 1 1 out
    if filename is not given, it's "./psi4_analytical.txt" by default.
*)

(* variable name : description *)

(* Begin of variables *)
(* m : the Teukolsky multiple order *)
(* a : wave amplitude *)
(* L : wave width *)
(* n : the order N in F formula *)
(* End of variables *)


(* Parsing arguments *)

argv = Rest@$ScriptCommandLine;
Check[
  {m, a, L, n} = Interpreter["Number"]/@argv[[;;4]];
  type = argv[[5]];
  If[Length@argv >= 6,
    filename = argv[[6]],
    filename = "./psi4_analytical.txt"
  ],
  Print["Error: Check your command line arguments!"];Abort[]
];

(* Set Assumptions *)
$Assumptions = aa > 0 && r > 0 && Pi > theta > 0 &&
   2 Pi >= phi >= 0 && t ~Element~ Reals &&
   x ~Element~ Reals && y ~Element~ Reals && z ~Element~ Reals;
(* Global Simplify time limit *)
tc = 60

(* Help functions *)
pSimplify[list_] :=
  ArrayReshape[
    ParallelMap[
      Simplify[#, TimeConstraint->tc]&, Flatten@list, Method->"FinestGrained"
    ],
    Dimensions@list
];

psi4analytical[g_, coordinates_, vecn_, vecmb_] :=
  Module[
    {dim = Length[coordinates], t1, t2,
      t3, t4, ig, dg, gamma, dgamma, RiemR, psi4},
    dg = D[g, {coordinates}];
    ig = Inverse[g] // pSimplify;
    t1 =
      AbsoluteTiming[
        gamma = pSimplify[
          ParallelTable[
            1/2 Plus @@ ((ig[[i, #]] (dg[[#, j, k]] + dg[[#, k, j]] -
                    dg[[k, j, #]])) &) /@ Range[dim],
            {i, 1, dim}, {j,1, dim}, {k, 1, dim}
          ]
        ];
      ][[1]];
    Print["Finish computing gamma, time cost:", t1, "s"];
    t2 =
      AbsoluteTiming[
        dgamma = pSimplify[
          D[gamma, {coordinates}]
        ];
      ][[1]];
    Print["Finish computing dgamma, time cost:", t2, "s"];
    t3 =
      AbsoluteTiming[
        RiemR = ParallelTable[
            dgamma[[l, i, k, j]] - dgamma[[l, j, k, i]] +
              Plus @@ ((gamma[[#, k, i]] gamma[[l, j, #]] -\
                gamma[[#, k, j]] gamma[[l, i, #]]) &) /@ Range[dim],
            {i, 1, dim}, {j, 1, dim}, {k, 1, dim}, {l, 1, dim}
        ];
      ][[1]];
    Print["Finish computing RiemR, time cost:" , t3, "s"];
    t4 =
      AbsoluteTiming[
        psi4 = ((((RiemR . g) . vecmb) . vecn) . vecmb) . vecn;
      ][[1]];
    Print["Finish computing psi4, time cost:" , t4, "s"];
    psi4
  ];

(* set metric *)
Switch[m,
  0,
    frr = 2 - 3 Sin[theta]^2;
    frtheta = -3 Cos[theta] Sin[theta];
    frphi = 0;
    fthetatheta1 = 3 Sin[theta]^2;
    fthetatheta2 = -1; fthetaphi = 0;
    fphiphi1 = -fthetatheta1;
    fphiphi2 = 3 Sin[theta]^2 - 1;,
  1,
    frr = Sin[2 theta] Cos[phi];
    frtheta = Cos[2 theta] Cos[phi];
    frphi = -Cos[theta] Sin[phi];
    fthetatheta1 = -frr;
    fthetatheta2 = 0;
    fthetaphi = -Sin[theta] Sin[phi];
    fphiphi1 = -fthetatheta1;
    fphiphi2 = -frr;,
  -1,
    frr = Sin[2 theta] Sin[phi];
    frtheta = Cos[2 theta] Sin[phi];
    frphi = Cos[theta] Cos[phi];
    fthetatheta1 = -frr;
    fthetatheta2 = 0;
    fthetaphi = Sin[theta] Cos[phi];
    fphiphi1 = -fthetatheta1;
    fphiphi2 = -frr;,
  2,
    frr = Sin[theta]^2 Cos[2 phi];
    frtheta = Sin[theta] Cos[theta] Cos[2 phi];
    frphi = -Sin[theta] Sin[2 phi];
    fthetatheta1 = (1 + Cos[theta]^2) Cos[2 phi];
    fthetatheta2 = -Cos[2 phi];
    fthetaphi = Cos[theta] Sin[2 phi];
    fphiphi1 = -fthetatheta1;
    fphiphi2 = Cos[theta]^2 Cos[2 phi];,
  -2,
    frr = Sin[theta]^2 Sin[2 phi];
    frtheta = Sin[theta] Cos[theta] Sin[2 phi];
    frphi = Sin[theta] Cos[2 phi];
    fthetatheta1 = (1 + Cos[theta]^2) Sin[2 phi];
    fthetatheta2 = -Sin[2 phi];
    fthetaphi = -Cos[theta] Cos[2 phi];
    fphiphi1 = -fthetatheta1; fphiphi2 = Cos[theta]^2 Sin[2 phi];,
  _,
    Print["Error: m must be one of {0,1,2,-1,-2}."];Abort[]
]

h = aa {
     {0, 0, 0, 0},
     {0, AA[t, r] frr, BB[t, r] frtheta r, BB[t, r] frphi r Sin[theta]},
     {0, BB[t, r] frtheta r,
        (CC[t, r] fthetatheta1 + AA[t, r] fthetatheta2) r^2,
        (AA[t, r] - 2 CC[t, r]) fthetaphi r^2 Sin[theta]},
     {0, BB[t, r] frphi r Sin[theta],
        (AA[t, r] - 2 CC[t, r]) fthetaphi r^2 Sin[theta],
        (CC[t, r] fphiphi1 + AA [t, r] fphiphi2) r^2 Sin[theta]^2}
    } // Simplify;
g = DiagonalMatrix[{-1, 1, r^2, r^2 Sin[theta]^2}] + h // Simplify;

rule1 = Switch[type,
  "in",
    {
      AA -> Function[{t, r},3 (F''[t + r]/r^3 - (3 F'[t + r])/r^4 + (3 F[t + r])/r^5)],
      BB -> Function[{t, r}, -(-(F'''[t + r]/r^2) + (3 F''[t + r])/r^3 - (6 F'[t + r])/r^4 + (6 F[t + r])/r^5)],
      CC -> Function[{t, r}, 1/4 (F''''[t + r]/r - (2 F'''[t + r])/r^2 + (9 F''[t + r])/r^3 - (21 F'[t + r])/r^4 + (21 F[t + r])/r^5)]
    },
  "out",
    {
      AA -> Function[{t, r}, 3 (F''[t - r]/r^3 + (3 F'[t - r])/r^4 + (3 F[t - r])/r^5)],
      BB -> Function[{t, r}, -(F'''[t - r]/r^2 + (3 F''[t - r])/r^3 + (6 F'[t - r])/r^4 + (6 F[t - r])/r^5)],
      CC -> Function[{t, r}, 1/4 (F''''[t - r]/r + (2 F'''[t - r])/r^2 + (9 F''[t - r])/r^3 + (21 F'[t - r])/r^4 + (21 F[t - r])/r^5)]
    },
  _,
    Print["Error: check wave type, can only be in or out"];Abort[]
]

rule2 = F -> Function[L^5 (#/L)^n E^(-(#^2/L^2))]

(* set tetrad *)
vecn = 1/2 {1, -1, 0, 0};
vecmb = 1/(Sqrt[2] r) {0, 0, 1, -(I/Sin[theta])};

(* Calculate psi4 *)
psi4 = psi4analytical[g, {t, r, theta, phi}, vecn, vecmb];
psi4 = psi4 /. {aa->a};
psi4 = psi4 /. rule1;
psi4 = psi4 /. rule2;

psi4 = psi4 /. {t->0};
(* psi4 = Simplify[psi4, TimeConstraint->tc]; *)

Put[psi4, filename]