src/app.py

"""
app.py  --  Streamlit UI for the AcademiQ PDF -> Summary -> Knowledge Graph project
==========================================================================
Owner: Aparna (UI layout, file upload, results display)

This file ONLY handles the interface. All the pipeline / model logic lives in
models.py, which currently returns demo placeholder data. When the real models
are ready, nothing in this file needs to change.

Run locally (Windows):
    py -m pip install -r requirements.txt
    py -m streamlit run app.py
"""

import html
import streamlit as st
import pandas as pd

import models  # our pipeline (placeholders for now)

# ---------------------------------------------------------------------------
# Page config  --  must be the first Streamlit call
# ---------------------------------------------------------------------------
st.set_page_config(
    page_title=" AcademiQ ",
    page_icon="📚",
    layout="wide",
)

# ---------------------------------------------------------------------------
# Session state: keep results so the tabs survive Streamlit's reruns
# ---------------------------------------------------------------------------
if "results" not in st.session_state:
    st.session_state.results = None


# ---------------------------------------------------------------------------
# Helper: turn text + entities into highlighted HTML
# (Iva)
# ---------------------------------------------------------------------------
def render_highlighted(text: str, entities: list) -> str:
    pieces = []
    cursor = 0
    for ent in entities:
        # plain text before this entity
        pieces.append(html.escape(text[cursor:ent["start"]]))
        color = models.LABEL_COLORS.get(ent["label"], models.LABEL_COLORS["ENTITY"])
        chunk = html.escape(text[ent["start"]:ent["end"]])
        pieces.append(
            f'<span style="background:{color};color:#102027;padding:2px 6px;'
            f'border-radius:6px;font-weight:600;" '
            f'title="{ent["label"]} ({ent["score"]})">{chunk}'
            f'<span style="font-size:0.7em;opacity:0.75;margin-left:4px;">'
            f'{ent["label"]}</span></span>'
        )
        cursor = ent["end"]
    pieces.append(html.escape(text[cursor:]))
    body = "".join(pieces)
    return (
        f'<div style="line-height:2.1;font-size:1.02rem;">{body}</div>'
    )


def label_legend() -> str:
    chips = []
    for label, color in models.LABEL_COLORS.items():
        if label == "ENTITY":
            continue
        chips.append(
            f'<span style="background:{color};color:#102027;padding:2px 10px;'
            f'border-radius:6px;font-size:0.8rem;font-weight:600;'
            f'margin-right:6px;">{label}</span>'
        )
    return '<div style="margin-bottom:12px;">' + "".join(chips) + "</div>"


# ---------------------------------------------------------------------------
# Sidebar: pipeline overview, settings, team credits
# ---------------------------------------------------------------------------
with st.sidebar:
    st.header("⚙️ Settings")

    demo_mode = st.toggle(
        "Demo mode (placeholder models)",
        value=True,
        help="ON: instant fake results, no model downloads. "
             "OFF: requires the real models wired into models.py.",
    )

    st.caption("Summary length (passed to the summariser)")
    max_len = st.slider("Max length", 60, 250, 130, step=10)
    min_len = st.slider("Min length", 10, 100, 30, step=5)

    st.divider()
    st.subheader("🔗 Pipeline")
    st.markdown(
        "1. **PDF → text** · PyMuPDF\n"
        "2. **Summary** · BART\n"
        "3. **Entities** · DistilBERT NER\n"
        "4. **Relations** · DistilBERT RE\n"
        "5. **Graph** · NetworkX"
    )

    st.divider()
    st.subheader("👥 Introducing our Team")
    st.markdown(
	"- **Jordan** — PDF ingestion and extraction using PyMuPDF\n"
	"- **Varsha** — Summarization BART-large\n"
	"- **Damir** - DistilBERT NER model\n"
	"- **Andy** - Handling Post-processing\n"
        "- **Aparna** - UI, upload, results\n"
        "- **Iva** - styling, NER highlighting, summary view\n"
	"-**Temirlan** - Integration and Testing"
    )

# ---------------------------------------------------------------------------
# Header
# ---------------------------------------------------------------------------
st.title("📚 AcademiQ Research Paper → Knowledge Graph")
st.markdown(
    "Upload an academic PDF. The app summarises it, pulls out key entities and "
    "their relationships, and builds a knowledge graph."
)

