#!/usr/bin/env python3
"""
fopen_pose_compare.py — Compare COLMAP camera poses with imagej-elphel ERS/GPS poses.

Reads:
  - COLMAP images.txt (IMAGE_ID QW QX QY QZ TX TY TZ CAMERA_ID NAME)
  - imagej-elphel *-INTERFRAME.corr-xml  (--xml, always present)
    OR  *-egomotion.csv                  (--ego, optional output file)
    --ego takes precedence if both are given (it has IMS/PIMU columns too).

Computes a 7-DOF Sim(3) alignment between the two POSITION sets (Umeyama 1991):
  COLMAP_center ≈ s * R_align * ego_pos + t_align

Then performs an ORIENTATION comparison two ways:
  A. Absolute: per-frame camera orientation in Sim(3)-aligned world frame vs ERS angles.
     Convention-dependent; printed raw to CSV for manual inspection.
  B. Relative (delta): inter-frame rotation change from COLMAP vs from ERS.
     Convention-independent — both must agree on how much the camera rotated between
     consecutive frames regardless of reference frame definition.

During a turning sequence, delta rotations are large and make sign/axis-swap
convention errors immediately visible.

Usage:
  python3 scripts/fopen_pose_compare.py \\
      --images  /path/to/sparse/0_txt/images.txt \\
      --xml     /path/to/*-INTERFRAME.corr-xml \\
      [--ego    /path/to/*-egomotion.csv] \\
      [--out    pose_compare.csv] \\
      [--plot]  \\
      [--skip   3 17 ...]  # pose row indices to skip (0-based)

Output:
  - Console: Sim(3) parameters, position RMSE (COLMAP units + physical metres),
             delta-orientation RMSE (deg), top-10 worst frames
  - CSV (--out): per-frame positions, position residuals, absolute COLMAP & ERS
                 orientation angles in Sim(3)-aligned world frame, delta-angle diffs
"""

import argparse
import sys
import re
import xml.etree.ElementTree as ET
import numpy as np
import csv

# ─────────────────────────── geometry helpers ────────────────────────────── #

def quat_to_R(qw, qx, qy, qz):
    """COLMAP quaternion (w,x,y,z) → 3×3 rotation matrix (world→camera)."""
    n = qw*qw + qx*qx + qy*qy + qz*qz
    if n < 1e-12:
        return np.eye(3)
    s = 2.0 / n
    R = np.array([
        [1 - s*(qy*qy + qz*qz),     s*(qx*qy - qw*qz),     s*(qx*qz + qw*qy)],
        [    s*(qx*qy + qw*qz), 1 - s*(qx*qx + qz*qz),     s*(qy*qz - qw*qx)],
        [    s*(qx*qz - qw*qy),     s*(qy*qz + qw*qx), 1 - s*(qx*qx + qy*qy)],
    ])
    return R


def R_to_axis_angle(R):
    """Rotation matrix → (axis 3-vec, angle_rad).  Returns angle≥0."""
    cos_a = np.clip((np.trace(R) - 1.0) / 2.0, -1.0, 1.0)
    angle = np.arccos(cos_a)
    if angle < 1e-10:
        return np.array([0., 0., 1.]), 0.0
    axis = np.array([R[2,1]-R[1,2], R[0,2]-R[2,0], R[1,0]-R[0,1]]) / (2.0*np.sin(angle))
    return axis / (np.linalg.norm(axis) + 1e-30), angle


def R_to_euler_zyx(R):
    """
    Rotation matrix → (yaw, pitch, roll) in degrees, ZYX / Tait-Bryan convention.
    yaw = rotation about world Z (azimuth)
    pitch = rotation about Y (tilt forward)
    roll  = rotation about X (bank)
    """
    sy = np.sqrt(R[0,0]**2 + R[1,0]**2)
    yaw   = np.degrees(np.arctan2( R[1,0],  R[0,0]))
    pitch = np.degrees(np.arctan2(-R[2,0],  sy))
    roll  = np.degrees(np.arctan2( R[2,1],  R[2,2]))
    return yaw, pitch, roll


