# -*- coding: utf-8 -*-
# @Time     : 2021/4/18 1:35
# @Author   : ufy
# @Email    : antarm@outlook.com / 549147808@qq.com
# @file     : classify.py
# @info     : 本模块用于根据提取特征，对切片进行癌和非癌的分类,本模块支持二维数据的处理

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import sklearn as sk
from sklearn.ensemble import RandomForestClassifier

try:
    # 老版本
    from sklearn import joblib
except:
    # 新版本
    import joblib
import os
import tensorflow as tf
from tensorflow.keras import layers
from xgboost import XGBClassifier
from tqdm import tqdm


def find_csv_in_dir(dirpath, allfiles=None):
    file_list = os.listdir(dirpath)
    for f in file_list:
        cur_f = os.path.join(dirpath, f)
        if os.path.isdir(cur_f):
            find_csv_in_dir(cur_f, allfiles)
        elif os.path.basename(f).split('.')[-1] == 'csv':
            allfiles.append(cur_f)


def load_data_csvs(filenames, target_cols=1, drop_id=True, id_cols=1, split_rate=0.7, shuffle=True):
    '''
    :param filename: csv数据文件名列表
    :param col_targets: 标签所在列的数量，可能是多标签,要求标签必须在最后几列
    :param drop_id: 是否删除id列
    :param id_endcols: id相关信息有多少列，默认前k列为id相关信息列
    :param split_rate: 换分比例，默认训练集占70%
    :return: 划分好的数据集(x_train,y_train),(x_test,y_test)
    '''
    data = []
    for filename in tqdm(filenames):
        print('load data:', filename)
        data_file = pd.read_csv(filename)
        data_file = np.array(data_file)
        if data == []:
            data = data_file
        else:
            data = np.vstack((data, data_file))
    if drop_id:
        data = data[:, id_cols:]
        data = data.astype('float32')

    if shuffle:
        # data = data[np.random.permutation(len(data))]
        np.random.shuffle(data)  # 功能与前行代码相同
    split_point = int(len(data) * split_rate)
    train_data = data[:split_point]
    test_data = data[split_point:]

    x_train, y_train = train_data[:, :-target_cols], train_data[:, -target_cols:]
    x_test, y_test = test_data[:, :-target_cols], test_data[:, -target_cols:]

    return (x_train, y_train.astype(int)), (x_test, y_test.astype(int))


