import numpy as np

import torch 
from torch import nn
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F

import pickle
import time
import math
import os
import csv
import glob

from sklearn import metrics
from sklearn.model_selection import KFold, StratifiedKFold

FILE_MODEL_TMP = "model_tmp.pkl"

MY_RANDOM_STATE = 5
torch.manual_seed(MY_RANDOM_STATE)

# NUMBER_EPOCHS = 20
# LEARNING_RATE = 1e-4

SAMPLE_LENGTH = 200
AVGPOOL1D_KERNEL_SIZE = 4
CONV1D_KERNEL_SIZE = 3
CONV1D_FEATURE_SIZE_BLOCK1 = 32
CONV1D_FEATURE_SIZE_BLOCK2 = 64
CONV1D_FEATURE_SIZE_BLOCK3 = 128

FULLY_CONNECTED_LAYER_SIZE = 256

# 这里不一样
MODEL_DIR = '../Train/model_layer1_seed' + str(MY_RANDOM_STATE)
if not os.path.exists(MODEL_DIR):
    os.makedirs(MODEL_DIR)

def one_hot(index, dimension):
    data = np.zeros((dimension))
    data[index] = 1
    return data

def load_text_file(file_text):
    with open(file_text) as f:
        lines = f.readlines()
        my_data = [line.strip().upper() for line in lines[1::2]]
        return my_data

my_dict = {'A': 0,
        'C': 1, 
        'G': 2,
        'T':3,
        'a':0,
        'c':1,
        'g':2,
        't':3}

class EnhancerDataset(Dataset):
    # X: list Enhancer sequence (每个序列200个字符)
    # Y: list label [0, 1]; 0: 负样本, 1: 正样本
    def __init__(self, X, Y):
        self.X = X
        self.Y = Y

    def __getitem__(self, index):
        label = self.Y[index]
        sample = self.X[index]
        
        values = np.zeros((4, SAMPLE_LENGTH))
        for i in range(SAMPLE_LENGTH):
            char_idx = my_dict[sample[i]]
            values[char_idx, i] = 1 
        
        values_one_mer = self.extract_1_mer(sample)
        #values = np.concatenate((values, values_one_mer), axis=0)
        values_two_mer = self.extract_2_mer(sample)
        #values = np.concatenate((values, values_two_mer), axis=0)
        values_three_mer = self.extract_3_mer(sample)
        #values = np.concatenate((values, values_three_mer), axis=0)
        values = np.concatenate((values, values_one_mer, values_two_mer, 
                        values_three_mer), axis=0)
        
        input = torch.from_numpy(values)
        return input, label
    
    def extract_1_mer(self, sample):
        my_count = {'A': 0.0, 'C': 0.0, 'G': 0.0, 'T': 0.0}        
        values = np.zeros((1, SAMPLE_LENGTH))
        for i in range(SAMPLE_LENGTH):
            my_count[sample[i]] += 1
        
        #for one_mer in my_count:
        #    print("one mer: ", one_mer, " : ", my_count[one_mer])
        
        for i in range(SAMPLE_LENGTH):
            values[0, i] = my_count[sample[i]] / SAMPLE_LENGTH;
        
        #print("values: ", values)    
        return values
    
    def extract_2_mer(self, sample):
        my_count = {'AA': 0.0, 'AC': 0.0, 'AG': 0.0, 'AT': 0.0,
                    'CA': 0.0, 'CC': 0.0, 'CG': 0.0, 'CT': 0.0,
                    'GA': 0.0, 'GC': 0.0, 'GG': 0.0, 'GT': 0.0,
                    'TA': 0.0, 'TC': 0.0, 'TG': 0.0, 'TT': 0.0} 
        values = np.zeros((2, SAMPLE_LENGTH))
        for i in range(SAMPLE_LENGTH - 1):
            two_mer = sample[i:i+2]
            #print("two_mer: ", two_mer)
            my_count[two_mer] += 1
        
        #for two_mer in my_count:
        #    print("two mer: ", two_mer, " : ", my_count[two_mer])
        
        values = np.zeros((2, SAMPLE_LENGTH))
        for i in range(1,SAMPLE_LENGTH-1):
            two_mer_left = sample[i-1:i+1]
            two_mer_right = sample[i:i+2]
            
            values[0, i] = my_count[two_mer_left] / (SAMPLE_LENGTH - 1);
            values[1, i] = my_count[two_mer_right] / (SAMPLE_LENGTH - 1);
        
        #print("values: ", values) 
        return values
    
    def extract_3_mer(self, sample):
        my_count = {}
                                        
        for firchCh in ['A', 'C', 'G', 'T']:
            for secondCh in ['A', 'C', 'G', 'T']:
                for thirdCh in ['A', 'C', 'G', 'T']:
                    three_mer = firchCh + secondCh + thirdCh
                    my_count[three_mer] = 0.0
        for i in range(SAMPLE_LENGTH - 2):
            three_mer = sample[i:i+3]
            #print("two_mer: ", two_mer)
            my_count[three_mer] += 1
                    
        values = np.zeros((1, SAMPLE_LENGTH))
        for i in range(1,SAMPLE_LENGTH-2):
            three_mer = sample[i-1:i+2]
            values[0, i] = my_count[three_mer] / SAMPLE_LENGTH;
                    
        return values
        
    def __len__(self):
        #return 100
        return len(self.X)

