"""Render an OpenFOAM internal.vtu result to PNGs with pyvista (off-screen)."""
import sys
import numpy as np
import pyvista as pv

pv.OFF_SCREEN = True
vtu = sys.argv[1] if len(sys.argv) > 1 else \
    r"C:\dev\ezflow_eval\of_sphere_re1000_vtk\of_sphere_re1000_2000\internal.vtu"
tag = sys.argv[2] if len(sys.argv) > 2 else "Re1000 (Cd=0.504)"
out_prefix = sys.argv[3] if len(sys.argv) > 3 else \
    r"C:\dev\ezflow_eval\of_sphere_re1000"

g = pv.read(vtu)
if "U" not in g.point_data:
    g = g.cell_data_to_point_data()
g["Umag"] = np.linalg.norm(np.asarray(g["U"]), axis=1)

# slice through the centre plane (z=0), focus on the region around the sphere
sl = g.slice(normal="z", origin=(0, 0, 0)).clip_box([-2, 6, -3, 3, -1, 1], invert=False)

# overlay the actual body silhouette from the exported boundary patch, if present
import os
body = None
bnd = os.path.join(os.path.dirname(vtu), "boundary.vtm")
if os.path.exists(bnd):
    try:
        mb = pv.read(bnd)
        merged = mb.combine() if hasattr(mb, "combine") else None
        if merged is not None and merged.n_points:
            body = merged.slice(normal="z", origin=(0, 0, 0)).clip_box(
                [-2, 6, -3, 3, -1, 1], invert=False)
    except Exception:
        body = None

for field, cmap, clim, title in [
    ("Umag", "turbo", [0.0, 1.5], f"OpenFOAM {tag}  |U|/Uinf"),
    ("p", "coolwarm", [-0.6, 0.6], f"OpenFOAM {tag}  kinematic p"),
]:
    p = pv.Plotter(off_screen=True, window_size=(1280, 720))
    p.add_mesh(sl, scalars=field, cmap=cmap, clim=clim,
               scalar_bar_args={"title": title.split("  ")[-1]})
    if body is not None:
        p.add_mesh(body, color="black", line_width=2)
    p.view_xy()
    p.add_text(title, font_size=9)
    png = f"{out_prefix}_{field}.png"
    p.screenshot(png)
    print("saved", png)