#!/usr/bin/env python3
"""Compare fixed-stationary and motion-capable ESKF models on UART logs.

The current board applications emit GPS and IMU in separate runs.  Their
firmware/host elapsed axes are therefore replayed from t=0 for this explicitly
stationary experiment.  This is useful for noise/model comparison, but it is
not a synchronized moving-platform fusion data set.
"""

from __future__ import annotations

import argparse
import csv
import math
import re
from dataclasses import dataclass
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

from fusion_eskf import (
    G,
    StationaryEskf,
    align_quat,
    geodetic_to_ecef,
    robust_sigma,
)


POSITION_RE = re.compile(
    r"position\s+lat=(?P<lat>[-+0-9.eE]+)\s+lon=(?P<lon>[-+0-9.eE]+)"
    r"(?:\s+height=(?P<height>[-+0-9.eE]+))?"
)
GDOP_RE = re.compile(r"(?:DOP\s+)?GDOP=(?P<gdop>[-+0-9.eE]+)")
TRACKED_RE = re.compile(r"tracked satellites=(?P<tracked>[0-9]+)")
IMU_RE = re.compile(
    r"(?P<t>[0-9]+\.[0-9]+)\s+s\s+\|\s+samples=\s*(?P<n>[0-9]+)\s+"
    r"total=\s*(?P<total>[0-9]+)\s+\|\s+Temp=\s*(?P<temp>[-+0-9.]+)\s+C\s+\|\s+"
    r"Accel\[g\]\s+X=\s*(?P<ax>[-+0-9.]+)\s+Y=\s*(?P<ay>[-+0-9.]+)\s+"
    r"Z=\s*(?P<az>[-+0-9.]+)\s+\|\s+Gyro\[dps\]\s+"
    r"X=\s*(?P<gx>[-+0-9.]+)\s+Y=\s*(?P<gy>[-+0-9.]+)\s+Z=\s*(?P<gz>[-+0-9.]+)"
)


@dataclass
class Gps:
    t: float
    unix_s: float
    lat: float
    lon: float
    height: float
    gdop: float
    tracked: int


@dataclass
class Imu:
    t: float
    unix_s: float
    samples: int
    accel: np.ndarray
    gyro: np.ndarray


def rows(path: Path):
    with path.open(newline="", encoding="utf-8") as stream:
        yield from csv.DictReader(stream)


def parse_gps(path: Path) -> list[Gps]:
    result: list[Gps] = []
    gdop = math.nan
    tracked = 0
    for row in rows(path):
        line = row["uart_line"]
        match = GDOP_RE.search(line)
        if match:
            gdop = float(match["gdop"])
        match = TRACKED_RE.search(line)
        if match:
            tracked = int(match["tracked"])
        match = POSITION_RE.search(line)
        if not match:
            continue
        height_text = match["height"]
        result.append(
            Gps(
                t=float(row["host_monotonic_s"]),
                unix_s=float(row["host_unix_s"]),
                lat=float(match["lat"]),
                lon=float(match["lon"]),
                height=float(height_text) if height_text is not None else 0.0,
                gdop=gdop,
                tracked=tracked,
            )
        )
    if len(result) < 4:
        raise ValueError(f"only {len(result)} GPS positions found in {path}")
    start = result[0].t
    for sample in result:
        sample.t -= start
    return result


def parse_imu(path: Path) -> list[Imu]:
    result: list[Imu] = []
    for row in rows(path):
        line = row["uart_line"]
        # A GDB restart can leave a few lines from the previous run in the
        # host UART buffer.  The firmware's capture banner is the unambiguous
        # beginning of the data set; discard anything seen before the latest
        # banner.
        if "Starting " in line and " s capture:" in line:
            result.clear()
            continue
        match = IMU_RE.search(line)
        if not match:
            continue
        result.append(
            Imu(
                t=float(match["t"]),
                unix_s=float(row["host_unix_s"]),
                samples=int(match["n"]),
                accel=G
                * np.array(
                    [float(match["ax"]), float(match["ay"]), float(match["az"])]
                ),
                gyro=np.deg2rad(
                    [float(match["gx"]), float(match["gy"]), float(match["gz"])]
                ),
            )
        )
    if len(result) < 4:
        raise ValueError(f"only {len(result)} IMU windows found in {path}")
    start = result[0].t
    for sample in result:
        sample.t -= start
    return result


