"""Phase 2: Baseline diagnostics — mean/RMS/vorticity + zone metrics for correction fields.

For each available scene type:
1. Mean/RMS/vorticity of dq_blk, dq_ctl, dq_tar
2. Three-zone spatial metrics
3. Figures saved to data/figures/

Usage:
    conda run -n pycuda_3_10 python correction_analysis/diagnose_corrections.py
"""
from __future__ import annotations

import json
import os
import sys
from typing import Optional

import numpy as np

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from CCD_analysis.configs import DATA_DIR, NX, NY, L0, CENTER_Y
from CCD_analysis.utils.resampling import load_aligned_fields
from CCD_analysis.correction_analysis.compute_correction_fields import (
    compute_correction, dict_to_field_matrix,
)

FIG_DIR = os.path.join(DATA_DIR, "figures")
os.makedirs(FIG_DIR, exist_ok=True)

# Scene types to process
SCENE_TYPES = [
    "illusion_0.75L",
    "illusion_1.0L",
    "illusion_1.5L",
    "steady_cloak",
    "karman_re100",
    "vortex_lamb",
    "vortex_taylor",
]


# ---------------------------------------------------------------------------
# Three-zone masks (unified geometry: pinball center at 613 px, sensors at 800 px)
# ---------------------------------------------------------------------------

def define_zones() -> dict:
    """Define three-zone masks for all scenes (unified geometry, 2026-06-28).

    All scenes now use the same pinball/sensor positions after unified collection.
    Zone ranges:
        near_body:   580-720 px (around pinball at x≈613)
        body_wake:   720-850 px (near wake downstream)
        sensor_zone: 780-850 px (around sensors at x=800)
    """
    zones = {}
    mask = np.zeros((NY, NX), dtype=bool)
    mask[:, 580:720] = True
    zones["near_body"] = mask

    mask = np.zeros((NY, NX), dtype=bool)
    mask[:, 720:850] = True
    zones["body_wake"] = mask

    mask = np.zeros((NY, NX), dtype=bool)
    mask[:, 780:850] = True
    zones["sensor_zone"] = mask

    return zones


# ---------------------------------------------------------------------------
# Field computation helpers
# ---------------------------------------------------------------------------

def mean_field(ux: np.ndarray, uy: np.ndarray) -> tuple:
    """Compute mean velocity field from snapshots."""
    return np.mean(ux, axis=0), np.mean(uy, axis=0)


