"""Generate comparison figures across all cloak & illusion scenarios.

All dq_ctl fields use unified geometry (pinball center at ~613px, sensors at ~800px),
set during GPU collection (configs.py UNIFIED coordinates).
Figures zoom into the region around the pinball/cylinder (x=300-1100) to exclude
boundary artifacts.

1. All-scenes panorama: steady_cloak, karman_re100, vortex_lamb, vortex_taylor,
   illusion_0.75L, illusion_1.0L, illusion_1.5L
2. Illusion-only comparison: 0.75L, 1.0L, 1.5L

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

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 DATA_DIR, NX, NY
from CCD_analysis.correction_analysis.compute_correction_fields import (
    compute_correction,
)

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

# Display crop region (pixels) — around pinball at x~613
CROP_X0, CROP_X1 = 300, 1100

# Scene groups
CLOAK_SCENES = ["steady_cloak", "karman_re100", "vortex_lamb", "vortex_taylor"]
ILLUSION_SCENES = ["illusion_0.75L", "illusion_1.0L", "illusion_1.5L"]
ALL_SCENES = CLOAK_SCENES + ILLUSION_SCENES

SCENE_LABELS = {
    "steady_cloak": "Steady Cloak",
    "karman_re100": "Karman Cloak",
    "vortex_lamb": "Vortex Lamb",
    "vortex_taylor": "Vortex Taylor",
    "illusion_0.75L": "Illusion 0.75L",
    "illusion_1.0L": "Illusion 1.0L",
    "illusion_1.5L": "Illusion 1.5L",
}

FIELD_METRICS = [
    ("ux_mean", r"mean $u_x$", "RdBu_r", True),
    ("uy_mean", r"mean $u_y$", "RdBu_r", True),
    ("rms", "RMS", "viridis", False),
    ("vorticity", r"$\omega_z$", "RdBu_r", True),
]


def compute_metrics(st: str) -> dict | None:
    """Load dq_ctl for a scene and compute metrics."""
    try:
        corr = compute_correction(st)
        dq = corr.get("dq_ctl")
        if dq is None:
            return None
        ux, uy = dq["ux"], dq["uy"]
        return {
            "ux_mean": np.mean(ux, axis=0),
            "uy_mean": np.mean(uy, axis=0),
            "rms": np.sqrt(np.std(ux, axis=0)**2 + np.std(uy, axis=0)**2),
            "vorticity": np.gradient(np.mean(uy, axis=0), axis=1)
                         - np.gradient(np.mean(ux, axis=0), axis=0),
        }
    except Exception as e:
        print(f"  SKIP {st}: {e}")
        return None


def crop_field(f: np.ndarray) -> np.ndarray:
    """Crop to display region (NY, NX_cropped)."""
    return f[:, CROP_X0:CROP_X1]


def plot_comparison(scene_list: str | list, name: str):
    """Generate a grid of dq_ctl metrics for selected scenes."""
    if isinstance(scene_list, str):
        scene_list = [scene_list]

    # Load all fields
    fields = {}
    for st in scene_list:
        print(f"  Loading {st}...", flush=True)
        m = compute_metrics(st)
        if m is not None:
            fields[st] = m

    n_scenes = len(fields)
    if n_scenes == 0:
        print("  No valid fields, skipping")
        return

    scene_names = list(fields.keys())
    n_rows = len(FIELD_METRICS)

    # Compute global vmax per metric from CROPPED fields
    metric_vmax = {}
    for mkey, _, _, _ in FIELD_METRICS:
        all_vals = np.concatenate(
            [abs(crop_field(fields[s][mkey])).ravel() for s in fields])
        vmax = float(np.percentile(all_vals[np.isfinite(all_vals)], 99.5))
        metric_vmax[mkey] = max(vmax, 1e-12)

    fig, axes = plt.subplots(n_rows, n_scenes,
                             figsize=(3.5 * n_scenes, 3.0 * n_rows))
    if n_rows == 1:
        axes = [axes]
    if n_scenes == 1:
        axes = [[a] for a in axes]

    nx_crop = CROP_X1 - CROP_X0
    extent = (CROP_X0, CROP_X1, 0, NY - 1)

    for row, (mkey, mlabel, cmap, symmetric) in enumerate(FIELD_METRICS):
        for col, sn in enumerate(scene_names):
            ax = axes[row][col]
            f = crop_field(fields[sn][mkey])
            vmax = metric_vmax[mkey]

            kwargs = {"cmap": cmap, "origin": "lower",
                      "aspect": "equal", "extent": extent}
            if symmetric:
                kwargs["vmin"] = -vmax
                kwargs["vmax"] = vmax
            else:
                kwargs["vmin"] = 0
                kwargs["vmax"] = vmax

            ax.imshow(f, **kwargs)
            ax.tick_params(left=False, right=False, labelleft=False,
                           bottom=False, top=False, labelbottom=False)

            if row == 0:
                ax.set_title(SCENE_LABELS.get(sn, sn), fontsize=10)
            if col == 0:
                ax.set_ylabel(mlabel, fontsize=10)

    plt.suptitle(f"dq_ctl: [{', '.join(SCENE_LABELS.get(s,s) for s in scene_names)}]",
                 fontsize=12, y=1.01)
    plt.tight_layout()
    path = os.path.join(FIG_DIR, name)
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {path}", flush=True)

    # Stats
    print(f"\n  --- RMS (cropped region) ---")
    for sn in scene_names:
        rms_crop = crop_field(fields[sn]["rms"])
        rms_val = float(np.sqrt(np.mean(rms_crop**2)))
        print(f"  {sn:22s}: RMS={rms_val:.6f}")


def main():
    print("=" * 60)
    print("Comparison: dq_ctl across all cloak & illusion scenes")
    print("=" * 60)

    # 1. All 7 scenes panorama
    print("\n--- All 7 scenes panorama ---")
    plot_comparison(ALL_SCENES, "corr_comparison_all_scenes.png")

    # 2. Illusion-only (3 diameters)
    print("\n--- Illusion-only comparison ---")
    plot_comparison(ILLUSION_SCENES, "corr_illusion_comparison_dqctl.png")

    # 3. Cloak-only (4 scenes) for reference
    print("\n--- Cloak-only comparison ---")
    plot_comparison(CLOAK_SCENES, "corr_cloak_comparison_dqctl.png")

    print("\nDone!")


if __name__ == "__main__":
    sys.exit(main())