def gps_to_enu(gps: list[Gps]):
    lat0 = float(np.median([x.lat for x in gps]))
    lon0 = float(np.median([x.lon for x in gps]))
    # The current level-2 firmware omits height, so this comparison is 2-D.
    h0 = float(np.median([x.height for x in gps]))
    origin = geodetic_to_ecef([lat0], [lon0], [h0])[0]
    xyz = geodetic_to_ecef(
        [x.lat for x in gps], [x.lon for x in gps], [x.height for x in gps]
    )
    lat = math.radians(lat0)
    lon = math.radians(lon0)
    rotation = np.array(
        [
            [-math.sin(lon), math.cos(lon), 0.0],
            [
                -math.sin(lat) * math.cos(lon),
                -math.sin(lat) * math.sin(lon),
                math.cos(lat),
            ],
            [
                math.cos(lat) * math.cos(lon),
                math.cos(lat) * math.sin(lon),
                math.sin(lat),
            ],
        ]
    )
    return (xyz - origin) @ rotation.T, (lat0, lon0, h0)


class HorizontalEskf(StationaryEskf):
    """The shared 15-state ESKF with a horizontal-only GPS update."""

    def gps_update_2d(self, position, gate=13.816):
        h = np.zeros((2, 15))
        h[:, 0:2] = np.eye(2)
        residual = np.asarray(position)[0:2] - self.p[0:2]
        r = self.gps_r[0:2, 0:2]
        innovation = h @ self.p_cov @ h.T + r
        mahalanobis = float(residual @ np.linalg.solve(innovation, residual))
        if not math.isfinite(mahalanobis) or mahalanobis > gate:
            self.rejected_gps += 1
            return False, mahalanobis
        self._update(residual, h, r)
        self.accepted_gps += 1
        return True, mahalanobis


def stationary_detector(accel, gyro, accel_limit=0.15, gyro_limit=math.radians(1.0)):
    return (
        abs(float(np.linalg.norm(accel)) - G) <= accel_limit
        and float(np.linalg.norm(gyro)) <= gyro_limit
    )


def run_model(gps, imu, enu, mode, calibration):
    accel, gyro, attitude, accel_sigma, gyro_sigma, gps_sigma, initial = calibration
    eskf = HorizontalEskf(initial, attitude, accel_sigma, gyro_sigma, gps_sigma)
    events = [(x.t, 0, i) for i, x in enumerate(imu)]
    events += [(x.t, 1, i) for i, x in enumerate(gps)]
    events.sort()
    current_accel = accel[0]
    current_gyro = gyro[0]
    last_t = 0.0
    detector_run = 0
    stationary_windows = 0
    zupt_updates = 0
    output = []

    for t, kind, index in events:
        eskf.propagate(current_accel, current_gyro, t - last_t)
        last_t = t
        if kind == 0:
            current_accel = accel[index]
            current_gyro = gyro[index]
            detected = stationary_detector(current_accel, current_gyro)
            detector_run = detector_run + 1 if detected else 0
            stationary_windows += int(detected)
            apply_constraint = mode == "fixed" or detector_run >= 3
            if apply_constraint:
                eskf.zero_velocity_update(sigma=0.03)
                eskf.zero_gyro_update(
                    current_gyro, np.maximum(gyro_sigma, math.radians(0.003))
                )
                zupt_updates += 1
        else:
            accepted, nis = eskf.gps_update_2d(enu[index])
            output.append(
                {
                    "t": t,
                    "raw": enu[index].copy(),
                    "fused": eskf.p.copy(),
                    "velocity": eskf.v.copy(),
                    "accepted": accepted,
                    "nis": nis,
                }
            )
    return eskf, output, stationary_windows, zupt_updates


def calibrate(gps, imu, enu):
    accel = np.vstack([x.accel for x in imu])
    gyro = np.vstack([x.gyro for x in imu])
    attitude = align_quat(np.median(accel, axis=0), [0.0, 0.0, 1.0])
    accel_sigma = np.array([robust_sigma(accel[:, i], 0.02) for i in range(3)])
    gyro_sigma = np.array(
        [robust_sigma(gyro[:, i], math.radians(0.003)) for i in range(3)]
    )
    gps_sigma = np.array(
        [robust_sigma(enu[:, i], 1.0 if i < 2 else 3.0) for i in range(3)]
    )
    center = np.median(enu[:, 0:2], axis=0)
    plausible = np.linalg.norm(enu[:, 0:2] - center, axis=1) < 10_000.0
    first = int(np.flatnonzero(plausible)[0]) if np.any(plausible) else 0
    initial = enu[first].copy()
    initial[2] = 0.0
    return accel, gyro, attitude, accel_sigma, gyro_sigma, gps_sigma, initial


