#!/usr/bin/env python3
"""Late-start initial data for the GRChombo PBHCosmo example from the exact LTB solution (e = 0 test, plan 6.12 stage 1).

At t_0 = f0 t_C(0) the growing-mode LTB solution for zeta = mu exp(-k^2 r^2/6) (pbhgr.ltb, Yoo units k = 1, H_i = 5)
is written in the synchronous comoving slice: dl^2 = R'^2/(1+2E) dr^2 + R^2 dOmega^2, K^r_r = -R'dot/R', K^th_th = -Rdot/R,
rho = m'/(4 pi R^2 R'), u^i = 0.  The slice is mapped to the isotropic radius varrho (conformally flat, psi^2 = R/varrho,
varrho -> a(t_0) r far away, so the grid coordinate is the physical radius at t_0 and chi -> 1 in the FLRW region).
Matter proxy: real scalar field with phi = 0, Pi = sqrt(2 rho) at t_0 (rho_sf = Pi^2/2 exactly, J_i = 0 exactly), so the
LTB geometry satisfies the constraints exactly.  De Broglie map: q = m/H_ent = 1.5 m t_H, sigma_ent = sqrt6/q.

Output (--out DIR): ltb_table.dat (uniform varrho grid: varrho chi K Arr Pi), params.txt for the example, setup.json.
"""
from __future__ import annotations
import argparse, json, os, sys
import numpy as np
from scipy.interpolate import CubicSpline
from scipy.integrate import cumulative_trapezoid

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".."))
from pbhgr.ltb import LTBGaussian

