#!/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 egomotion.csv (tab-separated, numeric-index rows)

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

Reports per-camera residuals and compares multiple egomotion position columns
(ERS-corrected x/y/z, IMS imsX/Y/Z, PIMU corrected pimuX-C/Y-C/Z-C).

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

Output:
  - Console: Sim(3) parameters, per-column RMSE, per-camera top-10 residuals
  - CSV (--out): per-frame residuals for all position columns
  - Optional 3D plot (--plot): COLMAP trajectory and aligned egomotion trajectories
"""

import argparse
import sys
import re
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."""
    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 colmap_center(qw, qx, qy, qz, tx, ty, tz):
    """COLMAP image pose → camera centre in world frame C = -R^T @ t."""
    R = quat_to_R(qw, qx, qy, qz)
    t = np.array([tx, ty, tz])
    return -R.T @ t


def umeyama_sim3(P, Q):
    """
    7-DOF Sim(3) alignment: find s, R, t such that  P ≈ s * R @ Q.T + t.

    P, Q : (N,3) float arrays — source (COLMAP) and target (egomotion)
    Returns (s, R, 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          # variance of source
    Sigma_PQ = (P0.T @ Q0) / n              # cross-covariance (3×3)

    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     = U @ np.diag(d) @ Vt
    s     = (lam * d).sum() / sigma2_Q
    t     = mu_P - s * R @ mu_Q
    return s, R, t, sigma2_Q


def apply_sim3(s, R, t, Q):
    """Apply Sim(3): P_est = s * R @ Q^T + t  (works on 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)}
    """
    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)
        frame_no = int(m.group(1)) if m else -1
        center   = colmap_center(qw, qx, qy, qz, tx, ty, tz)
        entries.append({'frame': frame_no, 'name': name, 'center': center})
        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 after splitting on tab):
      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: 0-based indices among valid rows to drop).
    """
    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 (IndexError, ValueError):
                    return float('nan')
            row = {
                'valid_idx':  valid_idx,
                'scene_idx':  scene_idx,
                'timestamp':  fg(1),
                # ERS-corrected pose
                'x':   fg(2),   'y':   fg(3),   'z':   fg(4),
                'az':  fg(5),   'tilt': fg(6),   'roll': fg(7),
                # IMS position
                'imsX': fg(98), 'imsY': fg(99),  'imsZ': fg(100),
                # PIMU corrected
                'pimuX_C': fg(110), 'pimuY_C': fg(111), 'pimuZ_C': fg(112),
            }
            rows.append(row)
            valid_idx += 1
    return rows


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

def match_sequences(colmap_entries, ego_rows):
    """
    Sequential 1:1 match between sorted COLMAP frames and egomotion rows.
    Returns list of (colmap_entry, ego_row) pairs.
    Warns if counts differ.
    """
    n_c = len(colmap_entries)
    n_e = 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)]


def compute_alignment(pairs, pos_key_x, pos_key_y, pos_key_z):
    """
    Given matched pairs and egomotion column names, compute Sim(3) and residuals.
    Returns (s, R, t, residuals_m) where residuals_m is (N,) array.
    Returns None if any position is NaN.
    """
    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
    Pv, Qv = P[valid], Q[valid]
    s, R, t, _ = umeyama_sim3(Pv, Qv)
    Q_aligned   = apply_sim3(s, R, t, Q)
    residuals   = np.linalg.norm(P - Q_aligned, axis=1)
    residuals[~valid] = float('nan')
    return s, R, t, residuals


