import numpy as np
import pandas as pd 
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder, LabelEncoder
from gensim.models import Word2Vec
from sklearn.model_selection import KFold
import numpy as np
from tqdm import tqdm
import os,logging,pickle,random
from sklearn import metrics as skmetrics
import warnings
import seaborn as sns
import matplotlib.pyplot as plt
# from torch import nn as nn
# from torch.nn import functional as F
# import torch,time,os
warnings.filterwarnings("ignore")
from tensorflow.keras.models import Model
from tensorflow.keras.regularizers import l2
from tensorflow.keras.constraints import max_norm
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras.layers import Input, Dense, Dropout, Flatten, Activation
from tensorflow.keras.layers import Conv1D, Add, MaxPooling1D, BatchNormalization
from tensorflow.keras.layers import Embedding, Bidirectional, GlobalMaxPooling1D
from tensorflow.compat.v1.keras.layers import CuDNNLSTM
import tensorflow as tf
gpus = tf.config.list_physical_devices(device_type='GPU')
cpus = tf.config.list_physical_devices(device_type='CPU')
print(gpus, cpus)
data = pd.read_csv("../input/protein-dataset/train.csv")
df_test = pd.read_csv("../input/protein-dataset/test_withid.csv")
data["sequence"] = data["xulie"].apply(lambda x : x.upper())
df_test["sequence"] = df_test["xulie"].apply(lambda x : x.upper())
data["seqLen"] = data["xulie"].apply(lambda x : len(x))
df_test["seqLen"] = df_test["xulie"].apply(lambda x : len(x))
df_test.head()
def integer_encoding(data):
    """
    - Encodes code sequence to integer values.
    - 20 common amino acids are taken into consideration
    and rest 4 are categorized as 0.
    """

    encode_list = []
    for row in data['sequence'].values:
        row_encode = []
        for code in row:
            row_encode.append(char_dict.get(code, 0))
        encode_list.append(np.array(row_encode))

    return encode_list
itemCounter = {}
for seq in data["sequence"]:
    for i in seq:
        itemCounter[i] = itemCounter.get(i,0)+1
codes = list(itemCounter.keys())
codes.sort()
codes
def create_dict(codes):
    char_dict = {}
    for index, val in enumerate(codes):
        char_dict[val] = index+1

    return char_dict

char_dict = create_dict(codes)

print(char_dict)
print("Dict Length:", len(char_dict))
max_length = 300
data_encode = integer_encoding(data)
data_pad = pad_sequences(data_encode, maxlen=max_length, padding='post', truncating='post')
data_ohe = to_categorical(data_pad, 23)
data_ohe.shape
"""
## 序列长度分布
"""
def plot_seq_count(df, data_name):
    sns.distplot(df['seqLen'].values)
    plt.title(f'Sequence char count: {data_name}')
    plt.grid(True)
plt.subplot(1, 2, 1)
plot_seq_count(data, 'Train')

plt.subplot(1, 2, 2)
plot_seq_count(df_test, 'Test')

plt.subplots_adjust(right=3.0)
plt.show()
"""
## 模型1 双向LSTM
"""
# 双向LSTM
def BiLSTM():  
    x_input = Input(shape=(max_length,))
    emb = Embedding(23, 128, input_length=max_length)(x_input)
    bi_rnn = Bidirectional(CuDNNLSTM(64, kernel_regularizer=l2(0.01), recurrent_regularizer=l2(0.01), bias_regularizer=l2(0.01)))(emb)
    x = Dropout(0.3)(bi_rnn)

    # softmax 
    x_output = Dense(245, activation='softmax')(x)

    model1 = Model(inputs=x_input, outputs=x_output)
    return model1
    

#model1.summary()
"""
## 模型2 ProtCNN
"""
def residual_block(data, filters, d_rate):
    """
    _data: input
    _filters: convolution filters
    _d_rate: dilation rate
    """

    shortcut = data

    bn1 = BatchNormalization()(data)
    act1 = Activation('relu')(bn1)
    conv1 = Conv1D(filters, 1, dilation_rate=d_rate, padding='same', kernel_regularizer=l2(0.001))(act1)

    #bottleneck convolution
    bn2 = BatchNormalization()(conv1)
    act2 = Activation('relu')(bn2)
    conv2 = Conv1D(filters, 3, padding='same', kernel_regularizer=l2(0.001))(act2)

    #skip connection
    x = Add()([conv2, shortcut])

    return x

# model

