import numpy as np
import gradio as gr
from transformers import AutoTokenizer,AutoModelForSequenceClassification
from transformers import set_seed
from torch.utils.data import Dataset,DataLoader
import torch
import torch.nn as nn
import numpy as np
import warnings
warnings.filterwarnings('ignore')
set_seed(4)  
device = "cpu"
model_checkpoint = "facebook/esm2_t6_8M_UR50D"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

def AMP(file):
    test_sequences = file
    class MyDataset(Dataset):
            def __init__(self,dict_data) -> None:
                super(MyDataset,self).__init__()
                self.data=dict_data
            def __getitem__(self, index):
                return self.data['text'][index]
            def __len__(self):
                return len(self.data['text'])
    def collate_fn(batch):
        max_len = 30
        pt_batch=tokenizer([b[0] for b in batch], max_length=max_len, padding="max_length",truncation=True, return_tensors='pt')
        return {'input_ids':pt_batch['input_ids'],
                'attention_mask':pt_batch['attention_mask']}
    test_dict = {"text":test_sequences}
    test_data=MyDataset(test_dict)
    test_dataloader=DataLoader(test_data,batch_size=len(test_sequences), collate_fn=collate_fn)
    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)
    model = MyModel()
    model.load_state_dict(torch.load("best_model.pth"))
    model = model.to(device)
    model.eval()
    out = []
    out_probability = []
    with torch.no_grad():
        for index, batch in enumerate(test_dataloader):
            batchs = {k: v for k, v in batch.items()}
            predict = model(batchs)
            # tsne_plot(output_feature,batchs)
            out_probability.extend(np.max(np.array(predict.cpu()),axis=1).tolist())
            # out_probability.extend(np.array(predict.cpu())[:, -1].tolist())
            test_argmax = np.argmax(predict.cpu(), axis=1).tolist()
            out.extend(test_argmax)

    return out, out_probability

iface = gr.Interface(fn=AMP,
                     inputs="text"(label="input sequence"), 
                    outputs= ["text"(label="class"), "text"(label="probability")])
iface.launch()