import numpy as np
from scipy.signal import welch
from collections import deque



class Beta_Calculate():
    def __init__(self, Threshold_value_low, Threshold_value_high, fs=250, win_len=5, config=None):
        self.Threshold_value_low = Threshold_value_low
        self.Threshold_value_high = Threshold_value_high
        self.fs = fs
        self.beta_result = []
        self.eegQueue = deque(maxlen=win_len)

    def calculate_all(self, data, fs, nperseg=1000):
        mean_x = np.mean(data, axis=-1, keepdims=True)
        data = data - mean_x
        freqs, psd = self.compute_psd_multichannel(data, fs, nperseg)
        beta_psd = np.sum(self.band_psd(freqs, psd, (13, 30)))
        alpha_psd = np.sum(self.band_psd(freqs, psd, (8, 13)))
        theta_psd = np.sum(self.band_psd(freqs, psd, (4, 8)))
        
        # print(f"[功率] β={beta_psd:.2f} | α={alpha_psd:.2f} | θ={theta_psd:.2f}")

        return beta_psd, alpha_psd, theta_psd

    def compute_psd_multichannel(self, data, fs=250, nperseg=1000):
        n_samples = data.shape[-1]
        if n_samples < nperseg:
            nperseg = n_samples
        
        noverlap = 500
        if noverlap >= nperseg:
            noverlap = int(nperseg / 2)

        if nperseg == 0:
            return np.array([]), np.zeros((data.shape[0], 0))

        freqs, psd = welch(data, fs=fs, nperseg=nperseg, noverlap=noverlap, axis=-1)
        return freqs, psd

    def band_psd(self, freqs, psd, band):
        idx = np.logical_and(freqs >= band[0], freqs <= band[1])
        return np.sum(psd[:, idx], axis=-1)


    def reset_queue(self):
        self.eegQueue.clear()


    def queueOpt(self, data):
        if data is None or data.size == 0:
            return None
        if len(self.eegQueue) < self.eegQueue.maxlen:
            self.eegQueue.append(data)
        else:
            self.eegQueue.append(data)

        if len(self.eegQueue) == self.eegQueue.maxlen:
            eegData = np.hstack([self.eegQueue[i] for i in range(len(self.eegQueue))])
            if eegData.size == 0:
                return None
            eegData -= np.mean(eegData, axis=-1, keepdims=True)

            beta_psd, alpha_psd, theta_psd = self.calculate_all(eegData, fs=self.fs, nperseg=1000)

            return (beta_psd)