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

def run_prediction(data_path, model_path, prediction_type="classification", output_path="results.csv"):
    command = [
        "python", model_path,
        "--data_path", data_path,
        "--checkpoint_path", model_path.replace("predict.py", "../models/fusion_model.pt" if "mod" in model_path else "../models/vanilla_model.pt"),
        "--prediction_type", prediction_type,
        "--save_dir", ".",
        "--preds_path", output_path
    ]
    subprocess.run(command, check=True)
    return pd.read_csv(output_path)

def predict_from_smiles(smiles, model_choice, dataset_choice):
    temp_id = uuid.uuid4().hex
    temp_csv = f"data/temp_{temp_id}.csv"
    out_csv = f"data/output_{temp_id}.csv"
    df = pd.DataFrame([{"smiles": smiles, "compound_name": "molecule"}])
    df.to_csv(temp_csv, index=False)

    model_path = "chemprop_mod/predict.py" if model_choice == "Fusion (GNN + Transformer)" else "chemprop/predict.py"
    
    try:
        predictions = run_prediction(temp_csv, model_path, output_path=out_csv)
        return predictions
    except subprocess.CalledProcessError as e:
        return f"Error: {str(e)}"

def predict_from_file(file_obj, model_choice, dataset_choice):
    temp_id = uuid.uuid4().hex
    file_path = f"data/uploaded_{temp_id}.csv"
    out_csv = f"data/output_file_{temp_id}.csv"
    
    with open(file_path, "wb") as f:
        f.write(file_obj.read())

    model_path = "chemprop_mod/predict.py" if model_choice == "Fusion (GNN + Transformer)" else "chemprop/predict.py"

    try:
        predictions = run_prediction(file_path, model_path, output_path=out_csv)
        return predictions
    except subprocess.CalledProcessError as e:
        return f"Error: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("## 🧪 Drug Property Prediction with Fusion Models")
    gr.Markdown("Choose prediction input type and compare Chemprop vs Fusion model")

    with gr.Tab("Single SMILES"):
        with gr.Row():
            smiles_input = gr.Textbox(label="Enter SMILES string")
            model_select = gr.Radio(["Vanilla Chemprop", "Fusion (GNN + Transformer)"], label="Model")
            dataset_select = gr.Dropdown(["BBBP", "BACE"], label="Dataset")
        predict_button = gr.Button("Predict")
        result_output = gr.Dataframe(label="Prediction Result")
        predict_button.click(fn=predict_from_smiles, inputs=[smiles_input, model_select, dataset_select], outputs=result_output)

    with gr.Tab("Upload File"):
        with gr.Row():
            file_input = gr.File(label="Upload CSV File")
            model_select_file = gr.Radio(["Vanilla Chemprop", "Fusion (GNN + Transformer)"], label="Model")
            dataset_select_file = gr.Dropdown(["BBBP", "BACE"], label="Dataset")
        predict_button_file = gr.Button("Predict")
        result_output_file = gr.Dataframe(label="Prediction Result")
        predict_button_file.click(fn=predict_from_file, inputs=[file_input, model_select_file, dataset_select_file], outputs=result_output_file)

demo.launch()