#!/usr/bin/env python3
"""Extract camera positions (timestamp, x, y, z) from lingbot-map NPZ poses.
Usage: poses_to_csv.py <input.npz> [--start_iso 2026-05-05T08:34:01] [--fps 1.0] > output.csv
"""
import argparse, sys
import numpy as np
from datetime import datetime, timezone

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('npz')
    ap.add_argument('--start_iso', default=None, help='ISO timestamp of frame 0 (e.g. 2026-05-05T08:34:01)')
    ap.add_argument('--fps', type=float, default=1.0, help='frames per second (default 1.0)')
    ap.add_argument('--label', default='', help='segment label for CSV col')
    args = ap.parse_args()
    
    data = np.load(args.npz, allow_pickle=True)
    keys = list(data.keys())
    # auto-detect poses key
    poses = None
    for k in ['poses', 'extrinsics', 'cam_poses', 'c2w']:
        if k in keys:
            poses = data[k]; break
    if poses is None:
        # take first 3D array
        for k in keys:
            arr = data[k]
            if arr.ndim == 3 and arr.shape[-2:] in [(3,4),(4,4)]:
                poses = arr; break
    if poses is None:
        sys.exit(f'No poses found in {args.npz} (keys: {keys})')
    
    # start timestamp
    if args.start_iso:
        try:
            t0 = datetime.fromisoformat(args.start_iso).replace(tzinfo=timezone.utc).timestamp()
        except Exception:
            t0 = 0.0
    elif 'start_ns' in keys:
        t0 = float(data['start_ns']) / 1e9
    else:
        t0 = 0.0
    
    fps = float(args.fps)
    if 'fps' in keys:
        try: fps = float(data['fps'])
        except: pass
    
    print('segment,frame_idx,timestamp_s,x,y,z')
    for i, P in enumerate(poses):
        if P.shape == (4,4):
            P = P[:3]
        R, t = P[:, :3], P[:, 3]
        # extrinsic = world→cam ; cam position world = -R^T t
        # but lingbot might save c2w directly; check determinant heuristic
        pos = -R.T @ t  # if extrinsic
        ts = t0 + i / fps
        print(f'{args.label},{i},{ts:.6f},{pos[0]:.6f},{pos[1]:.6f},{pos[2]:.6f}')

if __name__ == '__main__':
    main()