import torch
from DeepMFPP.predictor import predict
from DeepMFPP.data_helper import Data2EqlTensor
import gradio as gr
from Bio import SeqIO
device = "cuda" if torch.cuda.is_available() else "cpu"

def MFPs_classifier(file:str,threshold:float=0.5,top_k=0):
    data = []
    for record in SeqIO.parse(file.name, 'fasta'): # notebook测试时 for record in SeqIO.parse(file, 'fasta')
        data.append((record.id, str(record.seq))) 
    seqs,_ = Data2EqlTensor(data,51,AminoAcid_vocab='esm')
    model_weight_path = './weight/DeepMFPP-Best.pth'
    MFPs_pred = predict(seqs=seqs, data=data, model_path=model_weight_path, threshold=threshold,top_k=top_k,device=device)
    
    return MFPs_pred
"""
[['peptide_1', 'ARRRRCSDRFRNCPADEALCGRRRR'],
 {'ADP': 0.5739, 'BBP': 0.5358},
 ['peptide_2', 'FFHHIFRGIVHVGKTIHKLVTGT'],
 {'ADP': 0.5633, 'AIP': 0.5371},
 ['peptide_3', 'GLRKRLRKFRNKIKEKLKKIGQKIQGFVPKLAPRTDY'],
 {'AAP': 0.5418, 'ADP': 0.5346},
 ['peptide_4', 'FLGALWNVAKSVF'],
 {'ADP': 0.5684, 'BBP': 0.5619},
 ['peptide_5', 'KIKSCYYLPCFVTS'],
 {'ADP': 0.5862, 'BBP': 0.5636}]
"""

with gr.Blocks() as demo:
    gr.Markdown(" ## DeepMFPP")
    gr.Markdown("In this study, we developed a multi-functional peptides(MFPs) prediction model. The model was used to predict the \
                functional labels (21 categories of MFPs involved in our study) that peptide sequences have.")
    
    with gr.Tab("MFPs Prediction Model"):
        with gr.Row():
            with gr.Column(scale=2):
                input_fasta = gr.File()
            with gr.Column(scale=2):
                cutoff = gr.Slider(0, 1, step=0.01, value=0.5, interactive=True, label="Threshold")
                top_k = gr.Slider(0, 21, step=1, value=0, interactive=True, label="top_k")
                
                gr.Markdown("### Note")
                gr.Markdown("- Limit the number of input sequences to less than 128.")
                gr.Markdown("- If top_k is set to 0, the combination of all probability values (processed by the sigmoid function) \
                            larger than the threshold is returned by default. Otherwise, the specified top k predictions are returned")
                gr.Markdown("- The file should be the Fasta format.")
                gr.Markdown("- If the length of the sequence is less than 50, use 0 to fill to 50; if the length is over 50, \
                            We used only the first 25 amino acids of each N-terminal and C-terminal of the sequence for prediction.")
                image_button_MFPs = gr.Button("Submit")
        with gr.Column():
            # gr.Markdown(" ### Flip text or image files using this demo.")
            gr.Markdown("Note: The predicted probabilities are processed by sigmoid function")
            MFPs_output = gr.DataFrame(
                headers=["Sequence Id", "Sequence", "Categories and probabilities"],
                datatype=["str", "str", "str"],)
            
            
    image_button_MFPs.click(MFPs_classifier, inputs=[input_fasta, cutoff,top_k], outputs=MFPs_output)
    
    with gr.Accordion("Citation"):
        gr.Markdown("- GitHub: https://github.com/leonern/DeepMFPP-GitHub")

    with gr.Accordion("License"):
        gr.Markdown("- Released under the [MIT license](https://github.com/leonern/DeepMFPP-GitHub/blob/main/LICENSE). ")

    with gr.Accordion("Contact"):
        gr.Markdown("- If you have any questions, please file a Github issue or contact me at 107552103310@stu.xju.edu.cn")


demo.queue(4)
demo.launch() #share=True