"""
Per-trial figure generation for isogrid optimization.

Saves 4 PNG figures per sandbox per trial into the trial folder:
  {sandbox_id}_density.png    — density field heatmap η(x,y)         [after Brain]
  {sandbox_id}_mesh.png       — Gmsh triangulation overlaid on density [after Brain]
  {sandbox_id}_ribs.png       — final rib profile (pockets)           [after Brain]
  {sandbox_id}_stress.png     — Von Mises stress field (per-element)  [after NX solve]

The stress figure overlays FEA results onto the rib pattern so you can see
which triangles/pockets are over/under-stressed — the key diagnostic for tuning
density field parameters (η₀, α, β, R₀, s_min, s_max).

These PNGs are NEVER deleted by the retention policy — full history is preserved.

Usage:
    from plot_trial import plot_trial_figures, plot_stress_figures

    # After Brain (density, mesh, ribs):
    plot_trial_figures(sb_data, trial_dir)

    # After NX solve (stress):
    stress_fields = {
        "sandbox_1": {"nodes_xy": [...], "stress_values": [...], "n_elements": N},
        "sandbox_2": {...},
    }
    plot_stress_figures(sb_data, stress_fields, trial_dir, sigma_allow=100.6)
"""

from __future__ import annotations

import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")  # headless — required when NX session may own the display
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np

# Project root on path (run_optimization.py sets sys.path already)
sys.path.insert(0, str(Path(__file__).resolve().parents[4]))

from optimization_engine.isogrid.density_field import evaluate_density_grid


# ─── Resolution for density grid (mm).  5mm is fast enough for plotting. ────
_DENSITY_RESOLUTION = 5.0
_DPI = 150


def plot_trial_figures(sb_data: list[dict], trial_dir: Path) -> list[Path]:
    """
    Generate and save all figures for one trial.

    Parameters
    ----------
    sb_data : list of dicts, one per sandbox.
        Each dict must have keys:
            sandbox_id, geometry, params, triangulation, pockets, ribbed_plate
    trial_dir : Path
        The trial folder where PNGs will be saved.

    Returns
    -------
    List of Path objects for the files that were written.
    """
    written: list[Path] = []
    for sbd in sb_data:
        try:
            sb_id   = sbd["sandbox_id"]
            geom    = sbd["geometry"]
            params  = sbd["params"]
            tri     = sbd["triangulation"]
            pockets = sbd["pockets"]
            plate   = sbd["ribbed_plate"]

            written.append(_plot_density(geom, params, trial_dir / f"{sb_id}_density.png"))
            written.append(_plot_mesh(geom, params, tri, trial_dir / f"{sb_id}_mesh.png"))
            written.append(_plot_ribs(geom, pockets, plate, trial_dir / f"{sb_id}_ribs.png"))
        except Exception as exc:
            print(f"  [Plot] WARNING: could not save figures for {sbd.get('sandbox_id', '?')}: {exc}")

    return [p for p in written if p is not None]


# =============================================================================
# Figure 1 — Density heatmap
# =============================================================================

def _plot_density(geometry: dict, params: dict, out_path: Path) -> Path | None:
    """Save density field heatmap with boundary and hole outlines."""
    try:
        X, Y, eta = evaluate_density_grid(geometry, params, resolution=_DENSITY_RESOLUTION)
    except Exception:
        return None

    fig, ax = plt.subplots(figsize=(9, 5), dpi=_DPI)

    m = ax.pcolormesh(X, Y, eta, shading="auto", cmap="viridis", vmin=0.0, vmax=1.0)
    fig.colorbar(m, ax=ax, label="Density η  (0 = sparse · 1 = dense)", shrink=0.8)

    # Outer boundary
    _draw_boundary(ax, geometry["outer_boundary"], color="white", lw=1.5, alpha=0.85)

    # Holes
    for hole in geometry.get("holes", []):
        _draw_hole(ax, hole, color="#ff6b6b", lw=1.2)

    ax.set_aspect("equal")
    sb_id = geometry.get("sandbox_id", "?")
    ax.set_title(
        f"Density Field — {sb_id}\n"
        f"η₀={params['eta_0']:.2f}  α={params['alpha']:.2f}  β={params['beta']:.2f}  "
        f"γ_s={params['gamma_stress']:.2f}  R₀={params['R_0']:.0f}  R_e={params['R_edge']:.0f}",
        fontsize=8,
    )
    ax.set_xlabel("x (mm)", fontsize=8)
    ax.set_ylabel("y (mm)", fontsize=8)
    ax.tick_params(labelsize=7)
    fig.tight_layout()
    fig.savefig(out_path, dpi=_DPI, bbox_inches="tight")
    plt.close(fig)
    return out_path


