import os
import math
import time
import numpy as np
from typing import Union
from common.runtime import load_root_env

load_root_env(__file__)


def get_signal_length(freq):
    length = int(os.getenv('signal_length_' + str(freq)))
    return length


class Signal:
    """
    Класс сбора и предобработки сигнала.

    Атрибуты:
            length: Длина сигнала.
            signal: Массив, в который собираем сигнал.
    """

    def __init__(self, conv_method='average'):
        self.conv_method = conv_method
        self.signal = []
        self.signal_abs = []
        self.last_packet_ts = None
        self.metric_mode = os.getenv('signal_metric_mode', 'fft_top_bins').strip().lower()
        self.fft_top_bins = max(1, int(os.getenv('signal_fft_top_bins', '2048')))
        self.fft_window = os.getenv('signal_fft_window', 'hann').strip().lower()

    def get_signal(self):
        """
        Возвращает собранный сигнал.
        :return: Массив с сигналом.
        """
        return self.signal, self.signal_abs

    def get_last_packet_ts(self):
        return self.last_packet_ts

    def clear(self) -> None:
        """
        Очистить массив с сигналом после предобработки?
        :return: None
        """
        self.signal = []
        self.signal_abs = []
        self.last_packet_ts = None

    def _build_window(self, size: int) -> np.ndarray:
        if self.fft_window in {'', 'none', 'rect', 'rectangular'}:
            return np.ones(size, dtype=np.float32)
        if self.fft_window == 'hann':
            return np.hanning(size).astype(np.float32, copy=False)
        raise ValueError(f'unsupported fft window: {self.fft_window}')

    def _compute_iq_power(self, samples: np.ndarray, signal_abs: np.ndarray) -> float:
        if self.conv_method == 'max':
            return float(np.max(signal_abs * signal_abs))

        if self.metric_mode in {'fft', 'fft_top_bins', 'top_bins'}:
            window = self._build_window(samples.size)
            windowed = samples.astype(np.complex64, copy=False) * window
            spectrum = np.fft.fft(windowed)
            power_bins = (np.abs(spectrum) ** 2).astype(np.float32, copy=False)
            power_bins /= max(float(np.sum(window * window)), 1.0)

            bins_to_keep = min(self.fft_top_bins, power_bins.size)
            top_bins = np.partition(power_bins, power_bins.size - bins_to_keep)[-bins_to_keep:]
            return float(np.mean(top_bins))

        return float(np.mean(signal_abs * signal_abs))

    def signal_preprocessing(self, length) -> float:
        """
        Предобработка сигнала.

        :return: Число типа float - "характеристика сигнала".
        """
        samples = np.asarray(self.signal).ravel()[0:length]
        if samples.size == 0:
            return 0.0

        # Основной режим: считаем dBFS из IQ-вектора.
        if np.iscomplexobj(samples):
            i = samples.real.astype(np.float32, copy=False)
            q = samples.imag.astype(np.float32, copy=False)
            signal = np.array([i, q], dtype=np.float32)
            signal_abs = np.sqrt(i * i + q * q).astype(np.float32, copy=False)

            power = self._compute_iq_power(samples, signal_abs)

            result = 10.0 * math.log10(max(power, 1e-20))
            self.signal = signal
            self.signal_abs = signal_abs
            return result

        # Fallback: если на вход уже подали скалярную метрику, агрегируем как есть.
        scalar_samples = samples.astype(np.float32, copy=False)
        if self.conv_method == 'max':
            result = float(np.max(scalar_samples))
        else:
            result = float(np.median(scalar_samples))

        self.signal = scalar_samples
        self.signal_abs = np.abs(scalar_samples)
        return result

    def fill_signal(self, lvl, length) -> Union[int, float]:
        """
        Сбор сигнала в соответствующий массив. Если уже собран, то предобработка.
        :param lvl: Массив, без ограничения общности, с неизвестной длиной, содержащий сигнал.
        :param length:
        :return: 0 - если еще нет нужного количества сигнала, "характеристика" иначе.
        """
        if len(self.signal) <= length:
            y = np.array(lvl).ravel()
            self.signal = np.concatenate((self.signal, y), axis=None)
            return 0
        else:
            self.last_packet_ts = time.time()
            preproc_signal = self.signal_preprocessing(length)
            return preproc_signal


class SignalsArray:
    """
    Класс для сохранения медиан сигналов на частотах.
    Атрибуты:
            sig_array: Список для сохранения медиан.
            counter: Индикатор наполненности массива.
    """

    def __init__(self):
        self.sig_array = []
        self.sig_ts_array = []
        self.counter = 0

    def fill_sig_arr(self, metrica, packet_ts=None, num_chs=3):
        """
        Аппендим характеристику сигнала (метрику) в массив длиной num_chs.
        :param metrica: Характеристика сигнала (метрика).
        :param packet_ts: Время завершения набора пакета с SDR для текущей метрики.
        :param num_chs: Количество каналов на частоте.
        :return: Индекс канала внутри частоты и массив с характеристиками, если заполнен, иначе - пустой.
        """
        if num_chs:
            if self.counter < num_chs:
                self.sig_array.append(metrica)
                self.sig_ts_array.append(packet_ts)
                self.counter += 1
            if self.counter == num_chs:
                arr = self.sig_array
                arr_ts = self.sig_ts_array
                self.sig_array = []
                self.sig_ts_array = []
                self.counter = 0
                return num_chs - 1, arr, arr_ts
            else:
                return self.counter - 1, [], []
        else:
            return 0, [], []
