# -*- coding: utf-8 -*-
# @Author  : chq_N
# @Time    : 2020/8/26


from datetime import datetime

import numpy as np
from scipy.special import softmax
from sklearn.metrics import roc_auc_score, roc_curve
from sklearn.model_selection import StratifiedKFold

from model import init_model


def cross_val(X, y, seed):
    np.random.seed(seed)
    kf = StratifiedKFold(n_splits=5, shuffle=True)
    prob_list = list()
    auc_list = list()
    y_list = list()
    index_group = []
    for train_index, test_index in kf.split(X, y):
        index_group.append(test_index.tolist())

    for i in range(len(index_group)):
        fold = str(i) + '-5'
        test_idx = index_group[i]
        _val_i = i + 1
        if _val_i >= len(index_group): _val_i = 0
        val_idx = index_group[_val_i]
        _train_i = [_ for _ in range(len(index_group)) if _ != i and _ != _val_i]
        train_idx = []
        for _ in _train_i:
            train_idx += index_group[_]

        m = init_model(
            fold, X[train_idx], y[train_idx],
            X[val_idx], y[val_idx],
            X[test_idx], y[test_idx], )
        print('Training model:', fold)
        _time = datetime.now()
        m.fit()
        print('Training Time:', datetime.now() - _time)
        print('Testing model:', fold)

        test_corr = 0
        val_max = 0
        val_i = -1
        for j in range(500, 5001, 100):
            m.load_model(j)
            pred, label = m.transform('test', 24)
            pred = softmax(pred, axis=1)
            test_auc = roc_auc_score(label, pred[:, 1])
            pred, label = m.transform('val', 24)
            pred = softmax(pred, axis=1)
            val_auc = roc_auc_score(label, pred[:, 1])
            if val_auc > val_max:
                val_i = j
                val_max = val_auc
                test_corr = test_auc

        print('Best iter:', val_i, 'Best V auc:', val_max, 'Corr T auc:', test_corr)
        m.load_model(val_i)
        pred, label = m.transform('test', 24)
        pred = softmax(pred, axis=1)[:, 1]
        test_auc = roc_auc_score(label, pred)
        prob_list.append(pred)
        auc_list.append(test_auc)
        y_list.append(label)
    return np.concatenate(prob_list), np.concatenate(y_list), np.mean(auc_list), np.std(auc_list, ddof=1)


def detail_test(features, label, ppv_th=0.7):
    def get_sen_spe(pred, label):

        def criteria(x, th):
            return (x > th).astype('int')

        for j in range(0, 1000, 1):
            j = j / 1000
            TP = ((label == 1) * (criteria(pred, j) == 1))
            TN = ((label == 0) * (criteria(pred, j) == 0))
            FP = ((label == 0) * (criteria(pred, j) == 1))
            FN = ((label == 1) * (criteria(pred, j) == 0))
            sensitivity = TP.sum() / (TP.sum() + FN.sum() + 1e-9)
            specifity = TN.sum() / (TN.sum() + FP.sum() + 1e-9)
            ppv = TP.sum() / (TP.sum() + FP.sum() + 1e-9)
            npv = TN.sum() / (TN.sum() + FN.sum() + 1e-9)
            if ppv >= ppv_th:
                break
        return sensitivity, specifity, ppv, npv

    X_selected = features
    auc_all = list()
    auc_mean = list()
    auc_std = list()
    sensitivity = list()
    specificity = list()
    ppv = list()
    npv = list()
    tpr = list()
    for i in range(5):
        y_pred, _y, _auc_mean, _auc_std = cross_val(X_selected, label, i * 10)
        test_auc = roc_auc_score(_y, y_pred)
        _tpr = inter_auc(y_pred, _y)
        tpr.append(_tpr)
        auc_all.append(test_auc)
        auc_mean.append(_auc_mean)
        auc_std.append(_auc_std)
        _sen, _spe, _ppv, _npv = get_sen_spe(y_pred, _y)
        sensitivity.append(_sen)
        specificity.append(_spe)
        ppv.append(_ppv)
        npv.append(_npv)
    return tpr, auc_all, auc_mean, auc_std, sensitivity, specificity, ppv, npv


def inter_auc(y_pred, y):
    inter_fpr = np.linspace(0, 1, 1000)
    fpr, tpr, thresholds = roc_curve(y, y_pred)
    inter_tpr = np.interp(inter_fpr, fpr, tpr)
    inter_tpr[0] = 0.0
    inter_tpr[-1] = 1.0
    return inter_tpr