#!/usr/bin/env python3
"""Evaluate DreamWaQ checkpoints on a fixed MuJoCo stair benchmark.

The benchmark is independent from the training curriculum. Each trial must
climb ten box steps, cross the platform, descend ten steps, and finish upright.
Results are written to CSV and TensorBoard using the checkpoint iteration as
the global step.
"""

import argparse
import csv
import glob
import json
import math
import os
import re
import subprocess
import sys
from dataclasses import asdict, dataclass

import mujoco
import numpy as np
import onnxruntime as ort
from torch.utils.tensorboard import SummaryWriter

from dreamwaq_sim2sim_mujoco import (
    ACTION_SCALE,
    CLIP_ACTIONS,
    CLIP_TORQUES,
    DEFAULT_ANGLES,
    HISTORY_LEN,
    KD,
    KP,
    NUM_ACTIONS,
    NUM_OBS,
    body_pose,
    compute_obs,
    course_geometry,
    stage_for_x,
)


PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
XML_DIR = os.path.join(
    PROJECT_DIR, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls"
)
DEFAULT_RUNS_DIR = os.path.join(PROJECT_DIR, "runs", "go1-dreamwaq-walk", "rslrl")
CHECKPOINT_RE = re.compile(r"model_(\d+)\.pt$")


@dataclass
class TrialResult:
    success: bool
    reached_top: bool
    reached_finish: bool
    fell: bool
    timeout: bool
    completion_time_s: float
    max_x_m: float
    final_x_m: float
    final_tilt_deg: float
    final_stage: str


def checkpoint_iteration(path: str) -> int:
    match = CHECKPOINT_RE.search(path)
    if match is None:
        raise ValueError(f"checkpoint name must match model_N.pt: {path}")
    return int(match.group(1))


def latest_run_dir() -> str:
    runs = [path for path in glob.glob(os.path.join(DEFAULT_RUNS_DIR, "*")) if os.path.isdir(path)]
    if not runs:
        raise FileNotFoundError(f"no runs found under {DEFAULT_RUNS_DIR}")
    return max(runs, key=os.path.getmtime)


def discover_checkpoints(args) -> list[str]:
    if args.checkpoint:
        paths = [os.path.abspath(path) for path in args.checkpoint]
    else:
        run_dir = os.path.abspath(args.run_dir or latest_run_dir())
        paths = glob.glob(os.path.join(run_dir, "model_*.pt"))
    paths = [path for path in paths if os.path.isfile(path)]
    paths.sort(key=checkpoint_iteration)
    selected = []
    for path in paths:
        iteration = checkpoint_iteration(path)
        if iteration < args.min_iteration or iteration > args.max_iteration:
            continue
        if args.every > 0 and iteration % args.every != 0:
            continue
        selected.append(path)
    if not selected:
        raise FileNotFoundError("no checkpoints matched the requested iteration range")
    return selected


def export_checkpoint(checkpoint: str, output: str) -> None:
    command = [
        sys.executable,
        os.path.join(PROJECT_DIR, "scripts", "export_dreamwaq_onnx_new.py"),
        "--checkpoint",
        checkpoint,
        "--output",
        output,
    ]
    subprocess.run(command, cwd=PROJECT_DIR, check=True)


def load_stair_model(step_height_m: float) -> tuple[mujoco.MjModel, dict]:
    previous_cwd = os.getcwd()
    os.chdir(XML_DIR)
    try:
        model = mujoco.MjModel.from_xml_path(os.path.join(XML_DIR, "scene_stairs_box.xml"))
    finally:
        os.chdir(previous_cwd)

    up_ids = []
    down_ids = []
    for geom_id in range(model.ngeom):
        name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, geom_id) or ""
        if name.startswith("step_up_"):
            up_ids.append((int(name.rsplit("_", 1)[1]), geom_id))
        elif name.startswith("step_down_"):
            down_ids.append((int(name.rsplit("_", 1)[1]), geom_id))
    up_ids.sort()
    down_ids.sort()
    if not up_ids:
        raise RuntimeError("stairs_box has no ascending step geoms")

    num_steps = len(up_ids)
    for index, geom_id in up_ids:
        top = (index + 1) * step_height_m
        model.geom_pos[geom_id, 2] = top / 2.0
        model.geom_size[geom_id, 2] = top / 2.0
    for index, geom_id in down_ids:
        top = (num_steps - index - 1) * step_height_m
        model.geom_pos[geom_id, 2] = top / 2.0
        model.geom_size[geom_id, 2] = top / 2.0

    platform_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "platform")
    platform_top = num_steps * step_height_m
    model.geom_pos[platform_id, 2] = platform_top / 2.0
    model.geom_size[platform_id, 2] = platform_top / 2.0
    mujoco.mj_setConst(model, mujoco.MjData(model))
    course = course_geometry(model)
    if course is None:
        raise RuntimeError("failed to read stairs_box course geometry")
    return model, course