def ProtCNN():    
    #Input(shape=(100, 21))
    x_input = Input(shape=(max_length, 23))

    #initial conv
    conv = Conv1D(128, 1, padding='same')(x_input) 

    # per-residue representation
    res1 = residual_block(conv, 128, 2)
    res2 = residual_block(res1, 128, 3)

    x = MaxPooling1D(3)(res2)
    x = Dropout(0.5)(x)

    # softmax classifier
    x = Flatten()(x)
    x_output = Dense(245, activation='softmax', kernel_regularizer=l2(0.0001))(x)

    model2 = Model(inputs=x_input, outputs=x_output)
    return model2
#odel2.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

#odel2.summary()
load_model(self.saved_model, custom_objects={"f1_score ": f1_score })
from keras.models import load_model
def generate_submission(model_name):
    model = load_model(f'./model_{model_name}/model_{model_name}.h5', custom_objects={"f1": f1 })
    test_encode = integer_encoding(df_test)
    test_pad = pad_sequences(test_encode, maxlen=max_length, padding='post', truncating='post')
    test_ohe = to_categorical(test_pad, 23)
    if "LSTM" in model_name:
        y_pred = model.predict(test_pad,batch_size = 256)
    else:
        y_pred = model.predict(test_ohe,batch_size = 256)
    a = np.argmax(y_pred, axis=1)
    final_pred = le.inverse_transform(a.reshape(-1, 1))
    ans = []
    for i in list(final_pred):
        ans.append(i.strip()[0] + "." + i.strip()[1:])
    df_test["category_id"] = ans
    output = df_test[["sample_id", "category_id"]]
    output.to_csv(f"submission_{model_name}.csv", index = None)
le = LabelEncoder()
le.fit(data["label"])
y = le.transform(data["label"])
y = to_categorical(y)
#y[trainIdList].shape
"""
## 模型1训练（未进行模型保存和测试集推断）
"""
import os
def create_dir_not_exist(path):
    if not os.path.exists(path):
        os.mkdir(path)
create_dir_not_exist("./model_BiLSTM")
create_dir_not_exist("./model_ProtCNN")
from tensorflow.keras import backend as K

def f1(y_true, y_pred):
    def recall(y_true, y_pred):
        """Recall metric.

        Only computes a batch-wise average of recall.

        Computes the recall, a metric for multi-label classification of
        how many relevant items are selected.
        """
        true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
        possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))
        recall = true_positives / (possible_positives + K.epsilon())
        return recall

    def precision(y_true, y_pred):
        """Precision metric.

        Only computes a batch-wise average of precision.

        Computes the precision, a metric for multi-label classification of
        how many selected items are relevant.
        """
        true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
        predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))
        precision = true_positives / (predicted_positives + K.epsilon())
        return precision
    precision = precision(y_true, y_pred)
    recall = recall(y_true, y_pred)
    return 2*((precision*recall)/(precision+recall+K.epsilon()))
model_name = "BiLSTM"
test_encode = integer_encoding(df_test)
test_pad = pad_sequences(test_encode, maxlen=max_length, padding='post', truncating='post')
test_ohe = to_categorical(test_pad, 23)
if "LSTM" in model_name:
    y_pred = model.predict(test_pad,batch_size = 256)
else:
    y_pred = model.predict(test_ohe,batch_size = 256)
a = np.argmax(y_pred, axis=1)
final_pred = le.inverse_transform(a.reshape(-1, 1))
ans = []
for i in list(final_pred):
    ans.append(i.strip()[0] + "." + i.strip()[1:])
df_test["category_id"] = ans
output = df_test[["sample_id", "category_id"]]
output.to_csv(f"submission_{model_name}.csv", index = None)
kf = KFold(n_splits=30)
histories = []
le = LabelEncoder()
y = le.fit_transform(data["label"])
y = to_categorical(y)
max_length = 300 # 指定序列长度
model = BiLSTM()
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy', f1])
for fold_id, (trainIdList, validIdList) in enumerate(kf.split(data)):
    print(f"==============Fold{fold_id}==================")
    print("TRAIN:", len(trainIdList), "TEST:", len(validIdList))
    
    df_train = data.iloc[trainIdList]
    df_val = data.iloc[validIdList]
    df_train['seq_char_count']= df_train['sequence'].apply(lambda x: len(x))
    df_val['seq_char_count']= df_val['sequence'].apply(lambda x: len(x))
    
    train_encode = integer_encoding(df_train)
    val_encode = integer_encoding(df_val) 

    train_pad = pad_sequences(train_encode, maxlen=max_length, padding='post', truncating='post')
    val_pad = pad_sequences(val_encode, maxlen=max_length, padding='post', truncating='post')
    
    train_ohe = to_categorical(train_pad, 23)
    val_ohe = to_categorical(val_pad, 23)
    
    y_train = y[trainIdList]
    y_val = y[validIdList]

    es = EarlyStopping(monitor='val_loss', patience=50, verbose=1)
    history1 = model.fit(
        train_pad, y_train,
        epochs=300, batch_size=256,
        validation_data=(val_pad, y_val),
        callbacks=[es]
        )
    histories.append(history1)
