"""Regenerate dexterity figures from a saved dexterity.json (no GPU)."""
import json
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

out_dir = Path(sys.argv[1])
d = json.load(open(out_dir / "dexterity.json"))
report, curves = d["report"], d["curves"]
trajs, force_fig = d["trajectories"], d["force_fig"]
names = list(report)
GT = "#878c98"
palette = ["#3c7ec0", "#b54e7f", "#6f66bd", "#c26a1e"]
colors = {n: palette[i % len(palette)] for i, n in enumerate(names)}

fig, axes = plt.subplots(1, 4, figsize=(16, 3.4))
first = names[0]
for k in range(4):
    ax = axes[k]
    ax.plot(np.array(trajs[first]["gt"])[:, k], color=GT, ls="--", lw=2, label="GT")
    for n in names:
        ax.plot(np.array(trajs[n]["pred"])[:, k], color=colors[n], lw=1.8, label=n)
    ax.set_title(f"hand dim {trajs[first]['joints'][k]}", fontsize=10)
    ax.set_xlabel("step (20 Hz)")
axes[0].set_ylabel("normalized q"); axes[0].legend(fontsize=8)
fig.suptitle("Finger-joint open-loop trajectories (most-moving hand dims)")
fig.tight_layout(); fig.savefig(out_dir / "fig_finger_trajectories.png", dpi=150); plt.close(fig)

fig, ax = plt.subplots(figsize=(6, 3.6))
for n in names:
    ax.plot(curves[n], color=colors[n], lw=2, marker="o", ms=3, label=n)
ax.set_xlabel("chunk step"); ax.set_ylabel("hand MSE")
ax.set_title("Per-step hand error growth (mean over anchors)")
ax.legend(); fig.tight_layout()
fig.savefig(out_dir / "fig_per_step_error.png", dpi=150); plt.close(fig)

fig, axes = plt.subplots(1, len(names), figsize=(6.0 * len(names), 3.6), squeeze=False)
slot_x = 16 + np.arange(4) * 1.0
for ax, n in zip(axes[0], names):
    ff = force_fig[n]
    ax.plot(np.arange(16), np.array(ff["gt_history"])[:, 2], color=GT, lw=2,
            label="GT force (5 Hz history)")
    ax.plot(slot_x, np.array(ff["gt_slots"])[:, 2], color=GT, ls="--", marker="o",
            label="GT future slots")
    if ff["vq_recon"] is not None:
        ax.step(np.arange(16), np.array(ff["vq_recon"])[:, 2], color="#3c7ec0",
                where="mid", lw=1.6, label="VQ-discretized view")
    if ff["force_pred"] is not None:
        ax.plot(slot_x, np.array(ff["force_pred"])[:, 2], color=colors[n],
                marker="s", lw=1.8, label="predicted future force")
    ax.axvline(15.5, color="#cccccc", lw=1)
    ax.set_title(f"{n} · finger {ff['finger']} · Fz", fontsize=10)
    ax.set_xlabel("5 Hz frame (past | future)"); ax.legend(fontsize=7)
axes[0][0].set_ylabel("normalized force")
fig.suptitle("Continuous force vs discrete VQ bottleneck")
fig.tight_layout(); fig.savefig(out_dir / "fig_force_continuity.png", dpi=150); plt.close(fig)

metrics = ["mse_hand", "vel_mse_hand", "fine_energy_ratio_hand", "tactile_delta_hand"]
fig, axes = plt.subplots(1, len(metrics), figsize=(3.4 * len(metrics), 3.2))
for ax, metric in zip(axes, metrics):
    vals = [report[n][metric] for n in names]
    ax.bar(range(len(names)), vals, color=[colors[n] for n in names], width=0.55)
    ax.set_xticks(range(len(names)), names, fontsize=7)
    ax.set_title(metric, fontsize=9)
fig.suptitle("Dexterity metrics (hand dims)")
fig.tight_layout(); fig.savefig(out_dir / "fig_metrics.png", dpi=150); plt.close(fig)
print("replotted", out_dir)