# -*- coding: utf-8 -*-
"""
MDD(우울증) 딥러닝 분류기 - 신호 전처리

모델 제공사의 preprocessing.py(preprocess_subject)를 그대로 가져오되,
원본은 eeg/ppg 파일 경로(.txt)를 읽어서 처리하는 반면 여기서는 이미 메모리에
올라와 있는 raw 배열(EEG.py / PPG.py 가 만든 self.raw)을 바로 받는다.
그 외 필터 대역, 진폭 임계값, 에폭 길이/겹침, z-score 정규화, 패치 시각 계산은
원본과 동일하다.
"""
import numpy as np
from scipy.signal import resample

from mdd_function import (
    bandpass_filter,
    notch_filter,
    Epoching,
    GlobalThresholding,
    IndividualThresholding,
    get_common_indices,
    make_patch_times,
)

# ===== 모델이 학습된 조건과 동일한 상수 (원본 preprocessing.py 와 동일) =====
FS = 256                      # 모델이 학습된 샘플링 레이트
EPOCH_T = 8                   # 에폭 길이 (s)
FC_EEG = [4, 50]              # EEG 대역통과
FC_PPG = [0.5, 5]             # PPG 대역통과
FC_NOTCH = [59, 61]           # 전원 잡음 제거
THRESHOLD = [1, 100]          # GlobalThresholding 진폭 범위 (uV 기준)


def _resample_to(x, fs_in, fs_out):
    """(C, N) 신호를 fs_out 으로 리샘플링."""
    if fs_in == fs_out:
        return x
    n_out = int(round(x.shape[1] * fs_out / float(fs_in)))
    return resample(x, n_out, axis=1)


def _zscore_epochwise(x):
    """(N, C, L) 에서 각 (epoch, channel) 의 시간축 z-score."""
    m = x.mean(axis=-1, keepdims=True)
    s = x.std(axis=-1, keepdims=True) + 1e-6
    return (x - m) / s


def preprocess_subject_arrays(eeg_raw, ppg_raw, fs_eeg, fs_ppg,
                               eeg_patch_size, ppg_patch_size,
                               eeg_patch_stride, ppg_patch_stride,
                               epoch_t=EPOCH_T,
                               use_global_threshold=True, verbose=False):
    """
    한 측정(EEG 2채널 + PPG 1채널 raw 배열)을 모델 입력 형태로 전처리한다.

    Args:
        eeg_raw (np.ndarray) : (2, N) 원신호 (필터링 전, EEG.py 의 self.raw 와 동일한 raw 값)
        ppg_raw (np.ndarray) : (1, N) 원신호 (필터링 전, PPG.py 의 self.raw 와 동일한 raw 값)
        fs_eeg, fs_ppg (int) : 원신호 샘플링 레이트 (기기별로 다름)

    Returns
    -------
    dict 또는 None
        eeg  (N, 2, L_eeg)  z-score 된 EEG 에폭
        ppg  (N, 1, L_ppg)  z-score 된 PPG 에폭
        유효 에폭이 하나도 없으면 None.
    """
    stride_t = int(epoch_t / 2)              # 50% 겹침
    overlap = 1 - stride_t / epoch_t

    # ---- 1) 256 Hz 리샘플링 ----
    eeg_raw = _resample_to(np.asarray(eeg_raw, dtype=np.float64), fs_eeg, FS)
    ppg_raw = _resample_to(np.asarray(ppg_raw, dtype=np.float64), fs_ppg, FS)

    nyq = 0.5 * FS

    # ---- 2) 필터링 (notch -> bandpass) ----
    eeg_filt = bandpass_filter(
        notch_filter(eeg_raw, FC_NOTCH[0], FC_NOTCH[1], nyq, order=3),
        FC_EEG[0], FC_EEG[1], nyq, order=3)
    ppg_filt = bandpass_filter(
        notch_filter(ppg_raw, FC_NOTCH[0], FC_NOTCH[1], nyq, order=3),
        FC_PPG[0], FC_PPG[1], nyq, order=3)

    # ---- 3) 에폭 분할 ----
    eeg_epoch, t_eeg_mid_pre = Epoching(eeg_filt, epoch_t, FS, overlap)
    ppg_epoch, t_ppg_mid = Epoching(ppg_filt, epoch_t, FS, overlap)
    if eeg_epoch.shape[0] == 0 or ppg_epoch.shape[0] == 0:
        return None

    # ---- 4) 아티팩트 제거 ----
    if use_global_threshold:
        eeg_thres, idx_thres = GlobalThresholding(eeg_epoch, THRESHOLD)
        t_eeg_mid_thres = t_eeg_mid_pre[idx_thres]
    else:
        eeg_thres, t_eeg_mid_thres = eeg_epoch, t_eeg_mid_pre
    if eeg_thres.shape[0] == 0:
        if verbose:
            print("  [WARN] GlobalThresholding 이 에폭을 전부 제거함 (EEG 진폭 단위/스케일 확인 필요)")
        return None

    eeg_clean, idx_indiv = IndividualThresholding(eeg_thres)
    t_eeg_mid = t_eeg_mid_thres[idx_indiv]
    if eeg_clean.shape[0] == 0:
        return None

    # ---- 5) EEG/PPG 공통 에폭만 사용 ----
    idx_eeg, idx_ppg = get_common_indices(t_eeg_mid, t_ppg_mid)
    if len(idx_eeg) == 0:
        return None

    eeg_use = eeg_clean[idx_eeg, :, :]
    ppg_use = ppg_epoch[idx_ppg, :, :]

    # ---- 6) 정규화 ----
    eeg_use = _zscore_epochwise(eeg_use).astype(np.float32)
    ppg_use = _zscore_epochwise(ppg_use).astype(np.float32)

    # ---- 7) 패치 시각 벡터 (모델 time encoding 용, 참고용으로 계산만 함께 반환) ----
    l_eeg, l_ppg = eeg_use.shape[-1], ppg_use.shape[-1]
    eeg_t1 = make_patch_times(l_eeg, eeg_patch_size, FS, 0.0, eeg_patch_stride)
    ppg_t1 = make_patch_times(l_ppg, ppg_patch_size, FS, 0.0, ppg_patch_stride)
    eeg_times = np.tile(eeg_t1[None, :], (eeg_use.shape[0], 1)).astype(np.float32)
    ppg_times = np.tile(ppg_t1[None, :], (ppg_use.shape[0], 1)).astype(np.float32)

    return {
        'eeg': eeg_use,
        'ppg': ppg_use,
        'eeg_times': eeg_times,
        'ppg_times': ppg_times,
    }