"""Frequency-axis calibration helpers."""

from __future__ import annotations

from typing import Any, Mapping

import numpy as np

from rfg_adc_plotter.constants import SWEEP_FREQ_MAX_GHZ, SWEEP_FREQ_MIN_GHZ
from rfg_adc_plotter.types import SweepData


def recalculate_calibration_c(
    base_coeffs: np.ndarray,
    f_min: float = SWEEP_FREQ_MIN_GHZ,
    f_max: float = SWEEP_FREQ_MAX_GHZ,
) -> np.ndarray:
    """Recalculate coefficients while preserving sweep edges."""
    coeffs = np.asarray(base_coeffs, dtype=np.float64).reshape(-1)
    if coeffs.size < 3:
        out = np.zeros((3,), dtype=np.float64)
        out[: coeffs.size] = coeffs
        coeffs = out
    c0, c1, c2 = float(coeffs[0]), float(coeffs[1]), float(coeffs[2])
    x0 = float(f_min)
    x1 = float(f_max)
    y0 = c0 + c1 * x0 + c2 * (x0 ** 2)
    y1 = c0 + c1 * x1 + c2 * (x1 ** 2)
    if not (np.isfinite(y0) and np.isfinite(y1)) or y1 == y0:
        return np.asarray([c0, c1, c2], dtype=np.float64)
    scale = (x1 - x0) / (y1 - y0)
    shift = x0 - scale * y0
    return np.asarray(
        [
            shift + scale * c0,
            scale * c1,
            scale * c2,
        ],
        dtype=np.float64,
    )


CALIBRATION_C_BASE = np.asarray([0.0, 1.0, 0.025], dtype=np.float64)
CALIBRATION_C = recalculate_calibration_c(CALIBRATION_C_BASE)


def get_calibration_base() -> np.ndarray:
    return np.asarray(CALIBRATION_C_BASE, dtype=np.float64).copy()


def get_calibration_coeffs() -> np.ndarray:
    return np.asarray(CALIBRATION_C, dtype=np.float64).copy()


def set_calibration_base_value(index: int, value: float) -> np.ndarray:
    """Update one base coefficient and recalculate the working coefficients."""
    global CALIBRATION_C
    CALIBRATION_C_BASE[int(index)] = float(value)
    CALIBRATION_C = recalculate_calibration_c(CALIBRATION_C_BASE)
    return get_calibration_coeffs()


def calibrate_freqs(sweep: Mapping[str, Any]) -> SweepData:
    """Return a sweep copy with calibrated and resampled frequency axis."""
    freqs = np.asarray(sweep["F"], dtype=np.float64).copy()
    values = np.asarray(sweep["I"], dtype=np.float64).copy()
    coeffs = np.asarray(CALIBRATION_C, dtype=np.float64)
    if freqs.size > 0:
        freqs = coeffs[0] + coeffs[1] * freqs + coeffs[2] * (freqs * freqs)

    if freqs.size >= 2:
        freqs_cal = np.linspace(float(freqs[0]), float(freqs[-1]), freqs.size, dtype=np.float64)
        values_cal = np.interp(freqs_cal, freqs, values).astype(np.float64)
    else:
        freqs_cal = freqs.copy()
        values_cal = values.copy()

    return {
        "F": freqs_cal,
        "I": values_cal,
    }