#!/usr/bin/env python3
"""Coverage swath QC plot — project each frame footprint on ground.

Usage:
  python3 coverage_swath.py --traj-csv /tmp/dvl_loopclosed_GX039839.csv \
    --frames-dir /home/cosma/...AUV210/GX039839 \
    --altitude 1.5 --fov-h 122 --fov-v 80 --out /tmp/coverage_GX039839.png
"""
import argparse, csv, math
from pathlib import Path
import numpy as np
import cv2

def compute_qc(frame_path):
    """R<G-5 && R<B-5 underwater test, return ratio of bottom_visible-ish pixels."""
    img = cv2.imread(str(frame_path), cv2.IMREAD_COLOR)
    if img is None: return 0.0
    b, g, r = cv2.split(img)
    mask = (r < g.astype(int) - 5) & (r < b.astype(int) - 5)
    # contrast: stddev of gray channel
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    if gray.std() < 20: return 0.0  # turbid / blurry
    return float(mask.mean())

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--traj-csv', required=True)
    ap.add_argument('--frames-dir', required=True)
    ap.add_argument('--altitude', type=float, default=1.5)
    ap.add_argument('--fov-h', type=float, default=122.0)
    ap.add_argument('--fov-v', type=float, default=80.0)
    ap.add_argument('--out', required=True)
    ap.add_argument('--x-col', default='east_m_corr')
    ap.add_argument('--y-col', default='north_m_corr')
    ap.add_argument('--label', default='segment')
    ap.add_argument('--heading-csv', default=None, help='separate CSV with heading_deg per frame')
    ap.add_argument('--sample-every', type=int, default=10, help='draw every N frames')
    args = ap.parse_args()
    
    # Load trajectory
    rows = list(csv.DictReader(open(args.traj_csv)))
    # autodetect col names
    cols = rows[0].keys()
    if args.x_col not in cols: args.x_col = 'east_m' if 'east_m' in cols else 'x'
    if args.y_col not in cols: args.y_col = 'north_m' if 'north_m' in cols else 'y'
    print(f'[cov] {len(rows)} rows, x_col={args.x_col} y_col={args.y_col}', flush=True)
    
    # Load heading from a heading CSV (or assume 0)
    headings = {}
    if args.heading_csv:
        for r in csv.DictReader(open(args.heading_csv)):
            headings[int(r['frame_idx'])] = float(r['heading_deg'])
        print(f'[cov] {len(headings)} headings loaded', flush=True)
    elif 'heading_deg' in cols:
        for r in rows:
            headings[int(r['frame_idx'])] = float(r['heading_deg'])
    
    frames_dir = Path(args.frames_dir)
    # Footprint dimensions at altitude
    half_w = args.altitude * math.tan(math.radians(args.fov_h/2))
    half_h = args.altitude * math.tan(math.radians(args.fov_v/2))
    print(f'[cov] footprint at alt={args.altitude}m: {2*half_w:.2f}m × {2*half_h:.2f}m', flush=True)
    
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    from matplotlib.patches import Rectangle
    from matplotlib.transforms import Affine2D
    from matplotlib.collections import PatchCollection
    
    fig, axes = plt.subplots(1, 2, figsize=(20, 10))
    ax_cov, ax_traj = axes[0], axes[1]
    
    # Collect rectangles + colors
    rects = []
    colors = []
    qc_values = []
    xs = []; ys = []
    
    for r in rows[::args.sample_every]:
        fi = int(r['frame_idx'])
        x = float(r[args.x_col])
        y = float(r[args.y_col])
        xs.append(x); ys.append(y)
        hdg = headings.get(fi, 0.0)
        # QC
        fpath = frames_dir / f'frame_{fi+1:05d}.jpg'
        if fpath.exists():
            qc = compute_qc(fpath)
        else:
            qc = 0.0
        qc_values.append(qc)
        
        # Build rotated rectangle
        # rectangle in body frame: x = +/- half_w (right), y = +/- half_h (forward)
        # but in world: rotated by heading
        from matplotlib.patches import Polygon
        corners_body = np.array([
            [-half_w, -half_h],
            [+half_w, -half_h],
            [+half_w, +half_h],
            [-half_w, +half_h],
        ])
        # heading rotation (clockwise from north, so for math invert)
        th = math.radians(hdg)
        R = np.array([[math.cos(th), math.sin(th)],
                      [-math.sin(th), math.cos(th)]])
        corners_world = corners_body @ R.T + np.array([x, y])
        rects.append(corners_world)
        colors.append(qc)
    
    # Plot coverage
    from matplotlib.patches import Polygon
    for corners, qc in zip(rects, colors):
        poly = Polygon(corners, alpha=0.10, edgecolor='black', linewidth=0.1,
                       facecolor=plt.cm.RdYlGn(qc * 1.5 if qc < 0.7 else 1.0))
        ax_cov.add_patch(poly)
    
    ax_cov.plot(xs, ys, '-k', linewidth=0.5, alpha=0.5)
    ax_cov.plot(xs[0], ys[0], 'go', markersize=12, label='start')
    ax_cov.plot(xs[-1], ys[-1], 'r^', markersize=12, label='end')
    ax_cov.set_xlabel('East (m)'); ax_cov.set_ylabel('North (m)')
    ax_cov.set_title(f'Coverage swath — {args.label}\n{len(rects)} footprints @ alt={args.altitude}m FOV {args.fov_h}°×{args.fov_v}°\n(green=bottom visible, red=hors-eau/turbid)')
    ax_cov.set_aspect('equal'); ax_cov.legend(); ax_cov.grid(True, alpha=0.3)
    
    # Trajectory only with QC color
    sc = ax_traj.scatter(xs, ys, c=colors, cmap='RdYlGn', s=12, vmin=0, vmax=0.7)
    plt.colorbar(sc, ax=ax_traj, label='bottom_visible ratio')
    ax_traj.plot(xs[0], ys[0], 'go', markersize=12)
    ax_traj.plot(xs[-1], ys[-1], 'r^', markersize=12)
    ax_traj.set_xlabel('East (m)'); ax_traj.set_ylabel('North (m)')
    ax_traj.set_title(f'Trajectoire colorée par QC ({len(xs)} points)'); ax_traj.set_aspect('equal'); ax_traj.grid(True, alpha=0.3)
    
    plt.suptitle(f'Acquisition QC swath — {args.label}', fontsize=14)
    plt.tight_layout()
    plt.savefig(args.out, dpi=120, bbox_inches='tight')
    print(f'[plot] {args.out}', flush=True)

if __name__ == '__main__': main()