"""Smooth the corrected reference and project hand poses above the table."""
from pathlib import Path
import json, shutil
import numpy as np
import mujoco
from scipy.signal import savgol_filter

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "output/spider_dynamics_fix_20260915"
OUT = ROOT / "output/desktop_smooth_fix_20260915"
OUT.mkdir(exist_ok=True)
shutil.copy2(SRC / "reference_video_rate.npz", OUT / "reference_before_smoothing.npz")
shutil.copy2(SRC / "datasets/processed/current/l20/bimanual/boxes/scene_act.xml", OUT / "scene_act.xml")

z = np.load(SRC / "reference_video_rate.npz")
arrays = {k: z[k].copy() for k in z.files}
q = arrays["qpos"]
original = q.copy()
m = mujoco.MjModel.from_xml_path(str(OUT / "scene_act.xml"))
d = mujoco.MjData(m)
# Filter each hand component independently. The object six-DOF states are kept
# exactly as estimated; only hand states are regularized.
window, poly = (9, 2)
for sl in (slice(0, 27), slice(27, 54)):
    q[:, sl] = savgol_filter(q[:, sl], window_length=window, polyorder=poly, axis=0, mode="interp")

# Recompute the exact five linear-mimic joints from the MuJoCo equality rows.
for e in range(m.neq):
    j1, j2 = m.eq_obj1id[e], m.eq_obj2id[e]
    a1, a2 = m.jnt_qposadr[j1], m.jnt_qposadr[j2]
    p = m.eq_data[e, :5]
    q[:, a1] = sum(p[k] * q[:, a2] ** k for k in range(5))

floor_ids = [i for i in range(m.ngeom) if "_floor" in (m.geom(i).name or "")]
for i in floor_ids:
    m.geom_conaffinity[i] = 3
hand_geom = {
    "right": [i for i in range(m.ngeom) if m.geom_contype[i] == 1 and (m.geom(i).name or "").startswith("right_")],
    "left": [i for i in range(m.ngeom) if m.geom_contype[i] == 1 and (m.geom(i).name or "").startswith("left_")],
}
hand_vertices = {}
for side in ("right", "left"):
    hand_vertices[side] = []
    for i in hand_geom[side]:
        mid = int(m.geom_dataid[i])
        if mid < 0:
            continue
        # Collision meshes are convex approximations; use the original visual
        # mesh vertices to detect the visible hand crossing the table.
        start = int(m.mesh_vertadr[mid]); end = start + int(m.mesh_vertnum[mid])
        vertices = m.mesh_vert[start:end].copy()
        hand_vertices[side].append((i, vertices[::max(1, len(vertices)//128)]))

def min_floor_dist(side):
    floor_z = float(d.geom_xpos[floor_ids[0]][2])
    vals = []
    for i, vertices in hand_vertices[side]:
        world = vertices @ d.geom_xmat[i].reshape(3, 3).T + d.geom_xpos[i]
        vals.append(float(world[:, 2].min() - floor_z))
    return min(vals) if vals else 1.0

projected = []
for f in range(len(q)):
    d.qpos[:] = q[f]
    mujoco.mj_forward(m, d)
    for side, sl in (("right", slice(0, 27)), ("left", slice(27, 54))):
        # Binary search the smallest wrist-z lift giving 0.5 mm clearance.
        base = q[f].copy()
        d.qpos[:] = base; mujoco.mj_forward(m, d)
        if min_floor_dist(side) < 0.0005:
            lo, hi = 0.0, 0.15
            for _ in range(24):
                mid = (lo + hi) / 2
                trial = base.copy(); trial[sl.start + 2] += mid
                d.qpos[:] = trial; mujoco.mj_forward(m, d)
                if min_floor_dist(side) >= 0.0005: hi = mid
                else: lo = mid
            q[f, sl.start + 2] += hi
            projected.append((f, side, hi))

arrays["qpos"] = q
dt = float(np.median(np.diff(arrays["time"])))
arrays["qvel"] = np.gradient(q, dt, axis=0); arrays["qvel"][0] = 0
arrays["ctrl"] = q[:, m.jnt_qposadr[m.actuator_trnid[:, 0]]]
np.savez_compressed(OUT / "reference_video_rate.npz", **arrays)
report = {
    "frames": len(q), "filter": {"type": "Savitzky-Golay", "window": window, "polyorder": poly},
    "projected_hand_table_contacts": len(projected),
    "max_table_lift_mm": float(max((x[2] for x in projected), default=0) * 1000),
    "max_hand_change_mm": float(max(np.linalg.norm(q[:, :3]-original[:, :3], axis=1).max(), np.linalg.norm(q[:, 27:30]-original[:, 27:30], axis=1).max()) * 1000),
    "object_qpos_unchanged": bool(np.array_equal(q[:, -12:], original[:, -12:])),
    "all_finite": bool(np.isfinite(q).all()),
    "note": "Table projection uses the simulation floor plane and hand collision geoms; object pose remains fixed."
}
(OUT / "smoothing_table_report.json").write_text(json.dumps(report, indent=2))
print(json.dumps(report, indent=2))