#!/usr/bin/env python3
"""Test multiple focal lengths on DVL and compare trajectory drift."""
import subprocess, csv, sys, math
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

FRAMES = '/home/cosma/cosma-pipeline/data/20260505-Lepradet/frames/AUV210/GX039839'
START = '2026-05-05T08:33:41'
LABEL = 'GX039839'
SCRIPT = '/home/cosma/cosma-qc/scripts/dvl_optical_full.py'

# focal in px equivalent to W/2 / tan(fov/2). W=518
# fov 100°→f=217, 110°→f=185, 122°→f=143, 130°→f=121, 140°→f=94, 150°→f=70
# But focal in px doesn't directly map to FOV unless we use that conversion. We pass --fov-deg.
fovs = [90, 100, 110, 122, 135, 150]

results = []
for fov in fovs:
    out_csv = f'/tmp/sweep_fov{fov}.csv'
    print(f'[sweep] fov={fov}', flush=True)
    r = subprocess.run(['python3', SCRIPT, '--frames-dir', FRAMES, '--altitude', '1.5',
                        '--fov-deg', str(fov), '--fps', '1.0', '--start-iso', START,
                        '--label', LABEL, '--out', out_csv], capture_output=True, text=True, timeout=600)
    if r.returncode != 0:
        print(f'  FAIL: {r.stderr[-300:]}', flush=True)
        continue
    # parse last position + drift metrics
    rows = list(csv.DictReader(open(out_csv)))
    e = [float(r['east_m']) for r in rows]
    n = [float(r['north_m']) for r in rows]
    h = [float(r['heading_deg']) for r in rows]
    end_x, end_y = e[-1], n[-1]
    end_dist = math.sqrt(end_x**2 + end_y**2)
    path_len = sum(math.sqrt((e[i]-e[i-1])**2 + (n[i]-n[i-1])**2) for i in range(1, len(e)))
    bbox = (max(e)-min(e), max(n)-min(n))
    results.append({'fov': fov, 'csv': out_csv, 'end_x': end_x, 'end_y': end_y, 
                    'end_dist': end_dist, 'path_len': path_len, 'bbox': bbox, 'rows': rows,
                    'e_arr': e, 'n_arr': n, 'h_arr': h})
    print(f'  end=({end_x:.1f},{end_y:.1f}) dist={end_dist:.1f}m path={path_len:.1f}m bbox={bbox}', flush=True)

# Plot all trajectories
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
axes = axes.flatten()
for ax, res in zip(axes, results):
    ax.plot(res['e_arr'], res['n_arr'], '-b', linewidth=0.8)
    ax.plot(res['e_arr'][0], res['n_arr'][0], 'go', markersize=8)
    ax.plot(res['e_arr'][-1], res['n_arr'][-1], 'r^', markersize=8)
    ax.set_xlabel('East (m)'); ax.set_ylabel('North (m)')
    ax.set_title(f'FOV={res["fov"]}° bbox=({res["bbox"][0]:.0f}×{res["bbox"][1]:.0f})m\nend_dist={res["end_dist"]:.1f}m path={res["path_len"]:.0f}m')
    ax.set_aspect('equal'); ax.grid(True, alpha=0.3)
plt.suptitle(f'DVL focal sweep — {LABEL} (assume closed-loop survey → smaller end_dist=better)')
plt.tight_layout()
plt.savefig('/tmp/sweep_focal.png', dpi=110, bbox_inches='tight')
print('[plot] /tmp/sweep_focal.png', flush=True)

# Summary
print('\n=== Summary ===')
for r in sorted(results, key=lambda x: x['end_dist']):
    print(f"FOV={r['fov']}°: end_dist={r['end_dist']:.1f}m  path={r['path_len']:.0f}m  bbox={r['bbox'][0]:.0f}×{r['bbox'][1]:.0f}m")