"""Phase 4: Visualization — O_k heatmap, CCD modes, POD phase portraits, 1.5L special case.

Integrates 1.5L special-mechanism branch (no separate analyze_15L.py).

Usage:
    conda run -n pycuda_3_10 python src/CCD_analysis/scripts/visualize_ccd.py
"""
from __future__ import annotations

import json
import os
import sys
from collections import deque

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, SCENES
from CCD_analysis.utils.resampling import (
    compute_pod, compute_reduced_ccd, cumulative_energy,
    load_aligned_fields, make_force_obs,
    build_field_matrix, project_into_basis,
    detect_dominant_frequency, detect_cycle_stability,
)

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

CCD_Q = 6
N_PTS = 24
N_CYCLES = 4
NX_ = 1280
NY_ = 512

# -- helper: warp CCD directions back to physical space --
def warp(W: np.ndarray, modes: np.ndarray) -> np.ndarray:
    """Convert CCD weight vectors to physical modes."""
    return modes @ W


# ====================================================================
# Task 1: O_k heatmap (force_fy primary, from ccd_results.json)
# ====================================================================
def task_1():
    print("=== Task 1: O_k heatmap (force_fy) ===", flush=True)
    results_path = os.path.join(DATA_DIR, "ccd", "ccd_results.json")
    if not os.path.isfile(results_path):
        print("  SKIP: ccd_results.json not found", flush=True)
        return

    with open(results_path) as f:
        all_results = json.load(f)

    for r_label, r in [("r6", 6), ("r10", 10)]:
        diameters = [0.75, 1.0, 1.5]
        ov_matrix = np.full((3, 3), np.nan)

        for col_idx, diam in enumerate(diameters):
            tgt_name = f"target_cylinder_{diam}L"
            ill_name = f"illusion_{diam}L"

            # Build target-only POD basis and recompute CCD for O_k
            try:
                tgt_d = load_aligned_fields(tgt_name)
                ill_d = load_aligned_fields(ill_name)
                pin_d = load_aligned_fields("pinball")
            except FileNotFoundError:
                continue

            Q_tgt = build_field_matrix(tgt_d["ux"], tgt_d["uy"])
            mf, modes, _, coeffs = compute_pod(Q_tgt)
            modes_r = modes[:, :r]

            def get_ccd_w(name, data):
                a = project_into_basis(data["ux"], data["uy"], modes_r, mf)
                frc = data.get("forces")
                if frc is None:
                    return None
                y = make_force_obs(frc, name, mode="fy")
                W, _, _, _, _, _ = compute_reduced_ccd(a, y, Q_delay=CCD_Q)
                return W

            W_tgt = get_ccd_w(tgt_name, tgt_d)
            W_ill = get_ccd_w(ill_name, ill_d)
            W_pin = get_ccd_w("pinball", pin_d)

            def ov(Wa, Wb, k=0):
                if Wa is None or Wb is None:
                    return np.nan
                n = min(Wa.shape[1], Wb.shape[1])
                if k >= n:
                    return np.nan
                return float(abs(
                    Wa[:, k] / (np.linalg.norm(Wa[:, k]) + 1e-12) @
                    Wb[:, k] / (np.linalg.norm(Wb[:, k]) + 1e-12)
                ))

            ov_matrix[0, col_idx] = ov(W_tgt, W_ill)
            ov_matrix[1, col_idx] = ov(W_tgt, W_pin)
            ov_matrix[2, col_idx] = ov(W_ill, W_pin)

            for d in [tgt_d, ill_d, pin_d]:
                if d is not None:
                    pass  # No explicit close needed, gc will handle

        fig, ax = plt.subplots(figsize=(8, 6))
        im = ax.imshow(ov_matrix, cmap="viridis", vmin=0, vmax=1, aspect="auto")
        ax.set_xticks(range(3))
        ax.set_xticklabels(["0.75L", "1.0L", "1.5L"])
        ax.set_yticks(range(3))
        ax.set_yticklabels(["target-illusion", "target-pinball", "illusion-pinball"])

        for i in range(3):
            for j in range(3):
                v = ov_matrix[i, j]
                if not np.isnan(v):
                    ax.text(j, i, f"{v:.3f}", ha="center", va="center",
                            color="white" if v > 0.5 else "black", fontsize=12)

        # Annotate 1.5L as special mechanism
        ax.annotate("special mechanism", xy=(2.0, -0.15), fontsize=9,
                    ha="center", va="center", color="orange",
                    xycoords="axes fraction")

        plt.colorbar(im, label="O_1 (modal overlap)")
        plt.title(f"Force-CCD (SigmaFy) O_1 heatmap ({r_label})")
        plt.tight_layout()
        path = os.path.join(FIG_DIR, f"Ok_heatmap_fy_{r_label}.png")
        fig.savefig(path, dpi=150)
        plt.close(fig)
        print(f"  Saved: {path}", flush=True)


