import streamlit
import streamlit as st

# st.set_page_config(layout="wide")

from ring import LLM_GPT35
from ring import LLM_GPT4
from ring import LLM_GPT4o
from ring import LLM_LLAMA2
from ring import LLM_GEMINI
from ring import PREF_ORDER_LLMS


from ring import requires_openai_key
from ring import get_openai_api_key
from ring import generate_response
from ring import generate_panel_response

__version__ = "0.1.9"

col1, col2 = st.columns([3, 1])

with col1:
    st.title("Project Sydney")
    st.markdown("**OpenWorm LLM v%s** - work in progress!" % __version__)

with col2:
    st.image("images/OpenWormLogo.png")


tab_free, tab_panel, tab_pubs, tab_data, tab_corpus, tab_model, tab_about = st.tabs(
    [
        "Individual LLMs",
        "Panel discussion",
        "Publications",
        "Structured data",
        "OpenWorm Corpus",
        "Run model",
        "About",
    ]
)


with tab_free:

    st.markdown("**Ask individual LLMs questions**")

    with st.form("form_free"):

        text = st.text_area(
            "Ask a question related to _C. elegans_:",
            "What is the primary role of the C. elegans neuron AVBL?",
        )

        llm_ver = st.selectbox("Which LLM version should I use?", PREF_ORDER_LLMS)

        temperature = st.text_input("Temperature", value=0.1)

        only_celegans = st.checkbox(
            "Only answer questions related to _C. elegans_", value=True
        )

        submitted = st.form_submit_button("Submit")

        if requires_openai_key(llm_ver) and not get_openai_api_key():
            st.info("Please add your OpenAI API key to continue.")
        elif submitted:
            response = generate_response(text, llm_ver, temperature, only_celegans)

            st.info(response)

with tab_panel:

    st.markdown("**Get a consensus answer across multiple LLMs**")

    with st.form("form_panel"):

        text = st.text_area(
            "Ask a question related to _C. elegans_:",
            "What is the typical length of the worm C. elegans?",
        )

        temperature = st.text_input("Temperature", value=0.1)

        ###Start of Kaan Changes
        #Selectbox to determine which LLM becomes panel lead
        panel_lead = st.selectbox("Which LLM would you like to lead this panel?", PREF_ORDER_LLMS)

        submitted = st.form_submit_button("Submit")

        if requires_openai_key(llm_ver) and not get_openai_api_key():
            st.info("Please add your OpenAI API key to continue.")
        elif submitted:
            #Create new list of LLMs that does not include panel lead
            llm_panelists = [llm for llm in PREF_ORDER_LLMS if llm !=panel_lead] 
            response = generate_panel_response(
                text,
                llm_panelists=llm_panelists,
                llm_panel_chair=panel_lead,
                temperature=temperature,
            )

            st.info(response)

with tab_pubs:

    from publications import find_basis_paper

    st.markdown(
        "**Find literature related to _C. elegans._** Uses: [Semantic Scholar](https://www.semanticscholar.org)"
    )

    with st.form("form_pubs"):

        text = st.text_input("Enter a topic:", "C. elegans locomotion")

        num = st.text_input("Number of results", value=10)

        submitted = st.form_submit_button("Submit")

        if submitted:
            response = find_basis_paper(query=text, result_limit=num)

            st.info(response)


with tab_data:

    st.markdown("**Query structured datasets**")

    with st.form("form_data"):

        from datasources import DS_WORMNEUROATLAS
        from datasources import FORMATS
        from datasources import query_data_source

        text = st.text_area("Which neuron would you like to know about?", "AVBL")

        source = st.selectbox("Select data source to use:", (DS_WORMNEUROATLAS,))
        format = st.selectbox("Return format:", FORMATS)

        submitted = st.form_submit_button("Submit")

        if submitted:
            response = query_data_source(text, source, format)

            st.info(response)

with tab_corpus:

    from corpus_query import run_query

    st.markdown(
        "**Search a (small) corpus of _C. elegans_ literature** Uses: [paper-qa](https://github.com/whitead/paper-qa)"
    )

    with st.form("form_corpus"):

        text = st.text_input(
            "Enter a query:",
            "What types of ion channels are present in C. elegans neurons?",
        )

        llm_ver = st.selectbox(
            "Which LLM version should I use?", (LLM_LLAMA2, LLM_GPT4o)
        )

        submitted = st.form_submit_button("Submit")

        if submitted:
            response = run_query(query=text, llm_ver=llm_ver)

            st.info(response)

with tab_model:

    st.markdown("**Run a _C. elegans_ cell model**")

    with st.form("form_model"):

        from model import run_model
        from model import MODELS

        text = st.text_area("Current injection level:", "4.1 pA")
        model_to_sim = st.selectbox("Model to simulate", list(MODELS.keys()))

        submitted = st.form_submit_button("Submit")

        if submitted:
            response, traces, events = run_model(text, model_to_sim)
            st.info(response)
            try:

                import numpy as np
                import matplotlib.pyplot as plt

                fig, ax = plt.subplots()

                for key in sorted(traces.keys()):
                    if key != "t":
                        ts = traces["t"]
                        vs = traces[key]
                        ax.plot(ts, vs, label=key)
                        ax.legend()

                        plt.xlabel("Time (s)")
                        plt.ylabel("(SI units)")

                st.pyplot(fig)
            except Exception as e:
                st.info("There was a problem...\n\n%s" % e)

with tab_about:

    st.markdown('### About "Project Sydney"')
    st.markdown(
        "This is an initiative by the [OpenWorm project](https://openworm.org) to investigate the use of LLMs for interacting with scientific literature and structured datasets related to _C. elegans_."
    )
    st.markdown("**Work in progress!!** Subject to change/removal without notice!")