"""Closed-loop MuJoCo evaluation for the trained phase-imitation policy."""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
import sys

import mujoco
import numpy as np
import onnxruntime as ort


DANCE_DIR = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(DANCE_DIR))

from choreography import SHOW_SECONDS  # noqa: E402
from generate_scene import build_scene  # noqa: E402
from physics_show import (  # noqa: E402
    CONTROL_HZ,
    DECIMATION,
    DEFAULT_POLICY,
    PHYSICS_HZ,
    ContactDiagnostics,
    MultiDuckPolicy,
)


def phase_command(time_seconds: float, duration: float) -> np.ndarray:
    phase = (time_seconds % duration) / duration
    result = np.zeros(13, dtype=np.float32)
    result[0] = math.cos(2.0 * math.pi * phase)
    result[1] = math.sin(2.0 * math.pi * phase)
    return result


def run(args: argparse.Namespace) -> dict:
    output = args.output.resolve()
    output.parent.mkdir(parents=True, exist_ok=True)
    scene = output.parent / "phase_policy_eval_scene.xml"
    build_scene(scene, grid=1, spacing=1.5)
    model = mujoco.MjModel.from_xml_path(str(scene))
    model.opt.timestep = 1.0 / PHYSICS_HZ
    data = mujoco.MjData(model)
    indexing = MultiDuckPolicy(model, data, args.teacher_policy.resolve(), grid=1)
    indexing.reset()

    session = ort.InferenceSession(str(args.policy.resolve()))
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    metadata = session.get_modelmeta().custom_metadata_map
    action_scale = float(metadata["action_scale"])
    default_pose = np.array(
        [float(value) for value in metadata["default_joint_pos"].split(",")],
        dtype=np.float32,
    )
    if session.get_inputs()[0].shape[-1] != 61 or session.get_outputs()[0].shape[-1] != 14:
        raise ValueError("Phase policy must use the 61D -> 14D deployment contract")
    indexing.default_pose = default_pose
    indexing.data.ctrl[:] = default_pose
    indexing.last_actions.fill(0.0)
    mujoco.mj_forward(model, data)

    for _ in range(int(round(args.warmup * CONTROL_HZ))):
        command = phase_command(0.0, args.duration)
        observation = indexing.observation(0, command)
        action = session.run([output_name], {input_name: observation[None, :]})[0][0]
        indexing.last_actions[0] = action
        data.ctrl[indexing.actuator_ids[0]] = default_pose + action * action_scale
        for _ in range(DECIMATION):
            mujoco.mj_step(model, data)

    initial_xy = data.qpos[indexing.free_qpos[0]:indexing.free_qpos[0] + 2].copy()
    diagnostics = ContactDiagnostics(indexing, initial_xy[None, :])
    reference = np.load(args.reference.resolve(), allow_pickle=True)
    steps = int(round(args.duration * CONTROL_HZ))
    rollout_joint = np.empty((steps, 14), dtype=np.float32)
    rollout_action = np.empty((steps, 14), dtype=np.float32)
    observations = np.empty((steps, 61), dtype=np.float32)

    for index in range(steps):
        command = phase_command(index / CONTROL_HZ, args.duration)
        observation = indexing.observation(0, command)
        action = session.run([output_name], {input_name: observation[None, :]})[0][0]
        indexing.last_actions[0] = action
        data.ctrl[indexing.actuator_ids[0]] = default_pose + action * action_scale
        for _ in range(DECIMATION):
            mujoco.mj_step(model, data)
        diagnostics.update()
        observations[index] = observation
        rollout_action[index] = action
        rollout_joint[index] = data.qpos[indexing.joint_qpos[0]]

    reference_joint = reference["joint_pos"][:steps]
    joint_rmse = np.sqrt(np.mean((rollout_joint - reference_joint) ** 2, axis=0))
    report = diagnostics.report(args.duration)
    report["closed_loop_joint_rmse_rad"] = float(np.sqrt(np.mean((rollout_joint - reference_joint) ** 2)))
    report["closed_loop_joint_rmse_per_joint_rad"] = joint_rmse.round(7).tolist()
    report["policy_path"] = str(args.policy)
    report["reference_path"] = str(args.reference)
    report["passed"] = (
        not report["fallen_ducks"]
        and report["minimum_foot_liftoffs"] >= 2
        and report["maximum_airborne_fraction"] <= 0.20
        and np.isfinite(rollout_joint).all()
    )
    np.savez_compressed(
        output,
        time_seconds=np.arange(steps, dtype=np.float32) / CONTROL_HZ,
        joint_pos=rollout_joint,
        action=rollout_action,
        observation_61d=observations,
        reference_joint_pos=reference_joint,
    )
    output.with_suffix(".json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps(report, indent=2))
    if not report["passed"]:
        raise SystemExit("Closed-loop phase-policy validation failed")
    return report


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--policy", type=Path, required=True)
    parser.add_argument("--reference", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--teacher-policy", type=Path, default=DEFAULT_POLICY)
    parser.add_argument("--duration", type=float, default=SHOW_SECONDS)
    parser.add_argument("--warmup", type=float, default=1.5)
    run(parser.parse_args())


if __name__ == "__main__":
    main()