def centered_metrics(values):
    center = np.median(values, axis=0)
    radius = np.linalg.norm(values - center, axis=1)
    return {
        "rms": float(np.sqrt(np.mean(radius**2))),
        "p95": float(np.percentile(radius, 95)),
        "max": float(np.max(radius)),
        "std_e": float(np.std(values[:, 0], ddof=1)),
        "std_n": float(np.std(values[:, 1], ddof=1)),
    }


def write_csv(path, gps, fixed, adaptive):
    with path.open("w", newline="", encoding="utf-8") as stream:
        writer = csv.writer(stream)
        writer.writerow(
            [
                "t_s", "lat_deg", "lon_deg", "gdop", "tracked", "height_m",
                "gps_e_m", "gps_n_m", "fixed_e_m", "fixed_n_m",
                "adaptive_e_m", "adaptive_n_m", "fixed_speed_mps",
                "adaptive_speed_mps", "fixed_gps_accepted",
                "adaptive_gps_accepted", "fixed_nis", "adaptive_nis",
            ]
        )
        for sample, a, b in zip(gps, fixed, adaptive):
            writer.writerow(
                [
                    sample.t, sample.lat, sample.lon, sample.gdop, sample.tracked,
                    sample.height,
                    *a["raw"][0:2], *a["fused"][0:2], *b["fused"][0:2],
                    np.linalg.norm(a["velocity"]), np.linalg.norm(b["velocity"]),
                    int(a["accepted"]), int(b["accepted"]), a["nis"], b["nis"],
                ]
            )


def plot(path, fixed, adaptive):
    time = np.array([x["t"] for x in fixed])
    raw = np.vstack([x["raw"] for x in fixed])[:, 0:2]
    fixed_p = np.vstack([x["fused"] for x in fixed])[:, 0:2]
    adaptive_p = np.vstack([x["fused"] for x in adaptive])[:, 0:2]
    center = np.median(raw, axis=0)
    raw_radius = np.linalg.norm(raw - center, axis=1)
    fixed_radius = np.linalg.norm(fixed_p - np.median(fixed_p, axis=0), axis=1)
    adaptive_radius = np.linalg.norm(
        adaptive_p - np.median(adaptive_p, axis=0), axis=1
    )

    fig, axes = plt.subplots(2, 1, figsize=(11, 8), sharex=True)
    axes[0].plot(time, raw_radius, color="0.65", lw=0.8, label="Ham GPS")
    axes[0].plot(time, fixed_radius, color="tab:blue", label="Sabit model")
    axes[0].plot(time, adaptive_radius, color="tab:orange", label="Hareketli/adaptif model")
    axes[0].set_ylabel("Medyan merkezden yatay uzaklık [m]")
    fixed_speed = np.linalg.norm(
        np.vstack([x["velocity"] for x in fixed]), axis=1
    )
    adaptive_speed = np.linalg.norm(
        np.vstack([x["velocity"] for x in adaptive]), axis=1
    )
    axes[1].plot(time, fixed_speed, color="tab:blue", label="Sabit model")
    axes[1].plot(time, adaptive_speed, color="tab:orange", label="Hareketli/adaptif model")
    axes[1].set_ylabel("Tahmin edilen hız [m/s]")
    axes[1].set_xlabel("Geçen süre [s]")
    for axis in axes:
        axis.grid(alpha=0.25)
        axis.legend()
    fig.suptitle("Durağan kart: sabit ve hareket-kabiliyetli ESKF karşılaştırması")
    fig.tight_layout()
    fig.savefig(path, dpi=170)
    plt.close(fig)


