# 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
"""Second pass: density-contrast frames + horizontally averaged (y,t) profiles."""
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_out2"
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

frames = sorted(glob.glob(os.path.join(RUN, "sun_plt?????")))
if ONLY:
    frames = [f for f in frames if f.endswith(ONLY)]
print("frames:", len(frames), flush=True)

NX, NY = 384, 1152
LX, LY = 4.0, 12.0
dx, dy = LX/NX, LY/NY
ys = (np.arange(NY) + 0.5) * dy
GAM, H, RHO_B = 1.4, 2.0, 1000.0
p0 = 1.0 - (GAM-1)/GAM * ys/H
rho0 = np.where(p0 > 0, np.maximum(RHO_B*np.power(np.clip(p0, 0, None), 1/(GAM-1)), 1e-3), 1e-3)

prof_rho, prof_flux, prof_up, times = [], [], [], []

for f in frames:
    ds = yt.load(f)
    g = ds.covering_grid(level=0, left_edge=ds.domain_left_edge, dims=[NX, NY, 1])
    rho  = np.array(g["boxlib","density"][:,:,0])
    ymom = np.array(g["boxlib","ymom"][:,:,0])
    t = float(ds.current_time)
    tag = os.path.basename(f).replace("sun_plt", "")

    times.append(t)
    prof_rho.append(rho.mean(axis=0))
    prof_flux.append(ymom.mean(axis=0))
    prof_up.append(np.where(ymom > 0, ymom, 0.0).mean(axis=0))

    # density contrast vs the initial hydrostatic profile: this is what makes
    # transported material visible (it is near the floor density in absolute terms)
    contrast = np.log10(np.maximum(rho, 1e-12) / rho0[None, :])
    fig, ax = plt.subplots(figsize=(3.5, 6.0))
    im = ax.imshow(contrast.T, origin="lower", extent=[0, LX, 0, LY], aspect="auto",
                   cmap="RdBu_r", vmin=-1.2, vmax=1.2, interpolation="bilinear")
    ax.axhline(7.0, color="k", lw=0.7, ls=":", alpha=0.6)
    ax.axhline(2.0, color="#1a7f1a", lw=0.9, ls="--", alpha=0.85)
    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_{\\rm initial})$  ·  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"cont{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, "sunconv_profiles.npz"),
                        t=np.array(times), y=ys, rho0=rho0,
                        rho=np.array(prof_rho), flux=np.array(prof_flux),
                        up=np.array(prof_up))
    print("wrote sunconv_profiles.npz", flush=True)
print("PASS2_DONE", flush=True)