def ers_to_R(azimuth_rad, tilt_rad, roll_rad):
    """
    imagej-elphel ERS angles → rotation matrix.
    Convention (from OpticalFlow / ErsCorrection):
      azimuth : rotation about world Z  (positive = turn left when viewed from above)
      tilt    : rotation about X after azimuth  (forward tilt / pitch)
      roll    : rotation about Y after tilt     (lateral roll / bank)
    R maps world → camera (same convention as COLMAP for fair comparison).

    NOTE: the sign conventions here are a HYPOTHESIS until confirmed against the code.
    The raw angles are also written to the CSV for inspection.
    """
    ca, sa = np.cos(azimuth_rad), np.sin(azimuth_rad)
    ct, st = np.cos(tilt_rad),    np.sin(tilt_rad)
    cr, sr = np.cos(roll_rad),    np.sin(roll_rad)

    # Rz(az) * Rx(tilt) * Ry(roll)  — naive guess; verify against code
    Rz = np.array([[ ca, sa, 0],[-sa, ca, 0],[0, 0, 1]])
    Rx = np.array([[1, 0, 0],[0, ct, st],[0,-st, ct]])
    Ry = np.array([[cr, 0,-sr],[0, 1, 0],[sr, 0, cr]])
    return Rz @ Rx @ Ry


def colmap_center(R_cw, tx, ty, tz):
    """COLMAP pose → camera centre in world frame: C = -R^T @ t."""
    t = np.array([tx, ty, tz])
    return -R_cw.T @ t


def umeyama_sim3(P, Q):
    """
    7-DOF Sim(3) alignment: find s, R, t such that  P ≈ s * R @ Q + t.
    P : (N,3) COLMAP camera centres
    Q : (N,3) egomotion positions (metres)
    Returns (s, R_align, t, sigma2_Q)
    Umeyama (1991), IEEE TPAMI 13(4): 376-380.
    """
    assert P.shape == Q.shape and P.ndim == 2 and P.shape[1] == 3
    n = P.shape[0]
    mu_P = P.mean(axis=0)
    mu_Q = Q.mean(axis=0)
    P0   = P - mu_P
    Q0   = Q - mu_Q
    sigma2_Q = (Q0 ** 2).sum() / n
    Sigma_PQ = (P0.T @ Q0) / n
    U, lam, Vt = np.linalg.svd(Sigma_PQ)
    d = np.ones(3)
    if np.linalg.det(U) * np.linalg.det(Vt) < 0:
        d[2] = -1
    R_align = U @ np.diag(d) @ Vt
    s       = (lam * d).sum() / sigma2_Q
    t       = mu_P - s * R_align @ mu_Q
    return s, R_align, t, sigma2_Q


def apply_sim3(s, R, t, Q):
    """Apply Sim(3): P_est = s * R @ Q^T + t  (Nx3 array)."""
    return (s * (R @ Q.T)).T + t


# ──────────────────────────── file parsers ───────────────────────────────── #

def parse_images_txt(path):
    """
    Parse COLMAP images.txt.
    Returns list of dicts sorted by frame number (frame_NNNN.png):
      {'frame': int, 'name': str,
       'center': np.array(3),       # world-frame camera position
       'R_cw':   np.array(3,3),     # world→camera rotation
       'qw/qx/qy/qz': float }
    """
    entries = []
    with open(path) as f:
        lines = [l for l in f if not l.startswith('#')]
    i = 0
    while i < len(lines):
        line = lines[i].strip()
        if not line:
            i += 1
            continue
        parts = line.split()
        if len(parts) < 9:
            i += 1
            continue
        # IMAGE_ID QW QX QY QZ TX TY TZ CAMERA_ID NAME
        try:
            qw, qx, qy, qz = float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4])
            tx, ty, tz      = float(parts[5]), float(parts[6]), float(parts[7])
            name            = parts[9]
        except (ValueError, IndexError):
            i += 2
            continue
        m    = re.search(r'frame_(\d+)', name)
        fno  = int(m.group(1)) if m else -1
        R_cw = quat_to_R(qw, qx, qy, qz)
        entries.append({'frame': fno, 'name': name,
                        'center': colmap_center(R_cw, tx, ty, tz),
                        'R_cw': R_cw,
                        'qw': qw, 'qx': qx, 'qy': qy, 'qz': qz})
        i += 2  # skip POINTS2D line
    entries.sort(key=lambda e: e['frame'])
    return entries


