import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import pandas as pd
from transformers import set_seed
import torch
import torch.nn as nn
from collections import OrderedDict
import warnings
import gradio as gr
from tqdm import tqdm

warnings.filterwarnings('ignore')
set_seed(4)  
device = "cpu"
model_checkpoint = "esm2_t30_150M_UR50D"

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.bert = AutoModelForSequenceClassification.from_pretrained(model_checkpoint,num_labels=320)
        self.bn1 = nn.BatchNorm1d(256)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(64)
        self.relu = nn.ReLU()
        self.fc1 = nn.Linear(320,256)
        self.fc2 = nn.Linear(256,128)
        self.fc3 = nn.Linear(128,64)
        self.output_layer = nn.Linear(64,2)
        self.dropout = nn.Dropout(0)

    def forward(self,x):
        with torch.no_grad():
            bert_output = self.bert(input_ids=x['input_ids'].to(device),attention_mask=x['attention_mask'].to(device)) 
        output_feature = self.dropout(bert_output["logits"])
        output_feature = self.relu(self.bn1(self.fc1(output_feature)))
        output_feature = self.relu(self.bn2(self.fc2(output_feature)))
        output_feature = self.relu(self.bn3(self.fc3(output_feature)))
        output_feature = self.output_layer(output_feature)
        return torch.softmax(output_feature,dim=1)

def Kmers_funct(seq,num):
    for i in range(len(seq)):
        a = seq[i]
        l = []
        for index in range(len(a)):
            t = a[index:index + num]
            if (len(t)) == num:
                l.append(t)
    return l

def ACE(file):
    test_seq = file
    all = []
    seq_len = len(test_seq)
    if seq_len > 30:
        for j in range(2, 11):
            X = Kmers_funct([test_seq], j)
            all.extend(X)
    else:
        all.append(test_seq)
    model = MyModel()
    model.load_state_dict(torch.load("best_model.pth", map_location=torch.device('cpu')), strict=False)
    model = model.to(device)
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
    max_len = 30

    seq_all = []
    output_all = []
    probability_all = []
    for seq in tqdm(all):
        test_data = tokenizer(seq, max_length=max_len, padding="max_length",truncation=True, return_tensors='pt')
        out_probability = []
        with torch.no_grad():
            predict = model(test_data)
            out_probability.extend(np.max(np.array(predict.cpu()),axis=1).tolist())
            test_argmax = np.argmax(predict.cpu(), axis=1).tolist()
            id2str = {0:"non-ACE", 1:"ACE"}
            output = id2str[test_argmax[0]]
            probability = out_probability[0]
            seq_all.append(seq)
            output_all.append(output)
            probability_all.append(probability)

    summary = OrderedDict()
    summary['Seq'] = seq_all
    summary['Class'] = output_all
    summary['Probability'] = probability_all
    summary_df = pd.DataFrame(summary)
    summary_df.to_csv('output.csv', index=False)
    
    if seq_len > 30:
        out_text = "None"
        out_prob = "None"

    else:
        out_text = output
        out_prob = probability
        

    return 'output.csv', out_text, out_prob

with open("ACE.md", "r") as f:
    description = f.read()
iface = gr.Interface(fn=ACE,
                     title="🏹DeepACE",
                     inputs=gr.Textbox(show_label=False, placeholder="Enter peptide or protein", lines=4), 
                     outputs= ["file",gr.Textbox(show_label=False, placeholder="class", lines=1),gr.Textbox(show_label=False, placeholder="probability", lines=1)],
                     description=description)
iface.launch()