# import torch.nn as nn
# import torch.nn.functional as F
# import torch

# 定义模型类
class EnhancerCnnModel(nn.Module):
    def __init__(self):
        super(EnhancerCnnModel, self).__init__()
        
        # 第一种卷积+第一种BN层
        self.c1_1 = nn.Conv1d(8, CONV1D_FEATURE_SIZE_BLOCK1, CONV1D_KERNEL_SIZE, padding=1)
        self.c1_1bn = nn.BatchNorm1d(CONV1D_FEATURE_SIZE_BLOCK1)
        
        # 第一种卷积+第一种BN层
        self.c1_2 = nn.Conv1d(CONV1D_FEATURE_SIZE_BLOCK1, CONV1D_FEATURE_SIZE_BLOCK1, 
            CONV1D_KERNEL_SIZE, padding=1)
        self.c1_2bn = nn.BatchNorm1d(CONV1D_FEATURE_SIZE_BLOCK1)
        
        # 第一种卷积+第一种BN层
        self.c1_3 = nn.Conv1d(CONV1D_FEATURE_SIZE_BLOCK1, CONV1D_FEATURE_SIZE_BLOCK1, 
            CONV1D_KERNEL_SIZE, padding=1)    
        self.c1_3bn = nn.BatchNorm1d(CONV1D_FEATURE_SIZE_BLOCK1)
        
        # 池化层
        self.p1 = nn.MaxPool1d(AVGPOOL1D_KERNEL_SIZE)
        
        # 第二种卷积+第二种BN层
        self.c2_1 = nn.Conv1d(CONV1D_FEATURE_SIZE_BLOCK1, 
            CONV1D_FEATURE_SIZE_BLOCK2, 
            CONV1D_KERNEL_SIZE, padding=1)
        self.c2_1bn = nn.BatchNorm1d(CONV1D_FEATURE_SIZE_BLOCK2)
        
        # 第二种卷积+第二种BN层
        self.c2_2 = nn.Conv1d(CONV1D_FEATURE_SIZE_BLOCK2, CONV1D_FEATURE_SIZE_BLOCK2, 
            CONV1D_KERNEL_SIZE, padding=1)
        self.c2_2bn = nn.BatchNorm1d(CONV1D_FEATURE_SIZE_BLOCK2)
        
        # 第二种卷积+第二种BN层
        self.c2_3 = nn.Conv1d(CONV1D_FEATURE_SIZE_BLOCK2, CONV1D_FEATURE_SIZE_BLOCK2, 
            CONV1D_KERNEL_SIZE, padding=1)
        self.c2_3bn = nn.BatchNorm1d(CONV1D_FEATURE_SIZE_BLOCK2)
        
        # 池化层
        self.p2 = nn.MaxPool1d(AVGPOOL1D_KERNEL_SIZE)
        
        # 全连接层
        self.fc = nn.Linear(768, FULLY_CONNECTED_LAYER_SIZE)        
        # 输出层
        self.out = nn.Linear(FULLY_CONNECTED_LAYER_SIZE, 1)
        
        # 定义损失函数
        self.criterion = nn.BCELoss()
        # 定义优化器
        # self.optimizer = torch.optim.Adam(self.parameters(), lr=LEARNING_RATE)
    
    # 前向传播函数
    def forward(self, inputs):
        batch_size = inputs.size(0)
        # 将输入的形状从(batch_size x seq_len) 转换为 (batch_size x input_size x seq_len)，以适应卷积层的输入
        #inputs = inputs.transpose(1,2)
        #print("inputs size: ", inputs.size())        
        output = F.relu(self.c1_1bn(self.c1_1(inputs)))
        output = F.relu(self.c1_2bn(self.c1_2(output)))
        output = F.relu(self.c1_3bn(self.c1_3(output)))
        output = self.p1(output)
        #print("After p1: ", output.shape) 
        
        output = F.relu(self.c2_1bn(self.c2_1(output)))
        output = F.relu(self.c2_2bn(self.c2_2(output)))
        output = F.relu(self.c2_3bn(self.c2_3(output)))
        output = self.p2(output)
        #print("After p2: ", output.shape)
        
        output = output.view(batch_size, -1)
        #print("Reshape : ", output.shape)
        
        output = F.relu(self.fc(output))
        #print("After FC layer: ", output.shape)  
        
        output = torch.sigmoid(self.out(output))
        #print("Final output (After sigmoid): ", output.shape)
        #print("Final output: ", output)
        
        return output 

