#!/usr/bin/env python3
"""Stage 06 — absolute alignment of lingbot relative trajectory using MCAP IMU+USBL.
Input  : trajectory CSV (segment, frame_idx, timestamp_s, x, y, z) + MCAP bag dir
Output : absolute trajectory CSV (timestamp_s, lat, lon, alt, east_m, north_m, up_m, segment, lingbot_x, lingbot_y, lingbot_z)
+ plot PNG with absolute trajectory (spiral expected for lawnmower)

Method:
  - Parse MCAP : /mavros/imu/data + /mavros/global_position/global + /mavros/imu/static_pressure
  - Convert lat/lon → ENU meters (origin = first fix)
  - For each frame timestamp, find nearest GPS/IMU
  - Umeyama similarity transform : align lingbot (x,y,z) → (east,north,depth)
  - Output enhanced CSV + plot

Usage:
  python3 stage06_align_absolute.py --traj /tmp/auv213_full_trajectory.csv \
      --mcap-dir /mnt/ssd/20260505-Lepradet/raw_data/logs/SUB/bag/20260505_150717_AUV013/ \
      --out /tmp/auv213_absolute.csv --plot /tmp/auv213_absolute.png
"""
import argparse, csv, glob, os, sys, math
import numpy as np
from pathlib import Path
from datetime import datetime, timezone

def parse_mcap_bag(bag_dir):
    """Return dict of {topic: [(ts_ns, msg_dict)]}."""
    from rosbags.highlevel import AnyReader
    bag_paths = sorted(Path(bag_dir).glob('*.mcap'))
    if not bag_paths:
        sys.exit(f'No .mcap in {bag_dir}')
    print(f'[mcap] reading {len(bag_paths)} bag files', flush=True)
    topics_of_interest = {
        '/mavros/imu/data': 'Imu',
        '/mavros/global_position/global': 'NavSatFix',
        '/mavros/imu/static_pressure': 'FluidPressure',
    }
    data = {t: [] for t in topics_of_interest}
    with AnyReader(bag_paths) as reader:
        conns = [c for c in reader.connections if c.topic in topics_of_interest]
        print(f'[mcap] connections: {[(c.topic, c.msgtype) for c in conns]}', flush=True)
        for conn, ts_ns, raw in reader.messages(connections=conns):
            try:
                m = reader.deserialize(raw, conn.msgtype)
                if conn.topic == '/mavros/imu/data':
                    q = m.orientation
                    data[conn.topic].append((ts_ns, {'qx': q.x, 'qy': q.y, 'qz': q.z, 'qw': q.w}))
                elif conn.topic == '/mavros/global_position/global':
                    if not (math.isnan(m.latitude) or math.isnan(m.longitude)):
                        data[conn.topic].append((ts_ns, {'lat': m.latitude, 'lon': m.longitude, 'alt': m.altitude}))
                elif conn.topic == '/mavros/imu/static_pressure':
                    data[conn.topic].append((ts_ns, {'pressure_pa': m.fluid_pressure}))
            except Exception as e:
                continue
    for t in data:
        print(f'[mcap] {t}: {len(data[t])} samples', flush=True)
    return data

def latlon_to_enu(lat, lon, alt, lat0, lon0, alt0):
    """Local ENU meters (flat Earth around (lat0, lon0))."""
    R = 6378137.0
    dlat = math.radians(lat - lat0)
    dlon = math.radians(lon - lon0)
    east = R * dlon * math.cos(math.radians(lat0))
    north = R * dlat
    up = alt - alt0
    return east, north, up

def nearest_ts(arr, ts_target_ns):
    if not arr: return None
    idx = np.searchsorted([a[0] for a in arr], ts_target_ns)
    if idx == 0: return arr[0]
    if idx >= len(arr): return arr[-1]
    if abs(arr[idx][0] - ts_target_ns) < abs(arr[idx-1][0] - ts_target_ns):
        return arr[idx]
    return arr[idx-1]

