#!/usr/bin/env python3
"""Stage 06b — IMU gravity + depth alignment with arbitrary origin.

Method (no GPS/USBL):
  1. Read MCAP for /mavros/imu/data (orientation) + /mavros/imu/static_pressure (depth)
  2. For each frame ts, find nearest IMU+depth
  3. Compute rotation = inverse of IMU body-to-world quaternion at frame 0
     → rotates lingbot frame so World Up is +Z
  4. Scale Z so range matches real depth (pressure-derived)
  5. Place origin at (east0, north0, 0) arbitrary
  6. Output CSV + plot trajectory in local ENU (origin = user-given coords)

Usage:
  python3 stage06b_imu_depth_align.py \
    --traj /tmp/auv213_full_trajectory.csv \
    --mcap-dir /mnt/ssd/.../20260505_150717_AUV013/ \
    --east0 1000 --north0 5000 \
    --out /tmp/auv213_aligned.csv --plot /tmp/auv213_aligned.png
"""
import argparse, csv, math, sys
import numpy as np
from pathlib import Path

def quat_to_rot(qw, qx, qy, qz):
    """Quaternion to 3x3 rotation matrix (body→world for mavros convention)."""
    n = math.sqrt(qw*qw + qx*qx + qy*qy + qz*qz)
    if n < 1e-9: return np.eye(3)
    qw, qx, qy, qz = qw/n, qx/n, qy/n, qz/n
    return np.array([
        [1 - 2*(qy*qy + qz*qz), 2*(qx*qy - qz*qw),     2*(qx*qz + qy*qw)],
        [2*(qx*qy + qz*qw),     1 - 2*(qx*qx + qz*qz), 2*(qy*qz - qx*qw)],
        [2*(qx*qz - qy*qw),     2*(qy*qz + qx*qw),     1 - 2*(qx*qx + qy*qy)],
    ])

