"""Action-CCD mode 1 visualization for cloak scenes.

Action-CCD finds correction-field structures most correlated with cylinder
rotation speeds. For cloak scenes (steady/karman/vortex), this should reveal
the structures the controller directly modulates — clean velocity deficit
compensation and cylinder dipoles, excluding upstream disturbance structures.

Usage:
    conda run -n pycuda_3_10 python correction_analysis/visualize_action_ccd.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, L0
from CCD_analysis.utils.resampling import (
    compute_pod, cumulative_energy, e95_index, compute_reduced_ccd,
    unstack_velocity_modes,
)
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)

CLOAK_SCENES = ["steady_cloak", "vortex_lamb", "vortex_taylor"]
# karman_re100 excluded due to 72 vs 96 frame mismatch

R = 10
CCD_Q = 6
CROP_X0, CROP_X1 = 300, 1100


def main():
    print("=" * 60)
    print("Action-CCD Mode 1: Cloak scenes")
    print("=" * 60)

    for st in CLOAK_SCENES:
        print(f"\n--- {st} ---", flush=True)

        # Load correction fields
        corr = compute_correction(st)
        dq_ctl = corr.get("dq_ctl")
        if dq_ctl is None or dq_ctl.get("actions") is None:
            print(f"  SKIP: no dq_ctl or no actions")
            continue

        # Build snapshot matrix and compute POD
        Q = dict_to_field_matrix(dq_ctl)
        N = Q.shape[1]
        mf, modes, sv, coeffs = compute_pod(Q)
        e95 = e95_index(cumulative_energy(sv))
        print(f"  POD: E95={e95}, N_modes={len(sv)}")

        # Action-CCD: find structures correlated with actions
        a_r = coeffs[:R, :]
        actions = dq_ctl["actions"][:N].T  # (3, N)
        W, sigma, _, _, _, _ = compute_reduced_ccd(a_r, actions, Q_delay=CCD_Q)

        print(f"  Action-CCD: sigma[0]={sigma[0]:.4f}, sigma_top3={sigma[:3]}")

        # Reconstruct CCD mode 1 in physical space
        # z1 = W[:, 0] @ A_z → CCD temporal coefficient
        # CCD mode = sum over POD modes of (CCD direction weights * POD mode)
        w1 = W[:, 0] / (np.linalg.norm(W[:, 0]) + 1e-12)
        ccd_mode1 = modes[:, :R] @ w1  # (2*NX*NY,)

        # Unstack into ux, uy
        half = NX * NY
        ux_mode = ccd_mode1[:half].reshape(NY, NX)
        uy_mode = ccd_mode1[half:].reshape(NY, NX)

        # Plot mode 1: ux + uy + vorticity, cropped
        vor = np.gradient(uy_mode, axis=1) - np.gradient(ux_mode, axis=0)

        fig, axes = plt.subplots(1, 3, figsize=(14, 4))
        extent = (CROP_X0, CROP_X1, 0, NY - 1)

        # ux
        vmax_ux = max(abs(ux_mode).max(), 1e-12)
        axes[0].imshow(ux_mode[:, CROP_X0:CROP_X1], cmap="RdBu_r",
                       vmin=-vmax_ux, vmax=vmax_ux,
                       origin="lower", aspect="equal", extent=extent)
        axes[0].set_title(f"{st}: Action-CCD mode 1 ux")

        # uy
        vmax_uy = max(abs(uy_mode).max(), 1e-12)
        axes[1].imshow(uy_mode[:, CROP_X0:CROP_X1], cmap="RdBu_r",
                       vmin=-vmax_uy, vmax=vmax_uy,
                       origin="lower", aspect="equal", extent=extent)
        axes[1].set_title(f"{st}: Action-CCD mode 1 uy")

        # vorticity
        vmax_vor = max(np.percentile(abs(vor), 99), 1e-12)
        axes[2].imshow(vor[:, CROP_X0:CROP_X1], cmap="RdBu_r",
                       vmin=-vmax_vor, vmax=vmax_vor,
                       origin="lower", aspect="equal", extent=extent)
        axes[2].set_title(f"{st}: Action-CCD mode 1 vorticity")

        for ax in axes:
            ax.tick_params(left=False, right=False, labelleft=False,
                           bottom=False, top=False, labelbottom=False)

        plt.tight_layout()
        path = os.path.join(FIG_DIR, f"action_ccd_mode1_{st}.png")
        fig.savefig(path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        print(f"  Saved: {path}")

    print("\nDone!")


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