# -*- coding: utf-8 -*-
"""
MDD(우울증) 딥러닝 분류기 - 신호 전처리 보조 함수
(모델 제공사의 function.py 를 그대로 이식. 기존 python/ 안의 다른 모듈과 이름 충돌을
방지하기 위해 mdd_ 접두어를 붙였다.)
"""
from scipy.signal import butter, filtfilt
import numpy as np


def butter_bandpass(lowcut, highcut, nyq, order=3):
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return b, a


def bandpass_filter(data, lowcut, highcut, nyq, order=3):
    b, a = butter_bandpass(lowcut, highcut, nyq, order)
    y = filtfilt(b, a, data)
    return y


def butter_notch(lowcut, highcut, nyq, order=3):
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='bandstop')
    return b, a


def notch_filter(data, lowcut, highcut, nyq, order=3):
    b, a = butter_notch(lowcut, highcut, nyq, order)
    y = filtfilt(b, a, data)
    return y


def Epoching(data, EpochT, fs, overlap=0):
    [_, nSamp] = data.shape
    Samp_per_epoch = EpochT * fs
    step_size = int(Samp_per_epoch * (1 - overlap))
    nEpoch = int((nSamp - (Samp_per_epoch - step_size)) / step_size)
    Epoch = []
    t_epoch_mid = []
    for iEpoch in range(nEpoch):
        Epoch_pre = data[:, iEpoch * step_size:iEpoch * step_size + Samp_per_epoch]
        t_epoch_mid_pre = iEpoch * step_size + Samp_per_epoch / 2
        Epoch.append(Epoch_pre)
        t_epoch_mid.append(t_epoch_mid_pre)

    return np.array(Epoch), np.array(t_epoch_mid)


def GlobalThresholding(data, threshold):
    [thershold_lower, threshold_upper] = threshold
    idxEpoch4use = (np.max(np.max(abs(data), axis=2), axis=1) < threshold_upper) & (
        np.sum((np.max(abs(data), axis=2) > thershold_lower), axis=1) == 2)
    Epoch4use = data[idxEpoch4use, :, :]
    return Epoch4use, idxEpoch4use


def IndividualThresholding(data):
    idxEpochAmp4use = np.zeros((data.shape[0], data.shape[1]), dtype=bool)
    idxEpochDiff4use = np.zeros((data.shape[0], data.shape[1]), dtype=bool)
    data_diff = data[:, :, 1:] - data[:, :, :-1]
    for iChan in range(data.shape[1]):
        ampMax = np.max(abs(data[:, iChan, :]), axis=1)
        ampMax_mean = np.mean(ampMax, axis=0)
        ampMax_std = np.std(ampMax, axis=0)
        idxEpochAmp4use[:, iChan] = (ampMax < ampMax_mean + ampMax_std * 5) & (ampMax > ampMax_mean - ampMax_std * 5)
        diffMax = np.max(abs(data_diff[:, iChan, :]), axis=1)
        diffMax_mean = np.mean(diffMax, axis=0)
        diffMax_std = np.std(diffMax, axis=0)
        idxEpochDiff4use[:, iChan] = (diffMax < diffMax_mean + diffMax_std * 5) & (diffMax > diffMax_mean - diffMax_std * 5)
    idxEpoch4use = (np.sum(idxEpochAmp4use, axis=1) == data.shape[1]) & (np.sum(idxEpochDiff4use, axis=1) == data.shape[1])
    Epoch4use = data[idxEpoch4use, :, :]
    return Epoch4use, idxEpoch4use


def get_common_indices(A, B):
    A = np.array(A)
    B = np.array(B)
    common_values = np.intersect1d(A, B)
    A_indices = np.where(np.isin(A, common_values))[0]
    B_indices = np.where(np.isin(B, common_values))[0]
    return A_indices, B_indices


def make_patch_times(total_len_samples: int, patch_size: int, fs: int, start_time_s: float, stride_size: int):
    starts = np.arange(0, total_len_samples - patch_size + 1, stride_size)
    centers = starts + patch_size / 2.0
    times = start_time_s + centers / fs
    return times.astype(np.float32)  # (Lp,)