from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import numpy as np

from .impedance import ImpedanceCurve, MIN_IMPEDANCE_OHM


@dataclass(frozen=True)
class WayConfig:
    name: str
    channel_zero_based: int
    impedance_path: Path
    amp_gain_db: float

    @property
    def amp_gain_linear(self) -> float:
        return 10.0 ** (self.amp_gain_db / 20.0)


@dataclass
class VoltageSpectrumFrame:
    rms_v: float
    peak_v: float
    crest_factor: float
    crest_factor_db: float
    frequencies_hz: np.ndarray
    psd_v2_per_hz: np.ndarray
    v2_per_bin: np.ndarray


@dataclass
class PowerFrame:
    rms_input_v: float
    rms_amp_v: float
    total_p_w: float
    total_q_var: float
    total_s_va: float
    frequencies_hz: np.ndarray
    p_w_per_bin: np.ndarray
    q_var_per_bin: np.ndarray
    s_va_per_bin: np.ndarray


def one_sided_welch_voltage_spectrum(
    samples_v: np.ndarray,
    samplerate: float,
    window: np.ndarray,
    window_power_sum: float,
) -> tuple[np.ndarray, np.ndarray]:
    """Return one-sided voltage PSD and bin-integrated voltage squared.

    ``psd_v2_per_hz`` has units V^2/Hz. Multiplying by ``df`` gives
    bin-integrated RMS voltage squared in V^2; summing those bins recovers the
    mean-square AC voltage of the DC-removed block, within window statistics.
    """
    spectrum = np.fft.rfft(samples_v * window)
    psd_v2_per_hz = (np.abs(spectrum) ** 2) / (samplerate * window_power_sum)
    if len(psd_v2_per_hz) > 2:
        psd_v2_per_hz[1:-1] *= 2.0
    elif len(psd_v2_per_hz) == 2 and len(window) % 2:
        psd_v2_per_hz[1] *= 2.0
    v2_per_bin = psd_v2_per_hz * (samplerate / len(window))
    return psd_v2_per_hz, v2_per_bin


class VoltageSpectrumEstimator:
    """Hann/Welch voltage spectrum and input-side level metrics."""

    def __init__(
        self,
        samplerate: float,
        fft_size: int,
        input_volts_per_sample_unit: float,
        smoothing_alpha: float,
    ):
        if fft_size <= 0:
            raise ValueError("FFT size must be positive")
        if input_volts_per_sample_unit <= 0:
            raise ValueError("input voltage scale must be positive")
        if not 0.0 <= smoothing_alpha < 1.0:
            raise ValueError("smoothing alpha must be in [0, 1)")

        self.samplerate = float(samplerate)
        self.fft_size = int(fft_size)
        self.input_volts_per_sample_unit = float(input_volts_per_sample_unit)
        self.smoothing_alpha = float(smoothing_alpha)
        self.window = np.hanning(self.fft_size).astype(np.float64)
        self.window_power_sum = float(np.sum(self.window * self.window))
        self.frequencies_hz = np.fft.rfftfreq(self.fft_size, d=1.0 / self.samplerate)
        self._smoothed: VoltageSpectrumFrame | None = None

    def estimate(self, samples: np.ndarray) -> VoltageSpectrumFrame:
        if samples.shape[0] != self.fft_size:
            raise ValueError(f"expected {self.fft_size} samples, got {samples.shape[0]}")

        volts = samples.astype(np.float64, copy=False) * self.input_volts_per_sample_unit
        volts = volts - float(np.mean(volts))
        rms_v = float(np.sqrt(np.mean(volts * volts)))
        peak_v = float(np.max(np.abs(volts)))
        crest = peak_v / rms_v if rms_v > 0.0 else 0.0
        crest_db = 20.0 * np.log10(crest) if crest > 0.0 else float("-inf")

        psd, v2_per_bin = one_sided_welch_voltage_spectrum(
            volts,
            self.samplerate,
            self.window,
            self.window_power_sum,
        )
        psd[0] = 0.0
        v2_per_bin[0] = 0.0

        frame = VoltageSpectrumFrame(
            rms_v=rms_v,
            peak_v=peak_v,
            crest_factor=float(crest),
            crest_factor_db=float(crest_db),
            frequencies_hz=self.frequencies_hz,
            psd_v2_per_hz=psd,
            v2_per_bin=v2_per_bin,
        )
        return self._smooth(frame)

    def _smooth(self, frame: VoltageSpectrumFrame) -> VoltageSpectrumFrame:
        alpha = self.smoothing_alpha
        if self._smoothed is None or alpha <= 0.0:
            self._smoothed = frame
            return frame

        prev = self._smoothed
        smoothed = VoltageSpectrumFrame(
            rms_v=alpha * prev.rms_v + (1.0 - alpha) * frame.rms_v,
            peak_v=alpha * prev.peak_v + (1.0 - alpha) * frame.peak_v,
            crest_factor=alpha * prev.crest_factor + (1.0 - alpha) * frame.crest_factor,
            crest_factor_db=alpha * prev.crest_factor_db + (1.0 - alpha) * frame.crest_factor_db,
            frequencies_hz=frame.frequencies_hz,
            psd_v2_per_hz=alpha * prev.psd_v2_per_hz + (1.0 - alpha) * frame.psd_v2_per_hz,
            v2_per_bin=alpha * prev.v2_per_bin + (1.0 - alpha) * frame.v2_per_bin,
        )
        self._smoothed = smoothed
        return smoothed


