import os
import gradio as gr
from cactus.agent import Cactus
from huggingface_hub import InferenceClient


MODEL_API_KEY_MAP = {
    "gpt-3.5-turbo": "OPENAI_API_KEY",
    "gpt-4": "OPENAI_API_KEY",
    "gemini-pro": "GOOGLE_API_KEY",
    "claude-3-haiku-20240307": "ANTHROPIC_API_KEY",
    "claude-3-opus-20240229": "ANTHROPIC_API_KEY",
    "claude-3-sonnet-20240229": "ANTHROPIC_API_KEY",
}
TITLE = """CACTUS 🌵: Chemistry Agent Connecting Tool-Usage to Science"""
DESC = """
CACTUS is an innovative tool-augmented language model designed to assist researchers and chemists in various chemistry-related tasks. By integrating state-of-the-art language models with a suite of powerful cheminformatics tools, CACTUS provides an intelligent and efficient solution for exploring chemical space, predicting molecular properties, and accelerating drug discovery. Just as the cactus thrives in the harsh desert environment, adapting to limited resources and extreme conditions, CACTUS has been implemented by Pacific Northwest National Laboratory (PNNL) Scientists to navigate the complex landscape of chemical data and extract valuable insights.
"""
CITE = """**Citation** 

If you use CACTUS in your research, please cite our preprint: 
```bibtex
@article{mcnaughton2024cactus,
    title={CACTUS: Chemistry Agent Connecting Tool-Usage to Science},
    author={Andrew D. McNaughton and Gautham Ramalaxmi and Agustin Kruel and Carter R. Knutson and Rohith A. Varikoti and Neeraj Kumar},
    year={2024},
    eprint={2405.00972},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}
```"""


def agent_fn(prompt, model_name, api_key=None):
    env_key_name = MODEL_API_KEY_MAP.get(model_name)
    if env_key_name:
        if api_key:
            os.environ[env_key_name] = api_key
        else:
            return f"Error: API key not found for {model_name}. Please provide it."
    try:
        cactus_model = Cactus(model_name, model_type="api")
        response = cactus_model.run(prompt)
        return response
    except Exception as e:
        return f"Error: {e}"


with gr.Blocks(title=TITLE) as iface:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("# " + TITLE)

    gr.Markdown(DESC)

    with gr.Row():
        model_name_dropdown = gr.Dropdown(
            list(MODEL_API_KEY_MAP.keys()), label="Select Model"
        )
        api_key_input = gr.Textbox(label="Enter API Key", type="password")
        prompt_input = gr.Textbox(
            lines=2, placeholder="Enter your cheminformatics question..."
        )

    
    with gr.Row():  # Row for sample questions
        sample_questions = [
            "What is the molecular weight of the SMILES: CCO ?",
            "Does the SMILES: C1CCOC1 , trigger the PAINS filter?",
            "Identify the molecular weight, logP, and Topological Polar Surface Area (TPSA) of the molecule with the SMILES notation CC(=O)NC1=CC=C(O)C=C1",
        ]
        for question in sample_questions:
            gr.Button(question).click(
                lambda q=question: gr.update(value=q),
                None,
                prompt_input,
            )

    with gr.Row():
        submit_button = gr.Button("Submit", variant="primary")

    response_output = gr.Textbox(label="Agent Response")

    gr.Markdown(CITE)

    # Event Triggers
    model_name_dropdown.change(
        lambda x: gr.update(visible=MODEL_API_KEY_MAP.get(x) is not None),
        model_name_dropdown,
        api_key_input,
    )

    submit_button.click(
        agent_fn,
        [prompt_input, model_name_dropdown, api_key_input],
        response_output,
    )

iface.launch()