#!/usr/bin/env python3
"""Compare frozen historical and reproduction PPO trajectories without refitting."""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np

from SR_analysis.configs import get_scene
from SR_analysis.utils.cfd_interface import (
    LEGACY_DTW_VERSION,
    compute_similarity_illusion_legacy_v1,
    compute_similarity_karman_legacy_v1,
)
from SR_analysis.utils.provenance import atomic_write_json, hash_file

SCHEMA_VERSION = "sr-trajectory-equivalence-v1"
ARRAY_KEYS = ("sensors", "forces", "actions")
CHANNELS = {
    "sensors": ("upper_ux", "upper_uy", "center_ux", "center_uy", "lower_ux", "lower_uy"),
    "forces": ("front_fx", "front_fy", "upper_fx", "upper_fy", "lower_fx", "lower_fy"),
    "actions": ("front", "upper", "lower"),
}


def _load(path: Path) -> dict[str, np.ndarray]:
    with np.load(path, allow_pickle=False) as data:
        return {key: np.asarray(data[key]) for key in data.files}


def _dominant_frequency(values: np.ndarray) -> dict[str, Any]:
    centered = np.asarray(values, dtype=np.float64) - np.mean(values, axis=0, keepdims=True)
    spectrum = np.fft.rfft(centered, axis=0)
    power = np.abs(spectrum)
    if power.shape[0] <= 1:
        bins = np.zeros(centered.shape[1], dtype=int)
    else:
        bins = np.argmax(power[1:], axis=0) + 1
    return {
        "bin": bins.tolist(),
        "cycles_per_sample": (bins / max(centered.shape[0], 1)).tolist(),
        "phase_rad": np.angle(spectrum[bins, np.arange(centered.shape[1])]).tolist(),
    }


def _array_metrics(reference: np.ndarray, candidate: np.ndarray, names: Sequence[str]) -> dict[str, Any]:
    rows = min(len(reference), len(candidate))
    ref = np.asarray(reference[:rows], dtype=np.float64)
    cur = np.asarray(candidate[:rows], dtype=np.float64)
    if ref.shape != cur.shape:
        raise ValueError(f"aligned arrays differ in shape: {ref.shape} versus {cur.shape}")
    difference = cur - ref
    rmse = np.sqrt(np.mean(difference**2, axis=0))
    scale = np.ptp(ref, axis=0)
    nrmse = rmse / np.where(scale > 1e-12, scale, 1.0)
    return {
        "n_aligned": rows,
        "channels": list(names),
        "rmse": rmse.tolist(),
        "range_normalized_rmse": nrmse.tolist(),
        "max_abs_error": np.max(np.abs(difference), axis=0).tolist(),
        "reference_mean": np.mean(ref, axis=0).tolist(),
        "candidate_mean": np.mean(cur, axis=0).tolist(),
        "reference_std": np.std(ref, axis=0).tolist(),
        "candidate_std": np.std(cur, axis=0).tolist(),
        "reference_range": np.ptp(ref, axis=0).tolist(),
        "candidate_range": np.ptp(cur, axis=0).tolist(),
        "reference_frequency": _dominant_frequency(ref),
        "candidate_frequency": _dominant_frequency(cur),
    }


def _target_states(path: Path) -> np.ndarray:
    data = _load(path)
    if "target_states" in data:
        return np.asarray(data["target_states"])
    if "target" in data:
        return np.asarray(data["target"])
    raise KeyError(f"target NPZ has no target_states array: {path}")


def compare(
    scene: str,
    historical: Path,
    candidate: Path,
    historical_target: Path,
    candidate_target: Path,
    *,
    transient_count: int,
) -> dict[str, Any]:
    cfg = get_scene(scene)
    old = _load(historical)
    new = _load(candidate)
    missing = [key for key in ARRAY_KEYS if key not in old or key not in new]
    if missing:
        raise KeyError(f"trajectory pair lacks required arrays: {missing}")
    n = min(*(len(old[key]) for key in ARRAY_KEYS), *(len(new[key]) for key in ARRAY_KEYS))
    if not 0 <= transient_count < n:
        raise ValueError("transient_count must be non-negative and smaller than aligned trajectory length")
    segments = {
        "all_recorded_steps": slice(0, n),
        "initial_control_segment": slice(0, transient_count),
        "post_initial_segment": slice(transient_count, n),
    }
    segment_metrics: dict[str, Any] = {}
    for label, selection in segments.items():
        if selection.start == selection.stop:
            continue
        segment_metrics[label] = {
            key: _array_metrics(old[key][selection], new[key][selection], CHANNELS[key])
            for key in ARRAY_KEYS
        }
    old_target = _target_states(historical_target)
    new_target = _target_states(candidate_target)
    conv_len = int(cfg["conv_len"])
    similarity = (
        compute_similarity_illusion_legacy_v1
        if cfg["scene_id"] == "illusion"
        else compute_similarity_karman_legacy_v1
    )
    return {
        "schema_version": SCHEMA_VERSION,
        "scene": scene,
        "status": "diagnostic_only",
        "comparison_semantics": {
            "alignment": "same recorded index; no lag optimization",
            "transient_count": int(transient_count),
            "metric_algorithm": LEGACY_DTW_VERSION,
            "metric_window": {
                "target": f"indices {conv_len}:{2 * conv_len}",
                "state": f"last {conv_len} recorded samples",
            },
        },
        "sources": {
            "historical": {"path": str(historical), "sha256": hash_file(historical)},
            "candidate": {"path": str(candidate), "sha256": hash_file(candidate)},
            "historical_target": {"path": str(historical_target), "sha256": hash_file(historical_target)},
            "candidate_target": {"path": str(candidate_target), "sha256": hash_file(candidate_target)},
        },
        "lengths": {
            "historical": {key: len(old[key]) for key in ARRAY_KEYS},
            "candidate": {key: len(new[key]) for key in ARRAY_KEYS},
            "aligned": n,
        },
        "segments": segment_metrics,
        "legacy_dtw": {
            "historical_with_historical_target": float(similarity(old_target, old["sensors"][:n], conv_len)),
            "candidate_with_candidate_target": float(similarity(new_target, new["sensors"][:n], conv_len)),
            "algorithm": LEGACY_DTW_VERSION,
            "window_label": "legacy_reference_cycle_vs_last_recorded_cycle",
        },
        "interpretation": "Differences are reported, not auto-accepted; frequency phase is descriptive and is not a physical-delay estimate.",
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--scene", required=True)
    parser.add_argument("--historical", type=Path, required=True)
    parser.add_argument("--candidate", type=Path, required=True)
    parser.add_argument("--historical-target", type=Path, required=True)
    parser.add_argument("--candidate-target", type=Path, required=True)
    parser.add_argument("--transient-count", type=int, default=50)
    parser.add_argument("--output", type=Path, required=True)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    document = compare(
        args.scene,
        args.historical.resolve(),
        args.candidate.resolve(),
        args.historical_target.resolve(),
        args.candidate_target.resolve(),
        transient_count=args.transient_count,
    )
    atomic_write_json(args.output.resolve(), document)
    print(json.dumps(document, indent=2, sort_keys=True))
    return 0


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