model.save(f'./model_BiLSTM/model_BiLSTM.h5')
generate_submission("BiLSTM")
kf = KFold(n_splits=30)
histories = []
le = LabelEncoder()
y = le.fit_transform(data["label"])
y = to_categorical(y)
max_length = 300 # 指定序列长度
model = ProtCNN()
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy', f1])
for fold_id, (trainIdList, validIdList) in enumerate(kf.split(data)):
    print(f"==============Fold{fold_id}==================")
    print("TRAIN:", len(trainIdList), "TEST:", len(validIdList))
    
    
    df_train = data.iloc[trainIdList]
    df_val = data.iloc[validIdList]
    df_train['seq_char_count']= df_train['sequence'].apply(lambda x: len(x))
    df_val['seq_char_count']= df_val['sequence'].apply(lambda x: len(x))
    
    train_encode = integer_encoding(df_train)
    val_encode = integer_encoding(df_val) 

    train_pad = pad_sequences(train_encode, maxlen=max_length, padding='post', truncating='post')
    val_pad = pad_sequences(val_encode, maxlen=max_length, padding='post', truncating='post')
    
    train_ohe = to_categorical(train_pad, 23)
    val_ohe = to_categorical(val_pad, 23)
    
    y_train = y[trainIdList]
    y_val = y[validIdList]

    es = EarlyStopping(monitor='val_loss', patience=5, verbose=1)
    history1 = model.fit(
        train_ohe, y_train,
        epochs=20, batch_size=256,
        validation_data=(val_ohe, y_val),
        callbacks=[es]
        )
    histories.append(history1)
model.save(f'./model_ProtCNN/model_ProtCNN.h5')
generate_submission("ProtCNN")
# model = load_model(f'./model_BiLSTM/model_ProtCNN.h5')
# test_encode = integer_encoding(df_test)
# test_pad = pad_sequences(test_encode, maxlen=max_length, padding='post', truncating='post')
# test_ohe = to_categorical(test_pad, 23)
# y_pred = model.predict(test_ohe,batch_size = 256)
# a = np.argmax(y_pred, axis=1)
# final_pred = le.inverse_transform(a.reshape(-1, 1))
# ans = []
# for i in list(final_pred):
#     ans.append(i.strip()[0] + "." + i.strip()[1:])
# df_test["category_id"] = ans
# output = df_test[["sample_id", "category_id"]]
# output.to_csv("submission_pro.csv", index = None)
"""
## 模型2 ProtCNN
"""
# def residual_block(data, filters, d_rate):
#     """
#     _data: input
#     _filters: convolution filters
#     _d_rate: dilation rate
#     """

#     shortcut = data

#     bn1 = BatchNormalization()(data)
#     act1 = Activation('relu')(bn1)
#     conv1 = Conv1D(filters, 1, dilation_rate=d_rate, padding='same', kernel_regularizer=l2(0.001))(act1)

#     #bottleneck convolution
#     bn2 = BatchNormalization()(conv1)
#     act2 = Activation('relu')(bn2)
#     conv2 = Conv1D(filters, 3, padding='same', kernel_regularizer=l2(0.001))(act2)

#     #skip connection
#     x = Add()([conv2, shortcut])

#     return x

# # model

# def ProtCNN():    
#     #Input(shape=(100, 21))
#     x_input = Input(shape=(max_length, 23))

#     #initial conv
#     conv = Conv1D(128, 1, padding='same')(x_input) 

#     # per-residue representation
#     res1 = residual_block(conv, 128, 2)
#     res2 = residual_block(res1, 128, 3)

#     x = MaxPooling1D(3)(res2)
#     x = Dropout(0.5)(x)

#     # softmax classifier
#     x = Flatten()(x)
#     x_output = Dense(245, activation='softmax', kernel_regularizer=l2(0.0001))(x)

#     model2 = Model(inputs=x_input, outputs=x_output)
#     return model2
# #odel2.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# #odel2.summary()