#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Расчет параметров гироскопа по файлам МК.dat и Дрейф.dat.

МК.dat:
    0 столбец — время, с
    1 столбец — заданная скорость, °/с
    2 столбец — выходной сигнал гироскопа, В

Дрейф.dat:
    0 столбец — время, с
    1 столбец — выходной сигнал гироскопа, В
"""

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


DEFAULT_SCALE_FILE = "МК.dat"
DEFAULT_DRIFT_FILE = "Дрейф.dat"
DRIFT_AVG_WINDOW_SEC = 120.0
ALLAN_FIT_POWERS = (-2, -1, 0, 1)


# ============================================================
# ЧТЕНИЕ ДАННЫХ
# ============================================================

def read_columns(filename, columns):
    """Читает нужные числовые столбцы из .dat, пропуская заголовки."""
    rows = []

    with Path(filename).open("r", encoding="utf-8") as file:
        for line in file:
            parts = line.split()

            if not parts or parts[0].startswith("#"):
                continue

            try:
                rows.append([float(parts[i]) for i in columns])
            except (ValueError, IndexError):
                continue

    if not rows:
        raise ValueError(f"В файле {filename} не найдены числовые данные")

    data = np.asarray(rows, dtype=float)
    return data.T


# ============================================================
# ОБЩИЕ ВСПОМОГАТЕЛЬНЫЕ ФУНКЦИИ
# ============================================================

def mean_by_x(x, y):
    """Возвращает средние y для одинаковых значений x."""
    x_unique, inverse = np.unique(x, return_inverse=True)
    y_sum = np.bincount(inverse, weights=y)
    count = np.bincount(inverse)
    return x_unique, y_sum / count


def centered_moving_average(y, window_size):
    """Быстрое центрированное скользящее среднее без изменения длины массива."""
    if window_size <= 1 or len(y) < window_size:
        return y.copy()

    left = window_size // 2
    right = window_size - 1 - left
    padded = np.pad(y, (left, right), mode="symmetric")

    cumsum = np.cumsum(np.insert(padded, 0, 0.0))
    return (cumsum[window_size:] - cumsum[:-window_size]) / window_size


def voltage_to_rate(voltage_v, slope_mv_per_dps, intercept_mv):
    """Переводит напряжение в скорость по BFSL: U = a * omega + b."""
    voltage_mv = voltage_v * 1e3
    return (voltage_mv - intercept_mv) / slope_mv_per_dps


def drift_voltage_to_rate(voltage_v, scale_factor_mv_per_dps):
    """Переводит напряжение дрейфа в скорость, используя только масштабный коэффициент."""
    voltage_mv = voltage_v * 1e3
    zero_offset_mv = float(voltage_mv.mean())
    return voltage_mv / scale_factor_mv_per_dps, zero_offset_mv


# ============================================================
# МАСШТАБНЫЙ КОЭФФИЦИЕНТ
# ============================================================

def calc_scale_metrics(omega, voltage_v):
    voltage_mv = voltage_v * 1e3
    omega_unique, voltage_mean_mv = mean_by_x(omega, voltage_mv)

    # BFSL: U = a * omega + b
    slope, intercept = np.polyfit(omega_unique, voltage_mean_mv, 1)
    scale_factor = abs(slope)

    bfsl = slope * omega_unique + intercept
    full_scale = voltage_mean_mv.max() - voltage_mean_mv.min()
    nonlinearity = 100 * np.abs(voltage_mean_mv - bfsl) / full_scale

    pos = omega_unique >= 0
    neg = omega_unique <= 0

    if pos.sum() >= 2 and neg.sum() >= 2:
        k_pos, _ = np.polyfit(omega_unique[pos], voltage_mean_mv[pos], 1)
        k_neg, _ = np.polyfit(omega_unique[neg], voltage_mean_mv[neg], 1)
        k_avg = 0.5 * (abs(k_pos) + abs(k_neg))
        asymmetry = 100 * abs(abs(k_pos) - abs(k_neg)) / k_avg
    else:
        k_pos = k_neg = np.nan
        asymmetry = np.nan

    return {
        "scale_factor": float(scale_factor),
        "slope": float(slope),
        "intercept": float(intercept),
        "linearity_mean": float(nonlinearity.mean()),
        "linearity_max": float(nonlinearity.max()),
        "asymmetry": float(asymmetry),
        "zero_offset": float(-intercept / slope),
        "k_pos": float(k_pos),
        "k_neg": float(k_neg),
        "omega_unique": omega_unique,
        "voltage_mean_mv": voltage_mean_mv,
        "nonlinearity": nonlinearity,
    }


# ============================================================
# ДЕВИАЦИЯ АЛЛАНА И ЕЕ АППРОКСИМАЦИЯ
# ============================================================

def allan_deviation(data, sample_rate):
    """Девиация Аллана для равномерно дискретизированных данных.

    Для каждого времени усреднения tau сигнал делится на блоки.
    В каждом блоке считается средняя угловая скорость.
    Потом сравниваются соседние средние:
        sigma^2(tau) = 1/2 * mean((wbar[i+1] - wbar[i])^2)
    Девиация Аллана - это корень из этой дисперсии.
    """
    n = len(data)                                      # n: сколько всего отсчетов omega(t)
    m_values = []                                      # m: длина блока в отсчетах
    m = 1                                              # начинаем с tau = 1 / sample_rate

    while m <= n // 2:
        m_values.append(m)                             # это даст одну точку графика sigma(tau)
        m = int(np.ceil(m * 1.2))                      # следующий tau больше примерно на 20%

    tau = []                                           # ось X: времена усреднения tau, секунды
    adev = []                                          # ось Y: девиация Аллана sigma(tau)

    for m in m_values:
        block_count = n // m                           # число средних wbar_i для данного tau

        if block_count < 2:
            break                                      # иначе нет пары wbar[i+1] - wbar[i]

        # 1) Берем только целое число блоков, чтобы reshape сработал без остатка.
        used_data = data[:block_count * m]             # длина станет block_count * m

        # 2) Превращаем длинный ряд в таблицу: строка = один блок tau.
        blocks = used_data.reshape(block_count, m)     # форма массива: (число блоков, m)

        # 3) Среднее в строке - это wbar_i, средняя скорость на интервале tau.
        means = blocks.mean(axis=1)                    # получаем массив [wbar_1, wbar_2, ...]

        # 4) Соседние средние сравниваются: Delta_i = wbar_{i+1} - wbar_i.
        deltas = np.diff(means)                        # массив соседних разностей Delta_i

        # 5) Дисперсия Аллана: половина среднего квадрата этих разностей.
        variance = 0.5 * np.mean(deltas ** 2)          # sigma^2(tau) = 1/2 * <Delta_i^2>

        tau.append(m / sample_rate)                    # tau = m / fs, где fs = sample_rate
        adev.append(np.sqrt(variance))                 # sigma(tau) = sqrt(sigma^2(tau))

    return np.asarray(tau), np.asarray(adev)           # возвращаем точки графика Аллана


def fit_allan_variance(tau, adev, powers=ALLAN_FIT_POWERS, tau_min=0.33, tau_max=3500.0):
    """Аппроксимирует Allan variance: sigma^2(tau) = sum(c_i * tau^p_i).

    Разные степени tau соответствуют разным шумам:
    tau^-1 - белый шум угловой скорости / ARW,
    tau^0  - плато нестабильности смещения / BI,
    tau^1  - медленное случайное блуждание / RRW.
    """
    mask = (tau > tau_min) & (tau < tau_max)           # берем рабочий участок графика
    mask &= np.isfinite(adev) & (adev > 0)             # убираем NaN/inf и нулевые точки
    tau_fit = tau[mask]                                # X-точки для подгонки модели
    adev_fit = adev[mask]                              # Y-точки sigma(tau) для подгонки

    if len(tau_fit) < len(powers):
        return None                                    # мало точек для всех c_i

    y = adev_fit ** 2                                  # левая часть модели: sigma^2(tau)
    A = np.vstack([tau_fit ** p for p in powers]).T    # строки: [tau^p1, tau^p2, ...]

    # Взвешивание по логарифмической сетке, чтобы плотные участки не доминировали.
    log_tau = np.log10(tau_fit)                        # на графике tau обычно логарифмическая
    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)                        # шаги между соседними log10(tau)
    weights = np.sqrt(np.maximum(dlog, 1e-12))         # вес точки пропорционален ее log-интервалу

    Aw = A * weights[:, None]                          # растягиваем вес на все столбцы A
    yw = y * weights                                   # так же взвешиваем sigma^2(tau)
    coeffs, *_ = np.linalg.lstsq(Aw, yw, rcond=None)   # решаем A*c ~= y методом МНК
    coeffs = np.maximum(coeffs, 0.0)                   # физически мощность шума не отрицательна
    model = np.maximum(A @ coeffs, 0.0)                # восстановленная кривая sigma^2(tau)

    return {
        "tau": tau_fit,                                # tau, которые реально подгонялись
        "coeffs": coeffs,                              # найденные c_i при tau^p_i
        "model": model,                                # модельная Allan variance
        "powers": powers,                              # степени p_i для расшифровки c_i
    }


def noise_params_from_allan(tau, adev, allan_fit):
    """Оценивает ARW, Bias Instability и RRW.

    Основные оценки берутся из аппроксимации Allan variance.
    Для Bias Instability оставлен устойчивый запасной вариант через минимум
    девиации Аллана, если свободный член аппроксимации получился нулевым.
    """
    if allan_fit is None:
        return None                                  # аппроксимация не получилась

    coeff = dict(zip(allan_fit["powers"], allan_fit["coeffs"]))  # p -> c из модели

    c_arw = coeff.get(-1, 0.0)                         # коэффициент участка sigma ~ tau^-1/2
    c_bias = coeff.get(0, 0.0)                         # коэффициент горизонтального плато
    c_rrw = coeff.get(1, 0.0)                          # коэффициент участка sigma ~ tau^+1/2

    arw = np.sqrt(c_arw) * np.sqrt(3600)               # ARW переводим из секунд в часы

    if c_bias > 0:
        bias_instability = np.sqrt(c_bias) / 0.664 * 3600  # BI из плато Аллана
    else:
        bias_instability = adev.min() / 0.664 * 3600   # если плато не нашлось, берем минимум

    rrw = np.sqrt(3 * c_rrw) * (3600 ** 1.5)           # RRW переводим к часовой размерности

    return {
        "ARW": float(arw),                           # angle random walk
        "BI": float(bias_instability),               # bias instability
        "RRW": float(rrw),                           # rate random walk
    }


# ============================================================
# ДРЕЙФ
# ============================================================

def calc_drift_metrics(times, rates, zero_offset_mv=None):
    zero_offset = float(rates.mean())

    trend_slope, _ = np.polyfit(times, rates, 1)        # °/с²
    trend = float(trend_slope * 3600 ** 2)              # °/ч/ч

    sample_rate = 1.0 / np.mean(np.diff(times))
    tau, adev = allan_deviation(rates, sample_rate)
    allan_fit = fit_allan_variance(tau, adev)
    noise = noise_params_from_allan(tau, adev, allan_fit)

    return {
        "zero_offset": zero_offset,
        "zero_offset_mv": zero_offset_mv,
        "trend": trend,
        "tau": tau,
        "adev": adev,
        "allan_fit": allan_fit,
        "noise": noise,
    }


# ============================================================
# ГРАФИКИ
# ============================================================

def plot_scale(omega, voltage_v, metrics):
    voltage_mv = voltage_v * 1e3
    omega_line = np.linspace(omega.min(), omega.max(), 300)

    plt.figure(figsize=(10, 6))
    plt.plot(omega, voltage_mv, ".", alpha=0.25, label="Сырые измерения")
    plt.plot(metrics["omega_unique"], metrics["voltage_mean_mv"], "o", label="Средние точки")
    plt.plot(
        omega_line,
        metrics["slope"] * omega_line + metrics["intercept"],
        "-",
        linewidth=2,
        label="BFSL",
    )
    plt.title("Калибровочная характеристика")
    plt.xlabel("Угловая скорость, °/с")
    plt.ylabel("Напряжение, мВ")
    plt.grid(True)
    plt.legend()


def plot_nonlinearity(metrics):
    plt.figure(figsize=(10, 6))
    plt.plot(metrics["omega_unique"], metrics["nonlinearity"], "o-", label="Нелинейность BFSL")
    plt.axhline(metrics["linearity_mean"], linestyle="--", label=f"Средняя = {metrics['linearity_mean']:.3f} %")
    plt.axhline(metrics["linearity_max"], linestyle=":", label=f"Максимальная = {metrics['linearity_max']:.3f} %")
    plt.title("Нелинейность масштабного коэффициента")
    plt.xlabel("Угловая скорость, °/с")
    plt.ylabel("Нелинейность, %")
    plt.grid(True)
    plt.legend()


def plot_drift(times, rates, drift):
    dt = np.mean(np.diff(times))
    window_size = max(1, int(round(DRIFT_AVG_WINDOW_SEC / dt)))
    rates_avg = centered_moving_average(rates, window_size)

    plt.figure(figsize=(14, 8))
    plt.plot(times, rates, alpha=0.4, label="Дрейф, пересчёт U/K")
    plt.plot(times, rates_avg, linewidth=2, label="Скользящее среднее, 120 с")
    plt.title("Анализ дрейфа")
    plt.xlabel("Время, с")
    plt.ylabel("Угловая скорость, °/с")
    plt.grid(True)
    plt.legend()


def plot_allan(drift):
    plt.figure(figsize=(10, 7))
    plt.loglog(drift["tau"], drift["adev"] * 3600, "o-", label="Девиация Аллана")

    if drift["allan_fit"] is not None:
        fit = drift["allan_fit"]
        plt.loglog(fit["tau"], np.sqrt(fit["model"]) * 3600, "--", label="Аппроксимация")

    plt.title("Девиация Аллана")
    plt.xlabel("τ, с")
    plt.ylabel("σ(τ), °/ч")
    plt.grid(True, which="both", linestyle="--")
    plt.legend()


# ============================================================
# ВЫВОД
# ============================================================

def print_results(scale, drift):
    print("\n=== ПАРАМЕТРЫ МАСШТАБНОГО КОЭФФИЦИЕНТА ===")
    print(f"МК: {scale['scale_factor']:.6f} мВ/(°/с)")
    print(f"Наклон BFSL: {scale['slope']:.6f} мВ/(°/с)")
    print(f"Свободный член BFSL: {scale['intercept']:.6f} мВ")
    print(f"Нелинейность средняя: {scale['linearity_mean']:.6f} %")
    print(f"Нелинейность максимальная: {scale['linearity_max']:.6f} %")
    print(f"K_pos: {scale['k_pos']:.6f} мВ/(°/с)")
    print(f"K_neg: {scale['k_neg']:.6f} мВ/(°/с)")
    print(f"Несимметричность: {scale['asymmetry']:.6f} %")
    print(f"Смещение нуля по BFSL: {scale['zero_offset']:.6f} °/с")

    print("\n=== ПАРАМЕТРЫ ДРЕЙФА ===")
    if drift["zero_offset_mv"] is not None:
        print(f"Смещение нуля по Дрейф.dat: {drift['zero_offset_mv']:.6f} мВ")
    print(f"Смещение нуля: {drift['zero_offset']:.6f} °/с")
    print(f"Тренд: {drift['trend']:.6f} °/ч/ч")

    if drift["noise"] is None:
        print("Аппроксимация девиации Аллана не выполнена: недостаточно точек.")
    else:
        print(f"ARW: {drift['noise']['ARW']:.6f} °/√ч")
        print(f"Bias Instability: {drift['noise']['BI']:.6f} °/ч")
        print(f"RRW: {drift['noise']['RRW']:.6f} °/ч^(3/2)")


# ============================================================
# ПОЛНЫЙ АНАЛИЗ
# ============================================================

def analyze(scale_file, drift_file, show_plots=True):
    _, omega, scale_voltage = read_columns(scale_file, columns=(0, 1, 2))
    drift_time, drift_voltage = read_columns(drift_file, columns=(0, 1))

    scale = calc_scale_metrics(omega, scale_voltage)

    # Для дрейфа используем из МК.dat только масштабный коэффициент,
    # а смещение нуля находим напрямую по среднему сигналу из Дрейф.dat.
    drift_rate, drift_zero_offset_mv = drift_voltage_to_rate(
        voltage_v=drift_voltage,
        scale_factor_mv_per_dps=scale["scale_factor"],
    )

    drift = calc_drift_metrics(drift_time, drift_rate, drift_zero_offset_mv)

    print_results(scale, drift)

    plot_scale(omega, scale_voltage, scale)
    plot_nonlinearity(scale)
    plot_drift(drift_time, drift_rate, drift)
    plot_allan(drift)

    plt.tight_layout()

    if show_plots:
        plt.show()

    return scale, drift


# ============================================================
# ЗАПУСК
# ============================================================

def parse_args():
    parser = argparse.ArgumentParser(description="Обработка данных гироскопа")
    parser.add_argument("--scale", default=DEFAULT_SCALE_FILE, help="Файл МК.dat")
    parser.add_argument("--drift", default=DEFAULT_DRIFT_FILE, help="Файл Дрейф.dat")
    parser.add_argument("--no-show", action="store_true", help="Не открывать окна с графиками")
    return parser.parse_args()


def main():
    args = parse_args()
    analyze(args.scale, args.drift, show_plots=not args.no_show)


if __name__ == "__main__":
    main()