# # 训练一个 epoch 的函数
# def train_one_epoch(model, train_loader, learning_rate):
#     # 设置模型为训练模式
#     model.train()
#     # 更新学习率
#     for param_group in model.optimizer.param_groups:
#         param_group['lr'] = learning_rate
    
#     epoch_loss_train = 0.0
#     nb_train = 0
#     # 遍历训练集中的每一批数据
#     for i, data in enumerate(train_loader, 0):
#         # 获取输入和标签
#         inputs, labels = data   
#         #print("\ninputs: ", inputs.shape)
#         #print("labels: ", labels.shape)
        
#         inputs_length = inputs.size()[0]
#         nb_train += inputs_length
        
#         inputs = inputs.float()
#         labels = labels.float()
        
#         #######################
#         #  使用 GPU 进行计算  #
#         #######################
#         if torch.cuda.is_available():
#             inputs = Variable(inputs.cuda())
#             labels = Variable(labels.cuda())            
#         else:
#             inputs = Variable(inputs)
#             labels = Variable(labels)
        
#         # 清除参数的梯度
#         model.optimizer.zero_grad()
#         # 进行前向传播
#         outputs = model(inputs)
        
#         # 计算损失函数
#         loss = model.criterion(outputs, labels)
#         #print("loss: ", loss)
#         # 反向传播
#         loss.backward()
#         # 更新参数
#         model.optimizer.step()
        
#         # 累加损失
#         epoch_loss_train = epoch_loss_train + loss.item() * inputs_length
    
#     print("nb_train: ", nb_train)
#     # 计算平均损失
#     epoch_loss_train_avg = epoch_loss_train / nb_train
#     return epoch_loss_train_avg

