"""
gradio

"""

import gradio as gr
from helper import *
import scripts.script as script


def process_data(input_type, input_text, input_file):
    print(input_type)
    if input_type == "Text":
        if input_text:
            print(input_text)
            sequence = []
            for line in input_text.splitlines():
                if line.startswith(">"):
                    sequence.append("")
                else:
                    sequence[-1] += line.strip().upper()
            result = classify_sequence_type_length(sequence)
            script.run_argnet(input_text, "output.txt", result[0], result[1])
            df, pie_chart = script.view_stat("output.txt")
        else:
            result = "No input provided."
    else:
        if input_file:
            sequence = []
            input_text = open(input_file.name, "r").read()
            with open(input_file.name, "r") as f:
                for line in f:
                    if line.startswith(">"):
                        print(line.strip())
                        sequence.append("")
                    else:
                        print(line.strip().upper())
                        sequence[-1] += line.strip().upper()
            result = classify_sequence_type_length(sequence)
            script.run_argnet(input_text, "output.txt", result[0], result[1])
            df, pie_chart = script.view_stat("output.txt")
        else:
            result = "No input provided."

    return df, pie_chart


# Create the interface with tabs
with gr.Blocks() as whole_block:
    tab_selected = gr.State("Text")

    gr.HTML(
        """
    <center>
    <h1>ARGnet</h1>
    <p>A deep neural network for robust identification and annotation of Antibiotic Resistance genes</p>
    </center>
    """
    )

    with gr.Row():
        with gr.Column():
            gr.Markdown(
                """
            ## Input
            """
            )
            with gr.Tab("Text") as text_tab:
                gr.Markdown(
                    """
                ### Enter the sequence (fasta format)
                """
                )
                input_textbox = gr.Textbox(label="Sequence")
                input_textbox_2 = gr.Textbox(label="Sequence", visible=False)
                gr.Examples(
                    examples=[
                        ["Amino Acid Long Sequence (>51aa)"],
                        ["Amino Acid Short Sequence (30-50aa)"],
                        ["Nucleotide Long Sequence (>150nt)"],
                        ["Nucleotide Short Sequence (100-150nt)"],
                    ],
                    inputs=input_textbox,
                    fn=get_sequence_example,
                    cache_examples=True,
                    outputs=input_textbox,
                )
            with gr.Tab("File") as file_tab:
                gr.Markdown(
                    """
                ### Upload a FASTA file
                """
                )
                input_file = gr.File(label="Sequence File")
            submit_button = gr.Button("Submit")
        with gr.Column():
            gr.Markdown(
                """
            ## Output
            """
            )
            table = gr.Dataframe(
                headers=[
                    "Test ID",
                    "ARG Prediction",
                    "Resistance Category",
                    "Probability",
                ]
            )
            pie_chart = gr.Plot(container=True)

    text_tab.select(lambda: "Text", inputs=None, outputs=tab_selected)
    file_tab.select(lambda: "File", inputs=None, outputs=tab_selected)

    submit_button.click(
        fn=process_data,
        inputs=[tab_selected, input_textbox, input_file],
        outputs=[table, pie_chart],
    )

whole_block.launch()