
SetEnhancedTimes[False];

(******************************************************************************)
(* Options *)
(******************************************************************************)

createCode[derivOrder_, useJacobian_, evolutionTimelevels_, addMatter_] :=
Module[{},

prefix = "ML_";
suffix =
  ""
  <> If [useJacobian, "_MP", ""]
  <> If [derivOrder!=4, "_O" <> ToString[derivOrder], ""]
  (* <> If [evolutionTimelevels!=3, "_TL" <> ToString[evolutionTimelevels], ""] *)
  (* <> If [addMatter==1, "_M", ""] *)
  ;

ADMQuantities = prefix <> "ADMQuantities" <> suffix;

(******************************************************************************)
(* Derivatives *)
(******************************************************************************)

KD = KroneckerDelta;

derivatives =
{
  PDstandardNth[i_]    -> StandardCenteredDifferenceOperator[1,derivOrder/2,i],
  PDstandardNth[i_,i_] -> StandardCenteredDifferenceOperator[2,derivOrder/2,i],
  PDstandardNth[i_,j_] -> StandardCenteredDifferenceOperator[1,derivOrder/2,i] *
                          StandardCenteredDifferenceOperator[1,derivOrder/2,j],
  
(* PD: These come from my mathematica notebook
   "Upwind-Kranc-Convert.nb" that converts upwinding finite
   differencing operators generated by
   StandardUpwindDifferenceOperator into this form *)

  Switch[derivOrder,
    2,
    PDupwindNth[1] -> (dir[1]*(-3 + 4*shift[1]^dir[1] - 
                         shift[1]^(2*dir[1])))/(2*spacing[1]),
    4,
    PDupwindNth[1] -> (dir[1]*(-10 - 3/shift[1]^dir[1] + 18*shift[1]^dir[1] -
                         6*shift[1]^(2*dir[1]) + 
                         shift[1]^(3*dir[1])))/(12*spacing[1]),
    6,
    PDupwindNth[1] -> (dir[1]*(-35 + 2/shift[1]^(2*dir[1]) - 
                         24/shift[1]^dir[1] + 80*shift[1]^dir[1] -
                         30*shift[1]^(2*dir[1]) + 8*shift[1]^(3*dir[1]) -
                         shift[1]^(4*dir[1])))/(60*spacing[1]),
    8,
    PDupwindNth[1] -> (dir[1]*(-378 - 5/shift[1]^(3*dir[1]) +
                         60/shift[1]^(2*dir[1]) - 420/shift[1]^dir[1] +
                         1050*shift[1]^dir[1] - 420*shift[1]^(2*dir[1]) + 
                         140*shift[1]^(3*dir[1]) - 30*shift[1]^(4*dir[1]) +
                         3*shift[1]^(5*dir[1])))/(840*spacing[1])],
  Switch[derivOrder,
    2,
    PDupwindNth[2] -> (dir[2]*(-3 + 4*shift[2]^dir[2] - 
                         shift[2]^(2*dir[2])))/(2*spacing[2]),
    4,
    PDupwindNth[2] -> (dir[2]*(-10 - 3/shift[2]^dir[2] + 18*shift[2]^dir[2] -
                         6*shift[2]^(2*dir[2]) + 
                         shift[2]^(3*dir[2])))/(12*spacing[2]),
    6,
    PDupwindNth[2] -> (dir[2]*(-35 + 2/shift[2]^(2*dir[2]) - 
                         24/shift[2]^dir[2] + 80*shift[2]^dir[2] -
                         30*shift[2]^(2*dir[2]) + 8*shift[2]^(3*dir[2]) -
                         shift[2]^(4*dir[2])))/(60*spacing[2]),
    8,
    PDupwindNth[2] -> (dir[2]*(-378 - 5/shift[2]^(3*dir[2]) +
                         60/shift[2]^(2*dir[2]) - 420/shift[2]^dir[2] +
                         1050*shift[2]^dir[2] - 420*shift[2]^(2*dir[2]) + 
                         140*shift[2]^(3*dir[2]) - 30*shift[2]^(4*dir[2]) +
                         3*shift[2]^(5*dir[2])))/(840*spacing[2])],
  Switch[derivOrder,
    2,
    PDupwindNth[3] -> (dir[3]*(-3 + 4*shift[3]^dir[3] - 
                         shift[3]^(2*dir[3])))/(2*spacing[3]),
    4,
    PDupwindNth[3] -> (dir[3]*(-10 - 3/shift[3]^dir[3] + 18*shift[3]^dir[3] -
                         6*shift[3]^(2*dir[3]) + 
                         shift[3]^(3*dir[3])))/(12*spacing[3]),
    6,
    PDupwindNth[3] -> (dir[3]*(-35 + 2/shift[3]^(2*dir[3]) - 
                         24/shift[3]^dir[3] + 80*shift[3]^dir[3] -
                         30*shift[3]^(2*dir[3]) + 8*shift[3]^(3*dir[3]) -
                         shift[3]^(4*dir[3])))/(60*spacing[3]),
    8,
    PDupwindNth[3] -> (dir[3]*(-378 - 5/shift[3]^(3*dir[3]) +
                         60/shift[3]^(2*dir[3]) - 420/shift[3]^dir[3] +
                         1050*shift[3]^dir[3] - 420*shift[3]^(2*dir[3]) + 
                         140*shift[3]^(3*dir[3]) - 30*shift[3]^(4*dir[3]) +
                         3*shift[3]^(5*dir[3])))/(840*spacing[3])],

  (* TODO: make these higher order stencils *)
  PDonesided[1] -> dir[1] (-1 + shift[1]^dir[1]) / spacing[1],
  PDonesided[2] -> dir[2] (-1 + shift[2]^dir[2]) / spacing[2],
  PDonesided[3] -> dir[3] (-1 + shift[3]^dir[3]) / spacing[3]
};

PD  = PDstandardNth;
PDu = PDupwindNth;
PDo = PDonesided;



(******************************************************************************)
(* Tensors *)
(******************************************************************************)

(* Register the tensor quantities with the TensorTools package *)
Map [DefineTensor,
     {normal, tangentA, tangentB, dir,
      nn, nu, nlen, nlen2, su, vg,
      g, K, alpha, beta, detg, gu, G, R, trR, Km, trK,
      phi, gt, At, Xt, Xtn, alpha, A, beta, B, Atm, Atu, trA, Ats, trAts,
      ephi, detgt, gtu, ddetgt, dgtu, ddgtu, Gtl, Gtlu, Gt,
      Rt,
      Madm, Jadm,
      T00, T0, T, rho, S,
      x, r}];

Map [AssertSymmetricIncreasing,
     {gt[la,lb], At[la,lb], Rt[la,lb], T[la,lb]}];
AssertSymmetricIncreasing [G[ua,lb,lc], lb, lc];
AssertSymmetricIncreasing [Gtl[la,lb,lc], lb, lc];
AssertSymmetricIncreasing [Gt[ua,lb,lc], lb, lc];
Map [AssertSymmetricDecreasing,
     {gu[ua,ub], gtu[ua,ub], Atu[ua,ub]}];
AssertSymmetricDecreasing [dgtu[ua,ub,lc], ua, ub];
AssertSymmetricDecreasing [ddgtu[ua,ub,lc,ld], ua, ub];
AssertSymmetricIncreasing [ddgtu[ua,ub,lc,ld], lc, ld];

(* Use the CartGrid3D variable names *)
x1=x; x2=y; x3=z;

(* Use the TmunuBase variable names *)
T00=eTtt;
T01=eTtx; T02=eTty; T03=eTtz;
T11=eTxx; T12=eTxy; T22=eTyy; T13=eTxz; T23=eTyz; T33=eTzz;

(******************************************************************************)
(* Expressions *)
(******************************************************************************)

detgtExpr = Det [MatrixOfComponents [gt[la,lb]]];

(******************************************************************************)
(* Groups *)
(******************************************************************************)

SetGroupTimelevels[g_,tl_] = Join[g, {Timelevels -> tl}];

evolvedGroups = {};
evaluatedGroups =
  {SetGroupName [CreateGroupFromTensor [Madm    ], prefix <> "Madm"],
   SetGroupName [CreateGroupFromTensor [Jadm[la]], prefix <> "Jadm"]};
(*
evolvedGroups =
  {SetGroupName [CreateGroupFromTensor [Madm    ], prefix <> "Madm"],
   SetGroupName [CreateGroupFromTensor [Jadm[la]], prefix <> "Jadm"]};
evaluatedGroups = {};
*)

declaredGroups = Join [evolvedGroups, evaluatedGroups];
declaredGroups = Map [AddGroupExtra [#, Timelevels -> evolutionTimelevels] &, declaredGroups];

declaredGroupNames = Map [First, declaredGroups];



extraGroups =
  {{"grid::coordinates", {x, y, z, r}},
   {"ADMBase::metric",   {gxx, gxy, gxz, gyy, gyz, gzz}},
   {"ADMBase::curv",     {kxx, kxy, kxz, kyy, kyz, kzz}},
   {"ADMBase::lapse",    {alp}},
   {"ADMBase::dtlapse",  {dtalp}},
   {"ADMBase::shift",    {betax, betay, betaz}},
   {"ADMBase::dtshift",  {dtbetax, dtbetay, dtbetaz}},
   {"ML_BSSN::ML_log_confac", {phi}},
   {"ML_BSSN::ML_metric",     {gt11, gt12, gt13, gt22, gt23, gt33}},
   {"ML_BSSN::ML_Gamma",      {Xt1, Xt2, Xt3}},
   {"ML_BSSN::ML_trace_curv", {trK}},
   {"ML_BSSN::ML_curv",       {At11, At12, At13, At22, At23, At33}},
   {"ML_BSSN::ML_lapse",      {alpha}},
   {"ML_BSSN::ML_shift",      {beta1, beta2, beta3}},
   {"TmunuBase::stress_energy_scalar", {eTtt}},
   {"TmunuBase::stress_energy_vector", {eTtx, eTty, eTtz}},
   {"TmunuBase::stress_energy_tensor", {eTxx, eTxy, eTxz, eTyy, eTyz, eTzz}}
};



groups = Join [declaredGroups, extraGroups];

(******************************************************************************)
(* ADM quantities *)
(******************************************************************************)

ADMQuantitiesCalc =
{
  Name -> ADMQuantities<>"_calculate",
  Schedule -> Automatic,
  Where -> Interior,
  After -> "MoL_PostStep",
  Shorthands -> {detgt, gtu[ua,ub], dgtu[ua,ub,lc],
                 Gtl[la,lb,lc], Gtlu[la,lb,uc], Gt[ua,lb,lc],
                 Xtn[ua], Rt[la,lb], trRt,
                 Atm[ua,lb],
                 ephi,
                 rho, S[la]},
  Equations -> 
  {
    detgt          -> 1 (* detgtExpr *),
    gtu[ua,ub]     -> 1/detgt detgtExpr MatrixInverse [gt[ua,ub]],
    dgtu[ua,ub,lc] -> - gtu[ua,ud] gtu[ub,ue] PD[gt[ld,le],lc],
    Gtl[la,lb,lc]  -> 1/2
                      (PD[gt[lb,la],lc] + PD[gt[lc,la],lb] - PD[gt[lb,lc],la]),
    Gtlu[la,lb,uc] -> gtu[uc,ud] Gtl[la,lb,ld],
    Gt[ua,lb,lc]   -> gtu[ua,ud] Gtl[ld,lb,lc],

    (* The conformal connection functions calculated from the conformal metric,
       used instead of Xt where no derivatives of Xt are taken *)
    Xtn[ui] -> gtu[uj,uk] Gt[ui,lj,lk],

    (* PRD 62, 044034 (2000), eqn. (18) *)
    Rt[li,lj] -> - (1/2) gtu[ul,um] PD[gt[li,lj],ll,lm]
                 + (1/2) gt[lk,li] PD[Xt[uk],lj]
                 + (1/2) gt[lk,lj] PD[Xt[uk],li]
                 + (1/2) Xtn[uk] Gtl[li,lj,lk]
                 + (1/2) Xtn[uk] Gtl[lj,li,lk]
                 + (+ Gt[uk,li,ll] Gtlu[lj,lk,ul]
                    + Gt[uk,lj,ll] Gtlu[li,lk,ul]
                    + Gt[uk,li,ll] Gtlu[lk,lj,ul]),

    trRt -> gtu[ua,ub] Rt[la,lb],

    ephi -> IfThen [conformalMethod, 1/Sqrt[phi], Exp[phi]],

    Atm[ua,lb] -> gtu[ua,uc] At[lc,lb],

    (* Matter terms *)

    (* rho = n^a n^b T_ab *)
    rho -> addMatter
           1/alpha^2 (T00 - 2 beta[ui] T0[li] + beta[ui] beta[uj] T[li,lj]),

    (* S_i = -p^a_i n^b T_ab, where p^a_i = delta^a_i + n^a n_i *)
    S[li] -> addMatter (-1/alpha (T0[li] - beta[uj] T[li,lj])),

    (* ADM quantities *)
    (* See PRD 66, 084026 (2002) *)

    Madm -> 1/(16 Pi)
            (+ ephi^5 (+ 16 Pi addMatter rho
                       + Atm[ua,lb] Atm[ub,la]
                       - 2/3 trK^2)
             - gtu[ua,ub] Gt[uc,la,ld] Gtlu[lb,lc,ud]
             + (1 - ephi) trRt),

    Jadm[li] -> 1/(16 Pi) Eps[li,lj,uk] ephi^6
                (+ 2 Atm[uj,lk]
                 + 16 Pi x[uj] S[lk]
                 + 4/3 x[uj] PD[trK,lk]
                 - x[uj] dgtu[ul,um,lk] At[ll,lm])
  }
};

(******************************************************************************)
(* Implementations *)
(******************************************************************************)

inheritedImplementations =
  Join[{"ADMBase"},
       {prefix <> "BSSN" <> suffix},
       If [addMatter!=0, {"TmunuBase"}, {}]];

(******************************************************************************)
(* Parameters *)
(******************************************************************************)

intParameters =
{
  {
    Name -> conformalMethod,
    Description -> "Treatment of conformal factor",
    AllowedValues -> {{Value -> "0", Description -> "phi method"},
                      {Value -> "1", Description -> "W method"}},
    Default -> 0
  }
};

(******************************************************************************)
(* Construct the thorns *)
(******************************************************************************)

calculations =
{
  ADMQuantitiesCalc
};

CreateKrancThornTT [groups, ".", ADMQuantities,
  Calculations -> calculations,
  DeclaredGroups -> declaredGroupNames,
  PartialDerivatives -> derivatives,
  EvolutionTimelevels -> evolutionTimelevels,
  UseJacobian -> True,
  UseLoopControl -> True,
  InheritedImplementations -> inheritedImplementations,
  IntParameters -> intParameters
];

];



(******************************************************************************)
(* Options *)
(******************************************************************************)

(* These are the arguments to createCode:
   - derivative order: 2, 4, 6, 8, ...
   - useJacobian: False or True
   - timelevels: 2 or 3
     (keep this at 3; this is better chosen with a run-time parameter)
   - matter: 0 or 1
     (matter seems cheap; it should be always enabled)
*)

(* createCode[2, False, 3, 1]; *)
createCode[4, False, 3, 1];
(* createCode[4, True,  3, 1]; *)
