app.py

#!/usr/bin/env python3
"""MedGenesis AI — CPU-only, dual-LLM (OpenAI / Gemini)"""

# ────────────── FIX: create a writable Streamlit data dir ──────────────
import os, pathlib
os.environ["STREAMLIT_DATA_DIR"]                = "/tmp/.streamlit"
os.environ["XDG_STATE_HOME"]                    = "/tmp"
os.environ["STREAMLIT_BROWSER_GATHERUSAGESTATS"] = "false"
pathlib.Path("/tmp/.streamlit").mkdir(parents=True, exist_ok=True)
# ───────────────────────────────────────────────────────────────────────

import asyncio, re
from pathlib import Path

import streamlit as st
import pandas as pd
import plotly.express as px
from fpdf import FPDF
from streamlit_agraph import agraph

from mcp.orchestrator    import orchestrate_search, answer_ai_question
from mcp.workspace       import get_workspace, save_query
from mcp.knowledge_graph import build_agraph
from mcp.graph_metrics   import build_nx, get_top_hubs, get_density
from mcp.alerts          import check_alerts

ROOT = Path(__file__).parent
LOGO = ROOT / "assets" / "logo.png"

# ──────────────── small helpers ─────────────────
def _pdf(papers):
    pdf = FPDF(); pdf.add_page(); pdf.set_font("Arial", size=11)
    pdf.cell(200, 8, "MedGenesis AI – Results", ln=True, align="C"); pdf.ln(3)
    for i, p in enumerate(papers, 1):
        pdf.set_font("Arial", "B", 11)
        pdf.multi_cell(0, 7, f"{i}. {p['title']}")
        pdf.set_font("Arial", "", 9)
        pdf.multi_cell(0, 6, f"{p['authors']}\n{p['summary']}\n{p['link']}\n")
        pdf.ln(1)
    return pdf.output(dest="S").encode("latin-1")

def _sidebar_workspace():
    with st.sidebar:
        st.header("🗂️ Workspace")
        ws = get_workspace()
        if not ws:
            st.info("Run a search then press **Save** to populate this list.")
            return
        for i, item in enumerate(ws, 1):
            with st.expander(f"{i}. {item['query']}"):
                st.write(item["result"]["ai_summary"])

# ──────────────── Streamlit UI ──────────────────
def render_ui():
    st.set_page_config("MedGenesis AI", layout="wide")
    _sidebar_workspace()

    # header
    c1, c2 = st.columns([0.15, 0.85])
    with c1:
        if LOGO.exists():
            st.image(str(LOGO), width=105)
    with c2:
        st.markdown("## 🧬 **MedGenesis AI**")
        st.caption("Multi-source biomedical assistant · OpenAI / Gemini")

    llm = st.radio("LLM engine", ["openai", "gemini"], horizontal=True)
    query = st.text_input("Enter biomedical question",
                          placeholder="e.g. CRISPR glioblastoma therapy")

    # alert check
    if get_workspace():
        try:
            news = asyncio.run(check_alerts([w["query"] for w in get_workspace()]))
            if news:
                st.sidebar.subheader("🔔 New papers")
                for q, lst in news.items():
                    st.sidebar.write(f"**{q}** – {len(lst)} new")
        except Exception as e:
            st.sidebar.warning(f"Alert check failed: {e}")

    if st.button("Run Search 🚀") and query:
        with st.spinner("Collecting literature & biomedical data …"):
            res = asyncio.run(orchestrate_search(query, llm=llm))
        st.success(f"Completed with **{res['llm_used'].title()}**")

        tabs = st.tabs(["Results", "Genes", "Trials", "Graph",
                        "Metrics", "Visuals"])

        # Results -----------------------------------------------------
        with tabs[0]:
            for i, p in enumerate(res["papers"], 1):
                st.markdown(f"**{i}. [{p['title']}]({p['link']})**  *{p['authors']}*")
                st.write(p["summary"])

            col1, col2 = st.columns(2)
            with col1:
                st.download_button("CSV",
                                   pd.DataFrame(res["papers"]).to_csv(index=False),
                                   "papers.csv", "text/csv")
            with col2:
                st.download_button("PDF", _pdf(res["papers"]),
                                   "papers.pdf", "application/pdf")

            if st.button("💾 Save"):
                save_query(query, res)
                st.success("Saved to workspace")

            st.subheader("UMLS concepts")
            for c in res["umls"]:
                if c.get("cui"):
                    st.write(f"- **{c['name']}** ({c['cui']})")

            st.subheader("OpenFDA safety")
            for d in res["drug_safety"]:
                st.json(d)

            st.subheader("AI summary")
            st.info(res["ai_summary"])

        # Genes -------------------------------------------------------
        with tabs[1]:
            st.header("Gene / Variant signals")
            for g in res["genes"]:
                st.write(f"- **{g.get('name', g.get('geneid'))}** "
                         f"{g.get('description', '')}")
            if res["gene_disease"]:
                st.markdown("### DisGeNET links")
                st.json(res["gene_disease"][:15])
            if res["mesh_defs"]:
                st.markdown("### MeSH definitions")
                for d in res["mesh_defs"]:
                    if d:
                        st.write("-", d)

        # Trials ------------------------------------------------------
        with tabs[2]:
            st.header("Clinical trials")
            if not res["clinical_trials"]:
                st.info("No trials (rate-limited or none found).")
            for t in res["clinical_trials"]:
                st.markdown(f"**{t['NCTId'][0]}** – {t['BriefTitle'][0]}")
                st.write(f"Phase {t.get('Phase', [''])[0]} | "
                         f"Status {t['OverallStatus'][0]}")

        # Graph -------------------------------------------------------
        with tabs[3]:
            nodes, edges, cfg = build_agraph(res["papers"],
                                             res["umls"],
                                             res["drug_safety"])
            hl = st.text_input("Highlight node:", key="hl")
            if hl:
                pat = re.compile(re.escape(hl), re.I)
                for n in nodes:
                    n.color = "#f1c40f" if pat.search(n.label) else "#d3d3d3"
            agraph(nodes, edges, cfg)

        # Metrics -----------------------------------------------------
        with tabs[4]:
            G = build_nx([n.__dict__ for n in nodes],
                         [e.__dict__ for e in edges])
            st.metric("Density", f"{get_density(G):.3f}")
            st.markdown("**Top hubs**")
            for nid, sc in get_top_hubs(G):
                lab = next((n.label for n in nodes if n.id == nid), nid)
                st.write(f"- {lab}  {sc:.3f}")

        # Visuals -----------------------------------------------------
        with tabs[5]:
            years = [p["published"] for p in res["papers"] if p.get("published")]
            if years:
                st.plotly_chart(px.histogram(years, nbins=12,
                                             title="Publication Year"))

        st.markdown("---")
        follow = st.text_input("Ask follow-up:")
        if st.button("Ask AI"):
            ans = asyncio.run(answer_ai_question(follow,
                                                 context=query,
                                                 llm=llm))
            st.write(ans["answer"])

    else:
        st.info("Enter a question and press **Run Search 🚀**")

# entry-point
if __name__ == "__main__":
    render_ui()