# =============================================================================
# Figure 2 — Gmsh triangulation overlay
# =============================================================================

def _plot_mesh(geometry: dict, params: dict, triangulation: dict, out_path: Path) -> Path | None:
    """Save triangulation overlaid on a translucent density background."""
    vertices  = triangulation.get("vertices")
    triangles = triangulation.get("triangles")
    if vertices is None or len(vertices) == 0:
        return None

    try:
        X, Y, eta = evaluate_density_grid(geometry, params, resolution=_DENSITY_RESOLUTION)
    except Exception:
        eta = None

    fig, ax = plt.subplots(figsize=(9, 5), dpi=_DPI)

    # Density background (translucent)
    if eta is not None:
        ax.pcolormesh(X, Y, eta, shading="auto", cmap="viridis",
                      vmin=0.0, vmax=1.0, alpha=0.35)

    # Triangle edges
    if triangles is not None and len(triangles) > 0:
        ax.triplot(
            vertices[:, 0], vertices[:, 1], triangles,
            "k-", lw=0.35, alpha=0.75,
        )

    # Outer boundary
    _draw_boundary(ax, geometry["outer_boundary"], color="#00cc66", lw=1.5)

    # Holes (keepout rings)
    for hole in geometry.get("holes", []):
        _draw_hole(ax, hole, color="#4488ff", lw=1.2)
        # Keepout ring (d_keep × hole_radius)
        d_keep = params.get("d_keep", 1.2)
        r_hole = hole.get("radius", 0) or hole.get("diameter", 0) / 2.0
        if r_hole > 0:
            keepout = plt.Circle(
                hole["center"], r_hole * (1.0 + d_keep),
                color="#4488ff", fill=False, lw=0.8, ls="--", alpha=0.5,
            )
            ax.add_patch(keepout)

    ax.set_aspect("equal")
    n_tri = len(triangles) if triangles is not None else 0
    n_pts = len(vertices)
    sb_id = geometry.get("sandbox_id", "?")
    ax.set_title(
        f"Triangulation — {sb_id}  ({n_tri} triangles · {n_pts} vertices)\n"
        f"s_min={params['s_min']:.1f} mm  s_max={params['s_max']:.1f} mm",
        fontsize=8,
    )
    ax.set_xlabel("x (mm)", fontsize=8)
    ax.set_ylabel("y (mm)", fontsize=8)
    ax.tick_params(labelsize=7)
    fig.tight_layout()
    fig.savefig(out_path, dpi=_DPI, bbox_inches="tight")
    plt.close(fig)
    return out_path


# =============================================================================
# Figure 3 — Final rib profile
# =============================================================================

