#!/usr/bin/env python3
"""
merge_nav_usbl.py — Merge USBL decoded data with USV navigation log

Inputs:
  --usbl   : combined_usbl.csv (output of parse_kogger_usbl.py)
  --nav-dir: directory containing *_navigation_log.csv files
  --output : output CSV (default: combined_nav_usbl.csv)

Nav log format: timestamp,data,value  (long format)
  data=Lat → latitude_deg
  data=Lon → longitude_deg
  data=Heading → heading_deg

Interpolation: for each USBL timestamp, find nearest nav point within 1s window.
If multiple sessions, use session matching by timestamp overlap.

AUV position calculation:
  azimuth_deg from USBL is RELATIVE to USV heading (yaw) based on USBL hardware mounting.
  AUV bearing from North = (USV heading + azimuth_deg) mod 360
  Horizontal dist = dist_m * cos(elevation_deg * pi/180)
  
  Note: if azimuth is already absolute (referenced to North), do NOT add heading.
  Check: usbl_yaw in payload should match nav Heading if relative azimuth.
  We use RELATIVE convention (add USV heading) — documented here.

Geodetic forward (haversine):
  Using flat-earth approximation valid for distances < 500m:
  dlat = horiz_dist * cos(bearing) / R_earth
  dlon = horiz_dist * sin(bearing) / (R_earth * cos(lat))
  
  R_earth = 6371000 m
"""

import csv
import io
import sys
import math
import os

R_EARTH = 6371000.0  # meters


def parse_nav_log(nav_file):
    """
    Parse navigation_log.csv into:
    - nav_points: sorted list of (timestamp_str, lat, lon) from Lat+Lon entries
    - heading_series: sorted list of (timestamp_str, heading_deg) from Heading entries

    Lat/Lon and Heading have different timestamps so must be interpolated separately.
    Returns (nav_points, heading_series).
    """
    lat_by_ts = {}
    lon_by_ts = {}
    heading_series_raw = []

    with open(nav_file) as f:
        reader = csv.DictReader(f)
        for row in reader:
            data = row.get('data', '')
            ts = row.get('timestamp', '')
            try:
                val_raw = row.get('value', '') or ''
                if not val_raw:
                    continue
                val = float(val_raw)
            except (ValueError, TypeError):
                continue
            if data == 'Lat':
                lat_by_ts[ts] = val
            elif data == 'Lon':
                lon_by_ts[ts] = val
            elif data == 'Heading':
                heading_series_raw.append((ts, val))

    # Build nav_points: timestamps where both Lat and Lon appear (same ts)
    nav_points = []
    for ts in sorted(set(lat_by_ts.keys()) & set(lon_by_ts.keys())):
        nav_points.append((ts, lat_by_ts[ts], lon_by_ts[ts]))

    # heading_series sorted by timestamp
    heading_series = sorted(heading_series_raw, key=lambda x: x[0])

    return nav_points, heading_series


def ts_to_seconds(ts_str):
    """Convert '2026-03-24 09:29:05.230392' to float seconds since epoch (approx)."""
    # Simple: parse date+time, compute offset
    try:
        date_part, time_part = ts_str.strip().split(' ', 1)
        y, mo, d = date_part.split('-')
        parts = time_part.split(':')
        h, m = int(parts[0]), int(parts[1])
        s_str = parts[2]
        s = float(s_str)
        # Days since fixed epoch (don't need absolute, just relative diffs)
        total = (int(y)*365 + int(mo)*30 + int(d)) * 86400 + h*3600 + m*60 + s
        return total
    except Exception:
        return 0.0


def find_nearest(ts_sec, series_sec, series, max_gap=1.0):
    """Find nearest entry in sorted series within max_gap seconds."""
    best_idx = -1
    best_dt = float('inf')
    for i, (s_sec, entry) in enumerate(zip(series_sec, series)):
        dt = abs(s_sec - ts_sec)
        if dt < best_dt:
            best_dt = dt
            best_idx = i
        elif s_sec > ts_sec + max_gap:
            break
    if best_idx >= 0 and best_dt <= max_gap:
        return series[best_idx], best_dt
    return None, None


def geodetic_forward(lat_deg, lon_deg, bearing_deg, dist_m):
    """
    Compute destination point given start lat/lon, bearing (deg from North), distance (m).
    Flat-earth approximation valid for dist < 500m.
    """
    bearing_rad = math.radians(bearing_deg)
    lat_rad = math.radians(lat_deg)
    dlat = dist_m * math.cos(bearing_rad) / R_EARTH
    dlon = dist_m * math.sin(bearing_rad) / (R_EARTH * math.cos(lat_rad))
    auv_lat = lat_deg + math.degrees(dlat)
    auv_lon = lon_deg + math.degrees(dlon)
    return auv_lat, auv_lon