def parse_egomotion_csv(path, skip_rows=None):
    """
    Parse imagej-elphel egomotion.csv.
    Valid rows: tab-separated, ≥10 columns, col[0] is a numeric scene index.

    Columns of interest (0-indexed):
      0  scene_index
      1  timestamp
      2  x(m)      ERS-corrected position
      3  y(m)
      4  z(m)
      5  a(rad)    azimuth
      6  tilt(rad)
      7  roll(rad)
      98  imsX
      99  imsY
      100 imsZ
      110 pimuX-C  IMU corrected
      111 pimuY-C
      112 pimuZ-C

    Returns list of dicts in file order.
    skip_rows: set of 0-based valid-row indices to exclude.
    """
    skip_set = set(skip_rows or [])
    rows = []
    valid_idx = 0
    with open(path) as f:
        for line in f:
            line = line.rstrip('\n')
            parts = line.split('\t')
            if len(parts) < 10:
                continue
            try:
                scene_idx = int(parts[0])
            except ValueError:
                continue
            if valid_idx in skip_set:
                valid_idx += 1
                continue
            def fg(i):
                try:    return float(parts[i])
                except: return float('nan')
            rows.append({
                'valid_idx':  valid_idx,
                'scene_idx':  scene_idx,
                'timestamp':  fg(1),
                'x':   fg(2),   'y':   fg(3),   'z':   fg(4),
                'az':  fg(5),   'tilt': fg(6),  'roll': fg(7),
                'imsX': fg(98), 'imsY': fg(99), 'imsZ': fg(100),
                'pimuX_C': fg(110), 'pimuY_C': fg(111), 'pimuZ_C': fg(112),
            })
            valid_idx += 1
    return rows


def parse_interframe_xml(path, skip_rows=None):
    """
    Parse imagej-elphel *-INTERFRAME.corr-xml.
    Returns list of dicts sorted by timestamp — same order as COLMAP frame_NNNN
    (elphel_to_colmap.py assigns frame numbers by sorting timestamps).

    Fields match parse_egomotion_csv() for compatibility:
      scene_idx  — 0-based sequential index
      timestamp  — float seconds
      x/y/z      — position (m) relative to reference scene
      az/tilt/roll — orientation (rad)
      imsX/Y/Z, pimuX_C/Y_C/Z_C — NaN (not in XML)
    """
    skip_set = set(skip_rows or [])
    PREFIX    = "EYESIS_DCT_AUX."
    SCENE_KEY = PREFIX + "scenes_"

    tree    = ET.parse(path)
    root    = tree.getroot()
    entries = {e.get("key"): e.text for e in root.findall("entry")}

    raw = {}
    for key, value in entries.items():
        if not key.startswith(SCENE_KEY):
            continue
        ts_str = key[len(SCENE_KEY):]
        if ts_str.endswith("_dt") or ts_str.endswith("_d2t"):
            continue
        try:
            vals = [float(v) for v in value.split(",")]
        except (TypeError, ValueError):
            continue
        if len(vals) != 6:
            continue
        # "1763233715_106222"  →  1763233715.106222
        ts_float = float(ts_str.replace("_", ".", 1))
        raw[ts_float] = vals

    nan = float('nan')
    rows = []
    for i, (ts, vals) in enumerate(sorted(raw.items())):
        if i in skip_set:
            continue
        x, y, z, az, tilt, roll = vals
        rows.append({
            'scene_idx': i,
            'timestamp': ts,
            'x': x,   'y': y,   'z': z,
            'az': az, 'tilt': tilt, 'roll': roll,
            'imsX': nan, 'imsY': nan, 'imsZ': nan,
            'pimuX_C': nan, 'pimuY_C': nan, 'pimuZ_C': nan,
        })
    return rows


