"""Sanity check: compare raw force signals between target and illusion.

This script answers the most basic question:
  "Does the controller actually match the target's force signature?"

For each diameter (0.75L, 1.0L, 1.5L), loads the target cylinder's forces
and the illusion's total forces, plots them, and computes basic statistics.

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

Output: data/figures/sanity_force_{diam}L.png
         data/figures/sanity_force_report.json
"""
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)

DIAMETERS = [0.75, 1.0, 1.5]


def compute_cd_cl(forces: np.ndarray, scene_name: str) -> tuple:
    """Compute total Cd and Cl from raw forces.
    
    For target_cylinder: forces shape (N, 2), just [Fx, Fy]
    For illusion: forces shape (N, 6), need [F0+F2+F4, F1+F3+F5]
    """
    if "target_cylinder" in scene_name:
        return forces[:, 0], forces[:, 1]  # Fx, Fy
    else:
        return (forces[:, 0] + forces[:, 2] + forces[:, 4],
                forces[:, 1] + forces[:, 3] + forces[:, 5])


def run():
    print("=" * 60, flush=True)
    print("Sanity Check: Force Comparison", flush=True)
    print("=" * 60, flush=True)

    report = {}

    for diam in DIAMETERS:
        tgt_name = f"target_cylinder_{diam}L"
        ill_name = f"illusion_{diam}L"

        cfg_tgt = SCENES[tgt_name]
        cfg_ill = SCENES[ill_name]

        # Load target forces
        if cfg_tgt.get("source") == "open_loop":
            tgt_path = os.path.join(DATA_DIR, "target_cylinder", tgt_name, "sensors.npz")
            tgt_raw = np.load(tgt_path)
            tgt_forces = tgt_raw["forces"]  # (N, 2)
        else:
            tgt_path = os.path.join(DATA_DIR, "target_cylinder", tgt_name, "controlled.npz")
            tgt_raw = np.load(tgt_path)
            tgt_forces = tgt_raw["forces"]

        tgt_cd, tgt_cl = compute_cd_cl(tgt_forces, tgt_name)

        # Load illusion forces
        ill_raw = np.load(os.path.join(DATA_DIR, "illusion", ill_name, "controlled.npz"))
        ill_forces = ill_raw["forces"]  # (N, 6)
        ill_cd, ill_cl = compute_cd_cl(ill_forces, ill_name)

        # Align lengths
        n = min(len(tgt_cd), len(ill_cd))
        tgt_cd, tgt_cl = tgt_cd[:n], tgt_cl[:n]
        ill_cd, ill_cl = ill_cd[:n], ill_cl[:n]

        # Statistics
        cd_corr = float(np.corrcoef(tgt_cd, ill_cd)[0, 1])
        cl_corr = float(np.corrcoef(tgt_cl, ill_cl)[0, 1])
        cd_rmse = float(np.sqrt(np.mean((tgt_cd - ill_cd) ** 2)))
        cl_rmse = float(np.sqrt(np.mean((tgt_cl - ill_cl) ** 2)))
        cd_mean_tgt = float(np.mean(tgt_cd))
        cd_mean_ill = float(np.mean(ill_cd))
        cl_std_tgt = float(np.std(tgt_cl))
        cl_std_ill = float(np.std(ill_cl))

        print(f"\n--- {diam}L ---", flush=True)
        print(f"  Cd correlation: {cd_corr:.4f}")
        print(f"  Cl correlation: {cl_corr:.4f}")
        print(f"  Cd RMSE: {cd_rmse:.6f}")
        print(f"  Cl RMSE: {cl_rmse:.6f}")
        print(f"  Cd mean: target={cd_mean_tgt:.6f}, illusion={cd_mean_ill:.6f}")
        print(f"  Cl std:  target={cl_std_tgt:.6f}, illusion={cl_std_ill:.6f}")

        report[diam] = {
            "cd_correlation": cd_corr,
            "cl_correlation": cl_corr,
            "cd_rmse": cd_rmse,
            "cl_rmse": cl_rmse,
            "cd_mean_target": cd_mean_tgt,
            "cd_mean_illusion": cd_mean_ill,
            "cl_std_target": cl_std_tgt,
            "cl_std_illusion": cl_std_ill,
            "n_samples": n,
        }

        # Plot
        t = np.arange(n)
        fig, axes = plt.subplots(2, 2, figsize=(14, 8))

        # Cd time series
        ax = axes[0, 0]
        ax.plot(t, tgt_cd, "r-", alpha=0.7, label=f"Target {diam}L", linewidth=1)
        ax.plot(t, ill_cd, "b-", alpha=0.7, label=f"Illusion {diam}L", linewidth=1)
        ax.set_ylabel("Total Fx (lattice)")
        ax.set_title(f"{diam}L: Total Fx (Cd proxy)")
        ax.legend()
        ax.grid(True, alpha=0.3)

        # Cl time series
        ax = axes[0, 1]
        ax.plot(t, tgt_cl, "r-", alpha=0.7, label=f"Target {diam}L", linewidth=1)
        ax.plot(t, ill_cl, "b-", alpha=0.7, label=f"Illusion {diam}L", linewidth=1)
        ax.set_ylabel("Total Fy (lattice)")
        ax.set_title(f"{diam}L: Total Fy (Cl proxy)")
        ax.legend()
        ax.grid(True, alpha=0.3)

        # Cd scatter
        ax = axes[1, 0]
        ax.scatter(tgt_cd, ill_cd, s=2, alpha=0.5, c="purple")
        lim = max(np.abs(ax.get_xlim()).max(), np.abs(ax.get_ylim()).max())
        ax.plot([-lim, lim], [-lim, lim], "k--", alpha=0.3)
        ax.set_xlabel("Target Fx")
        ax.set_ylabel("Illusion Fx")
        ax.set_title(f"Cd scatter (r={cd_corr:.3f})")
        ax.set_aspect("equal")
        ax.grid(True, alpha=0.3)

        # Cl scatter
        ax = axes[1, 1]
        ax.scatter(tgt_cl, ill_cl, s=2, alpha=0.5, c="orange")
        lim = max(np.abs(ax.get_xlim()).max(), np.abs(ax.get_ylim()).max())
        ax.plot([-lim, lim], [-lim, lim], "k--", alpha=0.3)
        ax.set_xlabel("Target Fy")
        ax.set_ylabel("Illusion Fy")
        ax.set_title(f"Cl scatter (r={cl_corr:.3f})")
        ax.set_aspect("equal")
        ax.grid(True, alpha=0.3)

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

        tgt_raw.close()
        ill_raw.close()

    # Save report
    report_path = os.path.join(FIG_DIR, "sanity_force_report.json")
    with open(report_path, "w") as f:
        json.dump(report, f, indent=2)
    print(f"\nReport saved to {report_path}", flush=True)
    print("Done.", flush=True)


if __name__ == "__main__":
    run()