def parse_mcap(bag_dir):
    from rosbags.highlevel import AnyReader
    bags = sorted(Path(bag_dir).glob('*.mcap'))
    data = {'imu': [], 'press': []}
    with AnyReader(bags) as reader:
        for conn, ts_ns, raw in reader.messages(connections=[c for c in reader.connections if c.topic in ['/mavros/imu/data','/mavros/imu/static_pressure']]):
            m = reader.deserialize(raw, conn.msgtype)
            if conn.topic == '/mavros/imu/data':
                q = m.orientation
                data['imu'].append((ts_ns, [q.w, q.x, q.y, q.z]))
            else:
                data['press'].append((ts_ns, m.fluid_pressure))
    data['imu'].sort(); data['press'].sort()
    return data

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

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--traj', required=True)
    ap.add_argument('--mcap-dir', required=True)
    ap.add_argument('--east0', type=float, default=1000.0)
    ap.add_argument('--north0', type=float, default=5000.0)
    ap.add_argument('--out', required=True)
    ap.add_argument('--plot', default=None)
    args = ap.parse_args()

    rows = []
    with open(args.traj) as f:
        for r in csv.DictReader(f):
            rows.append({'segment': r['segment'], 'frame_idx': int(r['frame_idx']),
                'ts_s': float(r['timestamp_s']), 'p': np.array([float(r['x']), float(r['y']), float(r['z'])])})
    print(f'[traj] {len(rows)} rows', flush=True)

    data = parse_mcap(args.mcap_dir)
    print(f'[mcap] imu={len(data["imu"])} press={len(data["press"])}', flush=True)

    # 1. Gravity alignment using IMU at first frame ts
    ts0_ns = int(rows[0]['ts_s'] * 1e9)
    imu0 = nearest(data['imu'], ts0_ns)
    if imu0 is None:
        print('[err] no IMU sample', flush=True); sys.exit(1)
    R_body2world = quat_to_rot(*imu0[1])
    # Camera looks down on AUV → cam optical axis = body -Z. So cam Z (depth) in world = -R_body2world @ [0,0,1] = up direction inverted
    # In lingbot frame, frame 0 sets identity. So lingbot_R_world = R_body2world
    # World coords from lingbot coords: p_world = R_body2world @ p_lingbot
    # But this assumes lingbot world axes coincide with body frame at t=0
    R_world = R_body2world  # rough first-order approximation
    print(f'[gravity] IMU q0 = {imu0[1]} -> R_body2world =', flush=True)
    print(R_world, flush=True)

    # 2. Depth scale: match lingbot Z range to real depth range from pressure
    depths = []
    for r in rows:
        ts_ns = int(r['ts_s']*1e9)
        p = nearest(data['press'], ts_ns)
        if p:
            d = (p[1] - 101325.0) / (1025.0 * 9.81)
            depths.append(d)
        else:
            depths.append(None)
    valid_depths = [d for d in depths if d is not None and d > 0.1]
    if valid_depths:
        depth_min, depth_max = min(valid_depths), max(valid_depths)
        depth_range = depth_max - depth_min
        print(f'[depth] real range: {depth_min:.2f}m to {depth_max:.2f}m ({depth_range:.2f}m)', flush=True)
    else:
        depth_min = depth_max = depth_range = 0
        print('[depth] no valid pressure data', flush=True)

    # apply rotation (no scaling yet)
    rotated = np.array([R_world @ r['p'] for r in rows])
    z_min, z_max = rotated[:, 2].min(), rotated[:, 2].max()
    z_range = z_max - z_min
    print(f'[lingbot after rot] Z range: {z_min:.2f} to {z_max:.2f} ({z_range:.2f}m)', flush=True)

    # scale to match depth range (if both available and meaningful)
    if depth_range > 0.5 and z_range > 0.1:
        scale_z = depth_range / z_range
        print(f'[scale_z] {scale_z:.3f}', flush=True)
    else:
        scale_z = 1.0
        print(f'[scale_z] 1.0 (insufficient data)', flush=True)

    # Apply isotropic scale (since the camera tracking is consistent in 3 axes)
    rotated_scaled = rotated * scale_z

    # Translate origin to user-given (east, north, 0)
    east0, north0 = args.east0, args.north0
    final = rotated_scaled.copy()
    final[:, 0] += east0 - rotated_scaled[0, 0]
    final[:, 1] += north0 - rotated_scaled[0, 1]
    # Z aligned to depth_min if available else first frame
    if valid_depths:
        z_offset = -depth_min - final[0, 2]  # depth is positive down ; up = -depth
        final[:, 2] += z_offset

    with open(args.out, 'w', newline='') as f:
        w = csv.writer(f)
        w.writerow(['segment','frame_idx','timestamp_s','east_m','north_m','up_m','depth_real_m','lingbot_x','lingbot_y','lingbot_z'])
        for i, r in enumerate(rows):
            w.writerow([r['segment'], r['frame_idx'], f"{r['ts_s']:.6f}",
                f"{final[i,0]:.4f}", f"{final[i,1]:.4f}", f"{final[i,2]:.4f}",
                f"{depths[i]:.3f}" if depths[i] is not None else '',
                f"{r['p'][0]:.4f}", f"{r['p'][1]:.4f}", f"{r['p'][2]:.4f}"])
    print(f'[out] {args.out}', 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_d = 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 i, r in enumerate(rows):
            by_seg.setdefault(r['segment'], []).append((final[i], depths[i], r['ts_s']))
        t0 = rows[0]['ts_s']
        for seg, lst in by_seg.items():
            e = [v[0][0] for v in lst]; n = [v[0][1] for v in lst]; u = [v[0][2] for v in lst]
            d_real = [v[1] if v[1] is not None else 0 for v in lst]
            ts = [(v[2]-t0)/60 for v in lst]
            c = colors.get(seg, 'gray')
            ax_xy.plot(e, n, '-', color=c, label=f'{seg} ({len(lst)})', linewidth=1.2, alpha=0.8)
            ax_xy.plot(e[0], n[0], 'o', color=c, markersize=8)
            ax_xz.plot(e, u, '-', color=c, linewidth=1.2, alpha=0.8)
            ax_yz.plot(n, u, '-', color=c, linewidth=1.2, alpha=0.8)
            ax_d.plot(ts, d_real, '-', color=c, linewidth=1.2, alpha=0.8, label=seg)
        ax_xy.set_xlabel('East (m, origin=1000)'); ax_xy.set_ylabel('North (m, origin=5000)')
        ax_xy.set_title('Top view ENU — gravity+depth aligned'); ax_xy.set_aspect('equal'); ax_xy.legend(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('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('North-Up'); ax_yz.set_aspect('equal'); ax_yz.grid(True, alpha=0.3)
        ax_d.set_xlabel('Time (min)'); ax_d.set_ylabel('Depth real (m, from pressure)'); ax_d.set_title('Real depth from FluidPressure'); ax_d.legend(fontsize=8); ax_d.grid(True, alpha=0.3); ax_d.invert_yaxis()
        fig.suptitle('AUV213 — gravity+depth aligned (origin 1000,5000 ; from IMU q0 + pressure scale)')
        plt.tight_layout()
        plt.savefig(args.plot, dpi=130, bbox_inches='tight')
        print(f'[plot] {args.plot}', flush=True)

if __name__ == '__main__': main()
