#!/usr/bin/env python3

import argparse
import csv
import math
from pathlib import Path
from typing import Iterable, List, Optional, Sequence, Tuple


Sample = Tuple[float, float]


def parse_float(row: dict, key: str) -> Optional[float]:
    value = row.get(key)
    if value in ("", None):
        return None
    return float(value)


def load_series(csv_path: Path, key: str) -> List[Sample]:
    samples: List[Sample] = []
    with csv_path.open("r", encoding="utf-8", newline="") as f:
        reader = csv.DictReader(f)
        for row in reader:
            t = parse_float(row, "elapsed_s")
            y = parse_float(row, key)
            if t is None or y is None:
                continue
            samples.append((t, y))
    return samples


def fit_line(samples: Sequence[Sample]) -> Tuple[float, float]:
    n = len(samples)
    if n < 2:
        raise ValueError("need at least 2 valid samples")

    mean_t = sum(t for t, _ in samples) / n
    mean_y = sum(y for _, y in samples) / n
    s_tt = sum((t - mean_t) * (t - mean_t) for t, _ in samples)
    if s_tt == 0.0:
        raise ValueError("all timestamps are identical")
    s_ty = sum((t - mean_t) * (y - mean_y) for t, y in samples)
    slope = s_ty / s_tt
    intercept = mean_y - slope * mean_t
    return intercept, slope


def mean(values: Iterable[float]) -> float:
    values = list(values)
    if not values:
        raise ValueError("need at least 1 value")
    return sum(values) / len(values)


def stddev(values: Sequence[float]) -> float:
    if len(values) < 2:
        return 0.0
    mu = mean(values)
    return math.sqrt(sum((v - mu) * (v - mu) for v in values) / len(values))


def integrate(samples: Sequence[Sample]) -> float:
    if len(samples) < 2:
        return 0.0
    total = 0.0
    prev_t, prev_y = samples[0]
    for t, y in samples[1:]:
        dt = t - prev_t
        total += 0.5 * (prev_y + y) * dt
        prev_t, prev_y = t, y
    return total


def drift_rate_deg_per_min(samples: Sequence[Sample]) -> float:
    if len(samples) < 2:
        return 0.0
    total_time = samples[-1][0] - samples[0][0]
    if total_time <= 0.0:
        return 0.0
    yaw_drift_rad = integrate(samples)
    return math.degrees(yaw_drift_rad) / (total_time / 60.0)


def print_series_stats(name: str, samples: Sequence[Sample]) -> None:
    values = [y for _, y in samples]
    intercept, slope = fit_line(samples)
    print(f"{name}:")
    print(f"  samples={len(samples)}")
    print(f"  mean={mean(values):.12f} rad/s")
    print(f"  stddev={stddev(values):.12f} rad/s")
    print(f"  line_intercept={intercept:.12f} rad/s")
    print(f"  line_slope={slope:.12f} rad/s^2")
    print(f"  integrated_drift={integrate(samples):.12f} rad ({math.degrees(integrate(samples)):.6f} deg)")
    print(f"  drift_rate={drift_rate_deg_per_min(samples):.6f} deg/min")


def print_yaw_span(name: str, samples: Sequence[Sample]) -> None:
    if len(samples) < 2:
        return
    yaw_delta = samples[-1][1] - samples[0][1]
    total_time = samples[-1][0] - samples[0][0]
    print(f"{name}:")
    print(f"  start={samples[0][1]:.12f} rad")
    print(f"  end={samples[-1][1]:.12f} rad")
    print(f"  delta={yaw_delta:.12f} rad ({math.degrees(yaw_delta):.6f} deg)")
    if total_time > 0.0:
        print(f"  rate={math.degrees(yaw_delta) / (total_time / 60.0):.6f} deg/min")


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate gyro bias compensation quality from feedback.csv")
    parser.add_argument("csv_path", type=Path, help="path to feedback.csv")
    args = parser.parse_args()

    pre_bias = load_series(args.csv_path, "gyro_z_filtered_pre_bias")
    bias_model = load_series(args.csv_path, "gyro_z_bias_model")
    final_for_yaw = load_series(args.csv_path, "gyro_z_final_for_yaw")
    odom_yaw = load_series(args.csv_path, "odom_yaw_rad")
    imu_yaw = load_series(args.csv_path, "imu_yaw_rad")

    if pre_bias:
        print_series_stats("gyro_z_filtered_pre_bias", pre_bias)
        print()
    else:
        print("gyro_z_filtered_pre_bias: no valid samples")
        print()

    if bias_model:
        print_series_stats("gyro_z_bias_model", bias_model)
        print()
    else:
        print("gyro_z_bias_model: no valid samples")
        print()

    if final_for_yaw:
        print_series_stats("gyro_z_final_for_yaw", final_for_yaw)
        print()
    else:
        print("gyro_z_final_for_yaw: no valid samples")
        print()

    if odom_yaw:
        print_yaw_span("odom_yaw_rad", odom_yaw)
        print()

    if imu_yaw:
        print_yaw_span("imu_yaw_rad", imu_yaw)


if __name__ == "__main__":
    main()
