"""Deterministic, headless evaluation for Unitree Go2 velocity policies."""

from __future__ import annotations

import hashlib
import json
import os
import sys
from dataclasses import asdict, dataclass, field
from pathlib import Path
from statistics import fmean, pstdev
from typing import Literal

TRAINER_ROOT = Path(__file__).resolve().parents[1]
SERVICE_ROOT = TRAINER_ROOT.parent
for source_root in (TRAINER_ROOT, SERVICE_ROOT):
  if str(source_root) not in sys.path:
    sys.path.insert(0, str(source_root))

import torch
import tyro
import warp as wp

if not hasattr(wp, "context"):
  from warp._src import context as warp_context

  wp.context = warp_context  # type: ignore[attr-defined]

from mjlab.envs import ManagerBasedRlEnv
from mjlab.rl import MjlabOnPolicyRunner, RslRlVecEnvWrapper
from mjlab.tasks.registry import list_tasks, load_env_cfg, load_rl_cfg, load_runner_cls
from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
from mjlab.utils.torch import configure_torch_backends
from torch.utils.tensorboard import SummaryWriter

from tuning.schema import apply_reward_configuration, validate_configuration

SCENARIOS = (
  (0.0, 0.0, 0.0),
  (0.5, 0.0, 0.0),
  (1.0, 0.0, 0.0),
  (1.5, 0.0, 0.0),
  (0.0, 0.5, 0.0),
  (0.0, -0.5, 0.0),
  (0.0, 0.0, 0.5),
  (0.0, 0.0, -0.5),
  (0.8, 0.25, 0.35),
)
METRIC_NAMES = (
  "linear_velocity_rmse",
  "angular_velocity_rmse",
  "mean_action_acc",
  "orientation_error",
  "slip_velocity",
  "mechanical_power",
)


@dataclass(frozen=True)
class EvaluateConfig:
  checkpoint: str
  output: str
  reward_config: str | None = None
  num_envs: int = 256
  steps_per_seed: int = 1000
  seeds: tuple[int, ...] = field(default_factory=lambda: (101, 202, 303))
  device: str | None = None
  gpu_ids: list[int] | Literal["all"] | None = field(default_factory=lambda: [0])


def _sha256(path: Path) -> str:
  digest = hashlib.sha256()
  with path.open("rb") as stream:
    for chunk in iter(lambda: stream.read(1024 * 1024), b""):
      digest.update(chunk)
  return digest.hexdigest()


def _fixed_commands(command: torch.Tensor) -> torch.Tensor:
  values = torch.as_tensor(SCENARIOS, device=command.device, dtype=command.dtype)
  indexes = torch.arange(command.shape[0], device=command.device) % values.shape[0]
  command[:] = values[indexes]
  return command