# ──────────────────────────── matching ───────────────────────────────────── #

def match_sequences(colmap_entries, ego_rows):
    """Sequential 1:1 match; warn if counts differ."""
    n_c, n_e = len(colmap_entries), len(ego_rows)
    if n_c != n_e:
        print(f"[WARN] COLMAP frames={n_c}, egomotion rows={n_e}; "
              f"using first {min(n_c,n_e)} of each.", file=sys.stderr)
    n = min(n_c, n_e)
    return [(colmap_entries[i], ego_rows[i]) for i in range(n)]


# ──────────────────────── position alignment ─────────────────────────────── #

def compute_pos_alignment(pairs, pos_key_x, pos_key_y, pos_key_z):
    """
    Compute Sim(3) alignment and per-frame position residuals.
    Returns (s, R_align, t, residuals_colmap_units) or (None,…) on failure.
    """
    P = np.array([p[0]['center']                                       for p in pairs])
    Q = np.array([[p[1][pos_key_x], p[1][pos_key_y], p[1][pos_key_z]] for p in pairs])
    valid = np.isfinite(P).all(axis=1) & np.isfinite(Q).all(axis=1)
    if valid.sum() < 4:
        return None, None, None, None
    s, R_align, t, _ = umeyama_sim3(P[valid], Q[valid])
    Q_aligned = apply_sim3(s, R_align, t, Q)
    res = np.linalg.norm(P - Q_aligned, axis=1)
    res[~valid] = float('nan')
    return s, R_align, t, res


# ───────────────────────── orientation comparison ────────────────────────── #

def colmap_camera_to_world_R(R_cw, R_align):
    """
    Bring COLMAP world→camera rotation into the Sim(3)-aligned world frame.

    R_cw : world→camera  (COLMAP native)
    R_align : Sim(3) rotation  (ego → COLMAP world)
    → R_aligned_cw : world_gps → camera  (in GPS/ERS world frame)

    Camera optical axis (looking direction) in aligned world frame:
        optical_axis = R_aligned_cw.T @ [0,0,1]
    """
    # COLMAP_world = R_align @ EGO_world  (the Sim(3) maps ego coords to COLMAP coords)
    # R_cw acts in COLMAP_world, so in EGO_world frame:
    #   R_ego_cw = R_cw @ R_align
    return R_cw @ R_align