# ====================================================================
# Task 2: CCD mode 1 physical fields (target-only basis)
# ====================================================================
def task_2():
    print("=== Task 2: CCD mode 1 physical fields ===", flush=True)
    r = 6

    for diam in [0.75, 1.0]:
        tgt_name = f"target_cylinder_{diam}L"
        ill_name = f"illusion_{diam}L"

        try:
            tgt_d = load_aligned_fields(tgt_name)
            ill_d = load_aligned_fields(ill_name)
        except FileNotFoundError:
            continue

        # Build target-only POD basis
        Q_tgt = build_field_matrix(tgt_d["ux"], tgt_d["uy"])
        mf, modes, _, _ = compute_pod(Q_tgt)
        modes_r = modes[:, :r]

        for name, d_obj, label in [(tgt_name, tgt_d, "target"),
                                    (ill_name, ill_d, "illusion")]:
            a = project_into_basis(d_obj["ux"], d_obj["uy"], modes_r, mf)
            frc = d_obj.get("forces")
            if frc is None:
                continue
            y = make_force_obs(frc, name, mode="fy")
            W, _, _, _, _, _ = compute_reduced_ccd(a, y, Q_delay=CCD_Q)

            ccd_mode = warp(W[:, :1], modes_r)
            half = NX_ * NY_
            ux_m = ccd_mode[:half, 0].reshape(NY_, NX_)
            uy_m = ccd_mode[half:, 0].reshape(NY_, NX_)

            fig, axes = plt.subplots(1, 2, figsize=(14, 5))
            vmax = max(np.abs(ux_m).max(), np.abs(uy_m).max()) + 1e-12

            im0 = axes[0].imshow(ux_m, cmap="RdBu_r", vmin=-vmax, vmax=vmax,
                                 origin="lower", aspect="equal",
                                 extent=(0, NX_ - 1, 0, NY_ - 1))
            axes[0].set_title(f"{diam}L {label}: CCD mode 1 ux")
            plt.colorbar(im0, ax=axes[0])

            im1 = axes[1].imshow(uy_m, cmap="RdBu_r", vmin=-vmax, vmax=vmax,
                                 origin="lower", aspect="equal",
                                 extent=(0, NX_ - 1, 0, NY_ - 1))
            axes[1].set_title(f"{diam}L {label}: CCD mode 1 uy")
            plt.colorbar(im1, ax=axes[1])

            plt.tight_layout()
            path = os.path.join(FIG_DIR, f"ccd_mode1_fy_{diam}L_{label}.png")
            fig.savefig(path, dpi=150)
            plt.close(fig)
            print(f"  Saved: {path}", flush=True)

        # Mark end for this diameter
        del tgt_d, ill_d


