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

Input:  02_usbl_raw/<AUV>_nav_raw.json (or .csv)
Output: 03_usbl_filtered/<AUV>_nav_filtered.json + .csv

Steps:
1. Drop points with null lat/lon
2. MAD-3σ outlier removal on lat, lon, depth independently
3. Moving-average smoothing (window=5 by default)
4. Optional: simple 1D Kalman on each axis (KISS — no cross-covariance)

Usage:
    python3 03_usbl_filter.py /home/cosma/cosma-pipeline/20260505-Lepradet/02_usbl_raw/
    python3 03_usbl_filter.py /path/to/02_usbl_raw/ --auv AUV010 --sigma 2.5
"""
from __future__ import annotations

import argparse
import json
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

LOG_DIR = Path(os.environ.get("COSMA_PIPELINE_LOGS", "/home/cosma/cosma-pipeline/logs"))
LOG_DIR.mkdir(parents=True, exist_ok=True)


def mad_outlier_mask(arr: np.ndarray, sigma: float = 3.0) -> np.ndarray:
    """Returns boolean mask: True = keep (inlier). Uses MAD-based sigma."""
    if len(arr) < 4:
        return np.ones(len(arr), dtype=bool)
    median = np.median(arr)
    mad = np.median(np.abs(arr - median))
    if mad == 0:
        return np.ones(len(arr), dtype=bool)
    modified_z = 0.6745 * (arr - median) / mad
    return np.abs(modified_z) < sigma


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


def simple_kalman_1d(measurements: np.ndarray,
                     process_noise: float = 1e-4,
                     measurement_noise: float = 1e-2) -> np.ndarray:
    """Very simple 1D Kalman filter (scalar, no velocity state).
    KISS: just smooths, no cross-axis coupling.
    """
    n = len(measurements)
    filtered = np.zeros(n)
    x_est = measurements[0]
    p_est = 1.0

    for i, z in enumerate(measurements):
        # Predict
        p_pred = p_est + process_noise
        # Update
        K = p_pred / (p_pred + measurement_noise)
        x_est = x_est + K * (z - x_est)
        p_est = (1 - K) * p_pred
        filtered[i] = x_est

    return filtered


def filter_auv_nav(raw_path: Path, out_path: Path,
                   sigma: float = 3.0, window: int = 5,
                   use_kalman: bool = False) -> dict:
    """Filter nav for one AUV. Returns metrics dict."""
    data = json.loads(raw_path.read_text())
    points = data.get("points", [])
    auv_id = data.get("auv_id", raw_path.stem.replace("_nav_raw", ""))

    metrics = {
        "auv_id": auv_id,
        "points_before": len(points),
        "points_after": 0,
        "points_removed_null": 0,
        "points_removed_outlier": 0,
        "status": "ok",
    }

    if not points:
        metrics["status"] = "degraded"
        metrics["note"] = "no points to filter"
        _save_output(auv_id, [], out_path, metrics)
        return metrics

    # Step 1: Drop null lat/lon
    valid = [p for p in points if p.get("lat") is not None and p.get("lon") is not None]
    metrics["points_removed_null"] = len(points) - len(valid)

    if not valid:
        metrics["status"] = "degraded"
        metrics["note"] = "all points have null lat/lon (serial-only data)"
        print(f"  [03] {auv_id}: all null lat/lon — degraded (serial CSV source, no MCAP nav)")
        _save_output(auv_id, [], out_path, metrics)
        return metrics

    # Step 2: MAD outlier removal
    lats = np.array([p["lat"] for p in valid])
    lons = np.array([p["lon"] for p in valid])

    mask_lat = mad_outlier_mask(lats, sigma)
    mask_lon = mad_outlier_mask(lons, sigma)
    mask = mask_lat & mask_lon

    # Also filter depth if present
    depths = np.array([p.get("depth") or np.nan for p in valid])
    if not np.all(np.isnan(depths)):
        mask_depth = mad_outlier_mask(depths[~np.isnan(depths)], sigma)
        # Map back — only filter where we have depth
        depth_idx = np.where(~np.isnan(depths))[0]
        for i, keep in zip(depth_idx, mask_depth):
            if not keep:
                mask[i] = False

    filtered_points = [p for p, keep in zip(valid, mask) if keep]
    metrics["points_removed_outlier"] = int(np.sum(~mask))

    # Step 3: Sort by timestamp
    filtered_points.sort(key=lambda p: p["timestamp"])

    # Step 4: Smooth
    if len(filtered_points) >= window:
        filt_lats = moving_average(np.array([p["lat"] for p in filtered_points]), window)
        filt_lons = moving_average(np.array([p["lon"] for p in filtered_points]), window)
        for i, p in enumerate(filtered_points):
            p = dict(p)
            p["lat"] = float(filt_lats[i])
            p["lon"] = float(filt_lons[i])
            filtered_points[i] = p

    # Step 5: Optional Kalman
    if use_kalman and len(filtered_points) > 4:
        k_lats = simple_kalman_1d(np.array([p["lat"] for p in filtered_points]))
        k_lons = simple_kalman_1d(np.array([p["lon"] for p in filtered_points]))
        for i, p in enumerate(filtered_points):
            p = dict(p)
            p["lat"] = float(k_lats[i])
            p["lon"] = float(k_lons[i])
            filtered_points[i] = p

    metrics["points_after"] = len(filtered_points)
    if metrics["points_after"] < 5:
        metrics["status"] = "degraded"
        metrics["note"] = f"too few points after filter: {metrics['points_after']}"

    print(f"  [03] {auv_id}: {metrics['points_before']} → {metrics['points_after']} "
          f"(removed {metrics['points_removed_null']} null, {metrics['points_removed_outlier']} outliers)")

    _save_output(auv_id, filtered_points, out_path, metrics)
    return metrics


def _save_output(auv_id: str, points: list[dict], out_dir: Path, metrics: dict) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    json_out = out_dir / f"{auv_id}_nav_filtered.json"
    json_out.write_text(json.dumps({
        "auv_id": auv_id,
        "generated_at": now_iso(),
        "metrics": metrics,
        "points": points,
    }, indent=2, default=str))

    csv_out = out_dir / f"{auv_id}_nav_filtered.csv"
    with open(csv_out, "w") as f:
        f.write("timestamp,lat,lon,depth,heading,source\n")
        for p in points:
            f.write(f"{p['timestamp']},{p.get('lat','')},{p.get('lon','')},{p.get('depth','')},{p.get('heading','')},{p.get('source','')}\n")


def filter_mission(raw_dir: Path, out_dir: Path | None = None,
                   auv_filter: str | None = None,
                   sigma: float = 3.0, window: int = 5,
                   use_kalman: bool = False) -> list[dict]:
    if out_dir is None:
        out_dir = raw_dir.parent / "03_usbl_filtered"

    raw_files = sorted(raw_dir.glob("*_nav_raw.json"))
    if auv_filter:
        raw_files = [f for f in raw_files if auv_filter in f.name]

    all_metrics = []
    init_db()

    for raw_file in raw_files:
        out_file = out_dir / raw_file.name.replace("_raw", "_filtered")
        if out_file.exists():
            print(f"[03] {raw_file.stem}: output exists, skipping")
            existing = json.loads(out_file.read_text())
            all_metrics.append(existing.get("metrics", {}))
            continue

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

        # DB record
        with get_conn() as conn:
            mission_name = raw_dir.parent.name
            mission_row = conn.execute("SELECT id FROM missions WHERE name=?", (mission_name,)).fetchone()
            if mission_row:
                job_id = upsert_job(conn, mission_row["id"], m["auv_id"], "all", "03_usbl_filter",
                                    status="done" if m["status"] == "ok" else m["status"],
                                    output_path=str(out_dir))
                record_metric(conn, job_id, "usbl_points_before", value=m.get("points_before", 0))
                record_metric(conn, job_id, "usbl_points_after", value=m.get("points_after", 0),
                              pass_fail="pass" if m.get("points_after", 0) >= 5 else "fail")
                record_metric(conn, job_id, "usbl_points_removed_outlier",
                              value=m.get("points_removed_outlier", 0))

    return all_metrics


def main():
    ap = argparse.ArgumentParser(description="Stage 03 — Filter USBL navigation")
    ap.add_argument("raw_dir", type=Path, help="Directory with *_nav_raw.json files")
    ap.add_argument("--out", type=Path, default=None)
    ap.add_argument("--auv", type=str, default=None)
    ap.add_argument("--sigma", type=float, default=3.0, help="MAD sigma threshold")
    ap.add_argument("--window", type=int, default=5, help="Moving average window")
    ap.add_argument("--kalman", action="store_true", help="Apply simple Kalman smoothing")
    args = ap.parse_args()

    metrics = filter_mission(args.raw_dir, out_dir=args.out, auv_filter=args.auv,
                             sigma=args.sigma, window=args.window, use_kalman=args.kalman)

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


if __name__ == "__main__":
    main()