def evaluate(file_model, loader):
    # 定义一个 EnhancerCnnModel 的实例 model
    model = EnhancerCnnModel()
    
    # 如果 CUDA 可用，则将 model 移动到 GPU 上
    if torch.cuda.is_available():
        model.cuda()
    
    # 加载预训练的模型权重
    model.load_state_dict(torch.load(file_model))
    
    # 设置模型为评估模式
    model.eval()    
    
    epoch_loss = 0.0
    nb_samples = 0
    
    arr_labels = []
    arr_labels_hyp = []
    arr_prob = []
    
    # 遍历数据集
    for i, data in enumerate(loader, 0):
        # 获取输入和标签
        inputs, labels = data
        
        # inputs 的长度
        inputs_length = inputs.size()[0]
        nb_samples += inputs_length
        
        # 将标签添加到 arr_labels 列表中
        arr_labels += labels.squeeze(1).data.cpu().numpy().tolist()

        # 将 inputs 和 labels 转换为 float 类型
        inputs = inputs.float()
        labels = labels.float()
        
        # 如果 CUDA 可用，则将 inputs 和 labels 移动到 GPU 上
        if torch.cuda.is_available():
            inputs = Variable(inputs.cuda())
            labels = Variable(labels.cuda())
        else:
            inputs = Variable(inputs)
            labels = Variable(labels)
        
        # 前向传递
        outputs = model(inputs)
        
        # 计算损失使用模型 model 的损失函数来计算该模型在给定输出 outputs 和
        # 标签 labels 下的损失
        loss = model.criterion(outputs, labels)
        
        # 累加 epoch_loss
        epoch_loss = epoch_loss + loss.item() * inputs_length
        
        # 将模型的预测结果添加到 arr_prob 列表中
        arr_prob += outputs.squeeze(1).data.cpu().numpy().tolist()
    
    # 输出样本总数和平均损失
    print("nb_samples: ", nb_samples)
    epoch_loss_avg = epoch_loss / nb_samples
    print("epoch loss avg: ", epoch_loss_avg)
        
    # 输出预测概率和真实标签的列表长度
    print("arr_prob: ", len(arr_prob))
    print("arr_labels: ", len(arr_labels))
    
    # 计算 AUC 值
    auc = metrics.roc_auc_score(arr_labels, arr_prob)
    print("auc: ", auc)
    
    # 将预测概率转换为二分类标签，阈值为 0.5
    arr_labels_hyp = [int(prob > 0.5) for prob in arr_prob]
    arr_labels = [int(label) for label in arr_labels]
    
    # 计算各种分类指标
    acc, confusion_matrix, sensitivity, specificity, mcc = calculate_confusion_matrix(arr_labels, arr_labels_hyp)
    
    # 将结果存储到字典中
    result = {'epoch_loss_avg': epoch_loss_avg,
              'acc' : acc,
              'confusion_matrix' : confusion_matrix,
              'sensitivity' : sensitivity,
              'specificity' : specificity,
              'mcc' : mcc,
              'auc' : auc,
              'arr_prob': arr_prob,
              'arr_labels': arr_labels,
              'arr_labels_hyp':arr_labels_hyp
              }
    
    # 输出 accuracy
    print("acc: ", acc)
    
    # 返回结果字典
    return result

def calculate_confusion_matrix(arr_labels, arr_labels_hyp):
    # 初始化变量
    corrects = 0  # 正确分类的样本数
    confusion_matrix = np.zeros((2, 2))  # 混淆矩阵

    # 遍历所有样本
    for i in range(len(arr_labels)):
        # 更新混淆矩阵
        confusion_matrix[arr_labels_hyp[i]][arr_labels[i]] += 1

        # 统计正确分类的样本数
        if arr_labels[i] == arr_labels_hyp[i]:
            corrects = corrects + 1

    # 计算分类指标
    acc = corrects * 1.0 / len(arr_labels)  # 准确率
    specificity = confusion_matrix[0][0] / (confusion_matrix[0][0] + confusion_matrix[1][0])  # 特异度
    sensitivity = confusion_matrix[1][1] / (confusion_matrix[0][1] + confusion_matrix[1][1])  # 灵敏度
    tp = confusion_matrix[1][1]  # 真正例数
    tn = confusion_matrix[0][0]  # 真反例数
    fp = confusion_matrix[1][0]  # 假正例数
    fn = confusion_matrix[0][1]  # 假反例数
    mcc = (tp * tn - fp * fn ) / math.sqrt((tp + fp)*(tp + fn)*(tn + fp)*(tn + fn))  # MCC值

    # 返回分类指标
    return acc, confusion_matrix, sensitivity, specificity, mcc

