#!/usr/bin/env python3
"""Optical DVL + IMU heading correction.
Same as dvl_optical.py but rotates cam-frame flow to world frame using compass_hdg.

Usage:
  python3 dvl_optical_imu.py --frames-dir <dir> --bag-dir <auv_bags_dir> \
      --altitude 1.5 --fps 1.0 --start-iso ... --label ... \
      --out csv --plot png [--ref-csv ...]
"""
import argparse, csv, math, sys
from pathlib import Path
from datetime import datetime
import numpy as np
import cv2
from rosbags.highlevel import AnyReader

def load_heading(bag_dir, t_start, t_end):
    bags = sorted(Path(bag_dir).glob('*.mcap'))
    # filter empty
    bags = [b for b in bags if b.stat().st_size > 1000]
    headings = []  # list of (ts_s, heading_deg)
    for b in bags:
        try:
            with AnyReader([b]) as r:
                for conn, ts_ns, raw in r.messages(connections=[c for c in r.connections if c.topic == '/mavros/global_position/compass_hdg']):
                    t = ts_ns / 1e9
                    if t_start - 60 <= t <= t_end + 60:
                        m = r.deserialize(raw, conn.msgtype)
                        headings.append((t, m.data))
        except Exception as e:
            print(f'[warn] {b.name}: {e}', flush=True)
    headings.sort()
    return headings

