"""
Step 4 — Test on the 100 throws the network never saw.

    A. One step: from the TRUE state, how far off is the guess 10 ms later?
    B. Rollout:  from the true start only, feed each guess back in for 200 steps (2 s). Nothing corrects it.

The yardstick is "gravity only": exact free-fall physics with no floor.

Run:  python 04_evaluate.py [--tag hidden16_lr0.1_seed0]
Out:  results/<tag>/metrics.json, rollouts.npz, 04_rollouts.png
"""

import argparse, json
import numpy as np
import matplotlib.pyplot as plt
import ball
from model import load_network, predict_next

parser = argparse.ArgumentParser()
parser.add_argument("--tag", default="hidden16_lr0.1_seed0")
tag = parser.parse_args().tag
network = load_network(f"results/{tag}/network.pt")
data = np.load("data/throws.npz")
truth = data["states"][data["split"] == "test"]               # (100, 201, 4)
contact = data["contact"][data["split"] == "test"]            # (100, 200)


def gravity_only(s):
    """Newton with no floor: z += vz dt - g dt^2 / 2, vz -= g dt. x moves at constant vx."""
    dt, g = ball.MODEL_DT, 9.81
    return s + np.stack([s[..., 2] * dt, s[..., 3] * dt - g * dt**2 / 2, 0 * s[..., 2], -g * dt + 0 * s[..., 3]], -1)


def network_step(s):
    return predict_next(network, s)


# A. One step from the true state, scored separately in free flight and at the floor.
now, after = truth[:, :-1].reshape(-1, 4), truth[:, 1:].reshape(-1, 4)
floor = contact.reshape(-1)
one_step = {}
for name, step in [("network", network_step), ("gravity", gravity_only)]:
    miss = step(now) - after
    position_mm = np.linalg.norm(miss[:, :2], axis=1) * 1000
    velocity = np.linalg.norm(miss[:, 2:], axis=1)
    one_step[name] = {"flight_position_mm": float(np.sqrt(np.mean(position_mm[~floor] ** 2))),
                      "floor_velocity": float(np.sqrt(np.mean(velocity[floor] ** 2)))}
    print(f"{name:8s} free flight {one_step[name]['flight_position_mm']:7.3f} mm   "
          f"at the floor {one_step[name]['floor_velocity']:6.3f} m/s")

# B. Rollouts: each prediction becomes the next input.
rollouts = {}
for name, step in [("network", network_step), ("gravity", gravity_only)]:
    path = np.zeros_like(truth)
    path[:, 0] = truth[:, 0]
    for t in range(truth.shape[1] - 1):
        path[:, t + 1] = step(path[:, t])
    rollouts[name] = path
error_cm = {k: np.linalg.norm(v[:, :, :2] - truth[:, :, :2], axis=2) * 100 for k, v in rollouts.items()}
for k, e in error_cm.items():
    print(f"{k:8s} median rollout error: {np.median(e[:, 20]):6.1f} cm at 0.2 s, {np.median(e[:, 100]):6.1f} cm at 1 s")

json.dump({"one_step": one_step, "median_error_cm": {k: np.median(e, axis=0).round(3).tolist() for k, e in error_cm.items()}},
          open(f"results/{tag}/metrics.json", "w"))
np.savez_compressed(f"results/{tag}/rollouts.npz", truth=truth[:8], **{k: v[:8] for k, v in rollouts.items()})

fig, axes = plt.subplots(2, 2, figsize=(10, 6))
for i, ax in enumerate(axes.flat):
    ax.plot(truth[i, :, 0], truth[i, :, 1], lw=2.5, label="MuJoCo")
    ax.plot(rollouts["network"][i, :, 0], rollouts["network"][i, :, 1], "--", label="network rollout")
    ax.axhline(0, color="grey", lw=1); ax.set_ylim(-0.3, 1.8); ax.set_title(f"test throw {i}")
axes[0, 0].legend()
fig.tight_layout(); fig.savefig(f"results/{tag}/04_rollouts.png", dpi=120)