"""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

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:
    data = np.genfromtxt(path, delimiter=",", names=True, dtype=float)
    if data.ndim == 0:
        data = np.array([data])
    names = set(data.dtype.names or ())
    required = {"t", "gx", "gy", "gz", "ax", "ay", "az"}
    if not required.issubset(names):
        raise ValueError(f"IMU CSV must contain columns {sorted(required)}, got {sorted(names)}")
    t = np.asarray(data["t"], dtype=float).reshape(-1)
    gyro = np.column_stack([data["gx"], data["gy"], data["gz"]]).astype(float)
    acc = np.column_stack([data["ax"], data["ay"], data["az"]]).astype(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="")