# 将测试集的数据和标签分别存储在 dataset 字典的 "data_test" 和 "label_test" 键下
# dataset = {"data_test" : testset["data"], "label_test" : testset["label"]}

# 定义一个函数 check_dataset，用于检查数据集是否符合要求
def check_dataset(dataset):
    print("\n==> Checking dataset")  # 打印提示信息
    
    # 检查 data_test 数据是否符合要求
    data_test = dataset["data_test"]  # 获取 data_test 数据
    nb_error_samples = 0  # 初始化错误样本数量为 0
    for sample in data_test:  # 遍历每个样本
        if len(sample) != 200:  # 如果样本长度不为 200
            nb_error_samples += 1  # 错误样本数量加 1
    if nb_error_samples > 0:  # 如果存在错误样本
        print("data_teset error: ", nb_error_samples)  # 打印错误信息和错误样本数量
    else:  # 如果不存在错误样本
        print("data_test : OK!")  # 打印测试通过的信息

def read_test_file(filename):
    text_file = open(filename)  # 打开文件
    lines = text_file.readlines()  # 读取文件中所有行的内容到列表 lines 中
    my_data =[]  # 初始化空列表 my_data，用于存储处理后的数据
    
    for i in range(0, len(lines), 2):  # 每次跳两行，遍历每个样本
        # 读取第一行 DNA 序列名称并去除多余空格和换行符
        seq_name = lines[i].strip()
        # 拼接第二行 DNA 序列，并使用 strip() 方法去除多余空格和换行符，再转换为大写字母形式
        text = lines[i+1].strip().upper()
        my_data.append(text)  # 将处理后的数据添加到 my_data 列表中
    
    return my_data  # 返回处理后的数据列表 my_data

def prepare_test_data():
    print("\n ==> Loading test set")  # 打印提示信息
    
    # 加载增强子数据集
    data_strong = read_test_file('./Test/enhancer_200_process_rename_test.fasta')  # 调用函数 read_test_file 读取文件
    print("data_strong: ", len(data_strong))  # 打印加载的数据集大小
    #print("data_strong: ", data_strong)

    # data_weak = read_test_file('test_weak_enhancer.txt')
    # print("data_weak: ", len(data_weak))
    
    data_enhancer = data_strong # + data_weak
    print("data_enhancer: ", len(data_enhancer))
    
    # 加载非增强子数据集
    data_non_enhancer = read_test_file('./Test/noenhancer_200_process_rename_test.fasta')  # 调用函数 read_test_file 读取文件
    print("data_non_enhancer: ", len(data_non_enhancer))  # 打印加载的数据集大小
    
    # 构建测试集
    label_enhancer = np.ones((len(data_strong),1))  # 增强子数据集对应标签为 1
    label_non_enhancer = np.zeros((len(data_non_enhancer), 1))  # 非增强子数据集对应标签为 0
    
    data = np.concatenate((data_strong, data_non_enhancer))     # 将增强子数据集和非增强子数据集拼接在一起，构成测试集
    label = np.concatenate((label_enhancer, label_non_enhancer))
    
    testset = {"data" : data, "label" : label}  # 将测试集数据和标签存储在字典 testset 中
    
    return testset  # 返回测试集数据和标签字典 testset