def umeyama(src, dst):
    """Closed-form similarity transform (scale s, rotation R, translation t) minimizing ||s*R*src + t - dst||^2.
    src, dst: (N, 3) arrays."""
    src = np.asarray(src, dtype=np.float64)
    dst = np.asarray(dst, dtype=np.float64)
    mu_src = src.mean(axis=0)
    mu_dst = dst.mean(axis=0)
    src_c = src - mu_src
    dst_c = dst - mu_dst
    sigma_src = (src_c ** 2).sum() / len(src)
    cov = (dst_c.T @ src_c) / len(src)
    U, S, Vt = np.linalg.svd(cov)
    D = np.eye(3)
    if np.linalg.det(U @ Vt) < 0:
        D[2, 2] = -1
    R = U @ D @ Vt
    s = (S * np.diag(D)).sum() / sigma_src
    t = mu_dst - s * R @ mu_src
    return s, R, t

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--traj', required=True)
    ap.add_argument('--mcap-dir', required=True)
    ap.add_argument('--out', required=True)
    ap.add_argument('--plot', default=None)
    ap.add_argument('--min-fixes', type=int, default=10, help='min GPS fixes to attempt Umeyama')
    args = ap.parse_args()
    
    print(f'[traj] reading {args.traj}', flush=True)
    rows = []
    with open(args.traj) as f:
        reader = csv.DictReader(f)
        for r in reader:
            rows.append({
                'segment': r['segment'],
                'frame_idx': int(r['frame_idx']),
                'ts_s': float(r['timestamp_s']),
                'x': float(r['x']),
                'y': float(r['y']),
                'z': float(r['z']),
            })
    print(f'[traj] {len(rows)} rows', flush=True)
    
    data = parse_mcap_bag(args.mcap_dir)
    gps = data['/mavros/global_position/global']
    imu = data['/mavros/imu/data']
    press = data['/mavros/imu/static_pressure']
    
    if len(gps) < args.min_fixes:
        print(f'[warn] only {len(gps)} GPS fixes (need >= {args.min_fixes}), skipping Umeyama. Will output raw + IMU only.', flush=True)
    
    # ENU origin = first GPS fix
    if gps:
        lat0 = gps[0][1]['lat']; lon0 = gps[0][1]['lon']; alt0 = gps[0][1]['alt']
        print(f'[origin] lat0={lat0:.6f} lon0={lon0:.6f} alt0={alt0:.2f}', flush=True)
    else:
        lat0 = lon0 = alt0 = 0.0
    
    # Build per-frame absolute ENU using nearest GPS
    src_lingbot, dst_enu, gps_per_frame, imu_per_frame, depth_per_frame = [], [], [], [], []
    for r in rows:
        ts_ns = int(r['ts_s'] * 1e9)
        g = nearest_ts(gps, ts_ns)
        i = nearest_ts(imu, ts_ns)
        p = nearest_ts(press, ts_ns)
        # Always store IMU+depth
        imu_per_frame.append(i[1] if i else None)
        if p:
            depth_m = (p[1]['pressure_pa'] - 101325.0) / (1025.0 * 9.81)
            depth_per_frame.append(depth_m)
        else:
            depth_per_frame.append(None)
        gps_per_frame.append(g[1] if g else None)
        if g and abs(g[0] - ts_ns) < 2e9:  # GPS within 2s
            e, n, u = latlon_to_enu(g[1]['lat'], g[1]['lon'], g[1]['alt'], lat0, lon0, alt0)
            src_lingbot.append([r['x'], r['y'], r['z']])
            dst_enu.append([e, n, u])
    
    # Umeyama
    if len(src_lingbot) >= args.min_fixes:
        s, R, t = umeyama(src_lingbot, dst_enu)
        print(f'[umeyama] scale={s:.4f} translation={t}', flush=True)
        print(f'[umeyama] rotation_matrix=\n{R}', flush=True)
    else:
        s, R, t = 1.0, np.eye(3), np.zeros(3)
        print(f'[umeyama] insufficient pairs, identity transform', flush=True)
    
    # Apply transform
    out_rows = []
    for i_row, r in enumerate(rows):
        p_lingbot = np.array([r['x'], r['y'], r['z']])
        p_abs = s * R @ p_lingbot + t
        g = gps_per_frame[i_row]
        im = imu_per_frame[i_row]
        d = depth_per_frame[i_row]
        out_rows.append({
            'segment': r['segment'],
            'frame_idx': r['frame_idx'],
            'timestamp_s': r['ts_s'],
            'east_m': p_abs[0],
            'north_m': p_abs[1],
            'up_m': p_abs[2],
            'depth_m': d if d is not None else '',
            'gps_lat': g['lat'] if g else '',
            'gps_lon': g['lon'] if g else '',
            'imu_qw': im['qw'] if im else '',
            'imu_qx': im['qx'] if im else '',
            'imu_qy': im['qy'] if im else '',
            'imu_qz': im['qz'] if im else '',
            'lingbot_x': r['x'],
            'lingbot_y': r['y'],
            'lingbot_z': r['z'],
        })
    
    with open(args.out, 'w', newline='') as f:
        w = csv.DictWriter(f, fieldnames=list(out_rows[0].keys()))
        w.writeheader(); w.writerows(out_rows)
    print(f'[out] {args.out} {len(out_rows)} rows', flush=True)
    
    if args.plot:
        import matplotlib
        matplotlib.use('Agg')
        import matplotlib.pyplot as plt
        fig, axes = plt.subplots(2, 2, figsize=(14, 12))
        ax_xy, ax_xz, ax_yz, ax_dt = axes[0,0], axes[0,1], axes[1,0], axes[1,1]
        colors = {'GX020030':'tab:red','GX030030':'tab:blue','GX040030':'tab:green','GX050030':'tab:orange'}
        by_seg = {}
        for r in out_rows:
            by_seg.setdefault(r['segment'], []).append(r)
        # plot absolute trajectory
        for seg, rs in by_seg.items():
            e=[r['east_m'] for r in rs]; n=[r['north_m'] for r in rs]; u=[r['up_m'] for r in rs]; d=[r['depth_m'] if r['depth_m']!='' else 0 for r in rs]
            c = colors.get(seg, 'gray')
            ax_xy.plot(e, n, '-', color=c, label=seg, alpha=0.7, linewidth=1)
            ax_xy.plot(e[0], n[0], 'o', color=c, markersize=8)
            ax_xz.plot(e, u, '-', color=c, alpha=0.7, linewidth=1)
            ax_yz.plot(n, u, '-', color=c, alpha=0.7, linewidth=1)
            t0 = rs[0]['timestamp_s']
            ax_dt.plot([(r['timestamp_s']-t0)/60 for r in rs], d, '-', color=c, alpha=0.8, linewidth=1, label=seg)
        # also plot GPS fixes
        gps_e = [latlon_to_enu(g[1]['lat'], g[1]['lon'], g[1]['alt'], lat0, lon0, alt0)[0] for g in gps] if gps else []
        gps_n = [latlon_to_enu(g[1]['lat'], g[1]['lon'], g[1]['alt'], lat0, lon0, alt0)[1] for g in gps] if gps else []
        if gps_e:
            ax_xy.plot(gps_e, gps_n, 'k.', markersize=2, alpha=0.5, label='GPS fixes')
        ax_xy.set_xlabel('East (m)'); ax_xy.set_ylabel('North (m)'); ax_xy.set_title('Top view ENU')
        ax_xy.set_aspect('equal'); ax_xy.legend(loc='best', fontsize=8); ax_xy.grid(True, alpha=0.3)
        ax_xz.set_xlabel('East (m)'); ax_xz.set_ylabel('Up (m)'); ax_xz.set_title('Side East-Up'); ax_xz.set_aspect('equal'); ax_xz.grid(True, alpha=0.3)
        ax_yz.set_xlabel('North (m)'); ax_yz.set_ylabel('Up (m)'); ax_yz.set_title('Side North-Up'); ax_yz.set_aspect('equal'); ax_yz.grid(True, alpha=0.3)
        ax_dt.set_xlabel('Time (min)'); ax_dt.set_ylabel('Depth (m, from pressure)'); ax_dt.set_title('Depth over time (MCAP pressure)'); ax_dt.grid(True, alpha=0.3); ax_dt.legend(loc='best', fontsize=8); ax_dt.invert_yaxis()
        fig.suptitle('AUV213 trajectory — absolute (Umeyama lingbot → ENU)')
        plt.tight_layout()
        plt.savefig(args.plot, dpi=130, bbox_inches='tight')
        print(f'[plot] {args.plot}', flush=True)

if __name__ == '__main__':
    main()
