"""Run a policy (or a scripted action) and write a replay for the three.js bench.

  python rollout.py --policy ../runs/<run>/policy_best.json --out ../bench/replays/latest.json

The replay carries world poses per body, so the viewer needs no kinematics of
its own: it just moves the meshes it built from the same spec file.
"""

import argparse
import json
from pathlib import Path

import numpy as np

from env import RobotEnv
from ppo import Policy


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--spec", default=None)
    ap.add_argument("--policy", default=None, help="omit for a zero-action fall")
    ap.add_argument("--task", default=None, help="incentive to run under")
    ap.add_argument("--out", default=str(Path(__file__).parent.parent / "bench/replays/latest.json"))
    ap.add_argument("--seconds", type=float, default=12.0)
    ap.add_argument("--seed", type=int, default=123)
    ap.add_argument("--fps", type=int, default=50, help="replay frame rate")
    ap.add_argument("--no-dr", action="store_true")
    ap.add_argument("--stochastic", action="store_true")
    ap.add_argument("--cmd", type=float, nargs=2, default=None,
                    help="fix the command to vx wz instead of resampling")
    args = ap.parse_args()

    policy = Policy.load(args.policy) if args.policy else None
    spec_path = args.spec or (policy.meta.get("spec") if policy else None)
    if not spec_path:
        spec_path = str(Path(__file__).parent.parent / "robots/wheeled_biped.json")

    task = args.task or (policy.meta.get("task") if policy else None)
    # A checkpoint knows the goal-channel layout it was trained with. Honour it,
    # or the observation silently changes width under the policy.
    layout = policy.meta.get("obs_layout", "legacy") if policy else "union"
    env = RobotEnv(spec_path, seed=args.seed, task=task, layout=layout,
                   randomise=(False if args.no_dr else None))
    rng = np.random.default_rng(args.seed)
    obs = env.reset()
    if args.cmd:
        # Pin the command instead of letting the incentive resample it, so a
        # replay can show one specific manoeuvre.
        if not hasattr(env.task_obj, "cmd"):
            raise SystemExit(f"--cmd does not apply to the {env.task_name!r} incentive")
        env.task_obj.cmd = np.array(args.cmd, dtype=float)
        env.task_obj.on_step = lambda *a, **k: None

    replay = record(env, policy, obs, seconds=args.seconds, seed=args.seed,
                    fps=args.fps, deterministic=not args.stochastic,
                    policy_path=args.policy)

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    with open(out, "w") as f:
        json.dump(replay, f)

    s = replay["summary"]
    verdict = f"FELL at {s['survived_s']}s" if s["fell"] else f"survived {s['survived_s']}s"
    ok = " SUCCESS" if s["success"] else ""
    print(f"{out}  [{env.task_name}] {verdict}{ok}  "
          f"reward {s['total_reward']}  {len(replay['frames'])} frames")


def record(env, policy, obs, seconds=12.0, seed=0, fps=50, deterministic=True,
           policy_path=None, steps_trained=None):
    """Run one episode and return a replay dict.

    Shared by the CLI and by the trainer's periodic snapshots, so a snapshot
    taken mid-training is byte-compatible with a replay produced by hand.
    """
    rng = np.random.default_rng(seed)
    n_steps = int(seconds / env.control_dt)
    every = max(1, int(round((1.0 / fps) / env.control_dt)))

    frames = [env.frame()]
    series = {"t": [], "height": [], "up_z": [], "reward": [],
              "cmd_vx": [], "cmd_wz": [], "action": [], "tau": []}
    total = 0.0
    fell_at = None

    for i in range(n_steps):
        if policy:
            a, _, _ = policy.act(obs[None], rng, deterministic=deterministic)
            a = np.clip(a[0], -1, 1)
        else:
            a = np.zeros(env.act_dim)
        obs, r, fell, trunc, info = env.step(a)
        total += r

        if i % every == 0:
            frames.append(env.frame())
            series["t"].append(round(env.t, 3))
            series["height"].append(round(info["height"], 4))
            series["up_z"].append(round(info["up_z"], 4))
            series["reward"].append(round(r, 3))
            # Only some incentives issue a velocity command; the rest log zero
            # so the telemetry strip keeps a consistent shape.
            cmd = getattr(env.task_obj, "cmd", (0.0, 0.0))
            series["cmd_vx"].append(round(float(cmd[0]), 3))
            series["cmd_wz"].append(round(float(cmd[1]), 3))
            series["action"].append([round(float(v), 3) for v in a])
            # actual joint torques, so the viewer can colour each part by effort
            series["tau"].append([round(float(v), 4) for v in env.data.actuator_force])

        if fell:
            fell_at = round(env.t, 2)
            break

    # Carry the spec inside the replay. The bench used to rebuild the robot
    # from whatever robots/*.json currently says, so editing the geometry made
    # every existing replay render with the wrong limb lengths against frames
    # recorded from the old body. A replay has to be self-describing.
    spec_copy = {k: v for k, v in env.spec.items() if k != "_path"}

    return {
        "spec": Path(env.spec["_path"]).name,
        "spec_json": spec_copy,
        "task": env.task_name,
        "task_label": env.task_obj.label,
        "props": env.props,
        "spec_path": env.spec["_path"],
        "policy": policy_path,
        "steps_trained": steps_trained,
        "bodies": env.body_names(),
        "dt": round(every * env.control_dt, 4),
        "episode_s": env.max_steps * env.control_dt,
        "tau_limit": [round(float(x), 4) for x in env.tau_limit],
        "actuators": [a["name"] for a in env.spec.get("actuators", [])],
        "frames": frames,
        "series": series,
        "summary": {
            "survived_s": fell_at if fell_at is not None else round(env.t, 2),
            "fell": fell_at is not None,
            "total_reward": round(total, 1),
            "steps": i + 1,
            "randomised": bool(env.dr.get("enabled")),
            "success": bool(env.task_obj.success(env)),
        },
    }


if __name__ == "__main__":
    main()