def rms_field(ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
    """Compute RMS magnitude field."""
    ux_rms = np.std(ux, axis=0)
    uy_rms = np.std(uy, axis=0)
    return np.sqrt(ux_rms**2 + uy_rms**2)


def vorticity_field(ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
    """Compute mean z-vorticity from mean velocity field."""
    ux_m = np.mean(ux, axis=0)
    uy_m = np.mean(uy, axis=0)
    return np.gradient(uy_m, axis=1) - np.gradient(ux_m, axis=0)


# ---------------------------------------------------------------------------
# Zone metrics
# ---------------------------------------------------------------------------

def zone_metrics(dq: dict, zones: dict, label: str) -> dict:
    """Compute per-zone metrics for a correction field dict.

    Parameters
    ----------
    dq : dict with 'ux' (N, NY, NX), 'uy' (N, NY, NX)
    zones : dict of (NY, NX) boolean masks
    label : str for printing

    Returns
    -------
    metrics : dict with per-zone stats
    """
    if dq is None:
        print(f"  {label}: None, skipping zone metrics")
        return {}

    ux = dq["ux"]
    uy = dq["uy"]
    N = ux.shape[0]

    # Mean kinetic energy field (per snapshot, averaged)
    ke_field = 0.5 * np.mean(ux**2 + uy**2, axis=0)  # (NY, NX)

    # Vorticity field (from mean velocity)
    ux_m, uy_m = mean_field(ux, uy)
    vor = np.gradient(uy_m, axis=1) - np.gradient(ux_m, axis=0)
    enstrophy_field = vor**2

    metrics = {}
    total_ke = ke_field.sum()

    for zname, zmask in zones.items():
        n_pts = zmask.sum()
        if n_pts == 0:
            continue

        zone_ke = ke_field[zmask].mean()
        zone_enstrophy = enstrophy_field[zmask].mean()
        zone_ke_frac = ke_field[zmask].sum() / total_ke if total_ke > 0 else 0.0

        # Centreline asymmetry: ux mean above vs below centreline
        cy = int(CENTER_Y)
        y_indices = np.where(zmask.any(axis=1))[0]
        if len(y_indices) > 0:
            y_min, y_max = y_indices.min(), y_indices.max()
            above = zmask[y_min:cy, :].sum()
            below = zmask[cy:y_max, :].sum()
        else:
            above = below = 1

        mask_correction = f"_{label.replace(' ', '_')}"

        metrics[zname] = {
            "n_points": int(n_pts),
            "mean_KE": float(zone_ke),
            "mean_enstrophy": float(zone_enstrophy),
            "KE_fraction": float(zone_ke_frac),
        }

        print(f"    {zname:15s}: KE={zone_ke:.6e}, "
              f"ens={zone_enstrophy:.6e}, "
              f"KE_frac={zone_ke_frac:.4f}")

    return metrics


# ---------------------------------------------------------------------------
# Plotting helpers
# ---------------------------------------------------------------------------

def plot_mean_rms(dq: dict, label: str, prefix: str, zones: Optional[dict] = None):
    """Plot mean ux, mean uy, RMS magnitude for a correction field."""
    if dq is None:
        return

    ux = dq["ux"]
    uy = dq["uy"]
    ux_m, uy_m = mean_field(ux, uy)
    rms = rms_field(ux, uy)

    fig, axes = plt.subplots(1, 3, figsize=(18, 5))

    extent = (0, NX - 1, 0, NY - 1)

    # Mean ux
    vmax = max(abs(ux_m).max(), 1e-12)
    axes[0].imshow(ux_m, cmap="RdBu_r", vmin=-vmax, vmax=vmax,
                   origin="lower", aspect="equal", extent=extent)
    axes[0].set_title(f"{label}: mean ux")

    # Mean uy
    vmax = max(abs(uy_m).max(), 1e-12)
    axes[1].imshow(uy_m, cmap="RdBu_r", vmin=-vmax, vmax=vmax,
                   origin="lower", aspect="equal", extent=extent)
    axes[1].set_title(f"{label}: mean uy")

    # RMS magnitude
    axes[2].imshow(rms, cmap="viridis", origin="lower",
                   aspect="equal", extent=extent)
    axes[2].set_title(f"{label}: RMS magnitude")

    # Overlay zone boundaries if provided
    if zones is not None:
        # simple boundary: first/last column of each zone mask
        for zname, zmask in zones.items():
            for ax in axes:
                # Find leftmost and rightmost columns with True
                cols = np.where(zmask.any(axis=0))[0]
                if len(cols) > 1:
                    ax.axvline(cols[0], color="white", linewidth=0.5, alpha=0.5)
                    ax.axvline(cols[-1], color="white", linewidth=0.5, alpha=0.5)

    plt.tight_layout()
    path = os.path.join(FIG_DIR, f"{prefix}_{label.replace(' ', '_')}.png")
    fig.savefig(path, dpi=150)
    plt.close(fig)
    print(f"  Saved: {path}", flush=True)


def plot_vorticity(dq: dict, label: str, prefix: str):
    """Plot mean vorticity field."""
    if dq is None:
        return

    ux = dq["ux"]
    uy = dq["uy"]
    vor = vorticity_field(ux, uy)

    fig, ax = plt.subplots(figsize=(10, 4))
    vmax = max(np.percentile(abs(vor), 99), 1e-12)
    ax.imshow(vor, cmap="RdBu_r", vmin=-vmax, vmax=vmax,
              origin="lower", aspect="equal",
              extent=(0, NX - 1, 0, NY - 1))
    ax.set_title(f"{label}: mean vorticity")
    plt.tight_layout()
    path = os.path.join(FIG_DIR, f"{prefix}_vorticity_{label.replace(' ', '_')}.png")
    fig.savefig(path, dpi=150)
    plt.close(fig)
    print(f"  Saved: {path}", flush=True)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def run():
    print("=" * 60, flush=True)
    print("Phase 2: Baseline Diagnostics (correction fields)", flush=True)
    print("=" * 60, flush=True)

    zones = define_zones()

    all_metrics = {}

    for scene_type in SCENE_TYPES:
        print(f"\n{'=' * 60}", flush=True)
        print(f"Scene: {scene_type}", flush=True)
        print(f"{'=' * 60}", flush=True)

        try:
            corr = compute_correction(scene_type)
        except (FileNotFoundError, KeyError, AssertionError, ValueError) as e:
            print(f"  SKIP: {e}", flush=True)
            continue

        if corr["N"] == 0:
            print(f"  SKIP: no valid data (N=0)", flush=True)
            continue

        # Unified geometry — same zones for all scenes
        for dq_key, dq_label in [
            ("dq_blk", "dq_blk (pinball blockage)"),
            ("dq_ctl", "dq_ctl (control correction)"),
        ]:
            dq = corr.get(dq_key)
            if dq is None:
                continue

            prefix = f"corr_{scene_type}"
            plot_mean_rms(dq, dq_label, prefix, zones)
            plot_vorticity(dq, dq_label, prefix)

            print(f"  Zone metrics for {dq_label}:", flush=True)
            metrics = zone_metrics(dq, zones, dq_label)
            all_metrics[f"{scene_type}_{dq_key}"] = metrics

            # For scenes with a target, also plot dq_tar if available
            if dq_key == "dq_ctl" and corr.get("dq_tar") is not None:
                dq_tar = corr.get("dq_tar")
                if dq_tar is not None:
                    plot_mean_rms(dq_tar, "dq_tar (target correction)", prefix, zones)
                    plot_vorticity(dq_tar, "dq_tar (target correction)", prefix)

                    # dq_ctl vs dq_tar side-by-side comparison
                    fig, axes = plt.subplots(2, 2, figsize=(14, 8))
                    extent = (0, NX - 1, 0, NY - 1)

                    # Row 0: mean ux for dq_ctl and dq_tar
                    ux_ctl, _ = mean_field(dq["ux"], dq["uy"])
                    ux_tar, _ = mean_field(dq_tar["ux"], dq_tar["uy"])
                    vmax = max(abs(ux_ctl).max(), abs(ux_tar).max(), 1e-12)

                    axes[0, 0].imshow(ux_ctl, cmap="RdBu_r", vmin=-vmax, vmax=vmax,
                                      origin="lower", aspect="equal", extent=extent)
                    axes[0, 0].set_title("dq_ctl mean ux")

                    axes[0, 1].imshow(ux_tar, cmap="RdBu_r", vmin=-vmax, vmax=vmax,
                                      origin="lower", aspect="equal", extent=extent)
                    axes[0, 1].set_title("dq_tar mean ux")

                    # Row 1: RMS
                    rms_ctl = rms_field(dq["ux"], dq["uy"])
                    rms_tar = rms_field(dq_tar["ux"], dq_tar["uy"])
                    rmax = max(rms_ctl.max(), rms_tar.max(), 1e-12)

                    axes[1, 0].imshow(rms_ctl, cmap="viridis", vmin=0, vmax=rmax,
                                      origin="lower", aspect="equal", extent=extent)
                    axes[1, 0].set_title("dq_ctl RMS")

                    axes[1, 1].imshow(rms_tar, cmap="viridis", vmin=0, vmax=rmax,
                                      origin="lower", aspect="equal", extent=extent)
                    axes[1, 1].set_title("dq_tar RMS")

                    plt.suptitle(f"{scene_type}: dq_ctl vs dq_tar comparison")
                    plt.tight_layout()
                    path = os.path.join(FIG_DIR, f"corr_{scene_type}_ctl_vs_tar.png")
                    fig.savefig(path, dpi=150)
                    plt.close(fig)
                    print(f"  Saved: {path}", flush=True)

        # Steady cloak specific: dq_ctl + dq_blk check
        if scene_type == "steady_cloak":
            dq_b = corr.get("dq_blk")
            dq_c = corr.get("dq_ctl")
            if dq_b is not None and dq_c is not None:
                ux_b = np.mean(dq_b["ux"], axis=0)
                ux_c = np.mean(dq_c["ux"], axis=0)
                ux_cancel = ux_c + ux_b

                fig, axes = plt.subplots(1, 3, figsize=(18, 4))
                extent = (0, NX - 1, 0, NY - 1)
                vmax = max(abs(ux_b).max(), abs(ux_c).max(), abs(ux_cancel).max(), 1e-12)

                axes[0].imshow(ux_b, cmap="RdBu_r", vmin=-vmax, vmax=vmax,
                               origin="lower", aspect="equal", extent=extent)
                axes[0].set_title("dq_blk mean ux (blockage)")

                axes[1].imshow(ux_c, cmap="RdBu_r", vmin=-vmax, vmax=vmax,
                               origin="lower", aspect="equal", extent=extent)
                axes[1].set_title("dq_ctl mean ux (correction)")

                axes[2].imshow(ux_cancel, cmap="RdBu_r", vmin=-vmax, vmax=vmax,
                               origin="lower", aspect="equal", extent=extent)
                axes[2].set_title("dq_ctl + dq_blk (cancel test)")

                plt.suptitle(f"Steady cloak: cancellation test")
                plt.tight_layout()
                path = os.path.join(FIG_DIR, "steady_cloak_cancel_test.png")
                fig.savefig(path, dpi=150)
                plt.close(fig)
                print(f"  Saved: {path}", flush=True)

    # Save zone metrics
    metrics_path = os.path.join(DATA_DIR, "ccd", "zone_metrics.json")
    with open(metrics_path, "w") as f:
        json.dump(all_metrics, f, indent=2)
    print(f"\nZone metrics saved to {metrics_path}", flush=True)
    print("\nDone.", flush=True)


if __name__ == "__main__":
    run()