# ====================================================================
# Task 3: z_1(t) verification
# ====================================================================
def task_3():
    print("=== Task 3: z_1(t) verification ===", flush=True)

    for diam in [0.75, 1.0]:
        ill_name = f"illusion_{diam}L"
        tgt_name = f"target_cylinder_{diam}L"

        try:
            tgt_d = load_aligned_fields(tgt_name)
            ill_d = load_aligned_fields(ill_name)
        except FileNotFoundError:
            continue

        # Target-only POD basis
        Q_tgt = build_field_matrix(tgt_d["ux"], tgt_d["uy"])
        mf, modes, _, _ = compute_pod(Q_tgt)
        modes_r = modes[:, :6]

        a = project_into_basis(ill_d["ux"], ill_d["uy"], modes_r, mf)
        frc = ill_d.get("forces")
        if frc is None:
            continue
        y = make_force_obs(frc, ill_name, mode="fy")
        W, sig, _, z, _, _ = compute_reduced_ccd(a, y, Q_delay=CCD_Q)

        fig, axes = plt.subplots(2, 1, figsize=(12, 6))

        ax = axes[0]
        ax.plot(z[0, :], "b-", label="z_1(t)", alpha=0.8)
        ax.set_ylabel("CCD temporal coeff")
        ax.set_title(f"{diam}L illusion: Force-CCD (SigmaFy) z_1(t)")
        ax.legend()
        ax.grid(True, alpha=0.3)

        ax = axes[1]
        Nv = z.shape[1]
        y_norm = (y[0, :Nv] - np.mean(y[0, :Nv])) / (np.std(y[0, :Nv]) + 1e-12)
        z_norm = (z[0, :] - np.mean(z[0, :])) / (np.std(z[0, :]) + 1e-12)
        ax.plot(y_norm, "r-", label="norm SigmaFy", alpha=0.7)
        ax.plot(z_norm, "b--", label="norm z_1", alpha=0.7)
        ax.set_xlabel("Flat sample index")
        ax.set_ylabel("Normalized amplitude")
        ax.set_title(f"{diam}L: z_1 vs SigmaFy (normalized)")
        ax.legend()
        ax.grid(True, alpha=0.3)

        plt.tight_layout()
        path = os.path.join(FIG_DIR, f"z1_verification_fy_{diam}L.png")
        fig.savefig(path, dpi=150)
        plt.close(fig)
        print(f"  Saved: {path}", flush=True)


# ====================================================================
# Task 4: POD phase portraits (target-only basis)
# ====================================================================
def task_4():
    print("=== Task 4: POD phase portraits ===", flush=True)

    fig, axes = plt.subplots(1, 3, figsize=(15, 4))

    for idx, diam in enumerate([0.75, 1.0, 1.5]):
        if diam in [0.75, 1.0]:
            main_only = False
        else:
            main_only = False  # include 1.5L in phase portrait

        tgt_name = f"target_cylinder_{diam}L"
        ill_name = f"illusion_{diam}L"

        try:
            tgt_d = load_aligned_fields(tgt_name)
            ill_d = load_aligned_fields(ill_name)
            pin_d = load_aligned_fields("pinball")
        except FileNotFoundError:
            continue

        Q_tgt = build_field_matrix(tgt_d["ux"], tgt_d["uy"])
        mf, modes, _, _ = compute_pod(Q_tgt)
        modes_r = modes[:, :6]

        ax = axes[idx]
        colors = {"target": "red", "illusion": "blue", "pinball": "green"}

        for kind, d_obj, label in [("target", tgt_d, "target"),
                                    ("illusion", ill_d, "illusion"),
                                    ("pinball", pin_d, "pinball (unc)")]:
            if d_obj is None:
                continue
            a = project_into_basis(d_obj["ux"], d_obj["uy"], modes_r, mf)
            ax.plot(a[0, :], a[1, :], ".", color=colors[kind], markersize=3,
                    alpha=0.5, label=label if idx == 0 else "")

        ax.set_xlabel("a_1")
        ax.set_ylabel("a_2")
        title = f"{diam}L POD attractor"
        if diam == 1.5:
            title += " (special mechanism)"
        ax.set_title(title)
        ax.grid(True, alpha=0.3)
        ax.set_aspect("equal")
        if idx == 0:
            ax.legend(fontsize=8)

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