def run_trial(
    model: mujoco.MjModel,
    course: dict,
    session: ort.InferenceSession,
    speed_mps: float,
    timeout_s: float,
    lateral_offset_m: float,
    yaw_offset_rad: float,
) -> TrialResult:
    data = mujoco.MjData(model)
    mujoco.mj_resetData(model, data)
    spawn_x = -2.0
    spawn_z = 0.34
    data.qpos[0:3] = [spawn_x, lateral_offset_m, spawn_z]
    data.qpos[3:7] = [math.cos(yaw_offset_rad / 2.0), 0.0, 0.0,
                      math.sin(yaw_offset_rad / 2.0)]
    data.qpos[7:19] = DEFAULT_ANGLES
    mujoco.mj_forward(model, data)

    trunk_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "trunk")
    history = np.zeros((1, HISTORY_LEN, NUM_OBS), dtype=np.float32)
    history_initialized = False
    last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
    action = np.zeros(NUM_ACTIONS, dtype=np.float32)
    command = np.array([speed_mps, 0.0, 0.0], dtype=np.float32)
    actuator_joints = model.actuator_trnid[:, 0]
    decimation = 4
    physics_step = 0
    reached_top = False
    reached_finish = False
    fallen_since = None
    fell = False
    max_x = spawn_x
    final_stage = "approach"
    final_tilt = 0.0

    while data.time < timeout_s:
        if physics_step % decimation == 0:
            obs = compute_obs(model, data, command, last_action)
            if not history_initialized:
                history[:, -1] = obs
                history_initialized = True
            action = session.run(
                None,
                {"obs": obs.reshape(1, -1), "obs_history": history.reshape(1, -1)},
            )[0][0]
            if not np.all(np.isfinite(action)):
                fell = True
                break
            action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
            last_action = action.copy()
            if physics_step > 0:
                history[:, :-1] = history[:, 1:]
                history[:, -1] = obs

        target = np.clip(
            DEFAULT_ANGLES + action * ACTION_SCALE,
            model.jnt_range[actuator_joints, 0],
            model.jnt_range[actuator_joints, 1],
        )
        torque = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18]
        data.ctrl[:] = np.clip(torque, -CLIP_TORQUES, CLIP_TORQUES)
        mujoco.mj_step(model, data)
        physics_step += 1

        position, rpy = body_pose(data, trunk_id)
        max_x = max(max_x, float(position[0]))
        final_stage = stage_for_x(position[0], course)
        final_tilt = float(max(abs(rpy[0]), abs(rpy[1])))
        if (
            position[0] >= course["platform_start"] - 0.15
            and position[2] >= course["platform_top"] + 0.15
        ):
            reached_top = True
        if position[0] >= course["pass_x"]:
            reached_finish = True

        fallen = position[2] < 0.18 or final_tilt > 65.0
        if data.time > 1.0 and fallen:
            fallen_since = data.time if fallen_since is None else fallen_since
        else:
            fallen_since = None
        if fallen_since is not None and data.time - fallen_since >= 0.5:
            fell = True
            break

        if reached_top and reached_finish and position[2] < spawn_z + 0.25 and final_tilt < 35.0:
            return TrialResult(
                True, True, True, False, False, float(data.time), max_x,
                float(position[0]), final_tilt, final_stage,
            )

    position, _ = body_pose(data, trunk_id)
    return TrialResult(
        False,
        reached_top,
        reached_finish,
        fell,
        not fell and data.time >= timeout_s,
        float(data.time),
        max_x,
        float(position[0]),
        final_tilt,
        final_stage,
    )


