#!/usr/bin/env python3
"""Signal visualisation — obs/action timeseries, DTW, FFT spectrum.

Reads signals.npz and target.npz from scene output dirs.
Produces per-scene diagnostic plots.

Usage:
    conda run -n pycuda_3_10 python viz_signals.py [--scene re100] [--side train]
"""
from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path
from typing import Optional

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

_OUT_BASE = Path(__file__).resolve().parent / "output"


def log(msg: str) -> None:
    print(f"[viz_signals] {msg}", flush=True)


# ---------------------------------------------------------------------------
# Per-scene plot
# ---------------------------------------------------------------------------
def plot_scene_signals(scene_name: str, signals_path: Path,
                       target_path: Optional[Path],
                       out_dir: Path) -> None:
    """Full signal diagnostic for one scene."""
    if not signals_path.exists():
        log(f"  {scene_name}: signals.npz not found at {signals_path}, skipping")
        return

    data = np.load(signals_path, allow_pickle=True)
    sensors = data.get("sensors")
    forces = data.get("forces")
    actions = data.get("actions")
    rewards = data.get("rewards")

    fig = plt.figure(figsize=(16, 18))

    n_channels = 6
    ch_names_s = ["s0_ux", "s0_uy", "s1_ux", "s1_uy", "s2_ux", "s2_uy"]
    ch_names_f = ["F_fx", "F_fy", "T_fx", "T_fy", "B_fx", "B_fy"]

    # 1. Sensor timeseries
    ax1 = fig.add_subplot(4, 2, 1)
    if sensors is not None and sensors.shape[1] >= 6:
        for i in range(6):
            ax1.plot(sensors[:, i], lw=0.6, label=ch_names_s[i])
    ax1.set_title("Sensors (raw, 6 channels)")
    ax1.set_xlabel("Step"); ax1.legend(loc="upper right", fontsize=6, ncol=2)
    ax1.grid(alpha=0.3)

    # 2. Sensor uy channels only (for visual clarity)
    ax2 = fig.add_subplot(4, 2, 2)
    if sensors is not None and sensors.shape[1] >= 6:
        for idx, lbl in [(1, "s0_uy"), (3, "s1_uy"), (5, "s2_uy")]:
            ax2.plot(sensors[:, idx], lw=0.8, label=lbl)
    ax2.set_title("Sensors (uy channels)")
    ax2.set_xlabel("Step"); ax2.legend(fontsize=7)
    ax2.grid(alpha=0.3)

    # 3. Force timeseries
    ax3 = fig.add_subplot(4, 2, 3)
    if forces is not None and forces.shape[1] >= 6:
        for i in range(6):
            ax3.plot(forces[:, i], lw=0.6, label=ch_names_f[i])
    ax3.set_title("Forces (raw, 6 channels)")
    ax3.set_xlabel("Step"); ax3.legend(loc="upper right", fontsize=6, ncol=2)
    ax3.grid(alpha=0.3)

    # 4. Net force
    ax4 = fig.add_subplot(4, 2, 4)
    if forces is not None and forces.shape[1] >= 6:
        cd = (forces[:, 0] + forces[:, 2] + forces[:, 4]) / 3.0
        cl = (forces[:, 1] + forces[:, 3] + forces[:, 5]) / 3.0
        ax4.plot(cd, lw=1, label="Cd (net)", color="tab:red")
        ax4.plot(cl, lw=1, label="Cl (net)", color="tab:blue")
    ax4.set_title("Net force (Cd, Cl)")
    ax4.set_xlabel("Step"); ax4.legend(); ax4.grid(alpha=0.3)

    # 5. Actions
    ax5 = fig.add_subplot(4, 2, 5)
    if actions is not None and actions.shape[1] >= 3:
        ax5.plot(actions[:, 0], lw=0.8, label="aF (front)")
        ax5.plot(actions[:, 1], lw=0.8, label="aB (bottom)")
        ax5.plot(actions[:, 2], lw=0.8, label="aT (top)")
    ax5.set_title("Actions [-1, 1]")
    ax5.set_xlabel("Step"); ax5.legend(); ax5.grid(alpha=0.3)
    ax5.set_ylim(-1.05, 1.05)

    # 6. Reward (if available)
    ax6 = fig.add_subplot(4, 2, 6)
    if rewards is not None and len(rewards) > 0:
        ax6.plot(rewards, lw=1, color="tab:green", label="reward")
        ax6.axhline(y=np.mean(rewards[-180:]) if len(rewards) >= 180 else np.mean(rewards),
                    color="green", ls="--", lw=1, label="avg (tail)")
    ax6.set_title("Reward")
    ax6.set_xlabel("Step"); ax6.legend(); ax6.grid(alpha=0.3)

    # 7. FFT spectrum (controlled sensor uy vs target)
    ax7 = fig.add_subplot(4, 2, 7)
    if sensors is not None and sensors.shape[1] > 3:
        ax7, _ = _plot_fft(ax7, sensors[:, 3], label="Controlled s1_uy",
                           color="tab:blue")
        if target_path and target_path.exists():
            tgt = np.load(target_path, allow_pickle=True)
            tgt_data = tgt.get("target_states", tgt.get("sensors"))
            if tgt_data is not None and tgt_data.shape[1] > 3:
                _plot_fft(ax7, tgt_data[:, 3], label="Target s1_uy",
                          color="tab:red", linestyle="--")
    ax7.set_title("FFT: Center sensor uy")
    ax7.legend(fontsize=7); ax7.grid(alpha=0.3)

    # 8. Action histogram
    ax8 = fig.add_subplot(4, 2, 8)
    if actions is not None and actions.shape[1] >= 3:
        tail = slice(-180, None) if actions.shape[0] >= 180 else slice(None)
        ax8.hist(actions[tail, 0], bins=20, alpha=0.5, label="aF", color="tab:red")
        ax8.hist(actions[tail, 1], bins=20, alpha=0.5, label="aB", color="tab:orange")
        ax8.hist(actions[tail, 2], bins=20, alpha=0.5, label="aT", color="tab:green")
    ax8.set_title("Action distribution (tail 180 steps)")
    ax8.legend(fontsize=7); ax8.grid(alpha=0.3)

    fig.suptitle(f"{scene_name} - Signal Analysis", fontsize=14)
    fig.tight_layout()
    fig.savefig(out_dir / f"signals_{scene_name}.png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    log(f"  {scene_name}: signal plot saved")


def _plot_fft(ax, signal, label="", color="tab:blue", linestyle="-"):
    """Plot single-sided FFT magnitude spectrum."""
    n = len(signal)
    sig = np.asarray(signal, dtype=np.float64)
    sig = sig - np.mean(sig)
    fft = np.abs(np.fft.rfft(sig)) / n
    freqs = np.fft.rfftfreq(n, d=1)
    ax.plot(freqs[1:], fft[1:], lw=1, color=color, ls=linestyle, label=label)
    ax.set_xlim(0, 0.5)
    ax.set_xlabel("Frequency (1/step)")
    return ax, freqs, fft


# ---------------------------------------------------------------------------
# Cross-scene comparison: DTW / reward bars
# ---------------------------------------------------------------------------
def plot_cross_scene_summary(scene_dirs: dict, out_dir: Path) -> None:
    """Bar chart comparing key metrics across scenes."""
    import json

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    names = list(scene_dirs.keys())
    dtw_vals = []
    reward_vals = []
    for name, d in scene_dirs.items():
        m_path = d / "metrics.json"
        if m_path.exists():
            with open(m_path) as f:
                m = json.load(f)
            dtw_vals.append(m.get("dtw_sim_v5", m.get("sim_raw_mean", 0)))
            reward_vals.append(m.get("reward_mean", 0))
        else:
            dtw_vals.append(0)
            reward_vals.append(0)

    ax = axes[0]
    bars = ax.bar(names, dtw_vals, color="tab:blue", alpha=0.7)
    for b, v in zip(bars, dtw_vals):
        if v > 0:
            ax.text(b.get_x() + b.get_width() / 2, b.get_height() + 0.01,
                    f"{v:.3f}", ha="center", fontsize=8)
    ax.set_title("DTW Similarity (V5)")
    ax.set_ylim(0, 1.1)
    ax.tick_params(axis="x", rotation=30)
    ax.grid(axis="y", alpha=0.3)

    ax = axes[1]
    bars = ax.bar(names, reward_vals, color="tab:green", alpha=0.7)
    for b, v in zip(bars, reward_vals):
        if v > 0:
            ax.text(b.get_x() + b.get_width() / 2, b.get_height() + 0.01,
                    f"{v:.3f}", ha="center", fontsize=8)
    ax.set_title("Average Reward (tail 180)")
    ax.tick_params(axis="x", rotation=30)
    ax.grid(axis="y", alpha=0.3)

    fig.suptitle("Cross-Scene Summary")
    fig.tight_layout()
    fig.savefig(out_dir / "cross_scene_summary.png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    log(f"  Cross-scene summary saved")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> int:
    parser = argparse.ArgumentParser(description="Signal visualization")
    parser.add_argument("--scene", type=str, default=None)
    parser.add_argument("--side", type=str, default="all",
                        choices=["all", "train", "reproduce"])
    args = parser.parse_args()

    report_dir = _OUT_BASE / "reports" / "signal_plots"
    report_dir.mkdir(parents=True, exist_ok=True)

    for side_name, side_key in [("train", "train"), ("reproduce", "reproduce")]:
        if args.side != "all" and args.side != side_key:
            continue
        side_dir = _OUT_BASE / side_key
        if not side_dir.exists():
            continue

        scene_dirs_for_summary = {}
        for scene_dir in sorted(side_dir.iterdir()):
            if not scene_dir.is_dir():
                continue
            scene_name = scene_dir.name
            if args.scene and scene_name != args.scene and not scene_name.startswith(args.scene):
                continue

            sig_path = scene_dir / "signals.npz"
            tgt_paths = [scene_dir / "target.npz",
                         scene_dir / ".." / ".." / ".." / "calibrations" / "re100" / "target.npy"]
            tgt_found = None
            for tp in tgt_paths:
                if tp.exists():
                    tgt_found = tp
                    break

            plot_scene_signals(f"{side_key}_{scene_name}", sig_path, tgt_found,
                               report_dir)
            scene_dirs_for_summary[f"{side_key}_{scene_name}"] = scene_dir

        if scene_dirs_for_summary:
            plot_cross_scene_summary(scene_dirs_for_summary, report_dir)

    log(f"Reports in {report_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
