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

def mus_classifier(file,threshold):
    data = []
    for record in SeqIO.parse(file.name, 'fasta'):
        data.append((record.id, str(record.seq)))
    seqs,_ = Data2EqlTensor(data,40)
    mus_peptide_pred = predict(seqs,data, './weight-Mus/4.pth', threshold, device)
    return mus_peptide_pred

def homo_classifier(file,threshold):
    data = []
    for record in SeqIO.parse(file.name, 'fasta'):
        data.append((record.id, str(record.seq)))
    seqs,_ = Data2EqlTensor(data,40)
    homo_peptide_pred = predict(seqs, data, './weight-Homo/4.pth', threshold, device)
    return homo_peptide_pred
# {peptide_id:[Type:int(1->peptide,0->non-peptide)]}

with gr.Blocks() as demo:
    gr.Markdown(" ## DeepPD")
    gr.Markdown("In this study, we developed a peptide detectability prediction model. The model was used to predict the probability that an amino acid sequence is a peptide.")
    
    with gr.Tab("Prediction Model(Homo sapiens)"):
        with gr.Row():
            with gr.Column(scale=2):
                input_fasta_homo = gr.File()
            with gr.Column(scale=2):
                homo_cutoff = gr.Slider(0, 1, step=0.1, value=0.5, interactive=True, label="Threshold")
                gr.Markdown("### Note")
                gr.Markdown("- Limit the number of input sequences to less than 128.")
                gr.Markdown("- The file should be the Fasta format.")
                gr.Markdown("- We used only the first 20 amino acids of each N-terminal and C-terminal of the sequence for prediction.")
                image_button_homo = gr.Button("Submit")
        with gr.Column():
            # gr.Markdown(" ### Flip text or image files using this demo.")
            gr.Markdown("Note: the output scores indicates the probability of the input sequence to be predicted as a Peptide or a Non-Peptide.")
            frame_homo_output = gr.DataFrame(
                headers=["Sequence Id", "Sequence", "Probability of peptides", "Peptide"],
                datatype=["str", "str", "str", 'str'],)
            
    image_button_homo.click(homo_classifier, inputs=[input_fasta_homo, homo_cutoff], outputs=frame_homo_output)
    
    with gr.Tab("Prediction Model(Mus musculus)"):
        # cutoff = gr.Slider(0, 1, step=0.1, value=0.5, interactive=True)
        with gr.Row():
            with gr.Column(scale=2):
                input_fasta_mus = gr.File()
                # cutoff = gr.Slider(0, 1, step=0.1, value=0.5, interactive=True, label="threshold")
                # image_button = gr.Button("Submit")
            with gr.Column(scale=2):
                mus_cutoff = gr.Slider(0, 1, step=0.1, value=0.5, interactive=True, label="Threshold")
                gr.Markdown("### Note")
                gr.Markdown("- Limit the number of input sequences to less than 128.")
                gr.Markdown("- The file should be the Fasta format.")
                gr.Markdown("- We used only the first 20 amino acids of each N-terminal and C-terminal of the sequence for prediction.")
                image_button_mus = gr.Button("Submit")
        with gr.Column():
            # gr.Markdown(" ### Flip text or image files using this demo.")
            gr.Markdown("Note: the output scores indicates the probability of the input sequence to be predicted as a Peptide or a Non-Peptide.")
            frame_mus_output = gr.DataFrame(
                headers=["Sequence Id", "Sequence", "Probability of peptides", "Peptide"],
                datatype=["str", "str", "str", 'str'],)
    
    image_button_mus.click(mus_classifier, inputs=[input_fasta_mus, mus_cutoff], outputs=frame_mus_output)

    with gr.Accordion("Citation"):
        gr.Markdown("- GitHub: https://github.com/leonern/DeepPD")

    with gr.Accordion("License"):
        gr.Markdown("- Released under the [MIT license](https://github.com/leonern/DeepPD/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