from shiny import ui, reactive, render, App
from chemcrow.agents import ChemCrow
import shinyswatch

app_ui = ui.page_fluid(
    shinyswatch.theme.slate(),
    ui.panel_title("ChemCrow UI"),
    ui.p("An experiment with Shiny for Python and ChemCrow"),
    ui.br(),
    ui.row(
        ui.column(9, ui.input_text("prompt", label=None, width="100%")),
        ui.column(3, ui.input_action_button("chat", label="Chat", width="100%")),
    ),
    ui.help_text("Example 1: Propose a novel organicatalyst for enhancing carbon dioxide conversion in carbon capture and utilization processes."),
    ui.br(),
    ui.help_text("Example 2: What are the products of the reaction between 2-bromo-2-methylpropane and 4-(4-hydroxyphenyl)butan-2-one. Can this reaction run without problems?"),
    ui.output_ui("prompt_ui"),
    ui.output_ui("response_ui"),
    ui.br(),
    ui.hr(),
    ui.div(
        {"style": "align-items: center; display: flex; flex-direction: column; justify-content: center;"},
        ui.img(src="https://github.com/ur-whitelab/chemcrow-public/raw/main/assets/chemcrow_dark_thin.png", width="400px")
    ),
    ui.br(),
    ui.markdown(f'ChemCrow was [introduced](https://arxiv.org/abs/2304.05376) by Bran, Andres M., et al. "ChemCrow: Augmenting large-language models with chemistry tools." arXiv preprint arXiv:2304.05376 (2023). This tool is an extension of that work that puts the code into an interactive web app created by [James Wade](https://jameshwade.com) using [Shiny for Python](https://shiny.posit.co/py/). Find the code for the app [here](https://github.com/jameshwade/chemcrow) and the original code [here](https://github.com/ur-whitelab/chemcrow-public).')
)


def server(input, output, session):
    @reactive.Effect()
    def _():
        if input.chat():
            ui.update_text("prompt", value="")

    @output
    @render.ui
    @reactive.event(input.chat)
    def prompt_ui():
        list_ui = [ui.strong("Prompt"), ui.p(input.prompt())]
        return list_ui

    @output
    @render.ui
    @reactive.event(input.chat)
    def response_ui():
        ui.notification_show("Chatting with ChemCrow", type="message")
        chem_model = ChemCrow(model="gpt-4-0613", temp=0.1, verbose=True)
        response = chem_model.run(input.prompt())
        list_ui = [ui.strong("Thoughts"),
                   ui.markdown(response[0]),
                   ui.strong("Reasoning"),
                   ui.markdown(response[1]),
                   ui.strong("Answer"),
                   ui.markdown(response[2])]
        return list_ui


app = App(app_ui, server)