# 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
"""3D Castro RT: correct metrics + rho=1.5 isosurface renders. Runs ON the cluster."""
import glob, os, csv, sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LightSource, Normalize
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.ndimage import gaussian_filter
from skimage import measure
import yt
yt.set_log_level(50)

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

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

# A level-0 covering_grid MUST use that level's FULL cell count: dims*dx_level sets
# the extent covered, so dims=[64,64,128] silently returns only a sub-octant.
DIMS = [128, 128, 256]
DX = (0.5 / DIMS[0], 0.5 / DIMS[1], 1.0 / DIMS[2])
zs = (np.arange(DIMS[2]) + 0.5) * DX[2]

# Fixed z-window for the animation so the camera never jitters between frames.
ZLO, ZHI = 0.12, 0.80

def load_grid(f):
    ds = yt.load(f)
    g = ds.covering_grid(level=0, left_edge=ds.domain_left_edge, dims=DIMS)
    return ds, np.array(g["boxlib", "density"])

def render_iso(rho, t, path):
    # light smoothing suppresses marching-cubes stair-stepping on cell-aligned data
    rs = gaussian_filter(rho, sigma=0.8, mode="nearest")
    verts, faces, normals, _ = measure.marching_cubes(rs, level=1.5, spacing=DX)
    tri = verts[faces]
    zc = tri[:, :, 2].mean(axis=1)
    cols = plt.get_cmap("cividis")(Normalize(vmin=ZLO, vmax=ZHI)(zc))
    fn = normals[faces].mean(axis=1)
    fn /= (np.linalg.norm(fn, axis=1, keepdims=True) + 1e-12)
    shade = np.clip(0.62 + 0.38 * (fn @ np.array([-0.45, -0.45, 0.77])), 0.55, 1.0)
    cols[:, :3] = np.clip(cols[:, :3] * shade[:, None], 0, 1)

    fig = plt.figure(figsize=(5.2, 5.6))
    ax = fig.add_subplot(111, projection="3d")
    pc = Poly3DCollection(tri, facecolors=cols, edgecolors="none",
                          linewidths=0, antialiased=False)
    ax.add_collection3d(pc)
    ax.set_xlim(0, 0.5); ax.set_ylim(0, 0.5); ax.set_zlim(ZLO, ZHI)
    ax.set_box_aspect((0.5, 0.5, 0.72), zoom=1.32)
    ax.view_init(elev=16, azim=-58)
    ax.set_axis_off()
    if LABEL:
        ax.text2D(0.05, 0.93, f"$t$ = {t:.2f}", transform=ax.transAxes, size=12)
        ax.text2D(0.05, 0.045,
                  "Castro 3-D Rayleigh–Taylor\n$\\rho=1.5$ isosurface, $256^2{\\times}512$ effective",
                  transform=ax.transAxes, size=7.5, color="#555555")
    fig.subplots_adjust(0, 0, 1, 1)
    fig.savefig(path, dpi=140)
    plt.close(fig)
    return len(faces)

rows = []
for f in frames:
    ds, rho = load_grid(f)
    t = float(ds.current_time)
    tag = os.path.basename(f).replace("rt3d_plt", "")
    heavy = rho >= 1.5
    light = ~heavy
    kh = np.where(heavy.any(axis=(0, 1)))[0]
    kl = np.where(light.any(axis=(0, 1)))[0]
    spike = float(zs[kh.min()]) if kh.size else float("nan")
    bubble = float(zs[kl.max()]) if kl.size else float("nan")
    fbar = heavy.mean(axis=(0, 1))
    W = 4.0 * float(np.sum(fbar * (1 - fbar))) * DX[2]
    ntri = render_iso(rho, t, os.path.join(OUT, f"iso{SUFFIX}_{tag}.png"))
    rows.append({"t": t, "spike": spike, "bubble": bubble, "width": W,
                 "amp": 0.5 * (bubble - spike), "rho_min": float(rho.min()),
                 "rho_max": float(rho.max()), "n_tri": ntri,
                 "n_grids": int(ds.index.num_grids), "plotfile": os.path.basename(f)})
    print(f"done {tag} t={t:.3f} spike={spike:.4f} bubble={bubble:.4f} tri={ntri}", flush=True)

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