# 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
"""Render 3D Castro RT plotfiles ON the cluster; emit PNGs + a metrics CSV.

Only images and a small CSV leave the cluster - the plotfiles are multi-GB.
"""
import glob, os, sys, csv
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/rt3d"
OUT = "/public-workspace/castro/rt3d_out"
os.makedirs(OUT, exist_ok=True)

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

# ---------- metrics over every frame (cheap, uses covering_grid on level 0) ----------
NX = NY = 64
NZ = 128
zs = (np.arange(NZ) + 0.5) * (1.0 / NZ)
rows = []
for f in frames:
    ds = yt.load(f)
    g = ds.covering_grid(level=0, left_edge=ds.domain_left_edge, dims=[NX, NY, NZ])
    rho = np.array(g["boxlib", "density"])
    t = float(ds.current_time)
    heavy = rho >= 1.5
    light = ~heavy
    kz_heavy = np.where(heavy.any(axis=(0, 1)))[0]
    kz_light = np.where(light.any(axis=(0, 1)))[0]
    spike = zs[kz_heavy.min()] if kz_heavy.size else np.nan
    bubble = zs[kz_light.max()] if kz_light.size else np.nan
    fbar = heavy.mean(axis=(0, 1))
    W = 4.0 * float(np.sum(fbar * (1 - fbar))) * (1.0 / NZ)
    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_grids": int(ds.index.num_grids), "max_level": int(ds.index.max_level),
                 "plotfile": os.path.basename(f)})
    print("metrics", os.path.basename(f), f"t={t:.3f}", flush=True)

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)

# ---------- volume renderings: every frame, for an animation ----------
def volume_render(ds, path, sigma_clip=2.0):
    sc = yt.create_scene(ds, field=("boxlib", "density"), lens_type="perspective")
    src = sc[0]
    src.tfh.set_bounds((0.96, 2.09))
    src.tfh.set_log(False)
    src.tfh.grey_opacity = False
    tf = yt.ColorTransferFunction((0.96, 2.09))
    # emphasise the two fluids and the mixing band between them
    tf.add_gaussian(1.05, 0.0008, [0.18, 0.30, 0.62, 0.055])   # light fluid, blue
    tf.add_gaussian(1.50, 0.0040, [0.60, 0.60, 0.60, 0.030])   # mixing, grey
    tf.add_gaussian(1.95, 0.0010, [0.95, 0.85, 0.25, 0.075])   # heavy fluid, yellow
    src.set_transfer_function(tf)
    cam = sc.camera
    cam.resolution = (720, 900)
    cam.position = ds.arr([1.35, 1.15, 1.30], "code_length")
    cam.focus = ds.arr([0.25, 0.25, 0.47], "code_length")
    cam.north_vector = np.array([0.0, 0.0, 1.0])
    cam.set_width(ds.arr([1.5, 1.5, 1.5], "code_length"))
    sc.save(path, sigma_clip=sigma_clip, render=True)

for f in frames:
    ds = yt.load(f)
    t = float(ds.current_time)
    tag = os.path.basename(f).replace("rt3d_plt", "")
    png = os.path.join(OUT, f"vol_{tag}.png")
    try:
        volume_render(ds, png)
        print("volume", tag, f"t={t:.3f}", flush=True)
    except Exception as e:
        print("VOLFAIL", tag, type(e).__name__, e, flush=True)

# ---------- slice panels at selected times ----------
targets = [0.0, 0.8, 1.5, 2.0, 2.5]
times = np.array([r["t"] for r in rows])
sel = [frames[int(np.argmin(np.abs(times - tt)))] for tt in targets]
for f in sel:
    ds = yt.load(f)
    t = float(ds.current_time)
    tag = os.path.basename(f).replace("rt3d_plt", "")
    p = yt.SlicePlot(ds, "y", ("boxlib", "density"), center=[0.25, 0.25, 0.5])
    p.set_cmap(("boxlib", "density"), "cividis")
    p.set_zlim(("boxlib", "density"), 0.96, 2.09)
    p.set_log(("boxlib", "density"), False)
    p.annotate_timestamp(corner="upper_left", draw_inset_box=True)
    p.save(os.path.join(OUT, f"slice_y_{tag}.png"))
    print("slice", tag, f"t={t:.3f}", flush=True)

print("RENDER_DONE", flush=True)