"""Rebuild v4head prep (rot_centroids + z_range from grasp_site) and build a per-episode
calibration sidecar (K scaled to video res + T_w2c), reusing the exact Any4D conventions.
Keypoint = grasp_site via yam_fk @ C (action_format_v3). Runs in a4d-bw (no mujoco)."""
import argparse, glob, os, pickle
import numpy as np
from sklearn.cluster import KMeans

from custom.dataset.yam_fk import yam_fk
from custom.dataset.action_format_v3 import _FLANGE_TO_GRASP
import custom.dataset.action_format_v1 as v1
from custom.any4d.v4_head import rot6d_to_quat
import torch


def canon(q):
    return q if q[0] >= 0 else -q


def main(a):
    old = dict(np.load(a.old_prep))
    WM_W, WM_H = a.model_w, a.model_h
    eps = sorted(glob.glob(f'{a.raw}/[0-9][0-9][0-9][0-9]'))
    all_q, all_z = [], []
    ep_ids, Ks, Ts = [], [], []
    for ed in eps:
        ep = int(os.path.basename(ed))
        pkls = sorted(glob.glob(f'{ed}/lowdim/*.pkl'))
        if not pkls:
            continue
        fd0 = pickle.load(open(pkls[0], 'rb'))
        assert fd0.get('_scene_in_rbase', False), f'{ed} not scene_in_rbase'
        K_src = np.asarray(fd0['intrinsics']['scene_camera'], np.float64)
        W_src, H_src = fd0['scene_save_wh']
        T_w2c = np.linalg.inv(np.asarray(fd0['extrinsics']['scene_camera'], np.float64))
        K_wm = K_src.copy()
        K_wm[0, 0] *= WM_W / W_src; K_wm[0, 2] *= WM_W / W_src
        K_wm[1, 1] *= WM_H / H_src; K_wm[1, 2] *= WM_H / H_src
        ep_ids.append(ep); Ks.append(K_wm); Ts.append(T_w2c)
        for p in pkls:
            j = np.asarray(pickle.load(open(p, 'rb'))['joints'], np.float64)
            T = yam_fk(j[:6]) @ _FLANGE_TO_GRASP            # grasp_site pose
            rot6d = v1._matrix_to_rot6d(T[:3, :3])          # same packing as v3/head
            q = rot6d_to_quat(torch.from_numpy(rot6d[None]).float())[0].numpy()  # head's convention
            all_q.append(canon(q)); all_z.append(float(T[2, 3]))
    all_q = np.stack(all_q); all_z = np.array(all_z)
    print(f'{len(ep_ids)} episodes, {len(all_q)} frames')
    print(f'grasp_site z: [{all_z.min():.4f}, {all_z.max():.4f}]m  (old z_range {old["z_range"]})')

    # rot centroids: kmeans(256) on grasp_site quats (matches DINO n_rot_clusters=256)
    km = KMeans(n_clusters=a.n_rot, random_state=0, n_init=10, max_iter=100).fit(all_q)
    C = km.cluster_centers_
    C = (C / (np.linalg.norm(C, axis=1, keepdims=True) + 1e-8)).astype(np.float32)
    z_lo, z_hi = float(all_z.min()) - a.z_pad, float(all_z.max()) + a.z_pad

    prep = dict(old)
    prep['rot_centroids'] = C
    prep['z_range'] = np.array([z_lo, z_hi], np.float32)
    np.savez(a.out_prep, **prep)
    print(f'saved prep {a.out_prep}: z_range=[{z_lo:.4f},{z_hi:.4f}] centroids{C.shape}')

    # per-episode calibration sidecar
    np.savez(a.out_sidecar, ep_ids=np.array(ep_ids, np.int64),
             K_scene=np.stack(Ks).astype(np.float32),
             T_w2c_scene=np.stack(Ts).astype(np.float32),
             model_wh=np.array([WM_W, WM_H], np.int64))
    print(f'saved sidecar {a.out_sidecar}: {len(ep_ids)} episodes K/T @ {WM_W}x{WM_H}')
    print('BUILD_PREP_SIDECAR_OK')


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--raw', default=os.path.expanduser('~/any4d_work/raw/pickplace_70'))
    p.add_argument('--old_prep', default=os.path.expanduser('~/any4d_work/v4head_prep.npz'))
    p.add_argument('--out_prep', default=os.path.expanduser('~/any4d_work/v4head_prep_grasp.npz'))
    p.add_argument('--out_sidecar', default=os.path.expanduser('~/any4d_work/scene_calib_sidecar.npz'))
    p.add_argument('--model_w', type=int, default=320)
    p.add_argument('--model_h', type=int, default=176)
    p.add_argument('--n_rot', type=int, default=256)
    p.add_argument('--z_pad', type=float, default=0.01)
    main(p.parse_args())