import streamlit as st

# st.set_page_config(layout="wide")


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

from openworm_ai import __version__ as openworm_ai_version

import socket

__version__ = "0.4.6"


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

with col1:
    st.title("OpenWorm.ai")
    st.markdown(
        "Testing the use of LLMs and other AI technology for _C. elegans_ research and computational modelling of the worm. Part of the [OpenWorm](https://openworm.org) project."
    )
    st.markdown("**v%s** (openworm_ai **v%s**) - Note: this is a work in progress!" % (__version__, openworm_ai_version))

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


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


with tab_corpus:
    from corpus_query import run_query

    st.markdown("**Search a (small) corpus of _C. elegans_ literature**")

    with st.form("form_corpus"):
        text = st.text_area(
            "Enter a query:",
            "What are the main types of neurons in the C. elegans pharynx? Give me a list of the specific cells.",
        )

        submitted = st.form_submit_button("Submit")

        if submitted:
            response = run_query(query=text)

            st.info(response)

with tab_free:
    st.markdown("**Ask individual LLMs questions about _C. elegans_**")

    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)

        allow_all_topics = 'MBP' in socket.gethostname()

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

        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
        st.markdown(
            "<style> .stRadio> div {gap:30px;} </style>",
            unsafe_allow_html=True,
        )
        panel_lead = st.radio(
            "Which LLM would you like to lead this panel?",
            PREF_ORDER_LLMS,
            horizontal=True,
        )

        # Create checkboxes to determine which LLMs are included in the panel
        st.markdown(
            "<div style='text-align: left; font-size: 14px;'>Which LLMs would you like to attend the panel?</div>",
            unsafe_allow_html=True,
        )
        options = PREF_ORDER_LLMS  # [llm for llm in PREF_ORDER_LLMS if llm !=panel_lead] creates a list of options excluding the chosen panel lead
        selected_options = []

        # Put checkboxes into table arranged 4x2
        for row in range(2):  # 2 rows
            cols = st.columns(4)  # 4 columns in each row
            for col in range(4):  # Iterate over each column
                index = row * 4 + col
                if index < len(options):
                    with cols[col]:
                        selected = st.checkbox(options[index], key=options[index])
                        if selected:
                            selected_options.append(options[index])

        submitted = st.form_submit_button("Submit")

        if len(selected_options) == 1 and panel_lead in selected_options:
            st.markdown("Please choose a panelist that is not the panel lead.")
        elif len(selected_options) == 0:
            st.markdown("Please choose one or more panelists to join the panel lead.")
        elif 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 selected_options 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_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 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(
        "See also [https://github.com/openworm/openworm.ai](https://github.com/openworm/openworm.ai)."
    )
    st.markdown("**Work in progress!!** Subject to change/removal without notice!")