class PowerEstimator:
    """Windowed FFT/STFT electrical power estimator for one loudspeaker way."""

    def __init__(
        self,
        way: WayConfig,
        impedance: ImpedanceCurve,
        samplerate: float,
        fft_size: int,
        input_volts_per_sample_unit: float,
        smoothing_alpha: float,
    ):
        if fft_size <= 0:
            raise ValueError("FFT size must be positive")
        if input_volts_per_sample_unit <= 0:
            raise ValueError("input voltage scale must be positive")
        if not 0.0 <= smoothing_alpha < 1.0:
            raise ValueError("smoothing alpha must be in [0, 1)")

        self.way = way
        self.impedance = impedance
        self.samplerate = float(samplerate)
        self.fft_size = int(fft_size)
        self.input_volts_per_sample_unit = float(input_volts_per_sample_unit)
        self.smoothing_alpha = float(smoothing_alpha)

        self.window = np.hanning(self.fft_size).astype(np.float64)
        self.window_power_sum = float(np.sum(self.window * self.window))
        self.frequencies_hz = np.fft.rfftfreq(self.fft_size, d=1.0 / self.samplerate)
        self.df_hz = self.samplerate / self.fft_size
        self.z_bins = self.impedance.interpolate(self.frequencies_hz)

        self._smoothed: PowerFrame | None = None

    def estimate(self, samples: np.ndarray) -> PowerFrame:
        if samples.shape[0] != self.fft_size:
            raise ValueError(f"expected {self.fft_size} samples, got {samples.shape[0]}")

        dsp_v = samples.astype(np.float64, copy=False) * self.input_volts_per_sample_unit
        dsp_v = dsp_v - float(np.mean(dsp_v))
        rms_input_v = float(np.sqrt(np.mean(dsp_v * dsp_v)))

        _psd_v2_per_hz, dsp_v2_per_bin = one_sided_welch_voltage_spectrum(
            dsp_v,
            self.samplerate,
            self.window,
            self.window_power_sum,
        )
        amp_v2_per_bin = dsp_v2_per_bin * (self.way.amp_gain_linear ** 2)
        amp_v2_per_bin[0] = 0.0

        complex_power = amp_v2_per_bin / np.conj(self.z_bins)
        p_w = np.real(complex_power)
        q_var = np.imag(complex_power)
        s_va = amp_v2_per_bin / np.maximum(np.abs(self.z_bins), MIN_IMPEDANCE_OHM)

        frame = PowerFrame(
            rms_input_v=rms_input_v,
            rms_amp_v=float(rms_input_v * self.way.amp_gain_linear),
            total_p_w=float(np.sum(p_w)),
            total_q_var=float(np.sum(q_var)),
            total_s_va=float(np.sum(s_va)),
            frequencies_hz=self.frequencies_hz,
            p_w_per_bin=p_w,
            q_var_per_bin=q_var,
            s_va_per_bin=s_va,
        )
        return self._smooth(frame)

    def _smooth(self, frame: PowerFrame) -> PowerFrame:
        alpha = self.smoothing_alpha
        if self._smoothed is None or alpha <= 0.0:
            self._smoothed = frame
            return frame

        prev = self._smoothed
        smoothed = PowerFrame(
            rms_input_v=alpha * prev.rms_input_v + (1.0 - alpha) * frame.rms_input_v,
            rms_amp_v=alpha * prev.rms_amp_v + (1.0 - alpha) * frame.rms_amp_v,
            total_p_w=alpha * prev.total_p_w + (1.0 - alpha) * frame.total_p_w,
            total_q_var=alpha * prev.total_q_var + (1.0 - alpha) * frame.total_q_var,
            total_s_va=alpha * prev.total_s_va + (1.0 - alpha) * frame.total_s_va,
            frequencies_hz=frame.frequencies_hz,
            p_w_per_bin=alpha * prev.p_w_per_bin + (1.0 - alpha) * frame.p_w_per_bin,
            q_var_per_bin=alpha * prev.q_var_per_bin + (1.0 - alpha) * frame.q_var_per_bin,
            s_va_per_bin=alpha * prev.s_va_per_bin + (1.0 - alpha) * frame.s_va_per_bin,
        )
        self._smoothed = smoothed
        return smoothed
