#!/usr/bin/env python3
"""Stage 03 — Filter and smooth navigation trajectories.

Input:  ~/cosma-pipeline/data/<mission>/nav/<AUV>_<segment>.parquet
Output: ~/cosma-pipeline/data/<mission>/nav_filtered/<AUV>_<segment>.parquet

Steps:
1. Drop rows with null lat/lon OR lat==0 AND lon==0 (no GPS lock)
2. MAD-3σ outlier removal on lat, lon
3. Moving average smoothing (window 5s, KISS)
4. Depth from pressure: depth_m = (pressure_pa - 101325) / (1025 * 9.81)
5. Output: same columns + lat_smooth, lon_smooth, depth_m

Usage:
    python3 03_nav_filter.py /home/cosma/cosma-pipeline/data/20260505-Lepradet/nav/
    python3 03_nav_filter.py /path/nav/ --auv AUV013 --sigma 2.5
"""
from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).parent.parent))
from orchestrator.db import init_db, get_conn, upsert_job, record_metric, now_iso

PIPELINE_BASE = Path(os.environ.get("COSMA_PIPELINE_BASE", "/home/cosma/cosma-pipeline"))
RHO_SEA = 1025.0  # kg/m3
G = 9.81          # m/s2
P_ATM = 101325.0  # Pa


def mad_mask(arr: np.ndarray, sigma: float = 3.0) -> np.ndarray:
    """True = keep."""
    if len(arr) < 4:
        return np.ones(len(arr), dtype=bool)
    med = np.median(arr)
    mad = np.median(np.abs(arr - med))
    if mad == 0:
        return np.ones(len(arr), dtype=bool)
    return np.abs(0.6745 * (arr - med) / mad) < sigma


def moving_average(arr: np.ndarray, window: int = 5) -> np.ndarray:
    if len(arr) < window:
        return arr.copy()
    pad = window // 2
    padded = np.pad(arr, (pad, pad), mode="edge")
    return np.convolve(padded, np.ones(window) / window, mode="valid")[:len(arr)]


def filter_parquet(src: Path, dst_dir: Path, sigma: float = 3.0, window: int = 5) -> dict:
    import pandas as pd

    df = pd.read_parquet(src)
    auv_seg = src.stem
    metrics = {
        "file": src.name,
        "points_in": len(df),
        "points_out": 0,
        "status": "ok",
    }

    # Step 1: drop null/zero GPS
    has_lat = "lat" in df.columns and df["lat"].notna().any()
    if has_lat:
        mask_valid = df["lat"].notna() & df["lon"].notna() & (df["lat"] != 0) & (df["lon"] != 0)
        df_valid = df[mask_valid].copy()
    else:
        # No GPS — keep all rows for IMU/pressure
        df_valid = df.copy()
        metrics["degraded"] = True

    if len(df_valid) == 0:
        metrics["status"] = "degraded"
        metrics["note"] = "no valid GPS points"
        print(f"  [03] {auv_seg}: no valid GPS — saving as-is with depth calc only")
        df_out = df.copy()
    else:
        # Step 2: MAD outlier removal on lat/lon
        if has_lat and len(df_valid) >= 4:
            lats = df_valid["lat"].values
            lons = df_valid["lon"].values
            mask = mad_mask(lats, sigma) & mad_mask(lons, sigma)
            n_removed = int((~mask).sum())
            df_valid = df_valid[mask].copy()
            metrics["points_removed_outlier"] = n_removed
        else:
            metrics["points_removed_outlier"] = 0

        # Step 3: sort by timestamp
        if "ts_ns" in df_valid.columns:
            df_valid = df_valid.sort_values("ts_ns").reset_index(drop=True)

        # Step 4: smooth lat/lon
        if has_lat and len(df_valid) >= window:
            df_valid["lat_smooth"] = moving_average(df_valid["lat"].values, window)
            df_valid["lon_smooth"] = moving_average(df_valid["lon"].values, window)
        elif has_lat and len(df_valid) > 0:
            df_valid["lat_smooth"] = df_valid["lat"]
            df_valid["lon_smooth"] = df_valid["lon"]
        else:
            df_valid["lat_smooth"] = np.nan
            df_valid["lon_smooth"] = np.nan

        df_out = df_valid

    # Step 5: depth from pressure
    if "pressure_pa" in df_out.columns and df_out["pressure_pa"].notna().any():
        df_out["depth_m"] = (df_out["pressure_pa"] - P_ATM) / (RHO_SEA * G)
        df_out["depth_m"] = df_out["depth_m"].abs()  # negative when underwater (P < Patm)  # surface = 0
    else:
        df_out["depth_m"] = np.nan

    dst_dir.mkdir(parents=True, exist_ok=True)
    out_path = dst_dir / src.name
    df_out.to_parquet(out_path, index=False)

    metrics["points_out"] = len(df_out)
    removed_null = metrics["points_in"] - len(df_out) - metrics.get("points_removed_outlier", 0)
    metrics["points_removed_null"] = max(0, removed_null)
    print(f"  [03] {auv_seg}: {metrics['points_in']} → {metrics['points_out']} pts, "
          f"depth_m range=[{df_out['depth_m'].min():.1f}, {df_out['depth_m'].max():.1f}]"
          if df_out["depth_m"].notna().any() else
          f"  [03] {auv_seg}: {metrics['points_in']} → {metrics['points_out']} pts, no pressure")
    return metrics


def filter_mission(nav_dir: Path, auv_filter: str | None = None,
                   sigma: float = 3.0, window: int = 5) -> list[dict]:
    out_dir = nav_dir.parent / "nav_filtered"

    parquet_files = sorted(nav_dir.glob("*.parquet"))
    if auv_filter:
        parquet_files = [f for f in parquet_files if auv_filter in f.name]

    all_metrics = []
    init_db()

    for pf in parquet_files:
        out_file = out_dir / pf.name
        if out_file.exists():
            print(f"[03] {pf.stem}: cached")
            continue

        print(f"[03] Filtering {pf.name}...")
        m = filter_parquet(pf, out_dir, sigma=sigma, window=window)
        all_metrics.append(m)

        with get_conn() as conn:
            mission_name = nav_dir.parent.name
            mission_row = conn.execute("SELECT id FROM missions WHERE name=?", (mission_name,)).fetchone()
            if mission_row:
                auv_id = pf.stem.split("_")[0]
                job_id = upsert_job(conn, mission_row["id"], auv_id, "all", "03_nav_filter",
                                    status="done" if m.get("status") == "ok" else m.get("status", "done"),
                                    output_path=str(out_dir))
                record_metric(conn, job_id, "nav_points_filtered", value=m.get("points_out", 0),
                              pass_fail="pass" if m.get("points_out", 0) > 0 else "warn")

    return all_metrics


def main():
    ap = argparse.ArgumentParser(description="Stage 03 — Filter nav trajectories")
    ap.add_argument("nav_dir", type=Path, help="Directory with *.parquet from stage 02")
    ap.add_argument("--auv", type=str, default=None)
    ap.add_argument("--sigma", type=float, default=3.0)
    ap.add_argument("--window", type=int, default=5)
    args = ap.parse_args()

    metrics = filter_mission(args.nav_dir, auv_filter=args.auv,
                             sigma=args.sigma, window=args.window)

    print("\n=== Stage 03 summary ===")
    for m in metrics:
        print(f"  {m.get('file','?')}: {m.get('points_in',0)} → {m.get('points_out',0)} "
              f"[{m.get('status','?')}]")


if __name__ == "__main__":
    main()