def _evaluate_seed(task_id: str, cfg: EvaluateConfig, seed: int) -> dict[str, float]:
  torch.manual_seed(seed)
  env_cfg = load_env_cfg(task_id, play=False)
  agent_cfg = load_rl_cfg(task_id)
  env_cfg.seed = seed
  env_cfg.scene.num_envs = cfg.num_envs
  env_cfg.curriculum = {}
  env_cfg.observations["actor"].enable_corruption = False
  env_cfg.events.pop("push_robot", None)
  twist_cfg = env_cfg.commands["twist"]
  assert isinstance(twist_cfg, UniformVelocityCommandCfg)
  twist_cfg.heading_command = False
  twist_cfg.ranges.heading = None
  twist_cfg.rel_heading_envs = 0.0
  twist_cfg.rel_standing_envs = 0.0
  twist_cfg.resampling_time_range = (1.0e9, 1.0e9)
  if cfg.reward_config:
    reward_path = Path(cfg.reward_config).expanduser().resolve(strict=True)
    with reward_path.open(encoding="utf-8") as stream:
      apply_reward_configuration(env_cfg, validate_configuration(json.load(stream)))

  env = ManagerBasedRlEnv(cfg=env_cfg, device=cfg.device or "cuda:0")
  wrapped = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
  try:
    runner_cls = load_runner_cls(task_id) or MjlabOnPolicyRunner
    runner = runner_cls(wrapped, asdict(agent_cfg), log_dir=None, device=wrapped.device)
    runner.load(
      str(Path(cfg.checkpoint).expanduser().resolve(strict=True)),
      load_cfg={"actor": True},
      strict=True,
      map_location=str(wrapped.device),
    )
    policy = runner.get_inference_policy(device=str(wrapped.device))
    twist = wrapped.unwrapped.command_manager.get_term("twist")
    _fixed_commands(twist.command)
    obs = wrapped.get_observations()

    sums = {name: 0.0 for name in METRIC_NAMES}
    samples = 0
    terminations = 0
    completions = 0
    with torch.inference_mode():
      for _ in range(cfg.steps_per_seed):
        _fixed_commands(twist.command)
        obs = wrapped.get_observations()
        actions = policy(obs)
        obs, _rewards, _dones, _extras = wrapped.step(actions)
        manager = wrapped.unwrapped.metrics_manager
        for index, name in enumerate(manager.active_terms):
          if name in sums:
            sums[name] += float(torch.sum(manager._step_values[:, index]).item())
        samples += wrapped.num_envs
        terminated = wrapped.unwrapped.reset_terminated
        timed_out = wrapped.unwrapped.reset_time_outs
        terminations += int(torch.count_nonzero(terminated).item())
        completions += int(torch.count_nonzero(terminated | timed_out).item())
    result = {name: sums[name] / max(samples, 1) for name in METRIC_NAMES}
    result["fall_rate"] = terminations / max(completions, wrapped.num_envs)
    return result
  finally:
    wrapped.close()


def run_evaluation(task_id: str, cfg: EvaluateConfig) -> dict:
  configure_torch_backends()
  selected = cfg.gpu_ids
  if selected is None:
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    device = "cpu"
  else:
    if selected == "all":
      selected = [0]
    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, selected))
    device = cfg.device or "cuda:0"
  os.environ["MUJOCO_GL"] = "egl"
  cfg = EvaluateConfig(**{**asdict(cfg), "device": device})

  checkpoint = Path(cfg.checkpoint).expanduser().resolve(strict=True)
  output = Path(cfg.output).expanduser().resolve()
  output.parent.mkdir(parents=True, exist_ok=True)
  per_seed = [_evaluate_seed(task_id, cfg, seed) for seed in cfg.seeds]
  metrics = {
    key: fmean(seed_metrics[key] for seed_metrics in per_seed)
    for key in (*METRIC_NAMES, "fall_rate")
  }
  deviations = {
    key: pstdev(seed_metrics[key] for seed_metrics in per_seed)
    for key in (*METRIC_NAMES, "fall_rate")
  }
  result = {
    "protocolVersion": 1,
    "taskId": task_id,
    "checkpoint": checkpoint.name,
    "checkpointSha256": _sha256(checkpoint),
    "seeds": list(cfg.seeds),
    "numEnvs": cfg.num_envs,
    "stepsPerSeed": cfg.steps_per_seed,
    "scenarios": [list(value) for value in SCENARIOS],
    "metrics": metrics,
    "metricStd": deviations,
    "seedMetrics": [
      {"seed": seed, "metrics": values}
      for seed, values in zip(cfg.seeds, per_seed, strict=True)
    ],
  }
  output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
  writer = SummaryWriter(log_dir=str(output.parent / "evaluation-events"))
  try:
    for name, value in metrics.items():
      writer.add_scalar(f"Evaluation/{name}", value, 0)
  finally:
    writer.close()
  print("MUJOCO_EVALUATION " + json.dumps({"output": str(output), "metrics": metrics}))
  return result


def main() -> None:
  import mjlab.tasks  # noqa: F401
  import src.tasks  # noqa: F401

  chosen_task, remaining = tyro.cli(
    tyro.extras.literal_type_from_choices(list_tasks()),
    add_help=False,
    return_unknown_args=True,
    config=mjlab.TYRO_FLAGS,
  )
  args = tyro.cli(
    EvaluateConfig,
    args=remaining,
    prog=sys.argv[0] + f" {chosen_task}",
    config=mjlab.TYRO_FLAGS,
  )
  run_evaluation(chosen_task, args)


if __name__ == "__main__":
  main()