def compute_orientation_data(pairs, R_align):
    """
    For each frame compute:
      - COLMAP camera orientation in GPS-aligned world frame (Euler yaw/pitch/roll, deg)
      - ERS azimuth/tilt/roll (from egomotion, already in deg)
      - delta-rotation angle between consecutive frames (from COLMAP and from ERS)
      - delta-rotation angle difference |COLMAP_delta - ERS_delta| (deg)
      - cumulative rotation from COLMAP and ERS (both integrated from frame 0)

    Returns list of per-frame dicts; delta/cumulative values are NaN/0 for frame 0.

    Convention note for ers_to_R: uses Rz(az)*Rx(tilt)*Ry(roll) — hypothesis only.
    The delta rotation comparison is sign-convention-independent (rotation angles only).
    The cumulative comparison is also independent of the absolute reference orientation
    because both cumulative sums start from 0 at frame 0.
    """
    results = []

    # Pre-compute per-frame aligned rotation matrices
    R_aligned  = [colmap_camera_to_world_R(ce['R_cw'], R_align) for ce, _ in pairs]
    R_ers_list = []
    for _, er in pairs:
        if not any(np.isnan([er['az'], er['tilt'], er['roll']])):
            R_ers_list.append(ers_to_R(er['az'], er['tilt'], er['roll']))
        else:
            R_ers_list.append(None)

    cum_colmap = 0.0
    cum_ers    = 0.0

    for i, (ce, er) in enumerate(pairs):
        Ra = R_aligned[i]
        yaw_c, pitch_c, roll_c = R_to_euler_zyx(Ra)    # COLMAP cam in GPS frame

        row = {
            # COLMAP absolute orientation (in Sim3-aligned GPS world frame)
            'colmap_yaw_deg':   yaw_c,
            'colmap_pitch_deg': pitch_c,
            'colmap_roll_deg':  roll_c,
            # ERS absolute orientation (raw, converted to degrees)
            'ers_az_deg':   np.degrees(er['az'])   if np.isfinite(er['az'])   else float('nan'),
            'ers_tilt_deg': np.degrees(er['tilt']) if np.isfinite(er['tilt']) else float('nan'),
            'ers_roll_deg': np.degrees(er['roll']) if np.isfinite(er['roll']) else float('nan'),
            # Delta (inter-frame) values — filled in below
            'delta_colmap_deg':  float('nan'),
            'delta_ers_deg':     float('nan'),
            'delta_diff_deg':    float('nan'),
            # Cumulative rotation from frame 0
            'cum_colmap_deg':    0.0,
            'cum_ers_deg':       0.0,
            'cum_diff_deg':      float('nan'),
        }

        if i > 0:
            # COLMAP delta rotation  frame i-1 → frame i
            dR_c = Ra @ R_aligned[i-1].T
            _, angle_c = R_to_axis_angle(dR_c)
            dc = np.degrees(angle_c)
            row['delta_colmap_deg'] = dc
            cum_colmap += dc
            row['cum_colmap_deg']   = cum_colmap

            # ERS delta rotation  frame i-1 → frame i
            Re_cur  = R_ers_list[i]
            Re_prev = R_ers_list[i-1]
            if Re_cur is not None and Re_prev is not None:
                dR_e = Re_cur @ Re_prev.T
                _, angle_e = R_to_axis_angle(dR_e)
                de = np.degrees(angle_e)
                row['delta_ers_deg']  = de
                row['delta_diff_deg'] = abs(dc - de)
                cum_ers += de
                row['cum_ers_deg']    = cum_ers
                row['cum_diff_deg']   = abs(cum_colmap - cum_ers)
        else:
            row['cum_colmap_deg'] = 0.0
            row['cum_ers_deg']    = 0.0
            row['cum_diff_deg']   = 0.0

        results.append(row)
    return results


def print_orientation_report(orient_data, label="ERS x/y/z"):
    """Print delta-orientation RMSE, cumulative comparison, and worst frames."""
    diffs = np.array([r['delta_diff_deg'] for r in orient_data])
    valid = np.isfinite(diffs)
    if valid.sum() == 0:
        print("  (no orientation data)")
        return
    rmse = np.sqrt(np.nanmean(diffs[valid]**2))
    med  = np.nanmedian(diffs[valid])
    mx   = np.nanmax(diffs[valid])

    # Cumulative totals at end of sequence
    cum_c = orient_data[-1]['cum_colmap_deg']
    cum_e = orient_data[-1]['cum_ers_deg']
    cum_d = orient_data[-1]['cum_diff_deg']

    print(f"\n  Δ-orientation (|COLMAP_delta - ERS_delta|, per-frame):")
    print(f"    RMSE={rmse:.3f}°   median={med:.3f}°   max={mx:.3f}°")
    print(f"  Cumulative rotation over full sequence:")
    print(f"    COLMAP={cum_c:.2f}°   ERS={cum_e:.2f}°   |diff|={cum_d:.2f}°")
    idx = np.argsort(diffs)[::-1]
    print(f"  Top-5 worst per-frame deltas:")
    print(f"    {'Frame':>6}  {'COLMAP_Δ(°)':>12}  {'ERS_Δ(°)':>10}  {'diff(°)':>9}")
    for i in idx[:5]:
        r = orient_data[i]
        print(f"    {i:>6}  {r['delta_colmap_deg']:>12.4f}  "
              f"{r['delta_ers_deg']:>10.4f}  {r['delta_diff_deg']:>9.4f}")