def mastered_height(heights_cm: list[float], success_rates: list[float], threshold: float) -> float:
    mastered = 0.0
    for height, rate in sorted(zip(heights_cm, success_rates)):
        if rate < threshold:
            break
        mastered = height
    return mastered


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", action="append", help="model_N.pt; may be repeated")
    parser.add_argument("--run-dir", help="run directory; defaults to newest RSLRL run")
    parser.add_argument("--min-iteration", type=int, default=0)
    parser.add_argument("--max-iteration", type=int, default=10**9)
    parser.add_argument("--every", type=int, default=100, help="checkpoint interval")
    parser.add_argument("--heights-cm", default="4,6,8,10,12")
    parser.add_argument("--trials", type=int, default=5)
    parser.add_argument("--speed", type=float, default=0.5)
    parser.add_argument("--timeout", type=float, default=30.0)
    parser.add_argument("--success-threshold", type=float, default=0.8)
    parser.add_argument("--seed", type=int, default=5)
    parser.add_argument("--output-dir")
    parser.add_argument("--keep-onnx", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.trials < 1:
        raise ValueError("--trials must be at least 1")
    heights_cm = sorted({float(value) for value in args.heights_cm.split(",")})
    if not heights_cm or heights_cm[0] <= 0:
        raise ValueError("--heights-cm must contain positive values")
    checkpoints = discover_checkpoints(args)
    run_dir = os.path.dirname(checkpoints[-1])
    output_dir = os.path.abspath(args.output_dir or os.path.join(run_dir, "stairs_eval"))
    os.makedirs(output_dir, exist_ok=True)
    csv_path = os.path.join(output_dir, "results.csv")
    jsonl_path = os.path.join(output_dir, "trials.jsonl")
    writer = SummaryWriter(log_dir=os.path.join(output_dir, "tensorboard"))
    csv_exists = os.path.isfile(csv_path) and os.path.getsize(csv_path) > 0

    with open(csv_path, "a", newline="") as csv_file, open(jsonl_path, "a") as jsonl_file:
        fieldnames = [
            "iteration", "checkpoint", "step_height_cm", "trials", "success_rate",
            "reached_top_rate", "reached_finish_rate", "fall_rate",
            "median_completion_time_s", "mean_max_x_m",
        ]
        csv_writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
        if not csv_exists:
            csv_writer.writeheader()

        for checkpoint in checkpoints:
            iteration = checkpoint_iteration(checkpoint)
            onnx_path = os.path.join(output_dir, f"policy_{iteration}.onnx")
            export_checkpoint(checkpoint, onnx_path)
            session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
            rng = np.random.default_rng(args.seed + iteration)
            success_rates = []
            print(f"[Eval] iteration={iteration} checkpoint={checkpoint}")

            for height_cm in heights_cm:
                model, course = load_stair_model(height_cm / 100.0)
                results = []
                for trial in range(args.trials):
                    lateral = float(rng.uniform(-0.10, 0.10))
                    yaw = float(rng.uniform(-math.radians(5.0), math.radians(5.0)))
                    speed = float(args.speed * rng.uniform(0.9, 1.1))
                    result = run_trial(model, course, session, speed, args.timeout, lateral, yaw)
                    results.append(result)
                    jsonl_file.write(json.dumps({
                        "iteration": iteration,
                        "checkpoint": checkpoint,
                        "step_height_cm": height_cm,
                        "trial": trial,
                        "speed_mps": speed,
                        "lateral_offset_m": lateral,
                        "yaw_offset_rad": yaw,
                        **asdict(result),
                    }) + "\n")

                success_rate = float(np.mean([result.success for result in results]))
                top_rate = float(np.mean([result.reached_top for result in results]))
                finish_rate = float(np.mean([result.reached_finish for result in results]))
                fall_rate = float(np.mean([result.fell for result in results]))
                completed = [result.completion_time_s for result in results if result.success]
                median_time = float(np.median(completed)) if completed else float("nan")
                mean_max_x = float(np.mean([result.max_x_m for result in results]))
                success_rates.append(success_rate)

                suffix = f"{height_cm:g}cm"
                writer.add_scalar(f"Stairs/success_rate_{suffix}", success_rate, iteration)
                writer.add_scalar(f"Stairs/reached_top_rate_{suffix}", top_rate, iteration)
                writer.add_scalar(f"Stairs/reached_finish_rate_{suffix}", finish_rate, iteration)
                writer.add_scalar(f"Stairs/fall_rate_{suffix}", fall_rate, iteration)
                if completed:
                    writer.add_scalar(f"Stairs/completion_time_s_{suffix}", median_time, iteration)
                csv_writer.writerow({
                    "iteration": iteration,
                    "checkpoint": checkpoint,
                    "step_height_cm": height_cm,
                    "trials": args.trials,
                    "success_rate": success_rate,
                    "reached_top_rate": top_rate,
                    "reached_finish_rate": finish_rate,
                    "fall_rate": fall_rate,
                    "median_completion_time_s": median_time,
                    "mean_max_x_m": mean_max_x,
                })
                print(
                    f"  {height_cm:4.1f}cm success={success_rate:5.1%} "
                    f"top={top_rate:5.1%} finish={finish_rate:5.1%} "
                    f"fall={fall_rate:5.1%} max_x={mean_max_x:.2f}m"
                )

            mastered = mastered_height(heights_cm, success_rates, args.success_threshold)
            if len(heights_cm) == 1:
                auc = success_rates[0]
            else:
                auc = float(
                    np.trapezoid(success_rates, heights_cm)
                    / (heights_cm[-1] - heights_cm[0])
                )
            writer.add_scalar("Stairs/mastered_step_height_cm", mastered, iteration)
            writer.add_scalar("Stairs/success_auc", auc, iteration)
            writer.flush()
            csv_file.flush()
            jsonl_file.flush()
            print(f"  mastered={mastered:.1f}cm success_auc={auc:.3f}")
            if not args.keep_onnx:
                os.remove(onnx_path)

    writer.close()
    print(f"[Done] CSV: {csv_path}")
    print(f"[Done] TensorBoard: {os.path.join(output_dir, 'tensorboard')}")
    return 0


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