"""IMU adapters for the V1 standard intermediate format.

Accepted inputs
---------------
1. CSV with header:
   t,gx,gy,gz,ax,ay,az
   - ``t`` in seconds on the IMU clock
   - gyro in rad/s
   - accel in m/s^2

2. NPZ with arrays:
   t, gyro, acc
   shapes: (N,), (N,3), (N,3)
"""

from __future__ import annotations

import csv
from pathlib import Path

import numpy as np

from .contracts import ImuSeries


def load_imu_samples(path: Path | str) -> ImuSeries:
    """Load normalized IMU samples from CSV or NPZ."""

    source = Path(path)
    if not source.exists():
        raise FileNotFoundError(source)
    if source.suffix.lower() == ".csv":
        return _load_imu_csv(source)
    if source.suffix.lower() == ".npz":
        return _load_imu_npz(source)
    raise ValueError(f"unsupported IMU format '{source.suffix}' (use .csv or .npz)")


def _load_imu_csv(path: Path) -> ImuSeries:
    required_order = ["t", "gx", "gy", "gz", "ax", "ay", "az"]
    with path.open("r", encoding="utf-8-sig", newline="") as handle:
        header = next(csv.reader(handle), [])
    names = set(header)
    if not set(required_order).issubset(names):
        raise ValueError(f"IMU CSV must contain columns {sorted(required_order)}, got {sorted(names)}")
    usecols = [header.index(name) for name in required_order]
    data = np.loadtxt(path, delimiter=",", skiprows=1, usecols=usecols, ndmin=2)
    t = np.asarray(data[:, 0], dtype=float).reshape(-1)
    gyro = np.asarray(data[:, 1:4], dtype=float)
    acc = np.asarray(data[:, 4:7], dtype=float)
    order = np.argsort(t)
    return ImuSeries(t_s=t[order], gyro_rad_s=gyro[order], acc_m_s2=acc[order])


def _load_imu_npz(path: Path) -> ImuSeries:
    with np.load(path) as payload:
        keys = set(payload.files)
        if not {"t", "gyro", "acc"}.issubset(keys):
            raise ValueError(f"IMU NPZ must contain t, gyro, acc; got {sorted(keys)}")
        t = np.asarray(payload["t"], dtype=float).reshape(-1)
        gyro = np.asarray(payload["gyro"], dtype=float).reshape(-1, 3)
        acc = np.asarray(payload["acc"], dtype=float).reshape(-1, 3)
    order = np.argsort(t)
    return ImuSeries(t_s=t[order], gyro_rad_s=gyro[order], acc_m_s2=acc[order])


def save_imu_csv(path: Path | str, imu: ImuSeries) -> None:
    """Write IMU samples to the standard CSV format."""

    destination = Path(path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    array = np.column_stack([imu.t_s, imu.gyro_rad_s, imu.acc_m_s2])
    header = "t,gx,gy,gz,ax,ay,az"
    np.savetxt(destination, array, delimiter=",", header=header, comments="")