"""Phase 3: Generate vorticity images and verification reports.

For each case with fields_aligned.npz:
  1. Average the 96 snapshots (4 cycles x 24 phases) to get mean flow
  2. Compute vorticity from the mean flow
  3. Save vorticity image
  4. Also save the mean flow ux/uy fields for reference

Usage:
    python3 src/CCD_analysis/scripts/verify_cases.py

Output: data/figures/vorticity_{scene_name}.png
         data/figures/meanflow_{scene_name}.png
"""
from __future__ import annotations

import json
import os
import sys

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 SCENES, DATA_DIR

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

NX = 1280
NY = 512

# Cases with stable sampling periods
PERIODIC_CASES = [
    "pinball",
    "karman_re100",
    "illusion_0.75L",
    "illusion_1.0L",
    "illusion_1.5L",
    "target_cylinder_0.75L",
    "target_cylinder_1.0L",
    "target_cylinder_1.5L",
]

STEADY_CASES = [
    "steady_cloak",
    "target_channel",
]


def vorticity_from_uv(ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
    """Compute z-vorticity from velocity fields.

    ux, uy have shape (NX, NY) as returned by get_velocity_field.
    Returns omega_z with shape (NY, NX) for direct imshow use.
    """
    # Transpose to (NY, NX) for correct gradient axes
    ux_t = ux.T  # (NY, NX)
    uy_t = uy.T  # (NY, NX)
    # omega_z = duy/dx - dux/dy
    # gradient(uy_t, axis=1) = duy/dx  (axis=1 is x-direction in NY,NX layout)
    # gradient(ux_t, axis=0) = dux/dy  (axis=0 is y-direction in NY,NX layout)
    return np.gradient(uy_t, axis=1) - np.gradient(ux_t, axis=0)


def load_aligned_fields(scene_name: str, scene_id: str):
    """Load fields_aligned.npz or standard fields.npz."""
    d = os.path.join(DATA_DIR, scene_id, scene_name)
    paths = [
        os.path.join(d, "fields_aligned.npz"),
        os.path.join(d, "fields.npz"),
    ]
    for p in paths:
        if os.path.isfile(p):
            return np.load(p)
    return None


def plot_vorticity(omega: np.ndarray, title: str, path: str):
    """Save vorticity image with symmetric colorbar."""
    abs_o = np.abs(omega[np.isfinite(omega)])
    vmax = float(np.percentile(abs_o, 99.5)) if abs_o.size > 0 else 1.0
    if vmax <= 0:
        vmax = 1.0

    fig, ax = plt.subplots(figsize=(16, 5))
    im = ax.imshow(omega, origin="lower", aspect="equal", cmap="RdBu_r",
                   vmin=-vmax, vmax=vmax)
    ax.set_xlabel("x (lattice)")
    ax.set_ylabel("y (lattice)")
    ax.set_title(title)
    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label=r"$\omega_z$")
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)


def plot_meanflow(ux_mean: np.ndarray, uy_mean: np.ndarray, title: str, path: str):
    """Save mean flow magnitude image."""
    speed = np.sqrt(ux_mean**2 + uy_mean**2)
    vmax = float(np.percentile(speed, 99.5)) if speed.size > 0 else 1.0

    fig, axes = plt.subplots(1, 2, figsize=(20, 5))
    im0 = axes[0].imshow(ux_mean, origin="lower", aspect="equal", cmap="viridis")
    axes[0].set_title(f"{title}: ux mean")
    fig.colorbar(im0, ax=axes[0])

    im1 = axes[1].imshow(uy_mean, origin="lower", aspect="equal", cmap="viridis")
    axes[1].set_title(f"{title}: uy mean")
    fig.colorbar(im1, ax=axes[1])

    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)