# ──────────────────────────── reporting ───────────────────────────────────── #

def print_pos_alignment_report(label, s, R_align, t, residuals):
    """Print position alignment summary."""
    valid = np.isfinite(residuals)
    rmse_c = np.sqrt(np.nanmean(residuals[valid]**2))
    med_c  = np.nanmedian(residuals[valid])
    mx_c   = np.nanmax(residuals[valid])
    yaw, pitch, roll = R_to_euler_zyx(R_align)
    print(f"\n{'─'*60}")
    print(f"  {label}")
    print(f"{'─'*60}")
    print(f"  Scale:        {s:.6f}  (1 COLMAP unit = {1/s:.2f} m)")
    print(f"  Rotation:     yaw={yaw:+.2f}°  pitch={pitch:+.2f}°  roll={roll:+.2f}°")
    print(f"  Translation:  [{t[0]:+.4f},  {t[1]:+.4f},  {t[2]:+.4f}]  (COLMAP units)")
    print(f"  Pos residuals (COLMAP): RMSE={rmse_c:.4f}  median={med_c:.4f}  max={mx_c:.4f}")
    print(f"  Pos residuals (metres): RMSE={rmse_c/s:.2f} m  "
          f"median={med_c/s:.2f} m  max={mx_c/s:.2f} m")


def print_top_pos_residuals(pairs, residuals, s, n=10):
    """Top-n worst position residuals."""
    idx = np.argsort(residuals)[::-1]
    print(f"\n  Top-{n} worst position frames  (1 COLMAP = {1/s:.1f} m):")
    print(f"  {'Frame':>6}  {'Scene':>6}  {'res(COLMAP)':>12}  {'res(m)':>10}")
    for i in idx[:n]:
        c, e = pairs[i]
        print(f"  {c['frame']:>6}  {e['scene_idx']:>6}  {residuals[i]:>12.4f}  "
              f"{residuals[i]/s:>10.2f}")


# ─────────────────────────────── main ────────────────────────────────────── #

