import streamlit as st
from agent.gemini_agent import run_gemini_chat
from agent.clinical_nlp import umls_concept_lookup, bioportal_concept_lookup
from quantum.optimizer import optimize_treatment

st.set_page_config(page_title="Quantum Healthcare AI", layout="wide")

st.markdown("""
    <div style='display:flex;align-items:center'>
        <img src='https://cdn-icons-png.flaticon.com/512/891/891419.png' width='56' style='margin-right:15px;'/>
        <h1 style='display:inline'>Quantum-AI Healthcare Assistant</h1>
    </div>
    """, unsafe_allow_html=True)

st.write("**Ask anything about symptoms, conditions, diagnosis, or optimal care. Get instant LLM answers, clinical concept mapping, and even quantum-inspired optimization.**")

query = st.text_area("Describe symptoms, a clinical problem, or a medical question:", height=70)

cols = st.columns([1,1])
with cols[0]:
    get_ai = st.button("Get AI & Clinical Insight")
with cols[1]:
    run_quantum = st.button("Quantum Optimize Care Plan")

if get_ai and query.strip():
    with st.spinner("Gemini AI is reviewing your query..."):
        ai_response = run_gemini_chat(query)
    st.markdown(f"**Gemini Agent:**\n\n{ai_response}")

    st.markdown("---")
    with st.expander("🔎 UMLS Concept Mapping"):
        concepts = umls_concept_lookup(query)
        if concepts:
            for c in concepts:
                if 'error' in c:
                    st.error(c['error'])
                else:
                    st.write(f"**{c['name']}**  |  ID: `{c['ui']}`  |  Source: *{c['rootSource']}*")
        else:
            st.info("No UMLS concepts found.")

    with st.expander("🔬 BioPortal Ontology Lookup"):
        concepts = bioportal_concept_lookup(query)
        if concepts:
            for c in concepts:
                if 'error' in c:
                    st.error(c['error'])
                else:
                    st.write(f"**{c['prefLabel']}**  |  [Ontology]({c['ontology']})  |  [Concept ID]({c['id']})")
        else:
            st.info("No BioPortal concepts found.")

if run_quantum and query.strip():
    with st.spinner("Quantum optimizer working..."):
        plan = optimize_treatment(query)
    st.markdown("### 🧬 Optimized Care Plan")
    st.json(plan)

st.markdown("---")
st.caption("Powered by Gemini LLM, UMLS, BioPortal, and quantum-inspired algorithms. For research use only. No patient data is stored.")