# 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
"""Analyse the Castro stellar-convection / mass-ejection run ON the cluster.

Emits PNG frames + a metrics CSV; only small files leave the cluster.
"""
import glob, os, csv, 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_out"
os.makedirs(OUT, exist_ok=True)
LABEL = os.environ.get("CONV_LABEL", "1") != "0"
SUFFIX = os.environ.get("CONV_SUFFIX", "")

frames = sorted(glob.glob(os.path.join(RUN, "sun_plt?????")))
print("frames:", len(frames), flush=True)

# Uniform sampling grid: level-0 is 384x1152 over [0,4]x[0,12]
NX, NY = 384, 1152
LX, LY = 4.0, 12.0
dx, dy = LX/NX, LY/NY
ys = (np.arange(NY) + 0.5) * dy

# analytic initial (marginally stable) polytrope, for reference
GAM, H, RHO_B, GRAV = 1.4, 2.0, 1000.0, 2.0
prof0 = 1.0 - (GAM-1)/GAM * ys/H
rho0 = np.where(prof0 > 0, np.maximum(RHO_B*np.power(np.clip(prof0,0,None), 1/(GAM-1)), 1e-3), 1e-3)

def sample(f):
    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])
    xmom = np.array(g["boxlib","xmom"][:,:,0])
    T    = np.array(g["boxlib","Temp"][:,:,0])
    p    = np.array(g["boxlib","pressure"][:,:,0])
    return ds, rho, xmom, ymom, T, p

rows = []
for f in frames:
    ds, rho, xmom, ymom, T, p = sample(f)
    t = float(ds.current_time)
    vy = ymom / rho
    vx = xmom / rho

    # mass flux through horizontal cuts (per unit length in x)
    flux = (ymom.mean(axis=0)) * LX           # rho*vy averaged over x, times width
    up   = np.where(ymom > 0, ymom, 0).mean(axis=0) * LX
    dn   = np.where(ymom < 0, ymom, 0).mean(axis=0) * LX

    # how high has material been lifted: topmost row whose density exceeds
    # 10x the initial profile there (i.e. genuinely transported, not ambient)
    lifted = rho > np.maximum(10*rho0, 1e-2)[None, :]
    kk = np.where(lifted.any(axis=0))[0]
    h_lift = float(ys[kk.max()]) if kk.size else float("nan")

    # kinetic energy and peak upward speed
    ke = float((0.5*(xmom**2 + ymom**2)/rho).sum() * dx*dy)
    vmax = float(vy.max())
    mass = float(rho.sum() * dx*dy)

    # mass above the original atmosphere top (y=7): the "ejected" reservoir
    k7 = int(np.searchsorted(ys, 7.0))
    m_above = float(rho[:, k7:].sum() * dx*dy)
    f_out = float(flux[-1])          # flux through the open top boundary

    rows.append({"t": t, "mass": mass, "m_above7": m_above, "h_lift": h_lift,
                 "ke": ke, "vy_max": vmax, "flux_top": f_out,
                 "up_max": float(up.max()), "rho_min": float(rho.min()),
                 "plotfile": os.path.basename(f)})
    print(f"metrics {os.path.basename(f)} t={t:.3f} h_lift={h_lift:.3f} m>7={m_above:.4e} vmax={vmax:.3f}", flush=True)

with open(os.path.join(OUT, "sunconv_metrics.csv"), "w", newline="") as fh:
    w = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
    w.writeheader(); w.writerows(rows)
print("wrote sunconv_metrics.csv", flush=True)

# ---------- frames: density (log) with velocity structure ----------
for i, f in enumerate(frames):
    ds, rho, xmom, ymom, T, p = sample(f)
    t = float(ds.current_time)
    tag = os.path.basename(f).replace("sun_plt", "")
    fig, ax = plt.subplots(figsize=(3.6, 6.0))
    im = ax.imshow(np.log10(rho.T), origin="lower", extent=[0, LX, 0, LY],
                   aspect="auto", cmap="inferno", vmin=-3.0, vmax=3.0,
                   interpolation="bilinear")
    ax.axhline(7.0, color="w", lw=0.6, ls=":", alpha=0.55)
    ax.set_xticks([]); ax.set_yticks([])
    for s in ax.spines.values(): s.set_visible(False)
    if LABEL:
        ax.text(0.04, 0.975, f"$t$ = {t:.2f}", transform=ax.transAxes,
                color="w", size=10, va="top")
        ax.text(0.04, 0.015,
                "Castro: convective plumes ejecting material\n$\\log_{10}\\rho$, heated layer at $y=2$",
                transform=ax.transAxes, color="w", size=6.2, va="bottom", alpha=0.9)
    fig.subplots_adjust(0, 0, 1, 1)
    fig.savefig(os.path.join(OUT, f"conv{SUFFIX}_{tag}.png"), dpi=130)
    plt.close(fig)
    print("frame", tag, flush=True)

print("ANALYSIS_DONE", flush=True)