#!/usr/bin/env python3
"""Evaluate empirical PPO G-equivariance on recorded physical 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
from SR_analysis.stage_1_infer import (
    illusion_observation,
    load_existing_norm,
    model_path_for_scene,
    normalize_raw_observation,
)
from SR_analysis.utils.cfd_interface import load_ppo_model
from SR_analysis.utils.g_operator import apply_G_raw12, apply_G_target_force
from SR_analysis.utils.provenance import atomic_write_json, hash_file
from SR_analysis.utils.symmetry_dataset import empirical_equivariance_defect

SCHEMA_VERSION = "sr-ppo-g-audit-v1"


def _alpha(model: Any, observation: np.ndarray, cfg: dict[str, Any]) -> np.ndarray:
    normalized, _ = model.predict(observation, deterministic=True)
    normalized = np.asarray(normalized, dtype=np.float64).reshape(3)
    return normalized * float(cfg["action_scale"]) + np.asarray(cfg["action_bias"], dtype=np.float64)


def audit(scene: str, trajectory_path: 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_path, allow_pickle=False) as data:
        sensors = np.asarray(data["sensors"], dtype=np.float64)
        forces = np.asarray(data["forces"], dtype=np.float64)
        targets = np.asarray(data["target_forces"], dtype=np.float64) if "target_forces" in data else None
    raw = np.column_stack((sensors, forces))
    mirrored_raw = apply_G_raw12(raw)
    actions, mirrored_actions = [], []
    for index, (state, mirrored_state) in enumerate(zip(raw, mirrored_raw)):
        if cfg["scene_id"] == "illusion":
            if targets is None:
                raise ValueError("illusion symmetry audit requires target_forces")
            target = targets[index]
            observation = illusion_observation(state, norm, target)
            mirrored_observation = illusion_observation(
                mirrored_state, norm, apply_G_target_force(target),
            )
        else:
            observation = normalize_raw_observation(state, norm)
            mirrored_observation = normalize_raw_observation(mirrored_state, norm)
        actions.append(_alpha(model, observation, cfg))
        mirrored_actions.append(_alpha(model, mirrored_observation, cfg))
    return {
        "schema_version": SCHEMA_VERSION,
        "scene": scene,
        "status": "diagnostic_only",
        "g_contract": "raw state and target are mirrored before PPO normalization/evaluation",
        "trajectory": {"path": str(trajectory_path), "sha256": hash_file(trajectory_path)},
        "model": {"path": str(model_path), "sha256": hash_file(model_path)},
        "norm": {"path": str(norm_path), "sha256": hash_file(norm_path)},
        "defect": empirical_equivariance_defect(np.asarray(actions), np.asarray(mirrored_actions)),
        "interpretation": "A nonzero defect means hard-G fitting changes the empirical PPO policy and must be reported as physical symmetrization.",
    }


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)
    document = audit(args.scene, args.trajectory.resolve(), model_device=args.model_device)
    atomic_write_json(args.output.resolve(), document)
    print(json.dumps(document, indent=2, sort_keys=True))
    return 0


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