# legacy_test/core/comparator.py
"""Compare legacy test output against SR_analysis reference data.

Computes per-channel correlation, DTW similarity, RMS error, and spectral
comparison (FFT peak matching) between generated and reference signals.
"""

from __future__ import annotations

import os
import sys
from typing import Dict, Optional, Tuple

import numpy as np

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", ".."))
_SRC = os.path.join(_REPO, "src")
for p in [_REPO, _SRC]:
    if p not in sys.path:
        sys.path.insert(0, p)

from .dtw_metrics import calc_lag, calc_dtw_sim  # noqa: E402


def pearson_corr(x: np.ndarray, y: np.ndarray) -> float:
    """Pearson correlation coefficient between two 1-D arrays."""
    xm = x - x.mean()
    ym = y - y.mean()
    denom = np.sqrt((xm * xm).sum() * (ym * ym).sum())
    if denom < 1e-12:
        return 0.0
    return float((xm * ym).sum() / denom)


def channel_corr(ref: np.ndarray, gen: np.ndarray) -> np.ndarray:
    """Per-channel Pearson correlation.

    Args:
        ref: (N, D) reference array.
        gen: (N, D) generated array.

    Returns:
        (D,) correlation per channel.
    """
    n = min(ref.shape[0], gen.shape[0])
    ref = ref[:n]
    gen = gen[:n]
    return np.array([pearson_corr(ref[:, i], gen[:, i]) for i in range(ref.shape[1])])


def rms_error(ref: np.ndarray, gen: np.ndarray) -> float:
    """Root-mean-square error between two arrays."""
    n = min(ref.shape[0], gen.shape[0])
    ref = ref[:n]
    gen = gen[:n]
    return float(np.sqrt(np.mean((ref - gen) ** 2)))


def dtw_similarity(ref_sensors: np.ndarray, gen_sensors: np.ndarray,
                  conv_len: int = 30) -> float:
    """Compute DTW similarity across all sensor channels.

    Uses the same lag-compensated DTW as the legacy env reward:
    1. Compute lag from middle sensor (index 1) Uy component
    2. For all 6 channels, roll target by lag, compute DTW, average

    Args:
        ref_sensors: (N, 6) reference sensor data.
        gen_sensors: (N, 6) generated sensor data.
        conv_len: Convergence window length.

    Returns:
        Average DTW similarity in [0, 1].
    """
    n = min(ref_sensors.shape[0], gen_sensors.shape[0])
    target = np.asarray(ref_sensors[:n], dtype=np.float64)
    state = np.asarray(gen_sensors[:n], dtype=np.float64)

    id_sens = 1
    target_seq = target[conv_len:2 * conv_len, id_sens]
    state_seq = state[-conv_len:, id_sens]
    lag = calc_lag(target_seq, state_seq)

    similarities = 0.0
    for i in range(6):
        t_seq = np.roll(target[:, i], -lag)[conv_len:2 * conv_len]
        s_seq = state[-conv_len:, i]
        similarities += calc_dtw_sim(t_seq, s_seq)
    return float(similarities / 6.0)


