#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Построение трёх графиков девиации Аллана по данным HIGHRES_IMU из CSV:
1) гироскоп X/Y/Z на одном графике;
2) акселерометр X/Y/Z на одном графике;
3) аппроксимация девиации Аллана по оси X акселерометра.

Пример:
    python3 allan_three_plots.py highres_imu.csv --output-dir allan_out
"""

from __future__ import annotations

import argparse
import csv
import math
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


RAD_S_TO_DEG_H = 180.0 / math.pi * 3600.0
MS2_TO_MG = 1000.0 / 9.80665
ALLAN_FIT_POWERS = (-2, -1, 0, 1)


def read_highres_imu_csv(path: Path) -> dict[str, np.ndarray]:
    rows = []

    with path.open("r", encoding="utf-8", newline="") as csv_file:
        reader = csv.DictReader(csv_file)
        for row in reader:
            rows.append(
                (
                    float(row["HIGHRES_IMU.time_usec"]),
                    float(row["HIGHRES_IMU.xacc"]),
                    float(row["HIGHRES_IMU.yacc"]),
                    float(row["HIGHRES_IMU.zacc"]),
                    float(row["HIGHRES_IMU.xgyro"]),
                    float(row["HIGHRES_IMU.ygyro"]),
                    float(row["HIGHRES_IMU.zgyro"]),
                )
            )

    if not rows:
        raise RuntimeError("CSV-файл пуст или не содержит данных HIGHRES_IMU.")

    arr = np.asarray(rows, dtype=float)
    time_s = (arr[:, 0] - arr[0, 0]) / 1_000_000.0

    keep = np.concatenate(([True], np.diff(time_s) > 0))
    arr = arr[keep]
    time_s = time_s[keep]

    return {
        "time_s": time_s,
        "accel": arr[:, 1:4],
        "gyro": arr[:, 4:7],
    }


def estimate_sample_rate(time_s: np.ndarray) -> float:
    dt = np.diff(time_s)
    dt = dt[np.isfinite(dt) & (dt > 0)]
    if len(dt) == 0:
        raise ValueError("Не удалось оценить частоту дискретизации.")
    return 1.0 / np.median(dt)


def allan_deviation(signal: np.ndarray, fs: float, points_per_decade: int = 12):
    y = np.asarray(signal, dtype=float)
    y = y[np.isfinite(y)]
    n = len(y)
    if n < 10:
        raise ValueError("Слишком мало точек для расчёта девиации Аллана.")

    max_m = n // 2
    raw_m = np.logspace(0, np.log10(max_m), int(points_per_decade * np.log10(max_m)))
    m_values = np.unique(raw_m.astype(int))
    m_values = m_values[m_values >= 1]

    cumsum = np.concatenate(([0.0], np.cumsum(y)))
    taus = []
    adevs = []

    for m in m_values:
        if 2 * m >= n:
            continue
        moving_avg = (cumsum[m:] - cumsum[:-m]) / m
        diff = moving_avg[m:] - moving_avg[:-m]
        sigma2 = 0.5 * np.mean(diff * diff)
        if sigma2 > 0 and np.isfinite(sigma2):
            taus.append(m / fs)
            adevs.append(math.sqrt(sigma2))

    return np.asarray(taus), np.asarray(adevs)


def fit_allan_variance(
    tau: np.ndarray,
    adev: np.ndarray,
    powers: tuple[int, ...] = ALLAN_FIT_POWERS,
    tau_min=None,
    tau_max=None,
):
    """Аппроксимирует Allan variance: sigma^2(tau) = sum(c_i * tau^p_i)."""
    mask = (
        np.isfinite(tau)
        & np.isfinite(adev)
        & (tau > 0)
        & (adev > 0)
    )
    if tau_min is not None:
        mask &= tau >= tau_min
    if tau_max is not None:
        mask &= tau <= tau_max

    tau_fit = tau[mask]
    adev_fit = adev[mask]

    if len(tau_fit) < len(powers):
        raise ValueError("Слишком мало точек для аппроксимации.")

    y = adev_fit ** 2
    A = np.vstack([tau_fit ** p for p in powers]).T

    log_tau = np.log10(tau_fit)
    dlog = np.empty_like(log_tau)
    dlog[0] = log_tau[1] - log_tau[0] if len(log_tau) > 1 else 1.0
    dlog[1:] = np.diff(log_tau)
    weights = np.sqrt(np.maximum(dlog, 1e-12))

    Aw = A * weights[:, None]
    yw = y * weights
    coeffs, *_ = np.linalg.lstsq(Aw, yw, rcond=None)
    coeffs = np.maximum(coeffs, 0.0)
    model = np.maximum(A @ coeffs, 0.0)

    return {
        "tau": tau_fit,
        "coeffs": coeffs,
        "model": model,
        "powers": powers,
    }


def plot_xyz(out_path: Path, results: dict[str, tuple[np.ndarray, np.ndarray]], title: str, unit: str):
    plt.figure(figsize=(8, 5), dpi=160)
    for axis_name, (tau, adev) in results.items():
        plt.loglog(tau, adev, marker="o", markersize=3, linewidth=1.2, label=axis_name)
    plt.grid(True, which="both", alpha=0.35)
    plt.xlabel(r"$\tau$, c")
    plt.ylabel(f"Девиация Аллана, {unit}")
    plt.title(title)
    plt.legend()
    plt.tight_layout()
    plt.savefig(out_path)
    plt.close()


def plot_fit_x(out_path: Path, tau: np.ndarray, adev: np.ndarray, allan_fit: dict[str, np.ndarray]):
    plt.figure(figsize=(8, 5), dpi=160)
    plt.loglog(tau, adev, marker="o", markersize=3, linewidth=1.2, label="Акселерометр X")
    plt.loglog(
        allan_fit["tau"],
        np.sqrt(allan_fit["model"]),
        linestyle="--",
        linewidth=2,
        label=r"Аппроксимация $\sigma^2(\tau)=\sum c_i\tau^{p_i}$",
    )
    plt.grid(True, which="both", alpha=0.35)
    plt.xlabel(r"$\tau$, c")
    plt.ylabel("Девиация Аллана, mg")
    plt.title("Девиация Аллана акселерометра по оси X с аппроксимацией")
    plt.legend()
    plt.tight_layout()
    plt.savefig(out_path)
    plt.close()


def main():
    parser = argparse.ArgumentParser(description="Построение 3 графиков девиации Аллана по CSV HIGHRES_IMU")
    parser.add_argument(
        "csv",
        nargs="?",
        type=Path,
        default=Path("highres_imu.csv"),
        help="Путь к CSV-файлу HIGHRES_IMU",
    )
    parser.add_argument("--fit-tau-min", type=float, default=None, help="Начало участка аппроксимации, с")
    parser.add_argument("--fit-tau-max", type=float, default=None, help="Конец участка аппроксимации, с")
    parser.add_argument("--output-dir", type=Path, default=Path("."), help="Папка для графиков")
    args = parser.parse_args()

    args.output_dir.mkdir(parents=True, exist_ok=True)

    imu = read_highres_imu_csv(args.csv)
    fs = estimate_sample_rate(imu["time_s"])

    gyro_results = {}
    accel_results = {}
    for i, axis_name in enumerate(["X", "Y", "Z"]):
        gyro_results[axis_name] = allan_deviation(imu["gyro"][:, i] * RAD_S_TO_DEG_H, fs)
        accel_results[axis_name] = allan_deviation(imu["accel"][:, i] * MS2_TO_MG, fs)

    tau_x, adev_x = accel_results["X"]
    allan_fit_x = fit_allan_variance(tau_x, adev_x, tau_min=args.fit_tau_min, tau_max=args.fit_tau_max)

    gyro_path = args.output_dir / "allan_gyro_xyz.png"
    accel_path = args.output_dir / "allan_accel_xyz.png"
    fit_path = args.output_dir / "allan_accel_x_fit.png"

    plot_xyz(gyro_path, gyro_results, "Девиация Аллана гироскопа по осям X/Y/Z", "град/ч")
    plot_xyz(accel_path, accel_results, "Девиация Аллана акселерометра по осям X/Y/Z", "mg")
    plot_fit_x(fit_path, tau_x, adev_x, allan_fit_x)

    print(f"IMU-точек: {len(imu['time_s'])}")
    print(f"Длительность записи: {imu['time_s'][-1]:.2f} c")
    print(f"Оценка частоты дискретизации: {fs:.3f} Hz")
    coeff_text = ", ".join(
        f"c(tau^{power})={coeff:.6g}" for power, coeff in zip(allan_fit_x["powers"], allan_fit_x["coeffs"])
    )
    print(f"Аппроксимация Allan variance для акселерометра X: sigma^2(tau) = sum(c_i * tau^p_i)")
    print(f"Коэффициенты: {coeff_text}")
    print(f"Участок аппроксимации: tau = {allan_fit_x['tau'][0]:.6g} ... {allan_fit_x['tau'][-1]:.6g} c")
    print(f"Сохранено: {gyro_path}")
    print(f"Сохранено: {accel_path}")
    print(f"Сохранено: {fit_path}")


if __name__ == "__main__":
    main()