def haversine_dist(lat1, lon1, lat2, lon2):
    """Distance in meters between two lat/lon points."""
    phi1, phi2 = math.radians(lat1), math.radians(lat2)
    dphi = math.radians(lat2 - lat1)
    dlam = math.radians(lon2 - lon1)
    a = math.sin(dphi/2)**2 + math.cos(phi1)*math.cos(phi2)*math.sin(dlam/2)**2
    return 2 * R_EARTH * math.asin(math.sqrt(a))


def find_nav_file_for_session(usbl_file, nav_dir):
    """Match nav file by common timestamp prefix."""
    base = os.path.basename(usbl_file)
    # e.g. 2026-03-24_09-28-44_USV003_usbl.csv -> 2026-03-24_09-28-44_USV003
    prefix = base.replace('_usbl.csv', '')
    nav_candidate = os.path.join(nav_dir, prefix + '_navigation_log.csv')
    if os.path.exists(nav_candidate):
        return nav_candidate
    return None


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--usbl', required=True, help='combined_usbl.csv')
    parser.add_argument('--nav-dir', required=True, help='Directory with *_navigation_log.csv')
    parser.add_argument('--output', default='combined_nav_usbl.csv')
    parser.add_argument('--max-gap', type=float, default=1.0, help='Max timestamp gap in seconds')
    args = parser.parse_args()

    # Load USBL records
    usbl_records = []
    with open(args.usbl) as f:
        reader = csv.DictReader(f)
        for row in reader:
            usbl_records.append(row)

    print("USBL records loaded: %d" % len(usbl_records))

    # Group by source_file
    from collections import defaultdict
    by_source = defaultdict(list)
    for rec in usbl_records:
        by_source[rec.get('source_file', '')].append(rec)

    # Load nav files
    nav_data = {}  # source_file -> (nav_points, nav_points_sec, heading_series, heading_sec)
    nav_dir = args.nav_dir

    for source_file in by_source.keys():
        # Try to match nav file
        prefix = source_file.replace('_usbl.csv', '')
        nav_file = os.path.join(nav_dir, prefix + '_navigation_log.csv')
        if not os.path.exists(nav_file):
            # Try to find by scanning directory
            matched = None
            for fn in os.listdir(nav_dir):
                if prefix in fn and 'navigation_log' in fn:
                    matched = os.path.join(nav_dir, fn)
                    break
            if matched is None:
                print("WARNING: no nav file found for %s" % source_file)
                nav_data[source_file] = ([], [], [], [])
                continue
            nav_file = matched
        nav_points, heading_series = parse_nav_log(nav_file)
        nav_points_sec = [ts_to_seconds(pt[0]) for pt in nav_points]
        heading_sec = [ts_to_seconds(h[0]) for h in heading_series]
        nav_data[source_file] = (nav_points, nav_points_sec, heading_series, heading_sec)
        print("Nav loaded for %s: %d pos points, %d heading points" % (
            source_file, len(nav_points), len(heading_series)))

    # Process and write output
    output_rows = []
    stats_match = 0
    stats_nomatch = 0

    for source_file, records in by_source.items():
        nav_points, nav_points_sec, heading_series, heading_sec = nav_data.get(
            source_file, ([], [], [], []))

        for rec in records:
            ts_str = rec.get('Timestamp', '')
            ts_sec = ts_to_seconds(ts_str)

            dist_str = rec.get('Dist', '')
            azimuth_str = rec.get('Azimuth', '')
            elev_str = rec.get('Elev', '')
            snr_str = rec.get('SNR', '')

            try:
                dist = float(dist_str) if dist_str else float('nan')
                azimuth = float(azimuth_str) if azimuth_str else float('nan')
                elev = float(elev_str) if elev_str else float('nan')
                snr = float(snr_str) if snr_str else float('nan')
            except ValueError:
                dist, azimuth, elev, snr = float('nan'), float('nan'), float('nan'), float('nan')

            nav_pt, dt = find_nearest(ts_sec, nav_points_sec, nav_points, args.max_gap)
            hdg_pt, _ = find_nearest(ts_sec, heading_sec, heading_series, args.max_gap)

            if nav_pt is None:
                stats_nomatch += 1
                lat_usv, lon_usv, heading_usv = float('nan'), float('nan'), float('nan')
                auv_lat, auv_lon = float('nan'), float('nan')
            else:
                stats_match += 1
                _, lat_usv, lon_usv = nav_pt
                heading_usv = hdg_pt[1] if hdg_pt is not None else float('nan')

                # Calculate AUV absolute position
                # Azimuth from USBL is relative to USV heading (yaw convention)
                # AUV bearing from North = (USV Heading + azimuth_deg) mod 360
                if not (math.isnan(dist) or math.isnan(azimuth) or
                        math.isnan(lat_usv) or math.isnan(heading_usv)):
                    horiz_dist = dist * math.cos(math.radians(elev)) if not math.isnan(elev) else dist
                    abs_bearing = (heading_usv + azimuth) % 360
                    auv_lat, auv_lon = geodetic_forward(lat_usv, lon_usv, abs_bearing, horiz_dist)
                else:
                    auv_lat, auv_lon = float('nan'), float('nan')

            output_rows.append({
                'Timestamp': ts_str,
                'lat': '%.7f' % lat_usv if not math.isnan(lat_usv) else '',
                'lon': '%.7f' % lon_usv if not math.isnan(lon_usv) else '',
                'Heading': '%.2f' % heading_usv if not math.isnan(heading_usv) else '',
                'Dist': '%.4f' % dist if not math.isnan(dist) else '',
                'Azimuth': '%.4f' % azimuth if not math.isnan(azimuth) else '',
                'Elev': '%.4f' % elev if not math.isnan(elev) else '',
                'SNR': '%.4f' % snr if not math.isnan(snr) else '',
                'auv_lat': '%.7f' % auv_lat if not math.isnan(auv_lat) else '',
                'auv_lon': '%.7f' % auv_lon if not math.isnan(auv_lon) else '',
                'nav_dt_s': '%.3f' % dt if dt is not None else '',
            })

    print("\n=== Merge stats ===")
    print("Matched: %d  No nav match: %d" % (stats_match, stats_nomatch))

    with open(args.output, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['Timestamp', 'lat', 'lon', 'Heading', 'Dist', 'Azimuth', 'Elev', 'SNR',
                         'auv_lat', 'auv_lon', 'nav_dt_s'])
        for row in output_rows:
            writer.writerow([
                row['Timestamp'], row['lat'], row['lon'], row['Heading'],
                row['Dist'], row['Azimuth'], row['Elev'], row['SNR'],
                row['auv_lat'], row['auv_lon'], row['nav_dt_s']
            ])

    print("Output: %s (%d rows)" % (args.output, len(output_rows)))

    # Sample output
    if output_rows:
        print("\n=== Sample (first 5 rows) ===")
        print("Timestamp,lat,lon,Heading,Dist,Azimuth,Elev,SNR,auv_lat,auv_lon")
        for row in output_rows[:5]:
            print("%s,%s,%s,%s,%s,%s,%s,%s,%s,%s" % (
                row['Timestamp'], row['lat'], row['lon'], row['Heading'],
                row['Dist'], row['Azimuth'], row['Elev'], row['SNR'],
                row['auv_lat'], row['auv_lon']))

    # AUV position stats
    auv_lats = [float(r['auv_lat']) for r in output_rows if r['auv_lat']]
    auv_lons = [float(r['auv_lon']) for r in output_rows if r['auv_lon']]
    usv_lats = [float(r['lat']) for r in output_rows if r['lat']]
    usv_lons = [float(r['lon']) for r in output_rows if r['lon']]

    if auv_lats and usv_lats:
        print("\n=== Position stats ===")
        print("AUV lat range: %.6f - %.6f" % (min(auv_lats), max(auv_lats)))
        print("AUV lon range: %.6f - %.6f" % (min(auv_lons), max(auv_lons)))
        print("USV lat range: %.6f - %.6f" % (min(usv_lats), max(usv_lats)))
        
        # Average USV-AUV distance
        dists = []
        for r in output_rows:
            if r['lat'] and r['auv_lat']:
                d = haversine_dist(float(r['lat']), float(r['lon']),
                                   float(r['auv_lat']), float(r['auv_lon']))
                dists.append(d)
        if dists:
            print("Avg USV-AUV dist: %.2f m  (min=%.2f max=%.2f)" % (
                sum(dists)/len(dists), min(dists), max(dists)))


if __name__ == '__main__':
    main()