#!/usr/bin/env python3

import argparse
import csv
from datetime import datetime
from pathlib import Path


def load_samples(csv_path: Path):
    rows = []
    source_name = None
    base_time = None
    with csv_path.open("r", encoding="utf-8", newline="") as f:
        reader = csv.DictReader(f)
        fieldnames = reader.fieldnames or []
        if "gyro_z_filtered_pre_bias" in fieldnames:
            source_name = "gyro_z_filtered_pre_bias"
        elif "gyro_raw_z_rad_s" in fieldnames:
            source_name = "gyro_raw_z_rad_s"
        elif "imu_angular_velocity_z" in fieldnames:
            source_name = "imu_angular_velocity_z"
        else:
            raise ValueError("CSV does not contain gyro_z_filtered_pre_bias, gyro_raw_z_rad_s, or imu_angular_velocity_z")
        for row in reader:
            t = row.get("elapsed_s")
            z = row.get(source_name)
            if not z:
                continue
            if t:
                sample_t = float(t)
            else:
                wall_time = row.get("wall_time")
                if not wall_time:
                    continue
                parsed_time = datetime.fromisoformat(wall_time)
                if base_time is None:
                    base_time = parsed_time
                sample_t = (parsed_time - base_time).total_seconds()
            rows.append((sample_t, float(z)))
    return rows, source_name


def fit_line(samples):
    n = len(samples)
    if n < 2:
        raise ValueError("need at least 2 valid samples")

    sum_t = sum(t for t, _ in samples)
    sum_y = sum(y for _, y in samples)
    mean_t = sum_t / n
    mean_y = sum_y / 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 main():
    parser = argparse.ArgumentParser(description="Fit linear gyro z bias from a static feedback CSV")
    parser.add_argument("csv_path", type=Path, help="path to feedback.csv")
    args = parser.parse_args()

    samples, source_name = load_samples(args.csv_path)
    intercept, slope = fit_line(samples)

    print(f"source={source_name}")
    print(f"samples={len(samples)}")
    print(f"gyro_z_bias_intercept={intercept:.12f}")
    print(f"gyro_z_bias_slope={slope:.12f}")
    print(f"bias_model(t)= {intercept:.12f} + {slope:.12f} * t")


if __name__ == "__main__":
    main()