import gradio as gr
import os
import subprocess
import uuid
import pandas as pd

import sys

# Append chemprop and chemprop_updated folders explicitly
sys.path.append(os.path.abspath("chemprop"))
sys.path.append(os.path.abspath("chemprop-updated"))


# Helper to save SMILES string to a temporary CSV
def save_smiles_to_csv(smiles: str, temp_dir="temp_inputs"):
    os.makedirs(temp_dir, exist_ok=True)
    file_path = os.path.join(temp_dir, f"{uuid.uuid4().hex}.csv")
    df = pd.DataFrame({"smiles": [smiles]})
    df.to_csv(file_path, index=False)
    return file_path

# Core prediction logic
def predict(model_version, dataset_name, input_type, file=None, smiles=None):
    # Set paths
    if model_version == "Vanilla Chemprop":
        model_dir = "chemprop"
        model_path = f"model_weight/{dataset_name}/best_unbalanced.pt"
        script_path = "chemprop/chemprop/cli/predict.py"
    else:
        model_dir = "chemprop_updated"
        model_path = f"model_weight/{dataset_name}/best_bert_fusion.pt"
        script_path = "chemprop-updated/chemprop_updated/cli/predict.py"

    # Prepare input file
    if input_type == "Upload CSV":
        if file is None:
            return "Please upload a CSV file."
        input_path = file.name
    else:
        if not smiles:
            return "Please enter a SMILES string."
        input_path = save_smiles_to_csv(smiles)

    # Run prediction command
    cmd = [
        "python",script_path,
        "--test-path", input_path,
        "--model-paths", model_path,
        "--smiles-columns", "smiles"
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            return f"Error:\n{result.stderr}"
        return f"Prediction Output:\n{result.stdout}"
    except Exception as e:
        return f"Execution Failed: {str(e)}"

# Gradio UI setup
with gr.Blocks() as demo:
    gr.Markdown("## 🧪 Molecular Property Prediction using GNN and Transformers")

    model_version = gr.Radio(
        ["Vanilla Chemprop", "Updated Fusion Model"],
        label="Select Model Version"
    )
    dataset_name = gr.Radio(["BBBP", "ClinTox"], label="Select Dataset")

    input_type = gr.Radio(["Upload CSV", "Single SMILES"], label="Input Type")

    file_input = gr.File(file_types=[".csv"], label="Upload CSV", visible=True)
    smiles_input = gr.Textbox(label="Enter SMILES string", visible=False)

    def toggle_inputs(choice):
        return {
            file_input: gr.update(visible=(choice == "Upload CSV")),
            smiles_input: gr.update(visible=(choice == "Single SMILES"))
        }

    input_type.change(toggle_inputs, input_type, [file_input, smiles_input])

    predict_button = gr.Button("Predict")
    output = gr.Textbox(label="Output")

    predict_button.click(
        fn=predict,
        inputs=[model_version, dataset_name, input_type, file_input, smiles_input],
        outputs=output
    )

demo.launch()