#!/usr/bin/env python3
"""Verify recorded PPO actions by replaying the policy on causal recorded states."""
from __future__ import annotations

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

import numpy as np

from SR_analysis.configs import get_scene, model_path_for_scene
from SR_analysis.stage_1_infer import illusion_observation, load_existing_norm, normalize_raw_observation
from SR_analysis.utils.cfd_interface import load_ppo_model
from SR_analysis.utils.provenance import atomic_write_json, hash_file

SCHEMA_VERSION = "sr-policy-replay-parity-v1"


def replay(scene: str, trajectory: Path, *, model_device: str = "cpu") -> dict[str, Any]:
    cfg = get_scene(scene)
    norm, norm_path = load_existing_norm(scene, cfg)
    model_path = model_path_for_scene(scene)
    if model_path is None:
        raise FileNotFoundError(f"no PPO model configured for {scene}")
    model = load_ppo_model(model_path, device=model_device, s_dim=int(cfg["s_dim"]))
    with np.load(trajectory, allow_pickle=False) as data:
        sensors = np.asarray(data["sensors"], dtype=np.float64)
        forces = np.asarray(data["forces"], dtype=np.float64)
        actions = np.asarray(data["actions_norm" if "actions_norm" in data else "actions"], dtype=np.float64)
        targets = np.asarray(data["target_forces"], dtype=np.float64) if "target_forces" in data else None

    if len(actions) < 2:
        raise ValueError("policy replay requires at least two recorded actions")
    raw = np.column_stack((sensors, forces))
    predicted = []
    for index in range(1, len(actions)):
        state = raw[index - 1]
        if cfg["scene_id"] == "illusion":
            if targets is None:
                raise ValueError("Illusion policy replay requires target_forces")
            observation = illusion_observation(state, norm, targets[index])
        else:
            observation = normalize_raw_observation(state, norm)
        action, _ = model.predict(observation, deterministic=True)
        predicted.append(np.asarray(action, dtype=np.float64).reshape(3))

    predicted_array = np.asarray(predicted)
    recorded = actions[1:]
    residual = predicted_array - recorded
    max_abs = np.max(np.abs(residual), axis=0)
    rmse = np.sqrt(np.mean(residual**2, axis=0))
    tolerance = 2e-6
    return {
        "schema_version": SCHEMA_VERSION,
        "scene": scene,
        "status": "passed" if float(np.max(max_abs)) <= tolerance else "failed",
        "semantics": "recorded post-state i-1 is replayed to predict recorded normalized action i; first action is excluded because its pre-state is not stored",
        "tolerance": tolerance,
        "n_compared": int(len(recorded)),
        "action_layout": list(cfg["action_layout"]),
        "max_abs_error": max_abs.tolist(),
        "rmse": rmse.tolist(),
        "sources": {
            "trajectory": {"path": str(trajectory), "sha256": hash_file(trajectory)},
            "model": {"path": str(model_path), "sha256": hash_file(Path(model_path))},
            "norm": {"path": str(norm_path), "sha256": hash_file(norm_path)},
        },
        "interpretation": "This checks policy observation/action wiring without requiring separately initialized CFD trajectories to be pointwise identical.",
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--scene", required=True)
    parser.add_argument("--trajectory", type=Path, required=True)
    parser.add_argument("--model-device", default="cpu")
    parser.add_argument("--output", type=Path, required=True)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    report = replay(args.scene, args.trajectory.resolve(), model_device=args.model_device)
    atomic_write_json(args.output.resolve(), report)
    print(json.dumps(report, indent=2, sort_keys=True))
    return 0 if report["status"] == "passed" else 1


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