# ====================================================================
# Task 5: 1.5L special-mechanism diagnostics
# ====================================================================
def task_5():
    """1.5L special-mechanism analysis — raw diagnostics, action compactness, phase drift."""
    print("=== Task 5: 1.5L special-mechanism diagnostics ===", flush=True)

    SI = 800  # 1.5L sample interval

    try:
        ill_d = load_aligned_fields("illusion_1.5L")
        tgt_d = load_aligned_fields("target_cylinder_1.5L")
        pin_d = load_aligned_fields("pinball")
    except FileNotFoundError as e:
        print(f"  SKIP: {e}", flush=True)
        return

    sens_i = ill_d.get("sensors")
    forc_i = ill_d.get("forces")
    act_i = ill_d.get("actions")
    sens_t = tgt_d.get("sensors")
    forc_t = tgt_d.get("forces")

    # ---- Panel 5a: Raw diagnostics (sensors, forces, actions) ----
    print("  -- 5a: Raw time-series diagnostics", flush=True)
    n_plot = min(400, len(sens_i) if sens_i is not None else 0)
    t = np.arange(n_plot) * SI / 1000.0

    fig, axes = plt.subplots(3, 1, figsize=(14, 10))

    ax = axes[0]
    if sens_i is not None:
        for ch in range(6):
            ax.plot(t, sens_i[:n_plot, ch], label=f"ill_s{ch}", alpha=0.7)
    if sens_t is not None:
        ax.plot(t, sens_t[:n_plot, 3], "k--", label="target_s1_v", linewidth=2)
    ax.set_ylabel("Velocity (lattice)")
    ax.set_title("1.5L Sensors: Illusion vs Target")
    ax.legend(fontsize=7, ncol=3)
    ax.grid(True, alpha=0.3)

    ax = axes[1]
    if forc_i is not None:
        for ch in range(6):
            ax.plot(t, forc_i[:n_plot, ch], label=f"ill_F{ch}", alpha=0.7)
    if forc_t is not None:
        ax.plot(t, forc_t[:n_plot, 0], "k--", label="target_Fx", linewidth=2)
        ax.plot(t, forc_t[:n_plot, 1], "k:", label="target_Fy", linewidth=2)
    ax.set_ylabel("Force (lattice)")
    ax.set_title("1.5L Forces")
    ax.legend(fontsize=7, ncol=3)
    ax.grid(True, alpha=0.3)

    ax = axes[2]
    if act_i is not None:
        for ch in range(3):
            ax.plot(t, act_i[:n_plot, ch], label=f"Omega_{ch}")
    ax.set_xlabel("Time (T0 units)")
    ax.set_ylabel("Omega (normalised)")
    ax.set_title("1.5L Actions (DRL output, [-1, 1])")
    ax.legend()
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    path = os.path.join(FIG_DIR, "15L_raw_timeseries.png")
    fig.savefig(path, dpi=150)
    plt.close(fig)
    print(f"    Saved: {path}", flush=True)

    # ---- Panel 5b: Force-CCD compactness ----
    print("  -- 5b: Force-CCD compactness", flush=True)
    results_path = os.path.join(DATA_DIR, "ccd", "ccd_results.json")
    if os.path.isfile(results_path):
        with open(results_path) as f:
            all_res = json.load(f)
        # Report key action-CCD m80
        for r in [6, 8, 10]:
            key = f"1.5L_illusion_1.5L_action_r{r}"
            if key in all_res:
                print(f"    {key}: m80={all_res[key]['m80']}, "
                      f"sigma1={all_res[key]['sigma_top3'][0]:.4f}", flush=True)
            # Force-fy compactness
            key_f = f"1.5L_illusion_1.5L_force_fy_r{r}"
            if key_f in all_res:
                print(f"    {key_f}: m80={all_res[key_f]['m80']}, "
                      f"sigma1={all_res[key_f]['sigma_top3'][0]:.4f}", flush=True)

    # ---- Panel 5c: Windowed periodicity (phase drift) ----
    # Use raw (non-aligned) sensor data for sufficient window length
    print("  -- 5c: Windowed periodicity", flush=True)
    raw_path = os.path.join(DATA_DIR, "illusion", "illusion_1.5L", "controlled.npz")
    if os.path.isfile(raw_path):
        raw_d = np.load(raw_path)
        raw_sensors = raw_d["sensors"]
        raw_d.close()
    else:
        raw_sensors = sens_i  # fallback to aligned data

    if raw_sensors is not None and len(raw_sensors) > 200:
        signal = raw_sensors[:, 1]  # center sensor v
        window = 200
        stride = 20
        n_windows = (len(signal) - window) // stride

        cv_vals, T_vals, f_vals, t_centers = [], [], [], []
        for w in range(n_windows):
            seg = signal[w * stride:w * stride + window]
            cv_T, mean_T, _ = detect_cycle_stability(seg, SI)
            f_dom, T_dom, _ = detect_dominant_frequency(seg, SI)
            cv_vals.append(cv_T)
            T_vals.append(mean_T)
            f_vals.append(f_dom)
            t_centers.append((w * stride + window // 2) * SI / 1000)

        fig, axes = plt.subplots(3, 1, figsize=(14, 8), sharex=True)

        ax = axes[0]
        ax.plot(t_centers, cv_vals, "o-", markersize=3)
        ax.axhline(0.10, color="r", ls="--", label="strict gate")
        ax.axhline(0.12, color="orange", ls="--", label="relaxed gate")
        ax.set_ylabel("CV_T")
        ax.set_title("1.5L Windowed cycle stability (window=200 steps)")
        ax.legend()
        ax.grid(True, alpha=0.3)

        ax = axes[1]
        ax.plot(t_centers, T_vals, "o-", markersize=3, color="green")
        ax.set_ylabel("Mean period (steps)")
        ax.grid(True, alpha=0.3)

        ax = axes[2]
        ax.plot(t_centers, f_vals, "o-", markersize=3, color="purple")
        ax.set_xlabel("Time (T0 units)")
        ax.set_ylabel("Freq (1/step)")
        ax.grid(True, alpha=0.3)

        plt.tight_layout()
        path = os.path.join(FIG_DIR, "15L_windowed_periodicity.png")
        fig.savefig(path, dpi=150)
        plt.close(fig)
        print(f"    Saved: {path}", flush=True)

    # ---- Panel 5d: O(target, illusion) overlap bar ----
    print("  -- 5d: Force-CCD overlap summary", flush=True)
    fig, ax = plt.subplots(figsize=(6, 4))
    diam_labels = ["0.75L", "1.0L", "1.5L"]
    ov_vals = []
    results_path = os.path.join(DATA_DIR, "ccd", "ccd_results.json")
    if os.path.isfile(results_path):
        with open(results_path) as f:
            all_res = json.load(f)
        for diam in [0.75, 1.0, 1.5]:
            tgt = f"target_cylinder_{diam}L"
            ill = f"illusion_{diam}L"
            k_tgt = f"{diam}L_{tgt}_force_fy_r6"
            k_ill = f"{diam}L_{ill}_force_fy_r6"
            # Need W from the saved results — but we don't store W in json.
            # Instead, recompute overlap quickly from the raw data.
            try:
                td = load_aligned_fields(tgt)
                id_ = load_aligned_fields(ill)
                Qt = build_field_matrix(td["ux"], td["uy"])
                mf, modes, _, _ = compute_pod(Qt)
                modes6 = modes[:, :6]
                a_t = project_into_basis(td["ux"], td["uy"], modes6, mf)
                a_i = project_into_basis(id_["ux"], id_["uy"], modes6, mf)
                y_t = make_force_obs(td["forces"], tgt, mode="fy")
                y_i = make_force_obs(id_["forces"], ill, mode="fy")
                Wt, _, _, _, _, _ = compute_reduced_ccd(a_t, y_t, Q_delay=CCD_Q)
                Wi, _, _, _, _, _ = compute_reduced_ccd(a_i, y_i, Q_delay=CCD_Q)
                ov_val = float(abs(
                    Wt[:, 0] / (np.linalg.norm(Wt[:, 0]) + 1e-12) @
                    Wi[:, 0] / (np.linalg.norm(Wi[:, 0]) + 1e-12)
                ))
                ov_vals.append(ov_val)
            except Exception as e:
                print(f"      {diam}L overlap failed: {e}", flush=True)
                ov_vals.append(0.0)

    bars = ax.bar(diam_labels, ov_vals, color=["blue", "green", "orange"], alpha=0.7)
    for bar, v in zip(bars, ov_vals):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.02,
                f"{v:.3f}", ha="center", fontsize=11)
    ax.set_ylim(0, 1.1)
    ax.set_ylabel("O_1 (target-illusion)")
    ax.set_title("Force-CCD (SigmaFy) overlap comparison")
    ax.grid(True, alpha=0.3, axis="y")

    # Annotate 1.5L
    ax.annotate("special mechanism", xy=(2, 0.05), fontsize=9,
                ha="center", color="orange", fontweight="bold")

    plt.tight_layout()
    path = os.path.join(FIG_DIR, "15L_overlap_summary.png")
    fig.savefig(path, dpi=150)
    plt.close(fig)
    print(f"    Saved: {path}", flush=True)


# ====================================================================
# Task 6: Cross-diameter overlap (all illusions in 1.0L target basis)
# ====================================================================
def task_6():
    print("=== Task 6: Cross-diameter overlap (1.0L target basis) ===", flush=True)

    try:
        tgt_10 = load_aligned_fields("target_cylinder_1.0L")
    except FileNotFoundError:
        print("  SKIP: missing 1.0L target data", flush=True)
        return

    Q_10 = build_field_matrix(tgt_10["ux"], tgt_10["uy"])
    mf_10, modes_10, _, _ = compute_pod(Q_10)
    modes6 = modes_10[:, :6]

    W_cross = {}
    for diam in [0.75, 1.0, 1.5]:
        name = f"illusion_{diam}L"
        try:
            d = load_aligned_fields(name)
        except FileNotFoundError:
            continue
        a = project_into_basis(d["ux"], d["uy"], modes6, mf_10)
        frc = d.get("forces")
        if frc is None:
            continue
        y = make_force_obs(frc, name, mode="fy")
        W, _, _, _, _, _ = compute_reduced_ccd(a, y, Q_delay=CCD_Q)
        W_cross[diam] = W

    if len(W_cross) < 2:
        print("  SKIP: not enough illusions", flush=True)
        return

    diam_list = sorted(W_cross.keys())
    ov_mat = np.ones((len(diam_list), len(diam_list)))
    print("  Cross-diameter O_1 matrix (1.0L target-only basis, force_fy):")
    for i, da in enumerate(diam_list):
        for j, db in enumerate(diam_list):
            if i >= j:
                continue
            Wa, Wb = W_cross[da], W_cross[db]
            ov = float(abs(
                Wa[:, 0] / (np.linalg.norm(Wa[:, 0]) + 1e-12) @
                Wb[:, 0] / (np.linalg.norm(Wb[:, 0]) + 1e-12)
            ))
            ov_mat[i, j] = ov
            ov_mat[j, i] = ov
            print(f"    O({da}L, {db}L) = {ov:.4f}")

    fig, ax = plt.subplots(figsize=(6, 5))
    im = ax.imshow(ov_mat, cmap="viridis", vmin=0, vmax=1)
    ax.set_xticks(range(len(diam_list)))
    ax.set_xticklabels([f"{d}L" for d in diam_list])
    ax.set_yticks(range(len(diam_list)))
    ax.set_yticklabels([f"{d}L" for d in diam_list])
    for i in range(len(diam_list)):
        for j in range(len(diam_list)):
            v = ov_mat[i, j]
            ax.text(j, i, f"{v:.3f}", ha="center", va="center",
                    color="white" if v > 0.5 else "black")
    plt.colorbar(im, label="O_1")
    plt.title("Cross-diam force-CCD (1.0L target basis, SigmaFy)")
    plt.tight_layout()
    path = os.path.join(FIG_DIR, "cross_diameter_overlap_fy.png")
    fig.savefig(path, dpi=150)
    plt.close(fig)
    print(f"  Saved: {path}", flush=True)


# ====================================================================
# Main
# ====================================================================
def main():
    print("=" * 60, flush=True)
    print("Phase 4: Visualization (Round 5)", flush=True)
    print("=" * 60, flush=True)

    task_1()  # O_k heatmap
    task_2()  # CCD physical modes (0.75L, 1.0L)
    task_4()  # POD phase portraits (all diameters)
    task_3()  # z_1 verification (0.75L, 1.0L)
    task_5()  # 1.5L special mechanism
    task_6()  # Cross-diameter overlap

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


if __name__ == "__main__":
    main()
