# -*- coding: utf-8 -*-
"""
MDD(우울증) 딥러닝 분류기 - 모델 로드 및 예측
"""
import copy
import json
import os

import numpy as np
import torch

import directory
from mdd_model_arch import BHITransformer
from mdd_preprocessing import preprocess_subject_arrays

MDD_MODEL_DIR = os.path.join(directory.ROOT_DIR, 'model', 'mdd')
BACKBONE_DIR = os.path.join(MDD_MODEL_DIR, 'BHITransformer')
FT_HEAD_DIR = os.path.join(MDD_MODEL_DIR, 'FT_HEAD')

FS = 256

_DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

_backbone = None
_head = None
_mu = None
_sd = None
_meta = None


def _load_model():
    global _backbone, _head, _mu, _sd, _meta

    meta = json.load(open(os.path.join(FT_HEAD_DIR, 'meta.json'), encoding='utf-8'))
    hp = json.load(open(os.path.join(BACKBONE_DIR, 'hyperparameter.json'), encoding='utf-8'))

    backbone = BHITransformer(
        eeg_channels=2, ppg_channels=1,
        eeg_patch_size=int(FS * hp['eeg_patch_size']),
        ppg_patch_size=int(FS * hp['ppg_patch_size']),
        eeg_patch_stride=int(FS * hp['eeg_patch_stride']),
        ppg_patch_stride=int(FS * hp['ppg_patch_stride']),
        embed_dim=hp['embed_dim'], temporal_dim=hp['temporal_dim'],
        activation=hp['activation'], num_heads=hp['num_heads'],
        num_layers=hp['num_layers'], num_classes=2,
        pos_encoding_mode=hp['pos_encoding_mode'],
        time_min_s=0.0, time_max_s=hp['EpochT'],
        dropout=hp['dropout'], droppath=hp['droppath'],
        isPreLN=hp['isPreLN'], ablation_mode=hp['ablation_mode'])
    ckpt = os.path.join(BACKBONE_DIR, 'models', "bestModel_fold{:02d}.pt".format(meta['pretrained_fold']))
    backbone.load_state_dict(torch.load(ckpt, map_location='cpu'))
    backbone.to(_DEVICE).eval()

    head = copy.deepcopy(backbone.classifier)
    head.load_state_dict(torch.load(os.path.join(FT_HEAD_DIR, 'head.pt'), map_location='cpu'))
    head.to(_DEVICE).eval()

    sc = np.load(os.path.join(FT_HEAD_DIR, 'scaler.npz'))

    _backbone, _head, _mu, _sd, _meta = backbone, head, sc['mu'], sc['sd'], meta


@torch.no_grad()
def _embed(m, eeg, ppg):
    out = []
    for i in range(0, eeg.shape[0], 256):
        e, p = eeg[i:i + 256], ppg[i:i + 256]
        et, pt = m.eeg_embed(e), m.ppg_embed(p)
        B = et.shape[0]
        ta = (torch.arange(et.shape[1], device=e.device) * m.eeg_embed.patch_stride
              + m.eeg_embed.patch_size / 2) / FS
        tb = (torch.arange(pt.shape[1], device=p.device) * m.ppg_embed.patch_stride
              + m.ppg_embed.patch_size / 2) / FS
        ef = m.eeg_pos(et, times=ta.unsqueeze(0).expand(B, -1))
        pf = m.ppg_pos(pt, times=tb.unsqueeze(0).expand(B, -1))

        se, sp, ce, cp = ef, pf, ef, pf
        for b in m.self_attn_eeg:
            se = b(se)[0]
        for b in m.self_attn_ppg:
            sp = b(sp)[0]
        for b in m.cross_attn_e2p:
            ce = b(ce, pf)[0]
        for b in m.cross_attn_p2e:
            cp = b(cp, ef)[0]

        g = m._get_bhi_summary
        out.append(torch.cat([g(se), g(sp), g(ce), g(cp)], -1))
    return torch.cat(out).cpu().numpy()


@torch.no_grad()
def predict(eeg_ch1, eeg_ch2, ppg_raw, fs_eeg, fs_ppg, use_finetuned_head=True):
    """
    Args:
        ...
        use_finetuned_head (bool) : True(기본) = FT_HEAD(head.pt, bwave 데이터로 파인튜닝된
            분류기)를 사용. False = 백본에 원래 붙어있던(파인튜닝 이전) classifier를 그대로
            사용 - 파인튜닝 효과를 비교해보기 위한 옵션.
            ⚠ False일 때 나오는 값은 '우울 확률'로서의 의미가 검증되지 않았다 - 원본
            classifier가 애초에 우울증 분류용으로 학습된 게 아닐 수 있고, scaler.npz(mu/sd)도
            FT_HEAD 학습 시점 기준으로 계산된 값이라 원본 classifier 입력 분포와 안 맞을 수
            있음. 성능/점수 비교용으로만 참고할 것.
    """
    if _meta is None:
        _load_model()

    eeg_raw = np.vstack([np.asarray(eeg_ch1, dtype=np.float64),
                         np.asarray(eeg_ch2, dtype=np.float64)])
    ppg_raw = np.asarray(ppg_raw, dtype=np.float64)[None, :]

    seg = preprocess_subject_arrays(
        eeg_raw, ppg_raw, fs_eeg, fs_ppg,
        eeg_patch_size=_meta['eeg_patch_size'], ppg_patch_size=_meta['ppg_patch_size'],
        eeg_patch_stride=_meta['eeg_patch_stride'], ppg_patch_stride=_meta['ppg_patch_stride'],
        use_global_threshold=True, verbose=False)
    if seg is None:
        return None

    Z = _embed(_backbone,
               torch.from_numpy(seg['eeg']).float().to(_DEVICE),
               torch.from_numpy(seg['ppg']).float().to(_DEVICE))
    x = torch.from_numpy((Z - _mu) / _sd).float().to(_DEVICE)
    classifier_head = _head if use_finetuned_head else _backbone.classifier
    p = 1.0 - torch.softmax(classifier_head(x), 1)[:, _meta['hc_index']].cpu().numpy()
    return float(p.mean())