import gradio as gr

# Assuming 'predict_stability' is your function that predicts protein stability
def predict_stability(model_choice, organism_choice, pdb_file=None, sequence=None):
    # Dummy return for illustration; replace with your actual prediction logic
    return f"Predicted Stability using {model_choice} for {organism_choice}: Example Output"

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # PLTNUM: Protein LifeTime Neural Model
        **Predict the protein half-life from its sequence or PDB file.**
        """
    )
    
    gr.Image("https://github.com/sagawatatsuya/PLTNUM/blob/main/model-image.png?raw=true", label="Model Image")

    # Model and Organism selection in the same row to avoid layout issues
    with gr.Row():
        model_choice = gr.Radio(
            choices=["SaProt", "ESM-2"], 
            label="Select PLTNUM's base model.", 
            value="SaProt"
        )
        organism_choice = gr.Radio(
            choices=["Mouse", "Human"], 
            label="Select the target organism.", 
            value="Mouse"
        )
    
    with gr.Tabs():
        with gr.TabItem("Upload PDB File"):
            gr.Markdown("### Upload your PDB file:")
            pdb_file = gr.File(label="Upload PDB File")
            predict_button = gr.Button("Predict Stability")
            prediction_output = gr.Textbox(label="Stability Prediction", interactive=False)

            predict_button.click(fn=predict_stability, inputs=[model_choice, organism_choice, pdb_file], outputs=prediction_output)
        
        with gr.TabItem("Enter Protein Sequence"):
            gr.Markdown("### Enter the protein sequence:")
            sequence = gr.Textbox(
                label="Protein Sequence",
                placeholder="Enter your protein sequence here...",
                lines=8,
            )
            predict_button = gr.Button("Predict Stability")
            prediction_output = gr.Textbox(label="Stability Prediction", interactive=False)

            predict_button.click(fn=predict_stability, inputs=[model_choice, organism_choice, sequence], outputs=prediction_output)
    
    gr.Markdown(
        """
        ### How to Use:
        - **Select Model**: Choose between 'SaProt' or 'ESM-2' for your prediction.
        - **Select Organism**: Choose between 'Mouse' or 'Human'.
        - **Upload PDB File**: Choose the 'Upload PDB File' tab and upload your file.
        - **Enter Sequence**: Alternatively, switch to the 'Enter Protein Sequence' tab and input your sequence.
        - **Predict**: Click 'Predict Stability' to receive the prediction.
        """
    )
    
    gr.Markdown(
        """
        ### About the Tool
        This tool allows researchers and scientists to predict the stability of proteins using advanced algorithms. It supports both PDB file uploads and direct sequence input.
        """
    )

demo.launch()