PARAMS = """# PBHCosmo (GRChombo) — e = 0 late start from LTB, generated by ltb_to_grchombo.py
# mu = {mu}, t_0 = {f0} t_C(0) = {t0_tH:.4f} t_H, a(t_0) = {a0:.3f}, m t_H = {mtH}, q = {q:.1f}, sigma_ent = {sigma:.4f}
# Yoo units k = 1: t_H = {tH:.4f}, t_C(0) = {tC0:.4f}, M_H(t_k) = 0.75 t_H = {MH:.3f}; physical r_m at t_0 = {rm_phys:.3f}
# GRChombo time t = physical time - t_0; expected first AH (1D CMC map, cold): t = {t_hor_exp:.2f} (1.1203 t_C(0) - t_0)
verbosity = 0
chk_prefix = PBH_
plot_prefix = PBHp_
checkpoint_interval = {chk_int}
plot_interval = {plot_int}
num_plot_vars = 7
plot_vars = chi lapse K rho Pi Ham Mom
hdf5_subpath = "hdf5"
pout_subpath = "pout"
data_subpath = "data"
print_progress_only_to_rank_0 = 1

G_Newton = 1.
scalar_mass = {m:.8e}
ltb_table = {table}
ltb_rho_far = {rho_far:.10e}
lineout_num_points = {lineout}
tagging_center = {c:.6f} {c:.6f} {c:.6f}
tagging_radii = {tag_radii}
tagging_rho_thr = {tag_rho}
diag_interval = {diag_int}
gauge_K_ref = {gauge_K_ref}

N_full = {N}
L_full = {L:.6f}
max_level = {max_level}
regrid_interval = {regrid}
regrid_threshold = {regrid_thr}
max_box_size = {box}
min_box_size = {box}
tag_buffer_size = 3

isPeriodic = 1 1 1
vars_parity            = 0 0 4 6 0 5 0    0 0 4 6 0 5 0    0 1 2 3    0 1 2 3 1 2 3    0 0 0
vars_parity_diagnostic = 0 1 2 3 0 0 0 0 0
num_nonzero_asymptotic_vars = 5
nonzero_asymptotic_vars = chi h11 h22 h33 lapse
nonzero_asymptotic_values = 1.0 1.0 1.0 1.0 1.0

dt_multiplier = {dtm:.12e}
use_subcycling = {subcyc}
fixed_dt = {fixed_dt:.12e}
stop_time = {stop:.4f}
max_steps = {max_steps}
max_spatial_derivative_order = 4
nan_check = 1

lapse_advec_coeff = 1.0
lapse_coeff = {lapse_coeff}
lapse_power = 1.0
shift_advec_coeff = 0.0
shift_Gamma_coeff = 0.75
eta = {eta:.6e}
formulation = 1
kappa1 = 0.
kappa2 = 0.
kappa3 = 0.
covariantZ4 = 1
sigma = 0.3
min_chi = 1.e-50
min_lapse = 1.e-50

AH_activate = {ah}
AH_num_ranks = 4
AH_num_points_u = 30
AH_num_points_v = 50
AH_solve_interval = {ah_int}
AH_print_interval = 1
AH_track_center = false
AH_predict_origin = false
AH_level_to_run = 0
AH_start_time = {ah_start:.4f}
AH_give_up_time = -1.
AH_allow_re_attempt = 1
AH_max_fails_after_lost = -1
AH_verbose = 1
AH_initial_guess = {ah_guess:.4f}
"""


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mu", type=float, default=0.10)
    ap.add_argument("--f0", type=float, default=0.5, help="t_0 / t_C(0)")
    ap.add_argument("--mtH", type=float, default=30.0, help="scalar mass in units 1/t_H")
    ap.add_argument("--box", type=float, default=6.0, help="box side in units of the physical r_m at t_0")
    ap.add_argument("--N", type=int, default=64)
    ap.add_argument("--max_level", type=int, default=3)
    ap.add_argument("--tag_radii", type=str, default="2.0 0.8 0.35 0.18 0.09",
                    help="nested refinement spheres, level l+1 inside tag_radii[l] (units of the physical r_m at t_0)")
    ap.add_argument("--gauge_K_ref", type=int, default=1, help="1: K at the box corner (FLRW), 0: <K>")
    ap.add_argument("--subcycling", type=int, default=1, help="0: same fixed dt on all levels (larger constraint violation!)")
    ap.add_argument("--rho_thr", type=str, default="0 0 50 300",
                    help="density gate per level in units of rho_far(t_0): level l+1 only where rho > rho_thr[l] rho_far")
    ap.add_argument("--diag_interval", type=int, default=8, help="coarse steps between data_out/lineout outputs")
    ap.add_argument("--teclyn_levels", type=int, default=0, help="GRTeclyn params: amr.max_level (nested halves)")
    ap.add_argument("--teclyn_gauge", type=int, default=0, help="GRTeclyn params: 1 = weak 1+log + Gamma driver, 0 = geodesic")
    ap.add_argument("--teclyn_ah_interval", type=int, default=0, help="GRTeclyn params: AH search every N coarse steps (0 = off)")
    ap.add_argument("--lapse_coeff", type=float, default=0.3,
                    help="1+log coefficient: d_t alpha = -coeff alpha (K - fref K_far); 2 = standard (freezes the core)")
    ap.add_argument("--boxsize", type=int, default=16)
    ap.add_argument("--stop", type=float, default=0.9, help="stop at t_0 + stop * t_C(0)")
    ap.add_argument("--max_steps", type=int, default=10**7)
    ap.add_argument("--osc", type=float, default=0.2, help="m dt (0.2 = GRChombo's warning threshold)")
    ap.add_argument("--cfl", type=float, default=0.25)
    ap.add_argument("--n_table", type=int, default=16384)
    ap.add_argument("--n_r", type=int, default=20000, help="comoving grid points for the LTB slice")
    ap.add_argument("--r_max", type=float, default=8.0, help="table extent in r_m (comoving)")
    ap.add_argument("--ah", type=int, default=1)
    ap.add_argument("--ah_start", type=float, default=0.5, help="AH search from t_0 + ah_start * t_C(0)")
    ap.add_argument("--plot_every", type=float, default=0.05, help="plot interval in t_C(0)")
    ap.add_argument("--sigma_inj", type=float, default=0.0,
                    help="warm particles: 1D velocity dispersion (units of c) injected at t_inj as in the 1D code")
    ap.add_argument("--t_inj", type=float, default=1.0, help="injection time in t_H")
    ap.add_argument("--xdisp", type=float, default=2.0, help="dispersion only for shells with R/a < xdisp r_m at t_inj")
    ap.add_argument("--out", required=True)
    a = ap.parse_args()

    ltb = LTBGaussian(a.mu)
    tH, rm, ti = ltb.t_H, ltb.r_m, ltb.t_i
    tC0 = float(ltb.tC(1e-6))
    t0 = a.f0 * tC0
    a0 = (t0 / ti) ** (2.0 / 3.0)

    # --- LTB slice at t_0 on a fine comoving grid (dense near the centre, r in units of k^-1)
    r = np.linspace(1e-6, a.r_max * rm, a.n_r)          # uniform: spline second derivatives stay smooth
    R = np.array([ltb.R(t0, ri) for ri in r]); Rd = np.array([ltb.Rdot(t0, ri) for ri in r])
    E = ltb.E(r); mp = ltb.mp(r)
    sR, sRd = CubicSpline(r, R), CubicSpline(r, Rd)
    Rp, Rdp = sR(r, 1), sRd(r, 1)
    if np.any(Rp <= 0):
        raise SystemExit(f"shell crossing before t_0: min R' = {Rp.min():.3e}")
    rho = mp / (4 * np.pi * R**2 * Rp)
    Krr = -Rdp / Rp; Kth = -Rd / R
    K = Krr + 2 * Kth; Arr = (2.0 / 3.0) * (Krr - Kth)
    # --- warm particles: the 1D code injects an isotropic Maxwellian of width sigma_inj at t_inj (shell rest frame).
    # Collisionless evolution to t_0 conserves R w_t and (R'/sqrt(1+2E)) w_r for small sigma, so the local dispersions
    # at t_0 are sigma_t = sigma_inj R(t_inj)/R(t_0) and sigma_r = sigma_inj R'(t_inj)/R'(t_0); the table carries the ratios.
    sig_r = np.zeros_like(r); sig_t = np.zeros_like(r)
    if a.sigma_inj > 0:
        t_inj = a.t_inj * tH
        if not t_inj < t0:
            raise SystemExit("t_inj must be earlier than t_0")
        a_inj = (t_inj / ti) ** (2.0 / 3.0)
        Rj = np.array([ltb.R(t_inj, ri) for ri in r]); Rjp = CubicSpline(r, Rj)(r, 1)
        inside = Rj / a_inj < a.xdisp * rm
        sig_t = np.where(inside, Rj / R, 0.0); sig_r = np.where(inside, Rjp / Rp, 0.0)
    # --- isotropic radius: ln(varrho/(a0 r)) = -int_r^inf [R'/(R sqrt(1+2E)) - 1/r'] dr'
    f = Rp / (R * np.sqrt(1 + 2 * E)) - 1.0 / r
    tail = cumulative_trapezoid(f[::-1], r[::-1], initial=0.0)[::-1]      # = -int_r^{rmax} f  (rmax -> inf)
    varrho = a0 * r * np.exp(tail)
    psi2 = R / varrho; chi = 1.0 / psi2**2
    Pi = np.sqrt(2 * rho)
    # --- Hamiltonian constraint check in the isotropic form: -8 psi^-5 Lap psi + K^2 - K_ij K^ij = 16 pi rho
    psi = np.sqrt(psi2); sp = CubicSpline(varrho, psi)
    lap = sp(varrho, 2) + 2 * sp(varrho, 1) / varrho
    ham = -8 * lap / psi**5 + K**2 - (Krr**2 + 2 * Kth**2) - 16 * np.pi * rho
    ham_rel = np.abs(ham) / (16 * np.pi * rho)
    sel = varrho > 0.01 * a0 * rm
    # --- uniform table in varrho (cubic interpolation of the profiles; last row = FLRW far field)
    vr = np.linspace(0.0, varrho[-1], a.n_table)
    cols = [CubicSpline(varrho, q, extrapolate=True)(vr) for q in (chi, K, Arr, Pi)]
    chi_t, K_t, Arr_t, Pi_t = cols
    chi_t[0], Arr_t[0] = CubicSpline(varrho, chi)(0.0), 0.0
    os.makedirs(a.out, exist_ok=True)
    table = os.path.abspath(os.path.join(a.out, "ltb_table.dat"))
    hdr = (f"LTB slice mu={a.mu} t0={t0:.8e} (= {a.f0} tC0, tC0={tC0:.8e}, tH={tH:.8e}) a0={a0:.8e} k=1 H_i=5\n"
           f"columns: varrho(isotropic physical radius at t_0)  chi  K  A^rho_rho  Pi ; uniform spacing, last row = FLRW")
    np.savetxt(table, np.column_stack([vr, chi_t, K_t, Arr_t, Pi_t]), header=hdr, fmt="%.12e")
    # Vlasov (GRTeclyn PBHVlasov) table: rho (normal-observer dust density), u_r = 0 in the synchronous slice
    rho_t = 0.5 * Pi_t ** 2
    vcols = [vr, chi_t, K_t, Arr_t, rho_t, np.zeros_like(vr)]
    vhdr = hdr.replace("Pi", "rho  u_r")
    if a.sigma_inj > 0:
        vcols += [np.interp(vr, varrho, sig_r), np.interp(vr, varrho, sig_t)]   # linear: the cut at xdisp is sharp
        vhdr = vhdr.replace("rho  u_r", "rho  u_r  sigma_r/sigma_inj  sigma_t/sigma_inj")
    np.savetxt(os.path.join(a.out, "ltb_vlasov_table.dat"), np.column_stack(vcols), header=vhdr, fmt="%.12e")

    # --- grid and time-step choices
    rm_phys = float(np.interp(rm, r, varrho))
    L = a.box * rm_phys
    dx = L / a.N
    m = a.mtH / tH
    dt_osc = a.osc / m
    dt_cfl = a.cfl * dx
    dt = min(dt_osc, dt_cfl)
    dtm = dt / dx
    # oscillation-limited dt: advance all levels with the same dt (no subcycling) if the finest level's CFL allows
    dx_fine = dx / 2**a.max_level
    subcyc = 0 if (a.subcycling == 0 and dt <= a.cfl * dx_fine) else 1
    fixed_dt = dt if subcyc == 0 else -1.0
    if a.subcycling == 0 and subcyc == 1:
        print(f"WARNING: dt = {dt:.3g} exceeds the finest-level CFL step {a.cfl * dx_fine:.3g}; subcycling kept", file=sys.stderr)
    stop = a.stop * tC0
    q = 1.5 * a.mtH; sigma = np.sqrt(6.0) / q
    rho_far = 0.5 * Pi_t[-1] ** 2
    MH = 0.75 * tH
    t_hor_exp = 1.1203 * tC0 - t0
    plot_int = max(1, int(round(a.plot_every * tC0 / dt)))
    t0_tH_str = f"{t0 / tH:.4f}"
    plot_v = max(1, int(round(a.plot_every * tC0 / (a.cfl * dx / 2 ** a.max_level))))
    params = PARAMS.format(mu=a.mu, f0=a.f0, t0_tH=t0 / tH, a0=a0, mtH=a.mtH, q=q, sigma=sigma, tH=tH, tC0=tC0, MH=MH,
                           rm_phys=rm_phys, t_hor_exp=t_hor_exp, chk_int=10 * plot_int, plot_int=plot_int, m=m,
                           table=os.path.basename(table), rho_far=rho_far, lineout=3 * a.N, c=L / 2, N=a.N, L=L,
                           tag_radii=" ".join(f"{float(x) * rm_phys:.3f}" for x in a.tag_radii.split()[:max(1, a.max_level)]),
                           gauge_K_ref=a.gauge_K_ref, diag_int=a.diag_interval,
                           tag_rho=" ".join(f"{float(x) * rho_far:.6e}" for x in (a.rho_thr.split() + ["0"] * 8)[:max(1, a.max_level)]),
                           max_level=a.max_level, regrid="16 " * max(1, a.max_level), regrid_thr=50.0, box=a.boxsize,
                           dtm=dtm, subcyc=subcyc, fixed_dt=fixed_dt, stop=stop, max_steps=a.max_steps, eta=1.0 / MH, ah=a.ah,
                           ah_int=max(1, plot_int // 4), ah_start=a.ah_start * tC0, ah_guess=0.1 * rm_phys,
                           lapse_coeff=a.lapse_coeff)
    with open(os.path.join(a.out, "params.txt"), "w") as fh:
        fh.write(params)
    dt_v = a.cfl * dx / 2 ** a.max_level  # particle run: CFL-limited only (no oscillation constraint)
    grteclyn = f"""# PBHVlasov (GRTeclyn) — LTB e = 0 late start on particles, mu = {a.mu}, t_0 = {a.f0} t_C(0) = {t0_tH_str} t_H
# t_C(0) = {tC0:.6f}, t_H = {tH:.6f}, M_H = {MH:.4f}; GRTeclyn time t = physical time - t_0
amr.verbose = 1
amr.plot_int = {plot_v}
amr.check_int = {10 * plot_v}
amr.plot_vars = chi lapse K rho_p h11 h22 A11 A22 shift1
amr.n_cell = {a.N} {a.N} {a.N}
geometry.prob_extent = {L:.6f} {L:.6f} {L:.6f}
geometry.is_periodic = 1 1 1
amr.max_level = {a.teclyn_levels}
amr.regrid_int = -1
amr.max_grid_size = {a.boxsize}
amr.blocking_factor = 8
amr.n_error_buf = 0
geometry.center = {L / 2:.6f} {L / 2:.6f} {L / 2:.6f}
evolution.dt_multiplier = {a.cfl}
evolution.stop_time = {stop:.4f}
evolution.max_steps = {a.max_steps}
evolution.sigma = 0.3
# gauge: 0 = geodesic slicing (LTB comparison before the caustic), 1 = weak profile-referenced 1+log + Gamma driver
gauge.lapse_advec_coeff = {1.0 if a.teclyn_gauge else 0.0}
gauge.lapse_coeff = {a.lapse_coeff if a.teclyn_gauge else 0.0}
gauge.lapse_power = 1.0
gauge.shift_advec_coeff = 0.0
gauge.shift_Gamma_coeff = {0.75 if a.teclyn_gauge else 0.0}
gauge.eta = {1.0 / MH:.6e}
ccz4.formulation = 1
pbh_vlasov.init = table
pbh_vlasov.table = ltb_vlasov_table.dat
pbh_vlasov.n_per_dir = 2
pbh_vlasov.lattice_buffer = 4
# apparent-horizon search: every ah_interval coarse steps from t = ah_start (run time), guess radius in coordinates
pbh_vlasov.ah_interval = {a.teclyn_ah_interval}
pbh_vlasov.ah_start = {max(0.0, a.ah_start * tC0 - t0):.4f}
pbh_vlasov.ah_guess = {0.1 * rm_phys:.4f}
pbh_vlasov.ah_num_particles = 2000
pbh_vlasov.ah_max_iter = 3000
# particle proper time starts at the LTB time of the initial slice; ray expansion diagnostic every coarse step
pbh_vlasov.tau0 = {t0:.6f}
pbh_vlasov.theta_rays = 1
pbh_vlasov.theta_ray_n = 256
pbh_vlasov.theta_ray_dr = 0.5
pbh_vlasov.theta_profile_interval = 4
pbh_vlasov.max_disp_cells = 0.8
# warm particles (0 = cold): dispersion at injection; the table columns 7-8 carry sigma_r, sigma_t / sigma at t_0
pbh_vlasov.sigma = {a.sigma_inj}
pbh_vlasov.sigma_seed = 12345
# AMReX TimeIntegrator used by the AH finder flow (required parameters)
integration.type = RungeKutta
integration.rk.type = 4
ah_finder.eta = 3.0
ah_finder.c = 1.0
ah_finder.cfl_factor = 2.0
ah_finder.tolerance = 1e-4
ah_finder.r = 1.15
rho_line_extraction.enabled = true
rho_line_extraction.end = {L * (1 - 1.0 / a.N):.6f} {L / 2:.6f} {L / 2:.6f}
chi_line_extraction.enabled = true
chi_line_extraction.end = {L * (1 - 1.0 / a.N):.6f} {L / 2:.6f} {L / 2:.6f}
K_line_extraction.enabled = true
K_line_extraction.end = {L * (1 - 1.0 / a.N):.6f} {L / 2:.6f} {L / 2:.6f}
lapse_line_extraction.enabled = true
lapse_line_extraction.end = {L * (1 - 1.0 / a.N):.6f} {L / 2:.6f} {L / 2:.6f}
"""
    with open(os.path.join(a.out, "params_vlasov.txt"), "w") as fh:
        fh.write(grteclyn)
    setup = dict(mu=a.mu, f0=a.f0, t0=t0, tC0=tC0, tH=tH, ti=ti, a0=a0, rm=rm, rm_phys=rm_phys, L=L, N=a.N, dx=dx,
                 m=m, mtH=a.mtH, q=q, sigma_ent=sigma, dt=dt, dt_osc=dt_osc, dt_cfl=dt_cfl, dt_multiplier=dtm,
                 use_subcycling=subcyc, fixed_dt=fixed_dt, dx_finest=dx_fine, lapse_coeff=a.lapse_coeff,
                 n_steps_coarse=int(stop / dt), stop=stop, MH=MH, rho_far=rho_far, t_hor_expected_1D=t_hor_exp,
                 chi_centre=float(chi_t[0]), K_centre=float(K_t[0]), K_far=float(K_t[-1]), K_flrw=-2.0 / t0,
                 rho_centre_over_far=float(0.5 * Pi_t[0] ** 2 / rho_far), ham_rel_max=float(ham_rel[sel].max()),
                 table=table)
    with open(os.path.join(a.out, "setup.json"), "w") as fh:
        json.dump(setup, fh, indent=1)
    print(json.dumps(setup, indent=1))


if __name__ == "__main__":
    main()