# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Evaluate trained Go1 policy with specific velocity commands.

Tests three command patterns:
  1. Forward-Backward (前后往返): vx oscillates +1.0 <-> -1.0
  2. Left-Right (左右往返): vy oscillates +1.0 <-> -1.0
  3. Rotation (旋转): wz oscillates +1.0 <-> -1.0

For each pattern, the command toggles direction every N seconds.
Metrics (tracking errors, etc.) are logged to CSV files for analysis.
"""

import logging
import time
from pathlib import Path

import numpy as np
from absl import app, flags

from motrix_envs import registry as env_registry
from motrix_rl import registry, utils
from motrix_rl.skrl.jax import wrap_env
from motrix_rl.skrl.jax.train.ppo import Trainer as SkrlJaxTrainer

logger = logging.getLogger(__name__)

_ENV = flags.DEFINE_string("env", "go1-flat-terrain-walk", "The env to evaluate")
_POLICY = flags.DEFINE_string(
    "policy",
    None,
    "Path to policy checkpoint. Auto-discovers the latest best_agent.pickle if not specified.",
)
_DURATION = flags.DEFINE_float("duration", 4.0, "Seconds per command direction before toggling")
_OUTDIR = flags.DEFINE_string("outdir", None, "Output directory for CSV logs")
_RENDER = flags.DEFINE_bool("render", False, "Enable rendering (may crash on headless)")

# ── Helpers ─────────────────────────────────────────────────────────────


def _find_best_policy(env_name: str) -> Path:
    """Auto-discover the best SKRL policy for the given env."""
    base = Path(f"runs/{env_name}/skrl")
    if not base.exists():
        raise FileNotFoundError(f"No training runs found at {base}")

    runs = sorted([d for d in base.iterdir() if d.is_dir()], key=lambda d: d.stat().st_mtime, reverse=True)
    if not runs:
        raise FileNotFoundError(f"No training runs found at {base}")

    ckpt_dir = runs[0] / "checkpoints"
    best = list(ckpt_dir.glob("best_agent.*"))
    if best:
        return best[0]

    # Fallback: highest timestep
    ckpts = list(ckpt_dir.glob("agent_*.pickle"))
    if not ckpts:
        raise FileNotFoundError(f"No checkpoints found in {ckpt_dir}")

    def _ts(p):
        try:
            return int(p.stem.split("_")[1])
        except (IndexError, ValueError):
            return 0

    return max(ckpts, key=_ts)


def _save_log(outdir: Path, label: str, records: list[dict]) -> None:
    """Save records to a CSV file."""
    outdir.mkdir(parents=True, exist_ok=True)
    import csv

    path = outdir / f"{label}.csv"
    with open(path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=records[0].keys())
        writer.writeheader()
        writer.writerows(records)
    logger.info(f"Saved {len(records)} records → {path}")


# ── Command generators ──────────────────────────────────────────────────


def _command_forward_backward(t: float, period: float):
    """vx oscillates +1.0 / -1.0, vy=0, wz=0."""
    half = period / 2.0
    phase = (t % period) / half  # 0..1 forward, 1..2 backward
    vx = 1.0 if phase < 1.0 else -1.0
    return np.array([vx, 0.0, 0.0], dtype=np.float32)


def _command_left_right(t: float, period: float):
    """vy oscillates +1.0 / -1.0, vx=0, wz=0."""
    half = period / 2.0
    phase = (t % period) / half
    vy = 1.0 if phase < 1.0 else -1.0
    return np.array([0.0, vy, 0.0], dtype=np.float32)


def _command_rotation(t: float, period: float):
    """wz oscillates +1.0 / -1.0, vx=0, vy=0."""
    half = period / 2.0
    phase = (t % period) / half
    wz = 1.0 if phase < 1.0 else -1.0
    return np.array([0.0, 0.0, wz], dtype=np.float32)


# ── Main evaluation logic ───────────────────────────────────────────────


def _run_pattern(
    trainer: SkrlJaxTrainer,
    policy_path: str,
    pattern_name: str,
    command_fn,
    period: float,
    total_steps: int,
    ctrl_dt: float,
) -> list[dict]:
    """Run one command pattern and return tracking records."""
    logger.info(f"--- {pattern_name} ---")

    rlcfg = trainer._rlcfg
    env = env_registry.make(trainer._env_name, sim_backend=trainer._sim_backend, num_envs=1)
    env = wrap_env(env, enable_render=False)

    # Build fresh agent for this run
    models = trainer._make_model(env, rlcfg)
    ppo_cfg = rlcfg.runner.agent.to_dict()
    from motrix_rl.skrl.jax.train.ppo import _add_runtime_config
    _add_runtime_config(ppo_cfg, env)
    agent = trainer._make_agent(models, env, ppo_cfg, rlcfg.runner.memory)
    agent.load(policy_path)

    obs, info = env.reset()
    state = env._env.state
    records = []

    for step in range(total_steps):
        t = step * ctrl_dt

        # Override command in env state
        cmd = command_fn(t, period)
        state.info["commands"] = cmd.reshape(1, -1)

        # Recompute observation with the new command
        new_obs = env._env._get_obs(state.data, state.info)
        state = state.replace(obs=new_obs)
        obs = new_obs

        # Agent inference
        outputs = agent.act(obs, timestep=0, timesteps=0)
        actions = outputs[-1].get("mean_actions", outputs[0])
        obs, reward, terminated, truncated, info = env.step(actions)
        state = env._env.state

        # Collect tracking data
        lin_vel = env._env.get_local_linvel(state.data)[0]  # [vx, vy, vz] (body frame)
        gyro = env._env.get_gyro(state.data)[0]

        tracking_err_xy = np.linalg.norm(cmd[:2] - lin_vel[:2])
        tracking_err_yaw = abs(cmd[2] - gyro[2])

        records.append({
            "step": step,
            "time": round(t, 3),
            "cmd_vx": round(float(cmd[0]), 4),
            "cmd_vy": round(float(cmd[1]), 4),
            "cmd_wz": round(float(cmd[2]), 4),
            "actual_vx": round(float(lin_vel[0]), 4),
            "actual_vy": round(float(lin_vel[1]), 4),
            "actual_vz": round(float(lin_vel[2]), 4),
            "actual_wz": round(float(gyro[2]), 4),
            "tracking_err_xy": round(float(tracking_err_xy), 6),
            "tracking_err_yaw": round(float(tracking_err_yaw), 6),
            "reward": round(float(reward[0][0]), 6),
        })

        if step % 100 == 0:
            logger.info(
                f"  [{pattern_name}] step {step:5d}/{total_steps}  "
                f"cmd=[{cmd[0]:+.1f},{cmd[1]:+.1f},{cmd[2]:+.1f}]  "
                f"actual_v=[{lin_vel[0]:+.3f},{lin_vel[1]:+.3f},{gyro[2]:+.3f}]  "
                f"track_err_xy={tracking_err_xy:.4f}  "
                f"reward={float(reward[0][0]):.4f}"
            )

    env.close()
    return records


def main(argv):
    env_name = _ENV.value
    ctrl_dt = 0.01  # matches Go1WalkNpEnvCfg.ctrl_dt
    period = _DURATION.value * 2.0  # full cycle: forward + backward

    # Resolve policy path
    if _POLICY.present:
        policy_path = _POLICY.value
    else:
        policy_path = str(_find_best_policy(env_name))
    logger.info(f"Policy: {policy_path}")

    # Build trainer (used to construct models & agent)
    rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="jax")
    trainer = SkrlJaxTrainer(env_name, sim_backend=None, enable_render=False)
    trainer._rlcfg = rlcfg

    # Total steps per pattern = enough full cycles
    cycles = 3
    total_steps = int(cycles * period / ctrl_dt)

    # Output directory
    if _OUTDIR.present:
        outdir = Path(_OUTDIR.value)
    else:
        ts = time.strftime("%y-%m-%d_%H-%M-%S")
        outdir = Path(f"runs/{env_name}/eval_{ts}")
    logger.info(f"Output directory: {outdir}")

    patterns = [
        ("forward_backward", _command_forward_backward),
        ("left_right", _command_left_right),
        ("rotation", _command_rotation),
    ]

    summary = {}
    for label, cmd_fn in patterns:
        records = _run_pattern(
            trainer, policy_path, label, cmd_fn,
            period=period, total_steps=total_steps, ctrl_dt=ctrl_dt,
        )
        _save_log(outdir, label, records)

        # Summary stats (steady-state, skip first 2s for stabilization)
        warmup = int(2.0 / ctrl_dt)
        steady = records[warmup:]
        if steady:
            avg_err = np.mean([r["tracking_err_xy"] for r in steady])
            avg_reward = np.mean([r["reward"] for r in steady])
        else:
            avg_err, avg_reward = float("nan"), float("nan")
        summary[label] = {"avg_tracking_err_xy": avg_err, "avg_reward": avg_reward}
        logger.info(f"  [{label}] steady-state avg tracking_err_xy = {avg_err:.4f}, avg_reward = {avg_reward:.4f}")

    # Print summary
    print("\n" + "=" * 70)
    print("EVALUATION SUMMARY")
    print("=" * 70)
    for label, stats in summary.items():
        print(f"  {label:25s}  track_err_xy={stats['avg_tracking_err_xy']:.4f}  avg_reward={stats['avg_reward']:.4f}")
    print(f"\nDetailed CSV logs saved to: {outdir}")
    print("=" * 70)


if __name__ == "__main__":
    app.run(main)