#!/usr/bin/env python3
"""evaluate.py — run a trained policy on every body in the physical test set, in sim.
This produces the SIM side of the scaling curve; the real side comes from the robot.

  python3 evaluate.py runs/n100.pkl                # all 10 test bodies
  python3 evaluate.py runs/n100.pkl --body tripod  # one
Outputs runs/<tag>_eval.json: per body, success rate, velocity error, CoT proxy.
"""
import argparse, functools, json, pickle
import compat  # noqa: F401  (jax/brax API shim)
import jax, jax.numpy as jp, numpy as np
from brax.training.agents.ppo import networks as ppo_networks
from brax.training.acme import running_statistics
from brax.envs.wrappers import training as W
import env as E, generator as G

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("ckpt"); ap.add_argument("--body", default=None)
    ap.add_argument("--trials", type=int, default=32); ap.add_argument("--cmd", type=float, default=0.25)
    ap.add_argument("--steps", type=int, default=500, help="50 Hz steps per trial; 500 = 10 s")
    a = ap.parse_args()
    ck = pickle.load(open(a.ckpt, "rb")); params = ck["params"]
    env = E.ModBotEnv()
    hid = tuple(ck.get("hidden", (256, 256, 128)))
    nets = ppo_networks.make_ppo_networks(env.observation_size, env.action_size,
                                          preprocess_observations_fn=running_statistics.normalize,
                                          policy_hidden_layer_sizes=hid, value_hidden_layer_sizes=hid)
    make_policy = ppo_networks.make_inference_fn(nets)
    policy = jax.jit(make_policy(params, deterministic=True))

    bodies = [a.body] if a.body else list(G.TEST_SET)
    results = {}
    for name in bodies:
        rand = E.make_randomizer(env, mode=f"test:{name}")
        rfn = functools.partial(rand, rng=jax.random.split(jax.random.PRNGKey(0), a.trials))
        venv = W.wrap(env, episode_length=env.max_steps, randomization_fn=rfn)
        reset, step = jax.jit(venv.reset), jax.jit(venv.step)
        s = reset(jax.random.split(jax.random.PRNGKey(1), a.trials))
        s.info["cmd"] = jp.tile(jp.array([a.cmd, 0.0, 0.0]), (a.trials, 1))
        key = jax.random.PRNGKey(2)
        verr, torque, alive = [], [], np.ones(a.trials)
        for t in range(a.steps):
            key, k = jax.random.split(key)
            act, _ = policy(s.obs, k)
            s = step(s, act)
            alive *= (1 - np.array(s.metrics["fall"]))
            verr.append(np.array(1 - s.metrics["r_lin"])); torque.append(np.array(s.metrics["p_torque"]))
        results[name] = {"success": float(alive.mean()),
                         "vel_err_proxy": float(np.mean(verr)),
                         "torque_sq": float(np.mean(torque)),
                         "legs": G.TEST_SET[name].n_legs}
        print(f"{name:14s} legs={results[name]['legs']}  success {results[name]['success']:.2f}  "
              f"vel_err {results[name]['vel_err_proxy']:.3f}  torque² {results[name]['torque_sq']:.2f}")
    out = a.ckpt.replace(".pkl", "_eval.json")
    json.dump(results, open(out, "w"), indent=1); print("saved", out)

if __name__ == "__main__":
    main()