def plot_roc_curve(y_trues, y_pres, model_names=None, pos_label=1, issave=False, savename='roc.jpg',
                   analysis_filename='results_analysis.txt'):
    '''
    :param y_trues: 真实标签
    :param y_pres:  预测值
    :param pos_label: 正例标签
    :param issave: 是否保存结果
    :param savename: 保存结果图片的名
    :return: 绘制ROC曲线
    '''
    if model_names == None:
        model_names = [str(i + 1) for i in range(len(y_pres))]
    fp = open(analysis_filename, 'w', encoding='utf-8')
    for i in range(len(y_pres)):
        fpr, tpr, thresholds = sk.metrics.roc_curve(y_trues[i], y_pres[i], pos_label=pos_label)
        # print(fpr,'\n',tpr)
        roc_auc = sk.metrics.auc(fpr, tpr)
        # roc_auc = sk.metrics.roc_auc_score(y_trues[i], y_pres[i])
        print(model_names[i] + ':')
        y_pres_i = (y_pres[i] > 0.5).astype(int)
        print(sk.metrics.classification_report(y_trues[i], y_pres_i, digits=4))
        fp.write(model_names[i] + ':\n')
        fp.writelines(sk.metrics.classification_report(y_trues[i], y_pres_i, digits=4))
        fp.write('\n\n')
        plt.plot(fpr, tpr, label=model_names[i] + ' (AUC={0:.4f})'.format(roc_auc), lw=2)
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([-0.05, 1.05])
    plt.ylim([-0.05, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('ROC Curve')
    plt.legend(loc='lower right')

    if issave:
        plt.savefig(savename)  # 要在show之前保存

    plt.show()


def min_max_normal(x):
    '''
    :param x: 是一个二维数据
    :return: 一个单位化的数据
    '''
    min_col = np.min(x, axis=0)
    max_col = np.max(x, axis=0)
    return (x - min_col) / (max_col - min_col + 1e-7)


def softmax_temp(x, T=1):
    '''
    :param x: 带处理的数据
    :param T: 温度，默认温度为1
    :return: 根据温度，来进行softmax变换，并返回变换结果
    '''
    x = np.array(x)
    x_row_max = np.max(x, axis=-1)
    # x_row_max = x_row_max.reshape(list(x.shape)[:-1]+[1])
    x = x - x_row_max
    x_exp = np.exp(x / T)
    # x_exp_row_sum = x_exp.sum(axis=-1).reshape(list(x.shape)[:-1]+[1])
    x_exp_row_sum = x_exp.sum(axis=-1)
    softmax = x_exp / x_exp_row_sum
    return softmax


def MLP_softmax(x_train, y_train, model_savepath='MLP_softmax.h5', istraining=False, max_epochs=10000,
                batch_size=32, validation_split=0.3, verbose=1):
    if os.path.exists(model_savepath):
        model = tf.keras.models.load_model(model_savepath)
    else:
        istraining = True
        model = tf.keras.Sequential()
        model.add(layers.Dense(512, activation='relu', input_shape=(884,)))
        model.add(layers.Dropout(0.5))
        model.add(layers.Dense(256, activation='relu'))
        model.add(layers.Dense(128, activation='relu'))
        model.add(layers.Dense(64, activation='relu'))
        model.add(layers.Dropout(0.2))
        model.add(layers.Dense(2, activation='softmax'))

        model.summary()

    if istraining:
        print(y_train.shape)
        y_train_onehot = tf.one_hot(y_train, depth=2)
        print(y_train_onehot.shape)
        y_train_onehot = tf.reshape(y_train_onehot, (-1, 2))
        print(y_train_onehot.shape)

        model.compile(optimizer=tf.keras.optimizers.Adam(),
                      loss=tf.keras.losses.categorical_crossentropy,
                      metrics=['acc'])
        callback_list = [tf.keras.callbacks.EarlyStopping(monitor='acc', patience=8),
                         tf.keras.callbacks.ModelCheckpoint(filepath=model_savepath, monitor='val_loss',
                                                            save_best_only=True)]
        model.fit(x_train, y_train_onehot, batch_size=batch_size, epochs=max_epochs, validation_split=validation_split,
                  verbose=verbose, callbacks=callback_list)

    return model


def MLP_sigmoid(x_train, y_train, model_savepath='MLP_softmax.h5', istraining=False, max_epochs=10000,
                batch_size=32, validation_split=0.3, verbose=1):
    if os.path.exists(model_savepath):
        model = tf.keras.models.load_model(model_savepath)
    else:
        istraining = True
        model = tf.keras.Sequential()
        model.add(layers.Dense(512, activation='relu', input_shape=(884,)))
        model.add(layers.Dropout(0.5))
        model.add(layers.Dense(256, activation='relu'))
        model.add(layers.Dense(128, activation='relu'))
        model.add(layers.Dense(64, activation='relu'))
        model.add(layers.Dropout(0.2))
        model.add(layers.Dense(1, activation='sigmoid'))

        model.summary()

    if istraining:
        model.compile(optimizer=tf.keras.optimizers.Adam(),
                      loss=tf.keras.losses.binary_crossentropy,
                      metrics=['acc'])
        callback_list = [tf.keras.callbacks.EarlyStopping(monitor='acc', patience=8),
                         tf.keras.callbacks.ModelCheckpoint(filepath=model_savepath, monitor='val_loss',
                                                            save_best_only=True)]
        model.fit(x_train, y_train, batch_size=batch_size, epochs=max_epochs, validation_split=validation_split,
                  verbose=verbose, callbacks=callback_list)

    return model


def KNN(x_train, y_train, model_savepath='knn.pkl', istraining=False):
    if os.path.exists(model_savepath):
        model = joblib.load(model_savepath)
    else:
        istraining = True
        model = sk.neighbors.KNeighborsClassifier()
    if istraining:
        model.fit(x_train, y_train)
        joblib.dump(model, model_savepath)
    return model


def RandomForest(x_train, y_train, model_savepath='RandomForest.pkl', istraining=True):
    if os.path.exists(model_savepath):
        model = joblib.load(model_savepath)
    else:
        istraining = True
        model = RandomForestClassifier()
    if istraining:
        model.fit(x_train, y_train)
        joblib.dump(model, model_savepath)
    return model


def XGboost(x_train, y_train, model_savepath='xgboost.pkl', istraining=True):
    if os.path.exists(model_savepath):
        model = joblib.load(model_savepath)
    else:
        istraining = True
        model = XGBClassifier()
    if istraining:
        model.fit(x_train, y_train, verbose=True)
        # model.save_model(model_savepath)
        joblib.dump(model, model_savepath)
    return model


if __name__ == '__main__':
    allfiles = []
    find_csv_in_dir('../data/diagnosis/new_cancer/', allfiles)
    print(allfiles)