def plot_scatter(path, fixed, adaptive):
    raw = np.vstack([x["raw"] for x in fixed])[:, 0:2]
    fixed_p = np.vstack([x["fused"] for x in fixed])[:, 0:2]
    adaptive_p = np.vstack([x["fused"] for x in adaptive])[:, 0:2]
    center = np.median(raw, axis=0)
    radius = max(20.0, float(np.percentile(np.linalg.norm(raw - center, axis=1), 99)))
    fig, axis = plt.subplots(figsize=(8, 8))
    axis.scatter(raw[:, 0], raw[:, 1], s=7, alpha=0.25, color="0.4", label="Ham GPS")
    axis.plot(fixed_p[:, 0], fixed_p[:, 1], lw=1.4, color="tab:blue", label="Sabit model")
    axis.plot(adaptive_p[:, 0], adaptive_p[:, 1], lw=1.2, color="tab:orange", label="Hareketli/adaptif")
    axis.set_xlim(center[0] - 1.15 * radius, center[0] + 1.15 * radius)
    axis.set_ylim(center[1] - 1.15 * radius, center[1] + 1.15 * radius)
    axis.set_aspect("equal", adjustable="box")
    axis.set_xlabel("East [m]")
    axis.set_ylabel("North [m]")
    axis.grid(alpha=0.25)
    axis.legend()
    axis.set_title("Yatay saçılım (merkezi %99 görünüm)")
    fig.tight_layout()
    fig.savefig(path, dpi=170)
    plt.close(fig)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--gps", type=Path, required=True)
    parser.add_argument("--imu", type=Path, required=True)
    parser.add_argument("--output", type=Path, default=Path("output/model_compare"))
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)

    gps = parse_gps(args.gps)
    imu = parse_imu(args.imu)
    duration = min(gps[-1].t, imu[-1].t)
    gps = [x for x in gps if x.t <= duration]
    imu = [x for x in imu if x.t <= duration]
    enu, reference = gps_to_enu(gps)
    calibration = calibrate(gps, imu, enu)
    fixed_eskf, fixed, fixed_detected, fixed_zupt = run_model(
        gps, imu, enu, "fixed", calibration
    )
    adaptive_eskf, adaptive, adaptive_detected, adaptive_zupt = run_model(
        gps, imu, enu, "adaptive", calibration
    )

    raw = np.vstack([x["raw"] for x in fixed])[:, 0:2]
    fixed_p = np.vstack([x["fused"] for x in fixed])[:, 0:2]
    adaptive_p = np.vstack([x["fused"] for x in adaptive])[:, 0:2]
    raw_m = centered_metrics(raw)
    fixed_m = centered_metrics(fixed_p)
    adaptive_m = centered_metrics(adaptive_p)
    gps_steps = np.linalg.norm(np.diff(raw, axis=0), axis=1)
    imu_rate = len(imu) / max(imu[-1].t - imu[0].t, 1e-9)
    gps_rate = len(gps) / max(gps[-1].t - gps[0].t, 1e-9)
    lines = [
        "Fixed-stationary vs motion-capable ESKF",
        "========================================",
        f"Common replay duration: {duration:.3f} s",
        f"GPS: {len(gps)} samples, mean rate {gps_rate:.3f} Hz",
        f"IMU: {len(imu)} windows, mean rate {imu_rate:.3f} Hz, underlying samples {sum(x.samples for x in imu)}",
        "Timing: separate stationary captures replayed from t=0 (not hardware-synchronized)",
        "GPS update dimensions: 2-D horizontal (current level-2 firmware log has no height)",
        f"ENU reference: lat={reference[0]:.10f}, lon={reference[1]:.10f}",
        f"GPS max consecutive horizontal step: {np.max(gps_steps):.3f} m",
        f"Raw GPS centered horizontal RMS/P95/max: {raw_m['rms']:.3f} / {raw_m['p95']:.3f} / {raw_m['max']:.3f} m",
        f"Fixed model centered horizontal RMS/P95/max: {fixed_m['rms']:.3f} / {fixed_m['p95']:.3f} / {fixed_m['max']:.3f} m",
        f"Adaptive model centered horizontal RMS/P95/max: {adaptive_m['rms']:.3f} / {adaptive_m['p95']:.3f} / {adaptive_m['max']:.3f} m",
        f"Fixed GPS accepted/rejected: {fixed_eskf.accepted_gps}/{fixed_eskf.rejected_gps}",
        f"Adaptive GPS accepted/rejected: {adaptive_eskf.accepted_gps}/{adaptive_eskf.rejected_gps}",
        f"IMU stationary detector true: {adaptive_detected}/{len(imu)} ({100.0 * adaptive_detected / len(imu):.2f}%)",
        f"ZUPT updates fixed/adaptive: {fixed_zupt}/{adaptive_zupt}",
        f"Final speed fixed/adaptive: {np.linalg.norm(fixed_eskf.v):.6f} / {np.linalg.norm(adaptive_eskf.v):.6f} m/s",
        "Adaptive model uses a 3-window stationary decision; it does not assume zero velocity while motion is detected.",
        "No surveyed truth is available: scatter measures precision, not absolute position accuracy.",
    ]
    summary = "\n".join(lines) + "\n"
    (args.output / "summary.txt").write_text(summary, encoding="utf-8")
    write_csv(args.output / "model_comparison.csv", gps, fixed, adaptive)
    plot(args.output / "model_comparison_timeseries.png", fixed, adaptive)
    plot_scatter(args.output / "model_comparison_scatter.png", fixed, adaptive)
    print(summary, end="")


if __name__ == "__main__":
    main()