# Public copy: internal compute identifiers and local environment paths were removed or generalized. Scientific values are retained. Paths may need adapting and some original run inputs are not included.
#!/usr/bin/env python3
"""Third pass: density contrast against the ACTUAL t=0 profile, correct annotations."""
import glob, os, sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import yt
yt.set_log_level(50)

RUN = "/public-workspace/castro/runs/sunconv"
OUT = "/public-workspace/castro/sunconv_out4"
os.makedirs(OUT, exist_ok=True)
LABEL = os.environ.get("CONV_LABEL", "1") != "0"
SUFFIX = os.environ.get("CONV_SUFFIX", "")
ONLY = sys.argv[1] if len(sys.argv) > 1 else None

allf = sorted(glob.glob(os.path.join(RUN, "sun_plt?????")))
NX, NY = 384, 1152
LX, LY = 4.0, 12.0
ys = (np.arange(NY) + 0.5) * (LY/NY)

def grid(f, field):
    ds = yt.load(f)
    g = ds.covering_grid(level=0, left_edge=ds.domain_left_edge, dims=[NX, NY, 1])
    return ds, np.array(g["boxlib", field][:, :, 0])

# REFERENCE = the actual initial state, not an analytic polytrope. Castro's HSE
# integration + low_density_cutoff truncates the atmosphere near y=5, well below
# where the ideal polytrope formula would put it (y=7); using the analytic curve
# as the reference paints a spurious depletion band at t=0.
_, rho_ref = grid(allf[0], "density")
rho_ref_prof = rho_ref.mean(axis=0)
FLOOR = 1e-3
k_surf = np.where(rho_ref_prof > FLOOR*1.001)[0].max()
Y_SURF = float(ys[k_surf])
Y_EJECT = 7.0
print(f"initial surface at y={Y_SURF:.3f}; ejection threshold y={Y_EJECT}", flush=True)

frames = [f for f in allf if f.endswith(ONLY)] if ONLY else allf
print("frames:", len(frames), flush=True)

for f in frames:
    ds, rho = grid(f, "density")
    t = float(ds.current_time)
    tag = os.path.basename(f).replace("sun_plt", "")
    contrast = np.log10(np.maximum(rho, 1e-12) / np.maximum(rho_ref_prof, 1e-12)[None, :])

    fig, ax = plt.subplots(figsize=(3.5, 6.0))
    ax.imshow(contrast.T, origin="lower", extent=[0, LX, 0, LY], aspect="auto",
              cmap="RdBu_r", vmin=-4.0, vmax=4.0, interpolation="bilinear")
    ax.axhline(Y_SURF, color="k", lw=0.8, ls="-", alpha=0.55)
    ax.axhline(Y_EJECT, color="k", lw=0.7, ls=":", alpha=0.5)
    ax.axhline(2.0, color="#128012", lw=0.9, ls="--", alpha=0.9)
    ax.set_xticks([]); ax.set_yticks([])
    for s in ax.spines.values(): s.set_visible(False)
    if LABEL:
        ax.text(0.04, 0.977, f"$t$ = {t:.2f}", transform=ax.transAxes, size=10, va="top")
        ax.text(0.04, 0.013,
                "Castro: convective mass ejection\n$\\log_{10}(\\rho/\\rho_{t=0})$  ·  red = material added",
                transform=ax.transAxes, size=6.0, va="bottom", color="#222222")
    fig.subplots_adjust(0, 0, 1, 1)
    fig.savefig(os.path.join(OUT, f"c4{SUFFIX}_{tag}.png"), dpi=130)
    plt.close(fig)
    print("frame", tag, f"t={t:.3f}", flush=True)

if not ONLY:
    np.savez_compressed(os.path.join(OUT, "ref.npz"), y=ys, rho_ref=rho_ref_prof,
                        y_surf=Y_SURF, y_eject=Y_EJECT)
print("PASS3_DONE", flush=True)