# 将测试数据集转化为 EnhancerDataset 类型，并使用 DataLoader 加载数据集，
# 每次迭代使用 32 个样本，关闭数据集随机重排，使用 4 个工作进程
def testing():
    test_dataset = EnhancerDataset(testset["data"], testset["label"])
    test_loader = DataLoader(dataset=test_dataset, batch_size=32,                              
                              shuffle=False, num_workers=4)
    
    # 创建一个 CSV 文件，用于保存测试结果
    with open(testresult_fn, mode='w') as outfile:
        outfile = csv.writer(outfile, delimiter=',')
        outfile.writerow(['model_fn', 'Accuracy score', 'AUC score', 'Sensitivity', 'Specificity'])
        
    # 获取所有模型文件的文件名，并按字典序排序
    list_model_fn = sorted(glob.glob(MODEL_DIR+"/enhancer_*.pkl"))
    #print(list_model_fn)
    y_prob_mtx = []  # 初始化 y_prob_mtx，用于存储每个模型所预测的概率矩阵
    
    for model_fn in list_model_fn:  # 遍历每个模型
        print(model_fn)
        result = evaluate(model_fn, test_loader)  # 调前 evaluate 函数对该模型进行测试，并返回测试结果
        #print(result['arr_prob'])
        y_prob_mtx.append(result['arr_prob'])  # 将该模型预测的概率矩阵添加到 y_prob_mtx 列表中
        
        # 将该模型的测试结果写入 CSV 文件中
        with open(testresult_fn, mode='a') as outfile:
            outfile = csv.writer(outfile, delimiter=',')
            outfile.writerow([model_fn, result['acc'], result['auc'], 
                result['sensitivity'], result['specificity']])
    
    # 将 y_prob_mtx 转换为 numpy 的 ndarray 类型
    y_prob_mtx = np.array(y_prob_mtx)
    print("y_prob_mtx: ", y_prob_mtx.shape)
    #print("y_prob_mtx: ", y_prob_mtx)
    
    # 对 y_prob_mtx 进行集成，计算每个样本预测为正例的概率平均值，并将概率值大于 0.5 的样本预测为正例，其余样本预测为负例
    y_prob_ensemble = [np.mean(y_prob_mtx[:,col]) for col in range(np.size(y_prob_mtx, 1))] 
    y_pred_ensemble = [np.float(each > 0.5) for each in y_prob_ensemble]
    
    y_true = testset["label"]  # 获取测试集的真实标签
    
    # 计算集成模型的 ROC AUC、Accuracy、Sensitivity 和 Specificity 分数
    auc_score_ensemble = metrics.roc_auc_score(y_true, y_prob_ensemble)
    accuracy_score_ensemble = metrics.accuracy_score(y_true, y_pred_ensemble)
    
    cm = metrics.confusion_matrix(y_true, y_pred_ensemble)  # 计算混淆矩阵
    specificity_ensemble = cm[0,0]/(cm[0,0] + cm[0,1])  # 计算 Specificity
    sensitivity_ensemble = cm[1,1]/(cm[1,1] + cm[1,0])  # 计算 Sensitivity

    # 将集成模型的 ROC AUC、Accuracy、Sensitivity 和 Specificity 分数打印出来
    print("Accuracy score (Testing Set) = ", accuracy_score_ensemble)
    print("ROC AUC score  (Testing Set) = ", auc_score_ensemble)
    print("Sensitivity    (Testing Set) = ", sensitivity_ensemble)
    print("Specificity    (Testing Set) = ", specificity_ensemble)
    
    # 将集成模型的 ROC AUC、Accuracy、Sensitivity 和 Specificity 分数写入 CSV 文件中
    with open(testresult_fn, mode='a') as outfile:
        outfile = csv.writer(outfile, delimiter=',')
        outfile.writerow(["ensemble", accuracy_score_ensemble, auc_score_ensemble, sensitivity_ensemble, specificity_ensemble])

# if __name__== "__main__":    
#     testresult_fn = MODEL_DIR + "/test_result.csv"
    
#     testset = prepare_test_data()    
    
#     testing()