if demo_mode:
    st.info(
        "🧪 **Demo mode is ON** — results are placeholders so the interface "
        "runs instantly. Turn it off once the real models are connected in "
        "`models.py`.",
        icon="🧪",
    )

# ---------------------------------------------------------------------------
# Input section
# ---------------------------------------------------------------------------
st.subheader("1 · Upload")

col_up, col_sample = st.columns([3, 1])
with col_up:
    uploaded = st.file_uploader("Choose a PDF", type=["pdf"], label_visibility="collapsed")
with col_sample:
    use_sample = st.checkbox("Use sample paper", value=not bool(uploaded))

ready = bool(uploaded) or use_sample
analyze = st.button("🚀 Analyze paper", type="primary", disabled=not ready,
                    use_container_width=True)

# ---------------------------------------------------------------------------
# Run the pipeline
# ---------------------------------------------------------------------------
if analyze:
    # Step 0: get the raw text
    if uploaded is not None:
        raw_text = models.extract_text_from_pdf(uploaded.getvalue())
    else:
        raw_text = models.SAMPLE_TEXT

    with st.status("Running the pipeline...", expanded=True) as status:
        st.write("📝 Summarising...")
        summary_out = models.summarize(raw_text, max_len=max_len, min_len=min_len)
        summary = summary_out["summary"]

        st.write("🏷️ Extracting entities (NER)...")
        entities = models.ner(summary)

        st.write("🔗 Extracting relations...")
        relations = models.extract_relations(summary, entities)

        st.write("🕸️ Building knowledge graph...")
        graph_fig = models.build_graph(relations, entities)

        status.update(label="Done!", state="complete", expanded=False)

    st.session_state.results = {
        "raw_text": raw_text,
        "summary": summary,
        "metrics": summary_out["metrics"],
        "entities": entities,
        "relations": relations,
        "graph_fig": graph_fig,
    }

# ---------------------------------------------------------------------------
# Results
# ---------------------------------------------------------------------------
res = st.session_state.results
if res:
    st.subheader("2 · Results")

    # quick metrics row
    m1, m2, m3, m4 = st.columns(4)
    m1.metric("Entities", len(res["entities"]))
    m2.metric("Relations", len(res["relations"]))
    m3.metric("Graph nodes", len({r["head"] for r in res["relations"]} |
                                  {r["tail"] for r in res["relations"]}))
    m4.metric("Compression", res["metrics"]["Compression"])

    tab_text, tab_sum, tab_ent, tab_rel, tab_graph = st.tabs(
        ["📄 Extracted Text", "📝 Summary", "🏷️ Entities", "🔗 Relations", "🕸️ Knowledge Graph"]
    )

    with tab_text:
        st.caption(f"{len(res['raw_text'].split())} words extracted from the PDF")
        st.text_area("Raw text", res["raw_text"], height=320, label_visibility="collapsed")

    with tab_sum:
        st.markdown(f"> {res['summary']}")
        st.divider()
        cols = st.columns(len(res["metrics"]))
        for col, (k, v) in zip(cols, res["metrics"].items()):
            col.metric(k, v)
        st.caption("ROUGE scores are placeholders until evaluated against reference abstracts.")

    with tab_ent:
        st.markdown(label_legend(), unsafe_allow_html=True)
        st.markdown(render_highlighted(res["summary"], res["entities"]),
                    unsafe_allow_html=True)
        st.divider()
        df_ent = pd.DataFrame(res["entities"])[["text", "label", "score"]]
        df_ent.columns = ["Entity", "Type", "Confidence"]
        st.dataframe(df_ent, use_container_width=True, hide_index=True)

    with tab_rel:
        df_rel = pd.DataFrame(res["relations"])
        df_rel.columns = ["Head", "Relation", "Tail", "Confidence"]
        st.dataframe(df_rel, use_container_width=True, hide_index=True)
        st.caption("Each row is a (head → relation → tail) triple feeding the graph.")

    with tab_graph:
        st.pyplot(res["graph_fig"], use_container_width=True)
        st.caption("Nodes coloured by entity type; arrows show the extracted relations.")
else:
    st.caption("Upload a PDF (or tick *Use sample paper*) and press **Analyze**.")