def nearest_hdg(headings, t_target):
    if not headings: return None
    ts = [h[0] for h in headings]
    idx = np.searchsorted(ts, t_target)
    if idx == 0: return headings[0][1]
    if idx >= len(headings): return headings[-1][1]
    if abs(headings[idx][0] - t_target) < abs(headings[idx-1][0] - t_target):
        return headings[idx][1]
    return headings[idx-1][1]

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--frames-dir', required=True)
    ap.add_argument('--bag-dir', required=True)
    ap.add_argument('--altitude', type=float, default=1.5)
    ap.add_argument('--fov-deg', type=float, default=122.0)
    ap.add_argument('--fps', type=float, default=1.0)
    ap.add_argument('--start-iso', default='2026-05-05T00:00:00')
    ap.add_argument('--label', default='segment')
    ap.add_argument('--out', required=True)
    ap.add_argument('--plot', default=None)
    ap.add_argument('--ref-csv', default=None)
    args = ap.parse_args()
    
    frames = sorted(Path(args.frames_dir).glob('frame_*.jpg'))
    print(f'[dvl] {len(frames)} frames', flush=True)
    
    t0 = datetime.fromisoformat(args.start_iso).timestamp()
    t_end = t0 + len(frames) / args.fps
    
    # Load heading
    print(f'[hdg] loading from {args.bag_dir}', flush=True)
    headings = load_heading(args.bag_dir, t0, t_end)
    print(f'[hdg] {len(headings)} samples loaded, t range: {headings[0][0]:.0f}-{headings[-1][0]:.0f}', flush=True)
    
    W, H = 518, 294
    f = (W/2) / math.tan(math.radians(args.fov_deg/2))
    px_to_m = args.altitude / f
    print(f'[dvl] px_to_m={px_to_m:.5f}', flush=True)
    
    rows = []
    rows.append({'frame_idx': 0, 'ts_s': t0, 'heading_deg': nearest_hdg(headings, t0) or 0, 'flow_x_px': 0, 'flow_y_px': 0, 
                 'speed_mps': 0, 'east_m': 0, 'north_m': 0})
    
    prev = cv2.imread(str(frames[0]), cv2.IMREAD_GRAYSCALE)
    east_cum, north_cum = 0.0, 0.0
    
    for i in range(1, len(frames)):
        curr = cv2.imread(str(frames[i]), cv2.IMREAD_GRAYSCALE)
        if curr is None: continue
        t_frame = t0 + i / args.fps
        hdg = nearest_hdg(headings, t_frame) or 0
        
        flow = cv2.calcOpticalFlowFarneback(prev, curr, None, 0.5, 3, 21, 3, 5, 1.2, 0)
        fx_cam = np.median(flow[..., 0])
        fy_cam = np.median(flow[..., 1])
        
        # convert px → m in CAM frame (cam right = +X_cam, cam down = +Y_cam image coord)
        dx_cam = -fx_cam * px_to_m  # AUV moves opposite to flow
        dy_cam = -fy_cam * px_to_m
        
        # Apply heading rotation: cam +X_cam = body forward? assume cam frame Y axis = AUV forward
        # The downward camera: cam +Y_image = body forward typically (or -Y if mounted otherwise)
        # heading = degrees clockwise from North in body frame
        # World rotation: rotate body (dy_cam = forward, dx_cam = right) by heading angle from north
        hdg_rad = math.radians(hdg)
        # body forward (north when hdg=0) component:
        # body_forward_m = dy_cam (assuming cam Y_image = forward)
        # body_right_m = dx_cam
        body_forward = dy_cam  # may need sign flip depending on mounting; we'll see
        body_right = dx_cam
        # world East = forward*sin(hdg) + right*cos(hdg)
        # world North = forward*cos(hdg) - right*sin(hdg)
        de = body_forward * math.sin(hdg_rad) + body_right * math.cos(hdg_rad)
        dn = body_forward * math.cos(hdg_rad) - body_right * math.sin(hdg_rad)
        
        east_cum += de
        north_cum += dn
        speed_mps = math.sqrt(de**2 + dn**2) * args.fps
        
        rows.append({'frame_idx': i, 'ts_s': t_frame, 'heading_deg': hdg, 'flow_x_px': float(fx_cam), 'flow_y_px': float(fy_cam),
                     'speed_mps': speed_mps, 'east_m': east_cum, 'north_m': north_cum})
        prev = curr
        if i % 100 == 0:
            print(f'[dvl] {i}/{len(frames)} hdg={hdg:.1f}° flow=({fx_cam:.1f},{fy_cam:.1f}) pos=({east_cum:.2f},{north_cum:.2f})', flush=True)
    
    print(f'[dvl] done. Final ENU: ({east_cum:.2f}, {north_cum:.2f}) m', flush=True)
    
    with open(args.out, 'w', newline='') as ff:
        w = csv.DictWriter(ff, fieldnames=list(rows[0].keys()))
        w.writeheader(); w.writerows(rows)
    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_hdg, ax_speed, ax_cmp = axes[0,0], axes[0,1], axes[1,0], axes[1,1]
        e = [r['east_m'] for r in rows]
        n = [r['north_m'] for r in rows]
        ax_xy.plot(e, n, '-b', linewidth=1.2)
        ax_xy.plot(e[0], n[0], 'go', markersize=10, label='start')
        ax_xy.plot(e[-1], n[-1], 'r^', markersize=10, label='end')
        ax_xy.set_xlabel('East (m)'); ax_xy.set_ylabel('North (m)'); ax_xy.set_title('DVL + IMU heading trajectory')
        ax_xy.set_aspect('equal'); ax_xy.legend(); ax_xy.grid(True, alpha=0.3)
        
        hdgs = [r['heading_deg'] for r in rows]
        ax_hdg.plot(range(len(rows)), hdgs, '-c'); ax_hdg.set_xlabel('Frame'); ax_hdg.set_ylabel('Heading (deg)'); ax_hdg.set_title('Compass heading from MCAP'); ax_hdg.grid(True, alpha=0.3)
        
        speeds = [r['speed_mps'] for r in rows]
        ax_speed.plot(range(len(rows)), speeds, '-r'); ax_speed.set_xlabel('Frame'); ax_speed.set_ylabel('Speed m/s'); ax_speed.set_title('Speed over time'); ax_speed.grid(True, alpha=0.3)
        
        if args.ref_csv:
            try:
                with open(args.ref_csv) as fff:
                    refrows = [r for r in csv.DictReader(fff) if r.get('segment','')==args.label or r.get('label','')==args.label]
                rx = [float(r['x']) for r in refrows]
                ry = [float(r['y']) for r in refrows]
                ax_cmp.plot(e, n, '-b', linewidth=1.2, label='DVL+IMU', alpha=0.7)
                ax_cmp.plot(rx, ry, '-r', linewidth=1.2, label='lingbot', alpha=0.7)
                ax_cmp.set_xlabel('X/East'); ax_cmp.set_ylabel('Y/North'); ax_cmp.set_title('Comparison'); ax_cmp.set_aspect('equal'); ax_cmp.legend(); ax_cmp.grid(True, alpha=0.3)
            except Exception as e: print(f'[ref] {e}')
        else:
            ax_cmp.set_title('(no reference)')
        
        plt.suptitle(f'DVL+IMU heading — {args.label}')
        plt.tight_layout()
        plt.savefig(args.plot, dpi=130, bbox_inches='tight')
        print(f'[plot] {args.plot}', flush=True)

if __name__ == '__main__': main()