def fft_peak_match(ref_signal: np.ndarray, gen_signal: np.ndarray,
                   top_n: int = 3) -> Tuple[float, np.ndarray, np.ndarray]:
    """Compare FFT peak frequencies between reference and generated signals.

    Args:
        ref_signal: 1-D reference signal.
        gen_signal: 1-D generated signal.
        top_n: Number of top peaks to compare.

    Returns:
        (fraction_matched, ref_peaks, gen_peaks) where fraction_matched
        is the fraction of top_n ref peaks that have a matching gen peak
        within 10% frequency tolerance.
    """
    n = min(len(ref_signal), len(gen_signal))
    ref_spec = np.abs(np.fft.rfft(ref_signal[:n]))
    gen_spec = np.abs(np.fft.rfft(gen_signal[:n]))
    freqs = np.fft.rfftfreq(n, d=1)

    # Exclude DC (freq=0)
    mask = freqs > 0
    freqs_nz = freqs[mask]
    ref_amps = ref_spec[mask] if len(ref_spec) == len(freqs) else ref_spec[1:]
    gen_amps = gen_spec[mask] if len(gen_spec) == len(freqs) else gen_spec[1:]

    if len(freqs_nz) == 0:
        return 1.0, np.array([]), np.array([])

    ref_idx = np.argsort(ref_amps)[::-1][:top_n]
    gen_idx = np.argsort(gen_amps)[::-1][:top_n]

    ref_peaks = freqs_nz[ref_idx]
    gen_peaks = freqs_nz[gen_idx]

    matched = 0
    for rp in ref_peaks:
        if rp < 1e-12:
            matched += 1
            continue
        for gp in gen_peaks:
            if abs(rp - gp) / max(rp, 1e-12) < 0.10:
                matched += 1
                break

    return float(matched / max(top_n, 1)), ref_peaks, gen_peaks


def compare_scene(
    ref_dir: str,
    gen_sensors: np.ndarray,
    gen_forces: np.ndarray,
    gen_actions: np.ndarray,
    *,
    conv_len: int = 30,
    label: str = "",
) -> Dict:
    """Full comparison of generated signals against SR_analysis reference.

    Args:
        ref_dir: Path to SR_analysis scene directory.
        gen_sensors: (N, 6) generated sensor signals.
        gen_forces: (N, 6) generated force signals.
        gen_actions: (N, 3) generated action signals.
        conv_len: DTW convergence window length.
        label: Optional scene label for printing.

    Returns:
        dict with keys:
            sensor_corr: (6,) per-channel sensor correlation
            force_corr: (6,) per-channel force correlation
            action_corr: (3,) per-channel action correlation
            sensor_rms: scalar RMS error
            force_rms: scalar RMS error
            action_rms: scalar RMS error
            dtw_sim: scalar DTW similarity
            fft_match: fraction of FFT peaks matched (sensor channel 1)
            passed: bool — True if all metrics meet thresholds
    """
    from .io_helpers import load_reference_signals

    ref = load_reference_signals(ref_dir)

    s_corr = channel_corr(ref["sensors"], gen_sensors)
    f_corr = channel_corr(ref["forces"], gen_forces)
    a_corr = channel_corr(ref["actions"], gen_actions)

    s_rms = rms_error(ref["sensors"], gen_sensors)
    f_rms = rms_error(ref["forces"], gen_forces)
    a_rms = rms_error(ref["actions"], gen_actions)

    dtw_sim = dtw_similarity(ref["sensors"], gen_sensors, conv_len=conv_len)
    fft_match, _, _ = fft_peak_match(ref["sensors"][:, 1], gen_sensors[:, 1])

    # Primary threshold: DTW similarity (phase-invariant)
    passed = dtw_sim > 0.95

    result = {
        "sensor_corr": s_corr.tolist(),
        "force_corr": f_corr.tolist(),
        "action_corr": a_corr.tolist(),
        "sensor_rms": float(s_rms),
        "force_rms": float(f_rms),
        "action_rms": float(a_rms),
        "dtw_sim": float(dtw_sim),
        "fft_match": float(fft_match),
        "passed": passed,
    }

    # Print summary
    prefix = f"[{label}] " if label else ""
    print(f"{prefix}Sensor corr: {s_corr}")
    print(f"{prefix}Force  corr: {f_corr}")
    print(f"{prefix}Action corr: {a_corr}")
    print(f"{prefix}DTW sim: {dtw_sim:.4f}, FFT match: {fft_match:.2f}")
    print(f"{prefix}RMS — sens: {s_rms:.6f}, force: {f_rms:.6f}, action: {a_rms:.6f}")
    print(f"{prefix}{'PASS' if passed else 'FAIL'}")

    return result
