import os
import sys
import tempfile
import pandas as pd
import gradio as gr
from smi_ted_light.load import load_smi_ted  

# 1) Ajuste de paths para encontrar o inference
BASE_DIR = os.path.dirname(__file__)
INFERENCE_DIR = os.path.join(BASE_DIR, "smi-ted", "inference")
sys.path.append(INFERENCE_DIR)

# 2) Carregando o modelo
MODEL_DIR = os.path.join("smi-ted", "inference", "smi_ted_light")
model = load_smi_ted(
    folder=MODEL_DIR,
    ckpt_filename="smi-ted-Light_40.pt",
    vocab_filename="bert_vocab_curated.txt",
)

# 3) Função única que gera embedding E CSV
def gerar_embedding_e_csv(smiles: str):
    smiles = smiles.strip()
    if not smiles:
        # Se não digitou nada, retorna erro e esconde o botão de download
        return {"erro": "digite uma sequência SMILES primeiro"}, gr.update(visible=False)

    try:
        # Gera o embedding
        vetor = model.encode(smiles, return_torch=True)[0].tolist()
        # Cria DataFrame e escreve CSV num arquivo temporário
        df = pd.DataFrame([vetor])
        tmp = tempfile.NamedTemporaryFile(suffix=".csv", delete=False)
        df.to_csv(tmp.name, index=False)
        tmp.close()
        # Retorna: 1) JSON, 2) update para o File (path + visível)
        return vetor, gr.update(value=tmp.name, visible=True)
    except Exception as e:
        # Em caso de erro interno, mostra mensagem e esconde o botão
        return {"erro": str(e)}, gr.update(visible=False)

# 4) Montando a interface Blocks
with gr.Blocks() as demo:
    gr.Markdown(
        """
        ## SMI-TED Embedding Generator  
        Cole uma sequência SMILES e receba:
        1. O vetor embedding (768 floats) em JSON  
        2. Um botão para baixar esse vetor em CSV  
        """
    )

    with gr.Row():
        inp_smiles = gr.Textbox(label="SMILES", placeholder="Ex.: CCO")
        btn = gr.Button("Gerar Embedding")
    with gr.Row():
        out_json = gr.JSON(label="Embedding (lista de floats)")
        out_file = gr.File(label="Download do CSV", visible=False)

    # 5) Ligando o botão à função única (dois outputs)
    btn.click(
        fn=gerar_embedding_e_csv,
        inputs=inp_smiles,
        outputs=[out_json, out_file]
    )

if __name__ == "__main__":
    demo.launch()