def print_alignment_report(label, pairs, s, R, t, residuals):
    """Print summary stats for one alignment.

    residuals are in COLMAP units; dividing by s converts to physical meters
    (the GPS/egomotion coordinate frame).
    """
    valid = np.isfinite(residuals)
    # COLMAP units
    rmse_c = np.sqrt(np.nanmean(residuals[valid]**2))
    med_c  = np.nanmedian(residuals[valid])
    mx_c   = np.nanmax(residuals[valid])
    # Physical meters (GPS/ERS frame)
    rmse_m = rmse_c / s
    med_m  = med_c  / s
    mx_m   = mx_c   / s

    # Extract Euler angles from R (ZYX / yaw-pitch-roll, degrees)
    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]))

    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]:+.3f},  {t[1]:+.3f},  {t[2]:+.3f}]  (COLMAP units)")
    print(f"  Residuals (COLMAP): RMSE={rmse_c:.4f}  median={med_c:.4f}  max={mx_c:.4f}")
    print(f"  Residuals (meters): RMSE={rmse_m:.2f} m  median={med_m:.2f} m  max={mx_m:.2f} m")


def print_top_residuals(pairs, residuals, s, n=10):
    """Print the n frames with largest residuals (in COLMAP units and physical meters)."""
    idx = np.argsort(residuals)[::-1]
    print(f"\n  Top-{n} worst frames  (1 COLMAP unit = {1/s:.1f} m):")
    print(f"  {'Frame':>6}  {'Scene':>6}  {'res(COLMAP)':>12}  {'res(m)':>10}")
    for i in idx[:n]:
        c = pairs[i][0]
        e = pairs[i][1]
        print(f"  {c['frame']:>6}  {e['scene_idx']:>6}  {residuals[i]:>12.4f}  {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('--ego', required=True,
                    help='imagej-elphel *-egomotion.csv')
    ap.add_argument('--out', default=None,
                    help='Output CSV with per-frame residuals (optional)')
    ap.add_argument('--plot', action='store_true',
                    help='Show 3D trajectory plot (requires matplotlib)')
    ap.add_argument('--skip', type=int, nargs='*', default=[],
                    metavar='ROW',
                    help='0-based egomotion valid-row indices to skip (rare mismatches)')
    args = ap.parse_args()

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

    print(f"Loading egomotion:     {args.ego}")
    ego_rows = parse_egomotion_csv(args.ego, skip_rows=args.skip)
    print(f"  {len(ego_rows)} rows loaded (scene {ego_rows[0]['scene_idx']} … "
          f"{ego_rows[-1]['scene_idx']})")

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

    # ── Align each egomotion position column ───────────────────────────────
    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'),
    ]

    results = {}
    for label, kx, ky, kz in col_specs:
        s, R, t, res = compute_alignment(pairs, kx, ky, kz)
        if s is None:
            print(f"\n[SKIP] {label}: insufficient valid data")
            continue
        print_alignment_report(label, pairs, s, R, t, res)
        print_top_residuals(pairs, res, s, n=10)
        results[label] = (s, R, t, res)

    # ── CSV output ─────────────────────────────────────────────────────────
    if args.out:
        fieldnames = ['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']
        for label, _, _, _ in col_specs:
            if label in results:
                key = label.replace(' ', '_').replace('/', '_')
                fieldnames += [f'res_{key}(colmap_units)',
                               f'res_{key}(phys_m)',
                               f'aligned_{key}_X', f'aligned_{key}_Y', f'aligned_{key}_Z']

        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'],
                }
                for label, kx, ky, kz in col_specs:
                    if label not in results:
                        continue
                    s, R, t, res = results[label]
                    key = label.replace(' ', '_').replace('/', '_')
                    Q_i = np.array([[er[kx], er[ky], er[kz]]])
                    aligned = apply_sim3(s, R, t, Q_i)[0]
                    row[f'res_{key}(colmap_units)'] = res[i] if np.isfinite(res[i]) else ''
                    row[f'res_{key}(phys_m)']  = (res[i]/s) if np.isfinite(res[i]) else ''
                    row[f'aligned_{key}_X']   = aligned[0]
                    row[f'aligned_{key}_Y']   = aligned[1]
                    row[f'aligned_{key}_Z']   = aligned[2]
                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

        # COLMAP trajectory
        P = np.array([ce['center'] for ce, er 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 results:
                continue
            s, R, t, res = results[label]
            Q = np.array([[er[kx], er[ky], er[kz]] for ce, er in pairs])
            Qa = apply_sim3(s, R, 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:.2f} 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()