def _plot_ribs(geometry: dict, pockets: list, ribbed_plate, out_path: Path) -> Path | None:
    """Save final rib pattern — pockets (material removed) + rib plate outline."""
    fig, ax = plt.subplots(figsize=(9, 5), dpi=_DPI)

    # Outer boundary filled (light grey = material to start with)
    outer = np.array(geometry["outer_boundary"])
    ax.fill(outer[:, 0], outer[:, 1], color="#e8e8e8", zorder=0)
    _draw_boundary(ax, geometry["outer_boundary"], color="#00cc66", lw=1.5, zorder=3)

    # Pockets (material removed = pink/salmon)
    # pockets is list[dict] from generate_pockets() — each dict has a 'polyline' key
    for pocket in pockets:
        try:
            polyline = pocket.get("polyline", [])
            if len(polyline) < 3:
                continue
            coords = np.array(polyline)
            patch = mpatches.Polygon(coords, closed=True,
                                     facecolor="#ffaaaa", edgecolor="#cc4444",
                                     lw=0.5, alpha=0.85, zorder=2)
            ax.add_patch(patch)
        except Exception:
            pass

    # Bolt holes (not pocketed — solid keep zones)
    for hole in geometry.get("holes", []):
        _draw_hole(ax, hole, color="#2255cc", lw=1.2, zorder=4)

    ax.set_aspect("equal")
    sb_id = geometry.get("sandbox_id", "?")
    n_pockets = len(pockets)
    ax.set_title(
        f"Rib Profile — {sb_id}  ({n_pockets} pockets)\n"
        f"pink = material removed · grey = rib material · blue = bolt holes",
        fontsize=8,
    )
    ax.set_xlabel("x (mm)", fontsize=8)
    ax.set_ylabel("y (mm)", fontsize=8)
    ax.tick_params(labelsize=7)

    # Auto-fit to sandbox bounds
    boundary = np.array(geometry["outer_boundary"])
    margin = 5.0
    ax.set_xlim(boundary[:, 0].min() - margin, boundary[:, 0].max() + margin)
    ax.set_ylim(boundary[:, 1].min() - margin, boundary[:, 1].max() + margin)

    fig.tight_layout()
    fig.savefig(out_path, dpi=_DPI, bbox_inches="tight")
    plt.close(fig)
    return out_path


# =============================================================================
# Geometry helpers
# =============================================================================

def _draw_boundary(ax, outer_boundary, color, lw, alpha=1.0, zorder=2):
    """Draw a closed polygon boundary."""
    pts = np.array(outer_boundary)
    x = np.append(pts[:, 0], pts[0, 0])
    y = np.append(pts[:, 1], pts[0, 1])
    ax.plot(x, y, color=color, lw=lw, alpha=alpha, zorder=zorder)


def _draw_hole(ax, hole: dict, color, lw, zorder=3):
    """Draw a circular hole outline.

    Note: geometry_schema normalizes inner boundaries to dicts with key
    'diameter' (not 'radius'), so we use diameter/2.
    """
    cx, cy = hole["center"]
    # Normalized hole dicts have 'diameter'; raw dicts may have 'radius'
    r = hole.get("radius", 0) or hole.get("diameter", 0) / 2.0
    if r > 0:
        circle = plt.Circle((cx, cy), r, color=color, fill=False, lw=lw, zorder=zorder)
        ax.add_patch(circle)


# =============================================================================
# Figure 4 — Von Mises stress field (post-NX-solve)
# =============================================================================

def plot_stress_figures(
    sb_data: list[dict],
    stress_fields: dict,
    trial_dir: Path,
    sigma_allow: float = 100.6,
) -> list[Path]:
    """
    Generate and save stress heatmap figures for one trial.

    Must be called AFTER the NX solve (stress_fields comes from
    extract_sandbox_stress_field()).

    Parameters
    ----------
    sb_data : list of dicts (same as plot_trial_figures)
    stress_fields : dict keyed by sandbox_id
        Each value: {"nodes_xy": [...], "stress_values": [...], "n_elements": int}
    trial_dir : Path where PNGs are saved
    sigma_allow : allowable stress in MPa — shown as reference line on colorbar
    """
    written: list[Path] = []
    for sbd in sb_data:
        sb_id = sbd["sandbox_id"]
        sf = stress_fields.get(sb_id, {})
        if not sf.get("n_elements", 0):
            continue
        try:
            p = _plot_stress(
                geometry=sbd["geometry"],
                pockets=sbd["pockets"],
                stress_field=sf,
                sigma_allow=sigma_allow,
                out_path=trial_dir / f"{sb_id}_stress.png",
            )
            if p:
                written.append(p)
        except Exception as exc:
            print(f"  [Plot] WARNING: stress figure failed for {sb_id}: {exc}")
    return written