def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('--images', required=True, help='COLMAP sparse/0_txt/images.txt')
    ap.add_argument('--xml',    default=None,  help='imagej-elphel *-INTERFRAME.corr-xml (preferred)')
    ap.add_argument('--ego',    default=None,  help='imagej-elphel *-egomotion.csv (optional; overrides --xml if given)')
    ap.add_argument('--out',    default=None,  help='Output CSV (optional)')
    ap.add_argument('--plot',   action='store_true', help='3-D trajectory plot')
    ap.add_argument('--skip',   type=int, nargs='*', default=[],
                    metavar='ROW', help='0-based pose rows to skip')
    args = ap.parse_args()

    if not args.xml and not args.ego:
        ap.error("provide at least one of --xml or --ego")

    # ── Load ───────────────────────────────────────────────────────────────
    print(f"Loading COLMAP images: {args.images}")
    colmap_entries = parse_images_txt(args.images)
    print(f"  {len(colmap_entries)} frames  "
          f"(frame_{colmap_entries[0]['frame']:04d} … frame_{colmap_entries[-1]['frame']:04d})")

    if args.ego:
        print(f"Loading egomotion:     {args.ego}")
        ego_rows = parse_egomotion_csv(args.ego, skip_rows=args.skip)
        print(f"  {len(ego_rows)} rows  (scene {ego_rows[0]['scene_idx']} … {ego_rows[-1]['scene_idx']})")
    else:
        print(f"Loading INTERFRAME XML: {args.xml}")
        ego_rows = parse_interframe_xml(args.xml, skip_rows=args.skip)
        print(f"  {len(ego_rows)} rows  (ts {ego_rows[0]['timestamp']:.3f} … {ego_rows[-1]['timestamp']:.3f})")

    pairs = match_sequences(colmap_entries, ego_rows)
    print(f"  {len(pairs)} matched pairs")

    # ── Position alignment for each egomotion column set ───────────────────
    col_specs = [
        ('ERS x/y/z',            'x',       'y',       'z'      ),
        ('IMS imsX/Y/Z',         'imsX',    'imsY',    'imsZ'   ),
        ('PIMU pimuX-C/Y-C/Z-C', 'pimuX_C', 'pimuY_C', 'pimuZ_C'),
    ]

    pos_results  = {}    # label → (s, R_align, t, res)
    ori_results  = {}    # label → orient_data list

    for label, kx, ky, kz in col_specs:
        s, R_align, t, res = compute_pos_alignment(pairs, kx, ky, kz)
        if s is None:
            if not args.ego:
                continue   # IMS/PIMU are always NaN when using XML — no need to warn
            print(f"\n[SKIP] {label}: insufficient valid data")
            continue
        print_pos_alignment_report(label, s, R_align, t, res)
        print_top_pos_residuals(pairs, res, s, n=10)
        pos_results[label] = (s, R_align, t, res)

        # Orientation comparison uses R_align from this column's Sim(3)
        od = compute_orientation_data(pairs, R_align)
        ori_results[label] = od
        print_orientation_report(od, label)

    # ── CSV ────────────────────────────────────────────────────────────────
    if args.out:
        base_fields = [
            'frame', 'scene_idx', 'timestamp',
            'colmap_X', 'colmap_Y', 'colmap_Z',
            'ego_x', 'ego_y', 'ego_z',
            'ego_imsX', 'ego_imsY', 'ego_imsZ',
            'ego_pimuX_C', 'ego_pimuY_C', 'ego_pimuZ_C',
            # Raw ERS orientation
            'ers_az_rad', 'ers_tilt_rad', 'ers_roll_rad',
        ]
        extra_fields = []
        for label, _, _, _ in col_specs:
            if label not in pos_results:
                continue
            k = label.replace(' ', '_').replace('/', '_')
            extra_fields += [
                f'pos_res_{k}(colmap)', f'pos_res_{k}(m)',
                f'aligned_{k}_X', f'aligned_{k}_Y', f'aligned_{k}_Z',
                f'colmap_yaw_{k}(deg)', f'colmap_pitch_{k}(deg)', f'colmap_roll_{k}(deg)',
                f'delta_colmap_{k}(deg)', f'delta_ers_{k}(deg)', f'delta_diff_{k}(deg)',
                f'cum_colmap_{k}(deg)', f'cum_ers_{k}(deg)', f'cum_diff_{k}(deg)',
            ]
        fieldnames = base_fields + extra_fields

        with open(args.out, 'w', newline='') as f:
            w = csv.DictWriter(f, fieldnames=fieldnames)
            w.writeheader()
            for i, (ce, er) in enumerate(pairs):
                row = {
                    'frame':      ce['frame'],
                    'scene_idx':  er['scene_idx'],
                    'timestamp':  er['timestamp'],
                    'colmap_X':   ce['center'][0],
                    'colmap_Y':   ce['center'][1],
                    'colmap_Z':   ce['center'][2],
                    'ego_x':      er['x'],   'ego_y': er['y'],   'ego_z': er['z'],
                    'ego_imsX':   er['imsX'],'ego_imsY': er['imsY'],'ego_imsZ': er['imsZ'],
                    'ego_pimuX_C': er['pimuX_C'],
                    'ego_pimuY_C': er['pimuY_C'],
                    'ego_pimuZ_C': er['pimuZ_C'],
                    'ers_az_rad':   er['az'],
                    'ers_tilt_rad': er['tilt'],
                    'ers_roll_rad': er['roll'],
                }
                for label, kx, ky, kz in col_specs:
                    if label not in pos_results:
                        continue
                    s, R_align, t, res = pos_results[label]
                    k = label.replace(' ', '_').replace('/', '_')
                    Qi = np.array([[er[kx], er[ky], er[kz]]])
                    aligned = apply_sim3(s, R_align, t, Qi)[0]
                    row[f'pos_res_{k}(colmap)'] = res[i] if np.isfinite(res[i]) else ''
                    row[f'pos_res_{k}(m)']      = (res[i]/s) if np.isfinite(res[i]) else ''
                    row[f'aligned_{k}_X'] = aligned[0]
                    row[f'aligned_{k}_Y'] = aligned[1]
                    row[f'aligned_{k}_Z'] = aligned[2]
                    od = ori_results[label][i]
                    row[f'colmap_yaw_{k}(deg)']   = od['colmap_yaw_deg']
                    row[f'colmap_pitch_{k}(deg)']  = od['colmap_pitch_deg']
                    row[f'colmap_roll_{k}(deg)']   = od['colmap_roll_deg']
                    row[f'delta_colmap_{k}(deg)']  = od['delta_colmap_deg'] \
                                                     if np.isfinite(od['delta_colmap_deg']) else ''
                    row[f'delta_ers_{k}(deg)']     = od['delta_ers_deg'] \
                                                     if np.isfinite(od['delta_ers_deg']) else ''
                    row[f'delta_diff_{k}(deg)']    = od['delta_diff_deg'] \
                                                     if np.isfinite(od['delta_diff_deg']) else ''
                    row[f'cum_colmap_{k}(deg)']    = od['cum_colmap_deg']
                    row[f'cum_ers_{k}(deg)']       = od['cum_ers_deg']
                    row[f'cum_diff_{k}(deg)']      = od['cum_diff_deg'] \
                                                     if np.isfinite(od['cum_diff_deg']) else ''
                w.writerow(row)
        print(f"\nSaved: {args.out}")

    # ── 3-D plot ───────────────────────────────────────────────────────────
    if args.plot:
        try:
            import matplotlib
            matplotlib.use('TkAgg')
            import matplotlib.pyplot as plt
            from mpl_toolkits.mplot3d import Axes3D  # noqa: F401
        except ImportError:
            print("[WARN] matplotlib not available; skipping plot", file=sys.stderr)
            return

        P = np.array([ce['center'] for ce, _ in pairs])
        fig = plt.figure(figsize=(12, 8))
        ax  = fig.add_subplot(111, projection='3d')
        ax.plot(P[:,0], P[:,1], P[:,2], 'k-', lw=1, label='COLMAP', alpha=0.8)
        ax.scatter(P[0,0], P[0,1], P[0,2], c='k', s=40, marker='o')

        colors = ['tab:blue', 'tab:orange', 'tab:green']
        for (label, kx, ky, kz), col in zip(col_specs, colors):
            if label not in pos_results:
                continue
            s, R_align, t, res = pos_results[label]
            Q = np.array([[er[kx], er[ky], er[kz]] for _, er in pairs])
            Qa = apply_sim3(s, R_align, t, Q)
            rmse = np.sqrt(np.nanmean(res**2))
            ax.plot(Qa[:,0], Qa[:,1], Qa[:,2], '-', color=col, lw=1,
                    label=f'{label}  RMSE={rmse/s:.1f} m', alpha=0.7)
            ax.scatter(Qa[0,0], Qa[0,1], Qa[0,2], c=col, s=40, marker='o')

        ax.set_xlabel('X (COLMAP)')
        ax.set_ylabel('Y (COLMAP)')
        ax.set_zlabel('Z (COLMAP)')
        ax.legend(loc='upper left', fontsize=8)
        ax.set_title('Camera trajectory: COLMAP vs imagej-elphel (Sim3-aligned)')
        plt.tight_layout()
        plt.show()


if __name__ == '__main__':
    main()
