# CelerisLab/common/render.py
"""
Field rendering utilities (vorticity computation and visualization).

Separated from the streakline module: these are macroscopic field diagnostics,
not particle-tracking operations.
"""

from __future__ import annotations

import os
from typing import Optional, Sequence

import numpy as np

from ._types import CylinderSpec


def compute_vorticity(ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
    """Central-difference vorticity omega = d(uy)/dx - d(ux)/dy on interior."""
    ux_arr = np.asarray(ux, dtype=np.float64)
    uy_arr = np.asarray(uy, dtype=np.float64)
    vort = np.zeros_like(ux_arr, dtype=np.float64)
    dudy = 0.5 * (ux_arr[2:, 1:-1] - ux_arr[:-2, 1:-1])
    dvdx = 0.5 * (uy_arr[1:-1, 2:] - uy_arr[1:-1, :-2])
    vort[1:-1, 1:-1] = dvdx - dudy
    return vort


def render_vorticity_field(
    vort: np.ndarray,
    *,
    nx: int,
    ny: int,
    out_path: str,
    cylinders: Optional[Sequence[CylinderSpec]] = None,
    vmin: Optional[float] = None,
    vmax: Optional[float] = None,
    percentile: float = 99.0,
    minimal_axes: bool = True,
) -> dict:
    """Render vorticity with white at zero, green negative, purple positive."""
    try:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        from matplotlib.colors import LinearSegmentedColormap, TwoSlopeNorm
    except ImportError as exc:
        raise RuntimeError("render_vorticity_field requires matplotlib.") from exc

    field = np.asarray(vort, dtype=np.float64)
    vmax_field = float(np.nanmax(np.abs(field)))
    abs_pct = float(np.nanpercentile(np.abs(field), percentile))
    abs_max = max(abs_pct, 0.35 * vmax_field, 1e-6)
    if vmin is None:
        vmin = -abs_max
    if vmax is None:
        vmax = abs_max
    if vmin >= 0.0 or vmax <= 0.0:
        raise ValueError(f"vmin/vmax must bracket zero: vmin={vmin}, vmax={vmax}")

    cmap = LinearSegmentedColormap.from_list(
        "vort_green_white_purple",
        ["#1a9850", "#f7f7f7", "#762a83"],
        N=256,
    )
    norm = TwoSlopeNorm(vmin=float(vmin), vcenter=0.0, vmax=float(vmax))

    rgba = cmap(norm(field))
    if cylinders:
        yy, xx = np.mgrid[0:ny, 0:nx]
        for (cx, cy), radius in cylinders:
            mask = (xx - float(cx)) ** 2 + (yy - float(cy)) ** 2 <= float(radius) ** 2
            rgba[mask] = (0.0, 0.0, 0.0, 1.0)

    os.makedirs(os.path.dirname(os.path.abspath(out_path)), exist_ok=True)
    fig = plt.figure(figsize=(nx / 100.0, ny / 100.0), dpi=100)
    ax = fig.add_axes([0.0, 0.0, 1.0, 1.0])
    ax.imshow(rgba, origin="lower", interpolation="nearest", aspect="auto")
    if minimal_axes:
        ax.axis("off")
    fig.savefig(out_path, dpi=100, bbox_inches="tight", pad_inches=0)
    plt.close(fig)

    return {
        "image_path": out_path,
        "vmin": float(vmin),
        "vmax": float(vmax),
        "vort_min": float(np.nanmin(field)),
        "vort_max": float(np.nanmax(field)),
    }