def _plot_stress(
    geometry: dict,
    pockets: list,
    stress_field: dict,
    sigma_allow: float,
    out_path: Path,
) -> Path | None:
    """
    Von Mises stress heatmap overlaid with rib pocket outlines.

    Shows which triangles/pockets are over-stressed vs under-stressed.
    White pocket outlines make the rib pattern visible against the stress field.
    """
    nodes_xy = np.array(stress_field.get("nodes_xy", []))
    stress_vals = np.array(stress_field.get("stress_values", []))

    if len(nodes_xy) < 3:
        return None

    fig, ax = plt.subplots(figsize=(9, 5), dpi=_DPI)

    # Stress field — tricontourf when enough points, scatter otherwise
    cmap = "RdYlGn_r"   # green = low stress, red = high/overloaded
    vmax = max(float(np.max(stress_vals)), sigma_allow * 1.05)

    if len(nodes_xy) >= 6:
        from matplotlib.tri import Triangulation, LinearTriInterpolator
        try:
            triang = Triangulation(nodes_xy[:, 0], nodes_xy[:, 1])
            tc = ax.tricontourf(triang, stress_vals, levels=20,
                                cmap=cmap, vmin=0, vmax=vmax)
            cb = fig.colorbar(tc, ax=ax, label="Von Mises (MPa)", shrink=0.8)
        except Exception:
            sc = ax.scatter(nodes_xy[:, 0], nodes_xy[:, 1], c=stress_vals,
                            cmap=cmap, s=10, vmin=0, vmax=vmax, alpha=0.85)
            cb = fig.colorbar(sc, ax=ax, label="Von Mises (MPa)", shrink=0.8)
    else:
        sc = ax.scatter(nodes_xy[:, 0], nodes_xy[:, 1], c=stress_vals,
                        cmap=cmap, s=10, vmin=0, vmax=vmax, alpha=0.85)
        cb = fig.colorbar(sc, ax=ax, label="Von Mises (MPa)", shrink=0.8)

    # Mark σ_allow on colorbar
    cb.ax.axhline(y=sigma_allow, color="black", lw=1.2, ls="--")
    cb.ax.text(1.05, sigma_allow / vmax, f"σ_allow\n{sigma_allow:.0f}",
               transform=cb.ax.transAxes, fontsize=6, va="center", ha="left")

    # Rib pocket outlines (white) — so we can visually correlate stress with pockets
    for pocket in pockets:
        polyline = pocket.get("polyline", [])
        if len(polyline) >= 3:
            coords = np.array(polyline)
            patch = mpatches.Polygon(coords, closed=True,
                                     facecolor="none", edgecolor="white",
                                     lw=0.6, alpha=0.75, zorder=3)
            ax.add_patch(patch)

    # Outer boundary + holes
    _draw_boundary(ax, geometry["outer_boundary"], color="#00cc66", lw=1.5, zorder=4)
    for hole in geometry.get("holes", []):
        _draw_hole(ax, hole, color="#4488ff", lw=1.2, zorder=4)

    ax.set_aspect("equal")
    sb_id = geometry.get("sandbox_id", "?")
    n_el = stress_field.get("n_elements", 0)
    max_s = float(np.max(stress_vals))
    feasible = max_s <= sigma_allow

    status = "OK" if feasible else f"OVER by {max_s - sigma_allow:.1f} MPa"
    ax.set_title(
        f"Von Mises Stress — {sb_id}  ({n_el} elements)  [{status}]\n"
        f"max = {max_s:.1f} MPa  ·  σ_allow = {sigma_allow:.0f} MPa  "
        f"·  dashed line = limit",
        fontsize=8,
    )
    ax.set_xlabel("x (mm)", fontsize=8)
    ax.set_ylabel("y (mm)", fontsize=8)
    ax.tick_params(labelsize=7)

    boundary = np.array(geometry["outer_boundary"])
    margin = 5.0
    ax.set_xlim(boundary[:, 0].min() - margin, boundary[:, 0].max() + margin)
    ax.set_ylim(boundary[:, 1].min() - margin, boundary[:, 1].max() + margin)

    fig.tight_layout()
    fig.savefig(out_path, dpi=_DPI, bbox_inches="tight")
    plt.close(fig)
    return out_path