def process_periodic(scene_name: str):
    """Process a periodic case: average fields, compute vorticity, save images."""
    cfg = SCENES[scene_name]
    scene_id = cfg["scene_id"]

    print(f"\n=== {scene_name} ===", flush=True)

    fd = load_aligned_fields(scene_name, scene_id)
    if fd is None:
        print(f"  SKIP: no field data found", flush=True)
        return

    ux = fd["ux"]
    uy = fd["uy"]
    print(f"  fields shape: {ux.shape}", flush=True)

    # Time-average over all snapshots
    ux_mean = np.mean(ux, axis=0)
    uy_mean = np.mean(uy, axis=0)

    # Vorticity of mean flow
    omega = vorticity_from_uv(ux_mean, uy_mean)

    # Save images
    plot_vorticity(omega, f"{scene_name}: mean vorticity (4-cycle avg)",
                   os.path.join(FIG_DIR, f"vorticity_{scene_name}.png"))
    plot_meanflow(ux_mean, uy_mean, scene_name,
                  os.path.join(FIG_DIR, f"meanflow_{scene_name}.png"))
    print(f"  Saved vorticity and meanflow images", flush=True)

    # Also compute: instantaneous vorticity of each snapshot
    if ux.shape[0] >= 96:
        # Last cycle's snapshots for phase-resolved comparison
        cycle_start = 72  # third cycle start
        fig, axes = plt.subplots(4, 6, figsize=(24, 10))
        for p in range(24):
            row, col = p // 6, p % 6
            idx = cycle_start + p
            om = vorticity_from_uv(ux[idx], uy[idx])
            abs_o = np.abs(om[np.isfinite(om)])
            vm = float(np.percentile(abs_o, 99.5)) if abs_o.size > 0 else 1.0
            axes[row, col].imshow(om, origin="lower", aspect="equal",
                                  cmap="RdBu_r", vmin=-vm, vmax=vm)
            axes[row, col].set_title(f"phase {p}")
            axes[row, col].axis("off")
        plt.suptitle(f"{scene_name}: phase-resolved vorticity (1 cycle)")
        plt.tight_layout()
        fig.savefig(os.path.join(FIG_DIR, f"phase_vorticity_{scene_name}.png"),
                    dpi=150, bbox_inches="tight")
        plt.close(fig)
        print(f"  Saved phase-resolved vorticity", flush=True)

    fd.close()


def process_steady(scene_name: str):
    """Process steady cases (no periodic averaging needed)."""
    cfg = SCENES[scene_name]
    scene_id = cfg["scene_id"]

    print(f"\n=== {scene_name} (steady) ===", flush=True)

    fd = load_aligned_fields(scene_name, scene_id)
    if fd is None:
        print(f"  SKIP: no field data", flush=True)
        return

    ux = fd["ux"]
    uy = fd["uy"]

    ux_mean = np.mean(ux, axis=0)
    uy_mean = np.mean(uy, axis=0)
    omega = vorticity_from_uv(ux_mean, uy_mean)

    plot_vorticity(omega, f"{scene_name}: mean vorticity",
                   os.path.join(FIG_DIR, f"vorticity_{scene_name}.png"))
    plot_meanflow(ux_mean, uy_mean, scene_name,
                  os.path.join(FIG_DIR, f"meanflow_{scene_name}.png"))
    print(f"  Saved vorticity and meanflow images", flush=True)
    fd.close()


def main():
    print("=" * 60, flush=True)
    print("Phase 3: Case Verification", flush=True)
    print("=" * 60, flush=True)

    for name in PERIODIC_CASES:
        process_periodic(name)

    for name in STEADY_CASES:
        process_steady(name)

    # Generate summary report
    report = {
        "periodic_cases": PERIODIC_CASES,
        "steady_cases": STEADY_CASES,
        "figure_dir": FIG_DIR,
    }
    with open(os.path.join(DATA_DIR, "figures", "verification_report.json"), "w") as f:
        json.dump(report, f, indent=2)

    print(f"\nAll figures saved to {FIG_DIR}", flush=True)


if __name__ == "__main__":
    main()
