app.py

"""
EGM Research Assistant — RAG + Visualisations
Gradio 5.x | sentence-transformers → ChromaDB → Claude + Plotly via HTML
"""

import os, re, io, base64
import pandas as pd
import chromadb
from chromadb.utils import embedding_functions
from anthropic import Anthropic
import gradio as gr
from collections import Counter
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
COLLECTION  = "egm_studies"
MAX_TOKENS  = 2000
BATCH_SIZE  = 50

REGION_COLOURS = {
    "Sub-Saharan Africa":           "#2ecc71",
    "Latin America and Caribbean":  "#3498db",
    "East Asia and Pacific":        "#e74c3c",
    "South Asia":                   "#f39c12",
    "Middle East and North Africa": "#9b59b6",
    "Europe and Central Asia":      "#1abc9c",
    "Multi-continent":              "#95a5a6",
}

# ---------------------------------------------------------------------------
# Load dataset
# ---------------------------------------------------------------------------
print("Loading dataset...")
df = pd.read_csv("dataset.csv").fillna("")

# Normalise column name for PSE strategy (long name)
PSE_COL  = "pse_strategy_multi_select_separate_values_with_single_vertical_bar"
CONF_COL = "overall_how_much_confidence_do_you_have_in_the_methods_used_to_analyse_the_findings_relative_to_the_primary_question_addressed_in_the_review"

df["year_of_publication"] = (
    df["year_of_publication"].astype(str).str.replace(r"\.0$", "", regex=True)
)
RECORDS = df.to_dict("records")
TOTAL   = len(RECORDS)
print(f"Loaded {TOTAL} studies.")

# ---------------------------------------------------------------------------
# Build embedding text — rich multi-field document
# ---------------------------------------------------------------------------
def make_doc(r):
    parts = [
        r.get("title", ""),
        "Authors: "                + r.get("authors", ""),
        "Intervention: "           + r.get("intervention", ""),
        "Intervention detail: "    + r.get("intervention_description", "")[:200],
        "Intervention category: "  + r.get("intervention_category", ""),
        "Intervention group (PSE): "+ r.get("interv_PSE", ""),
        "Outcome: "                + r.get("outcome", ""),
        "Outcome category: "       + r.get("outcome_categories_assigned", ""),
        "Outcome group (PSE): "    + r.get("outc_PSE", ""),
        "Theme: "                  + r.get("primary_theme", ""),
        "Sub-theme: "              + r.get("sub_primary_theme", ""),
        "Sector: "                 + r.get("sector_name", ""),
        "SDG: "                    + r.get("un_sustainable_development_goal", ""),
        "Evaluation design: "      + r.get("evaluation_design", ""),
        "Evaluation method: "      + r.get("evaluation_method", ""),
        "PSE strategy: "           + r.get(PSE_COL, ""),
        "Country: "                + r.get("country", ""),
        "Region: "                 + r.get("region", ""),
        "Income level: "           + r.get("income_level", ""),
        "GPI fragility rank: "     + str(r.get("gpi_rank", "") or ""),
        "Equity focus: "           + r.get("equity_focus", ""),
        "Population: "             + r.get("population_description", "")[:150],
        "Record type: "            + r.get("record_type", ""),
        "Method eval: "            + r.get("method_eval", ""),
        "Review confidence: "      + r.get(CONF_COL, ""),
        "Keywords: "               + r.get("keywords", ""),
        "Other topics: "           + r.get("other_topics", ""),
        r.get("abstract", "")[:600],
    ]
    return " | ".join(p for p in parts if p.strip() and not p.strip().endswith(": "))


DOCS  = [make_doc(r) for r in RECORDS]
IDS   = [str(i) for i in range(TOTAL)]
METAS = [
    {
        "title":   r["title"][:200],
        "authors": r["authors"][:100],
        "year":    r["year_of_publication"],
        "country": r["country"][:80],
        "url":     r.get("url_link", "")[:200] if str(r.get("url_link","")).startswith("http") else "",
    }
    for i, r in enumerate(RECORDS)
]

# ---------------------------------------------------------------------------
# ChromaDB vector index
# ---------------------------------------------------------------------------
print("Building vector index...")
ef     = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=EMBED_MODEL)
chroma = chromadb.Client()

try:
    chroma.delete_collection(COLLECTION)
except Exception:
    pass

col = chroma.create_collection(
    name=COLLECTION,
    embedding_function=ef,
    metadata={"hnsw:space": "cosine"},
)

for start in range(0, TOTAL, BATCH_SIZE):
    end = min(start + BATCH_SIZE, TOTAL)
    col.add(ids=IDS[start:end], documents=DOCS[start:end], metadatas=METAS[start:end])
    print(f"  Indexed {end}/{TOTAL} studies")

print("Vector index ready ✓")

# ---------------------------------------------------------------------------
# Anthropic client + system prompt
# ---------------------------------------------------------------------------
anthropic_client = Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY", ""))

SYSTEM = f"""You are a research synthesis assistant for an Evidence Gap Map (EGM) \
containing {TOTAL} studies on private sector engagement and development interventions.

STRICT RULES:
1. Answer ONLY using the studies explicitly provided in each message. Never use outside knowledge.
2. Only state what is explicitly present in the study data — do not generalise or assume.
3. Respect the diversity of study designs: studies may be Experimental, Quasi-experimental, 
   or Qualitative. Never describe the whole dataset as quantitative or experimental.
4. Every factual claim must have an inline [N] citation to a specific provided study.
5. If evidence is insufficient, say so and describe what the retrieved studies do cover.

RESPONSE FORMAT:
**Overview** — 2–3 sentences summarising what the evidence shows, noting study design diversity.
**Key Findings** — bullets grouped by theme/intervention, each with [N] citations.
**Study Designs & Methods** — note the methodological mix (RCTs, DiD, qualitative, etc.)
**Geographic & Context Coverage** — regions, income levels, GPI fragility rank (lower = more fragile), fragile/conflict-affected states if relevant.
**Equity & Population Focus** — note gender, equity dimensions, and target populations if relevant.
**Gaps & Caveats** — what is not well evidenced in the retrieved studies.
**Key Takeaways** — 2–3 concise conclusions from the provided studies only.

Use **bold** for key terms. Use [N] inline for every claim."""

# ---------------------------------------------------------------------------
# Retrieval + context builder
# ---------------------------------------------------------------------------
def retrieve(query, top_k):
    top_k = min(int(top_k), TOTAL)
    results = col.query(query_texts=[query], n_results=top_k,
                        include=["metadatas", "distances"])
    studies = []
    for i, doc_id in enumerate(results["ids"][0]):
        rec = RECORDS[int(doc_id)].copy()
        rec["_rank"] = i + 1
        studies.append(rec)
    return studies


def build_context(studies):
    """
    Two-tier context:
    - All studies: structured fields including PSE groups, keywords, GPI rank
    - <=80 studies: also include abstract
    """
    include_abstract = len(studies) <= 80
    parts = []
    for i, s in enumerate(studies, 1):
        line = (
            f"[{i}] \"{s['title']}\" | "
            f"{s['authors']} ({s['year_of_publication']}) | "
            f"{s['country']} | {s['region']} | Income: {s['income_level']} | "
            f"GPI rank: {s.get('gpi_rank','')} | "
            f"{s['record_type']} | {s['method_eval']} | "
            f"{s['evaluation_design']} | {s['evaluation_method']} | "
            f"Intervention: {s['intervention']} [{s['intervention_category']}] | "
            f"PSE intervention group: {s.get('interv_PSE','')} | "
            f"Outcome: {s['outcome']} [{s['outcome_categories_assigned']}] | "
            f"PSE outcome group: {s.get('outc_PSE','')} | "
            f"Theme: {s['primary_theme']} > {s['sub_primary_theme']} | "
            f"Sector: {s['sector_name']} | Equity: {s['equity_focus']} | "
            f"Keywords: {s.get('keywords','')} | Other topics: {s.get('other_topics','')} | "
            f"Confidence: {s[CONF_COL]}"
        )
        if include_abstract and s.get("abstract","").strip():
            line += f"\n  Abstract: {s['abstract'][:250]}"
        parts.append(line)
    return "\n".join(parts)




# ---------------------------------------------------------------------------
# Charts rendered as HTML (avoids gr.Plot version conflicts)
# ---------------------------------------------------------------------------
CHART_BG = "#f8f9fa"

def detect_chart_topics(query):
    """Return exactly 1-2 most relevant chart types for the question."""
    q = query.lower()

    # --- SR appraisal check FIRST — if any SR keyword present, always show sr_appraisal ---
    sr_triggers = [
        "systematic review", "critical appraisal", "appraisal", "review quality",
        "reliability of review", "reliability of the review", "srr", "amstar", "prisma",
        "cochrane", "meta-analysis", "narrative synthesis", "review methodology",
        "confidence in the methods", "rate the reliability",
    ]
    if any(w in q for w in sr_triggers):
        return {"sr_appraisal"}

    # --- Non-SR scoring ---
    scores = {
        "interventions": 0,
        "outcomes":      0,
        "geo":           0,
        "methods":       0,
        "themes":        0,
        "timeline":      0,
    }

    # Remove "systematic review" from method_words — handled above
    geo_words     = ["country","countries","region","africa","asia","latin","europe","pacific",
                     "middle east","where","location","fragil","conflict","gpi","income",
                     "low income","middle income","global","south asia","sub-saharan"]
    method_words  = ["method","design","rct","randomis","experiment","qualitative","quasi",
                     "evaluation","evidence type","how studied","quantitative","approach",
                     "study type","study design"]
    theme_words   = ["theme","sector","health","education","finance","economic","social",
                     "environment","energy","urban","rural","topic","area"]
    interv_words  = ["intervention","program","programme","activity","what works","pse","ppp",
                     "cash transfer","voucher","training","tvet","microfinance","credit",
                     "matchmaking","subsidy","capacity","financing"]
    outcome_words = ["outcome","impact","effect","result","employment","income","job","market",
                     "growth","food","nutrition","wellbeing","welfare","earnings"]
    time_words    = ["year","recent","latest","trend","over time","timeline",
                     "2020","2021","2022","2023","2024","when","publish"]

    for w in geo_words:     scores["geo"]           += q.count(w)
    for w in method_words:  scores["methods"]       += q.count(w)
    for w in theme_words:   scores["themes"]        += q.count(w)
    for w in interv_words:  scores["interventions"] += q.count(w)
    for w in outcome_words: scores["outcomes"]      += q.count(w)
    for w in time_words:    scores["timeline"]      += q.count(w)

    ranked = sorted(scores.items(), key=lambda x: -x[1])
    top    = [k for k, v in ranked if v > 0][:2]

    if not top:
        top = ["interventions", "geo"]
    elif len(top) == 1:
        pairs = {
            "interventions": "outcomes",
            "outcomes":      "interventions",
            "geo":           "interventions",
            "methods":       "interventions",
            "themes":        "interventions",
            "timeline":      "interventions",
        }
        top.append(pairs[top[0]])

    return set(top[:2])


def make_single_chart(kind, studies):
    """Make one matplotlib chart and return as HTML img tag."""
    NAVY   = "#0a2342"
    BLUE   = "#3498db"
    RED    = "#e74c3c"
    AMBER  = "#f39c12"
    PURPLE = "#9b59b6"
    GREEN  = "#2ecc71"
    GREY   = "#95a5a6"
    BG     = "#f8f9fa"

    fig, ax = plt.subplots(figsize=(9, 4), facecolor=BG)
    ax.set_facecolor(BG)
    for spine in ax.spines.values():
        spine.set_color("#e2e8f0")

    if kind == "interventions":
        ivs = []
        for s in studies:
            for iv in s.get("interv_PSE","").split("|"):
                iv = iv.strip()
                if iv and iv.lower() not in ("","not applicable","nan"):
                    ivs.append(iv)
        if not ivs:
            for s in studies:
                iv = s.get("intervention_category","").strip()
                if iv and iv.lower() not in ("","not applicable","nan"):
                    ivs.append(iv)
        counts = Counter(ivs).most_common(10)
        if not counts: return None
        labels = [x[0][:40] for x in counts][::-1]
        values = [x[1] for x in counts][::-1]
        bars = ax.barh(labels, values, color=BLUE, edgecolor="white", height=0.6)
        for bar, val in zip(bars, values):
            ax.text(bar.get_width()+0.1, bar.get_y()+bar.get_height()/2,
                    str(val), va="center", ha="left", fontsize=8, color="#444")
        ax.set_xlabel("Studies", fontsize=9, color="#666")
        ax.set_xlim(0, max(values)*1.15)
        ax.set_title("PSE Intervention Groups in retrieved studies", fontsize=11, fontweight="bold", color=NAVY, pad=8)

    elif kind == "outcomes":
        ocs = []
        for s in studies:
            for oc in s.get("outc_PSE","").split("|"):
                oc = oc.strip()
                if oc and oc.lower() not in ("","not applicable","nan"):
                    ocs.append(oc)
        if not ocs:
            for s in studies:
                for oc in s.get("outcome_categories_assigned","").split("|"):
                    oc = oc.strip()
                    if oc and oc.lower() not in ("","not applicable","nan"):
                        ocs.append(oc)
        counts = Counter(ocs).most_common(10)
        if not counts: return None
        labels = [x[0][:40] for x in counts][::-1]
        values = [x[1] for x in counts][::-1]
        colors = [GREEN if i % 2 == 0 else "#1abc9c" for i in range(len(labels))]
        bars = ax.barh(labels, values, color=colors, edgecolor="white", height=0.6)
        for bar, val in zip(bars, values):
            ax.text(bar.get_width()+0.1, bar.get_y()+bar.get_height()/2,
                    str(val), va="center", ha="left", fontsize=8, color="#444")
        ax.set_xlabel("Studies", fontsize=9, color="#666")
        ax.set_xlim(0, max(values)*1.15)
        ax.set_title("PSE Outcome Groups in retrieved studies", fontsize=11, fontweight="bold", color=NAVY, pad=8)

    elif kind == "geo":
        regions = []
        for s in studies:
            r = s.get("region","").split("|")[0].strip()
            if r and r.lower() not in ("","not applicable","nan"):
                regions.append(r)
        counts = Counter(regions)
        if not counts: return None
        rl = list(counts.keys())
        rv = list(counts.values())
        rcols = [REGION_COLOURS.get(r, GREY) for r in rl]
        wedges, _, autotexts = ax.pie(
            rv, labels=None, colors=rcols, autopct="%1.0f%%",
            startangle=90, wedgeprops=dict(width=0.55, edgecolor="white"),
            pctdistance=0.75,
        )
        for t in autotexts:
            t.set_fontsize(9); t.set_color("white")
        ax.legend(wedges, [r[:28] for r in rl], loc="center left",
                  bbox_to_anchor=(1.0, 0.5), fontsize=8, frameon=False)
        ax.set_title("Regional distribution of retrieved studies", fontsize=11, fontweight="bold", color=NAVY, pad=8)

    elif kind == "methods":
        designs = [s.get("evaluation_design","").strip() for s in studies
                   if s.get("evaluation_design","").strip() not in ("","Not applicable","nan")]
        counts = Counter(designs)
        if not counts: return None
        dl = list(counts.keys()); dv = list(counts.values())
        dc = {"Experimental":RED, "Quasi-experimental":BLUE, "Qualitative":GREEN}
        dcols = [dc.get(d, GREY) for d in dl]
        wedges2, _, autotexts2 = ax.pie(
            dv, labels=None, colors=dcols, autopct="%1.0f%%",
            startangle=90, wedgeprops=dict(width=0.55, edgecolor="white"),
            pctdistance=0.75,
        )
        for t in autotexts2:
            t.set_fontsize(9); t.set_color("white")
        ax.legend(wedges2, dl, loc="center left",
                  bbox_to_anchor=(1.0, 0.5), fontsize=8, frameon=False)
        ax.set_title("Study designs in retrieved studies", fontsize=11, fontweight="bold", color=NAVY, pad=8)

    elif kind == "themes":
        themes = []
        for s in studies:
            t = s.get("primary_theme","").strip()
            if t and t.lower() not in ("","not applicable","nan"):
                themes.append(t)
        counts = Counter(themes).most_common(8)
        if not counts: return None
        labels = [x[0][:35] for x in counts][::-1]
        values = [x[1] for x in counts][::-1]
        palette = ["#0a2342","#1a3a5c","#1e5280","#3498db","#5dade2","#85c1e9","#aed6f1","#d6eaf8"]
        colors  = [palette[i % len(palette)] for i in range(len(labels))]
        bars = ax.barh(labels, values, color=colors, edgecolor="white", height=0.6)
        for bar, val in zip(bars, values):
            ax.text(bar.get_width()+0.1, bar.get_y()+bar.get_height()/2,
                    str(val), va="center", ha="left", fontsize=8, color="#444")
        ax.set_xlabel("Studies", fontsize=9, color="#666")
        ax.set_xlim(0, max(values)*1.15)
        ax.set_title("Themes in retrieved studies", fontsize=11, fontweight="bold", color=NAVY, pad=8)

    elif kind == "timeline":
        years = []
        for s in studies:
            try: years.append(int(float(s.get("year_of_publication",""))))
            except: pass
        if not years: return None
        yc = sorted(Counter(years).items())
        ax.plot([x[0] for x in yc], [x[1] for x in yc],
                color=RED, linewidth=2.5, marker="o", markersize=5)
        ax.fill_between([x[0] for x in yc], [x[1] for x in yc], alpha=0.12, color=RED)
        ax.set_xlabel("Year", fontsize=9, color="#666")
        ax.set_ylabel("Studies", fontsize=9, color="#666")
        ax.set_title("Publication years of retrieved studies", fontsize=11, fontweight="bold", color=NAVY, pad=8)

    elif kind == "sr_appraisal":
        plt.close(fig)
        _N="#0a2342"; _G="#2ecc71"; _A="#f39c12"; _R="#e74c3c"; _BG2="#f8f9fa"

        # Always use ALL SR studies from the full dataset for the appraisal chart
        sr_studies = [r for r in RECORDS if r.get("record_type","") == "srr"]

        if not sr_studies:
            fig2, ax2 = plt.subplots(figsize=(8, 3), facecolor=_BG2)
            ax2.text(0.5, 0.5, "No systematic reviews found in dataset",
                     ha="center", va="center", transform=ax2.transAxes, color="#aaa", fontsize=11)
            ax2.axis("off")
            return fig_to_html(fig2)

        n_sr = len(sr_studies)

        # All appraisal criteria
        criteria = {
            "Reliability of review":         "sr_reliability",
            "Confidence: analysis methods":  "sr_confidence_analysis",
            "Confidence: inclusion methods": "sr_confidence_inclusion",
            "Quality/bias criteria used":    "sr_quality_criteria",
            "Inclusion criteria reported":   "sr_inclusion_criteria",
            "Comprehensive search":          "sr_comprehensive_search",
            "Selection bias avoided":        "sr_selection_bias",
            "Independent screening":         "sr_independent_screening",
            "Independent RoB assessment":    "sr_independent_rob",
            "Independent data extraction":   "sr_data_extraction",
            "Language bias avoided":         "sr_language_bias",
            "Grey literature included":      "sr_grey_literature",
            "Relevant databases searched":   "sr_databases_searched",
        }

        YES_VALS     = {"yes","high","low risk","adequate"}
        PARTIAL_VALS = {"partially","partial","moderate","unclear","can't tell"}
        NO_VALS      = {"no","low","high risk","inadequate","not applicable"}

        labels_list, yes_c, partial_c, no_c, missing_c = [], [], [], [], []
        for label, col in criteria.items():
            vals = [str(s.get(col,"")).strip().lower() for s in sr_studies]
            y = sum(1 for v in vals if v in YES_VALS)
            p = sum(1 for v in vals if v in PARTIAL_VALS)
            n = sum(1 for v in vals if v in NO_VALS)
            m = n_sr - y - p - n
            labels_list.append(label)
            yes_c.append(y); partial_c.append(p)
            no_c.append(n);  missing_c.append(m)

        fig2, ax2 = plt.subplots(figsize=(11, 6), facecolor=_BG2)
        ax2.set_facecolor(_BG2)
        y_pos = list(range(len(labels_list)))

        ax2.barh(y_pos, yes_c,     color=_G,        edgecolor="white", label="Yes / High",  height=0.6)
        ax2.barh(y_pos, partial_c, color=_A,        edgecolor="white", label="Partial",      height=0.6,
                 left=yes_c)
        ax2.barh(y_pos, no_c,      color=_R,        edgecolor="white", label="No / Low",     height=0.6,
                 left=[y+p for y,p in zip(yes_c, partial_c)])
        ax2.barh(y_pos, missing_c, color="#d5d8dc", edgecolor="white", label="Not recorded", height=0.6,
                 left=[y+p+n for y,p,n in zip(yes_c, partial_c, no_c)])

        ax2.set_yticks(y_pos)
        ax2.set_yticklabels(labels_list, fontsize=9, color="#333")
        ax2.set_xlabel(f"Number of systematic reviews (n={n_sr})", fontsize=9, color="#666")
        ax2.set_xlim(0, n_sr + 0.5)
        ax2.set_xticks(range(0, n_sr + 1))
        ax2.set_title(f"Critical Appraisal of {n_sr} Systematic Reviews",
                      fontsize=12, fontweight="bold", color=_N, pad=10)
        ax2.legend(loc="lower right", fontsize=8, frameon=False, ncol=2)
        for spine in ax2.spines.values(): spine.set_color("#e2e8f0")
        ax2.tick_params(colors="#666")
        plt.tight_layout()
        return fig_to_html(fig2)


    ax.tick_params(labelsize=8, colors="#444")
    plt.tight_layout()
    return fig_to_html(fig)


def fig_to_html(fig):
    """Convert matplotlib figure to base64 PNG embedded as HTML img tag."""
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=130, bbox_inches="tight",
                facecolor=fig.get_facecolor())
    plt.close(fig)
    buf.seek(0)
    b64 = base64.b64encode(buf.read()).decode()
    return f'<img src="data:image/png;base64,{b64}" style="width:100%;border-radius:8px;margin-top:8px;">'


def make_chart_html(studies, query=""):
    """Generate exactly 1-2 charts most relevant to the query."""
    topics = detect_chart_topics(query)  # always returns exactly 2

    html_parts = []
    for topic in ["interventions","outcomes","geo","methods","themes","timeline","sr_appraisal"]:
        if topic in topics:
            try:
                chart = make_single_chart(topic, studies)
                if chart:
                    html_parts.append(chart)
            except Exception as e:
                html_parts.append(
                    f"<p style='color:#e74c3c;font-size:12px;'>Chart error ({topic}): {e}</p>"
                )

    if not html_parts:
        return EMPTY_HTML

    label_map = {
        "interventions": "Interventions",
        "outcomes": "Outcomes",
        "geo": "Geography",
        "methods": "Study designs",
        "themes": "Themes",
        "timeline": "Timeline",
        "sr_appraisal": "SR Critical Appraisal",
    }
    topic_labels = " + ".join(label_map[t] for t in topics if t in label_map)
    header = (
        f'''<div style="font-family:Inter,sans-serif;font-size:12px;color:#888;
        padding:4px 0 8px;">📊 <b style="color:#0a2342;">{topic_labels}</b>
        — based on {len(studies)} retrieved studies</div>'''
    )
    return header + "".join(html_parts)


EMPTY_HTML = """<div style="height:180px;display:flex;align-items:center;justify-content:center;
background:#f8f9fa;border-radius:8px;color:#adb5bd;font-size:14px;font-family:Inter,sans-serif;">
Ask a question — charts will show the profile of studies retrieved for that specific question</div>"""

# ---------------------------------------------------------------------------
# Chat — Gradio 5.x messages format
# ---------------------------------------------------------------------------
def chat(user_message, history, top_k):
    if not user_message.strip():
        yield history, EMPTY_HTML
        return

    studies  = retrieve(user_message, top_k)
    context  = build_context(studies)
    augmented = (
        f"RETRIEVED STUDIES ({len(studies)} of {TOTAL} total):\n\n{context}\n\n---\n"
        f"Using ONLY the studies above, answer:\n{user_message}"
    )

    api_messages = []
    for turn in history:
        role    = turn.get("role", "")
        content = turn.get("content", "")
        if isinstance(content, list):
            content = " ".join(b.get("text","") if isinstance(b,dict) else str(b) for b in content)
        if role in ("user","assistant") and content:
            api_messages.append({"role": role, "content": content})
    api_messages.append({"role": "user", "content": augmented})

    history = history + [
        {"role": "user",      "content": user_message},
        {"role": "assistant", "content": ""},
    ]

    partial = ""
    try:
        stream_ctx = anthropic_client.messages.stream(
            model="claude-sonnet-4-6",
            max_tokens=MAX_TOKENS,
            system=SYSTEM,
            messages=api_messages,
        )
        with stream_ctx as stream:
            for text in stream.text_stream:
                partial += text
                history[-1]["content"] = partial
                yield history, EMPTY_HTML
    except Exception as e:
        err = str(e)
        if "rate_limit" in err or "429" in err:
            history[-1]["content"] = "⚠️ **Rate limit reached.** Too many studies were sent at once. Please reduce the **Studies retrieved** slider to 20–25 and try again."
        else:
            history[-1]["content"] = f"⚠️ **API error:** {err[:300]}"
        yield history, EMPTY_HTML
        return

    # Append reference list
    cited = sorted(set(int(n) for n in re.findall(r"\[(\d+)\]", partial)))
    if cited:
        refs = ["\n\n---\n**Studies cited:**"]
        for n in cited:
            if 1 <= n <= len(studies):
                s   = studies[n-1]
                url = s.get("url_link","")
                lnk = f"[{s['title']}]({url})" if url and url.startswith("http") else s["title"]
                refs.append(
                    f"**[{n}]** {lnk}  \n"
                    f"*{s['authors']} ({s['year_of_publication']}) · "
                    f"{s['country']} · {s['evaluation_design']}*"
                )
        history[-1]["content"] += "\n".join(refs)

    try:
        chart_html = make_chart_html(studies, query=user_message)
    except Exception as e:
        chart_html = f"<p style='color:#e74c3c;font-family:sans-serif'>Chart could not be generated: {e}</p>"
    yield history, chart_html


def clear_chat():
    return [], EMPTY_HTML

# ---------------------------------------------------------------------------
# Gradio 5.x UI
# ---------------------------------------------------------------------------

# ---------------------------------------------------------------------------
# Dataset Summary — direct column counts, no RAG/LLM needed
# ---------------------------------------------------------------------------

# Columns available for summary, with friendly labels
SUMMARY_COLS = {
    "Intervention Category":       "intervention_category",
    "PSE Intervention Group":      "interv_PSE",
    "Outcome Category":            "outcome_categories_assigned",
    "PSE Outcome Group":           "outc_PSE",
    "Primary Theme":               "primary_theme",
    "Sub-theme":                   "sub_primary_theme",
    "Sector":                      "sector_name",
    "Evaluation Design":           "evaluation_design",
    "Method / Study Type":         "method_eval",
    "Region":                      "region",
    "Country":                     "country",
    "Income Level":                "income_level",
    "GPI Fragility Rank":          "gpi_rank",
    "Equity Focus":                "equity_focus",
    "Record Type":                 "record_type",
    "Publication Type":            "publication_type",
    "SDG":                         "un_sustainable_development_goal",
    "Funding Agency":              "program_funding_agency",
    "Implementation Agency":       "implementation_agencies",
    "Critical Appraisal of SRs":    "sr_appraisal_summary",
}

SUMMARY_COLOURS = [
    "#3498db","#2ecc71","#e74c3c","#f39c12","#9b59b6",
    "#1abc9c","#e67e22","#34495e","#e91e63","#00bcd4",
    "#8bc34a","#ff5722","#607d8b","#795548","#9c27b0",
]

def make_summary_chart(col_label):
    col = SUMMARY_COLS.get(col_label)
    if not col:
        return "<p>Select a column to summarise.</p>"

    # Virtual key — delegate to SR appraisal chart directly
    if col == "sr_appraisal_summary":
        para_white = "Critical appraisal ratings across all <b style=\"color:#ffffff;font-weight:700;\">13 systematic reviews</b> in the dataset, assessed across 13 quality criteria. Bars show <b style=\"color:#ffffff;font-weight:700;\">Yes/High</b> (green), <b style=\"color:#ffffff;font-weight:700;\">Partial</b> (amber), <b style=\"color:#ffffff;font-weight:700;\">No/Low</b> (red), and not recorded (grey). Most reviews are rated <b style=\"color:#ffffff;font-weight:700;\">Low</b> on reliability and confidence."
        desc_html = f"""<div style="background:linear-gradient(135deg,#0a2342 0%,#1a3a5c 100%);border-radius:10px;padding:18px 22px;margin-bottom:14px;font-family:Inter,sans-serif;font-size:14px;line-height:1.8;box-shadow:0 4px 16px rgba(10,35,66,0.25);"><div style="font-size:12px;font-weight:700;color:rgba(255,255,255,0.65);margin-bottom:8px;text-transform:uppercase;letter-spacing:0.8px;">Critical Appraisal of Systematic Reviews &nbsp;·&nbsp; 13 studies &nbsp;·&nbsp; 13 criteria</div><span style="color:#ffffff;display:block;">{para_white}</span></div>"""
        chart = make_single_chart("sr_appraisal", [])
        return desc_html + (chart or "<p>No SR data found.</p>")

    series = df[col].replace("", pd.NA).dropna()
    if series.empty:
        return f"<p>No data available for <b>{col_label}</b>.</p>"

    all_vals = []
    for v in series:
        for part in str(v).split("|"):
            part = part.strip()
            if part and part.lower() not in ("not applicable", "nan", "n/a", "multi-country"):
                all_vals.append(part)

    counts   = Counter(all_vals).most_common(20)
    if not counts:
        return f"<p>No data available for <b>{col_label}</b>.</p>"

    total    = sum(v for _, v in counts)
    n_cats   = len(Counter(all_vals))
    coverage = len(series)
    missing  = TOTAL - coverage
    top1     = counts[0]
    top2     = counts[1] if len(counts) > 1 else None
    top3     = counts[2] if len(counts) > 2 else None
    top3_sum = sum(v for _, v in counts[:3])

    # Full paragraph description
    cat_word = "category" if n_cats == 1 else "categories"
    rem_word = "category" if n_cats - 3 == 1 else "categories"
    para = (
        f"Across {coverage} studies with available data, "
        f"<b>{col_label}</b> spans <b>{n_cats} distinct {cat_word}</b>. "
        f"The most represented is <b>{top1[0]}</b> with {top1[1]} studies ({top1[1]/total*100:.1f}% of all values)"
    )
    if top2:
        para += f", followed by <b>{top2[0]}</b> ({top2[1]} studies, {top2[1]/total*100:.1f}%)"
    if top3:
        para += f" and <b>{top3[0]}</b> ({top3[1]} studies, {top3[1]/total*100:.1f}%)"
    para += f". Together, the top 3 categories account for <b>{top3_sum/total*100:.0f}%</b> of all values"
    if n_cats > 3:
        para += f", while the remaining {n_cats - 3} {rem_word} collectively represent {(total - top3_sum)/total*100:.0f}%"
    para += ". "
    if missing > 0:
        para += f"Note: {missing} of {TOTAL} studies have no value recorded for this field."
    else:
        para += f"Data is complete — all {TOTAL} studies have a value for this field."

    # Replace <b> tags with white inline spans so browser defaults don't override
    para_white = para.replace("<b>", '<b style="color:#ffffff;font-weight:700;">').replace("</b>", "</b>")
    desc_html = f"""<div style="background:linear-gradient(135deg,#0a2342 0%,#1a3a5c 100%);
        border-radius:10px;padding:18px 22px;margin-bottom:14px;
        font-family:Inter,sans-serif;font-size:14px;line-height:1.8;
        box-shadow:0 4px 16px rgba(10,35,66,0.25);">
        <div style="font-size:12px;font-weight:700;color:rgba(255,255,255,0.65);
            margin-bottom:8px;text-transform:uppercase;letter-spacing:0.8px;">
            {col_label} &nbsp;·&nbsp; {n_cats} categories &nbsp;·&nbsp; {coverage} studies
        </div>
        <span style="color:#ffffff;display:block;">{para_white}</span>
    </div>"""

    # GPI Rank — histogram
    if col == "gpi_rank":
        gpi_vals = []
        for v in series:
            try:
                fv = float(str(v).replace(",",""))
                if fv > 0:
                    gpi_vals.append(fv)
            except: pass
        if not gpi_vals:
            return desc_html + "<p style='color:#888;'>No GPI rank data available.</p>"
        fig2, ax = plt.subplots(figsize=(10, 4), facecolor="#f8f9fa")
        ax.set_facecolor("#f8f9fa")
        ax.hist(gpi_vals, bins=20, color="#9b59b6", edgecolor="white", linewidth=0.5)
        ax.set_xlabel("GPI Rank  (← more fragile | less fragile →)", fontsize=10, color="#666")
        ax.set_ylabel("Number of studies", fontsize=10, color="#666")
        ax.tick_params(labelsize=9, colors="#444")
        for spine in ax.spines.values(): spine.set_color("#e2e8f0")
        plt.tight_layout()
        return desc_html + fig_to_html(fig2)

    # World map for Country
    if col == "country":
        try:
            import geopandas as gpd
        except ImportError:
            # Fallback: bar chart of top countries
            country_counts = Counter(all_vals)
            top_c = country_counts.most_common(25)
            fig2, ax = plt.subplots(figsize=(10, 8), facecolor="#f8f9fa")
            ax.set_facecolor("#f8f9fa")
            c_labels = [x[0][:25] for x in top_c][::-1]
            c_values = [x[1] for x in top_c][::-1]
            c_colors = ["#0a2342" if v == max(c_values) else
                        "#1a5276" if v >= max(c_values)*0.5 else
                        "#3498db" if v >= max(c_values)*0.2 else
                        "#85c1e9" for v in c_values]
            bars = ax.barh(c_labels, c_values, color=c_colors, edgecolor="white", height=0.7)
            for bar, val in zip(bars, c_values):
                ax.text(bar.get_width() + 0.1, bar.get_y() + bar.get_height()/2,
                        str(val), va="center", ha="left", fontsize=8, color="#444")
            ax.set_xlabel("Studies", fontsize=10, color="#666")
            ax.set_title("Top 25 Countries by Study Count", fontsize=11,
                         fontweight="bold", color="#0a2342", pad=8)
            ax.tick_params(labelsize=8, colors="#444")
            for spine in ax.spines.values(): spine.set_color("#e2e8f0")
            plt.tight_layout()
            return desc_html + fig_to_html(fig2)

    # Bar chart for all other columns
    bar_labels = [x[0][:45] for x in counts][::-1]
    bar_values = [x[1] for x in counts][::-1]
    bar_pcts   = [f"{v/total*100:.1f}%" for v in bar_values]
    bar_colors = [SUMMARY_COLOURS[i % len(SUMMARY_COLOURS)] for i in range(len(bar_labels))]

    fig2, ax = plt.subplots(figsize=(10, max(4, len(bar_labels) * 0.5 + 1)),
                            facecolor="#f8f9fa")
    ax.set_facecolor("#f8f9fa")
    bars = ax.barh(bar_labels, bar_values, color=bar_colors, edgecolor="white", height=0.65)
    for bar, val, pct in zip(bars, bar_values, bar_pcts):
        ax.text(bar.get_width() + 0.1, bar.get_y() + bar.get_height()/2,
                f"{val}  ({pct})", va="center", ha="left", fontsize=8, color="#444")
    ax.set_xlabel("Number of studies", fontsize=10, color="#666")
    ax.set_xlim(0, max(bar_values) * 1.25)
    ax.tick_params(labelsize=9, colors="#444")
    for spine in ax.spines.values(): spine.set_color("#e2e8f0")
    plt.tight_layout()
    return desc_html + fig_to_html(fig2)

# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
# Evidence Gap Analysis
# ---------------------------------------------------------------------------

INTERVENTIONS = [
    "Capacity development of private, public, or hybrid sector actors",
    "Public-private partnerships (PPPs)",
    "Multicomponent",
    "Financing with private sector actors",
    "Financing of private sector actors",
    "Matchmaking and consulting",
    "Structured dialogue",
]
OUTCOMES = ["Work-Labor/Social","Jobs and Income","Market/Sector Growth",
            "Coordination/Reforms","Environmental","Energy/Extractives"]
SDGS = [
    "Decent Work and Economic Growth","Good Health and Well-being","Quality Education",
    "No Poverty","Industry, Innovation and Infrastructure","Reduced Inequality",
    "Gender Equality","Zero Hunger","Climate Action",
]
INTERV_PSE = [
    "Technical and vocational training",
    "Capacity strengthening to enhance private sector engagement",
    "Service management and delivery PPP",
    "Direct financing of private sector actors",
    "Payments by results",
    "Matchmaking",
    "Infrastructure development and financing",
    "Advisory support",
    "Multi-stakeholder platforms",
    "Blended finance",
]
OUTC_PSE = [
    "Employment","Income","Food security, health, and nutrition","Education",
    "Domestic business growth and market development",
    "Training, knowledge, and technology transfer",
    "Social factors","Productivity gains","Financial inclusion","International trade",
]

INT_SHORT = {
    "Capacity development of private, public, or hybrid sector actors": "Capacity Dev.",
    "Public-private partnerships (PPPs)": "PPPs",
    "Multicomponent": "Multicomponent",
    "Financing with private sector actors": "Financing (with)",
    "Financing of private sector actors": "Financing (of)",
    "Matchmaking and consulting": "Matchmaking",
    "Structured dialogue": "Structured Dialogue",
}
SDG_SHORT = {
    "Decent Work and Economic Growth": "SDG8: Decent Work",
    "Good Health and Well-being": "SDG3: Health",
    "Quality Education": "SDG4: Education",
    "No Poverty": "SDG1: No Poverty",
    "Industry, Innovation and Infrastructure": "SDG9: Industry",
    "Reduced Inequality": "SDG10: Inequality",
    "Gender Equality": "SDG5: Gender",
    "Zero Hunger": "SDG2: Zero Hunger",
    "Climate Action": "SDG13: Climate",
}


def build_matrix(row_keys, col_keys, row_col, col_col, split_col=False, dff=None):
    if dff is None:
        dff = df
    mat = pd.DataFrame(0, index=row_keys, columns=col_keys)
    for _, row in dff.iterrows():
        rval = str(row.get(row_col, "")).strip()
        if not rval or rval not in row_keys:
            # also try splitting row value
            for rp in rval.split("|"):
                rp = rp.strip()
                if rp in row_keys:
                    cvals = str(row.get(col_col, "")).split("|") if split_col else [str(row.get(col_col,""))]
                    for cv in cvals:
                        cv = cv.strip()
                        if cv in col_keys:
                            mat.loc[rp, cv] += 1
            continue
        cvals = str(row.get(col_col, "")).split("|") if split_col else [str(row.get(col_col,""))]
        for cv in cvals:
            cv = cv.strip()
            if cv in col_keys:
                mat.loc[rval, cv] += 1
    return mat


def make_gap_heatmap(matrix, row_labels, col_labels, title, height=420):
    """Render gap heatmap as PNG using matplotlib."""
    z    = matrix.values
    rows = [row_labels.get(r, r) for r in matrix.index]
    cols = [col_labels.get(c, c) for c in matrix.columns]

    nrows, ncols = z.shape
    fig, ax = plt.subplots(figsize=(max(10, ncols * 1.4), max(5, nrows * 0.8)),
                           facecolor="#f8f9fa")
    ax.set_facecolor("#f8f9fa")

    # Custom colormap: red (0) → amber → yellow → blue → dark navy
    import matplotlib.colors as mcolors
    cmap = mcolors.LinearSegmentedColormap.from_list("egm", [
        "#e74c3c", "#f39c12", "#f1c40f", "#3498db", "#0a2342"
    ])
    norm = mcolors.Normalize(vmin=0, vmax=max(z.max(), 1))

    im = ax.imshow(z, cmap=cmap, norm=norm, aspect="auto")

    # Annotations
    max_val = z.max() or 1
    for i in range(nrows):
        for j in range(ncols):
            val = z[i, j]
            txt = "GAP" if val == 0 else str(int(val))
            col = "white" if (val == 0 or val / max_val > 0.3) else "#2c3e50"
            ax.text(j, i, txt, ha="center", va="center",
                    fontsize=9, color=col, fontweight="bold" if val == 0 else "normal")

    ax.set_xticks(range(ncols))
    ax.set_xticklabels([c[:30] for c in cols], rotation=40, ha="right", fontsize=8, color="#444")
    ax.set_yticks(range(nrows))
    ax.set_yticklabels(rows, fontsize=8, color="#444")
    ax.set_title(title, fontsize=12, fontweight="bold", color="#0a2342", pad=10)

    cbar = fig.colorbar(im, ax=ax, fraction=0.03, pad=0.02)
    cbar.set_label("Studies", fontsize=9, color="#666")
    cbar.ax.tick_params(labelsize=8)

    for spine in ax.spines.values():
        spine.set_visible(False)
    plt.tight_layout()
    return fig


def make_gap_summary(matrix, row_labels, col_labels, type_label):
    total_cells   = matrix.size
    gap_cells     = int((matrix == 0).sum().sum())
    strong_cells  = int((matrix >= 10).sum().sum())
    limited_cells = int(((matrix >= 1) & (matrix < 5)).sum().sum())

    flat = matrix.stack().reset_index()
    flat.columns = ["row", "col", "count"]
    flat = flat[flat["count"] > 0].sort_values("count", ascending=False)

    top_combo = ""
    if not flat.empty:
        t = flat.iloc[0]
        top_combo = (
            f"The best-evidenced combination is "
            f"<b>{row_labels.get(t['row'], t['row'])}</b> × "
            f"<b>{col_labels.get(t['col'], t['col'])}</b> ({int(t['count'])} studies). "
        )

    gaps = [(row_labels.get(r,""), col_labels.get(c,""))
            for r in matrix.index for c in matrix.columns if matrix.loc[r,c] == 0]
    gap_examples = "; ".join(f"{r} × {c}" for r, c in gaps[:3])

    parts = type_label.split(" × ")
    dim2  = parts[1] if len(parts) > 1 else type_label

    para = (
        f"The matrix covers <b>{total_cells} possible combinations</b> "
        f"({len(matrix.index)} rows × {len(matrix.columns)} {dim2} values). "
        f"<b>{gap_cells} combinations ({gap_cells/total_cells*100:.0f}%) have no evidence</b> — "
        f"clear gaps where research is needed. "
        f"<b>{strong_cells} combinations ({strong_cells/total_cells*100:.0f}%)</b> have 10+ studies. "
        f"{top_combo}"
        + (f"Key gaps include: {gap_examples}." if gaps else "")
    )

    badges = (
        f'''<div style="display:flex;gap:10px;margin-top:12px;flex-wrap:wrap;">
        <span style="background:#e74c3c;color:white;border-radius:6px;padding:5px 12px;font-size:12px;font-weight:600;">🔴 {gap_cells} Gaps</span>
        <span style="background:#e67e22;color:white;border-radius:6px;padding:5px 12px;font-size:12px;font-weight:600;">🟡 {limited_cells} Limited (&lt;5)</span>
        <span style="background:#1a5276;color:white;border-radius:6px;padding:5px 12px;font-size:12px;font-weight:600;">🔵 {strong_cells} Strong (10+)</span>
        </div>'''
    )

    para_white = para.replace("<b>", '<b style="color:#ffffff;font-weight:700;">').replace("</b>", "</b>")
    return f'''<div style="background:linear-gradient(135deg,#0a2342 0%,#1a3a5c 100%);
        border-radius:10px;padding:18px 22px;margin-bottom:14px;
        font-family:Inter,sans-serif;font-size:14px;line-height:1.8;
        box-shadow:0 4px 16px rgba(10,35,66,0.25);">
        <div style="font-size:12px;font-weight:700;color:rgba(255,255,255,0.65);
            margin-bottom:8px;text-transform:uppercase;letter-spacing:0.8px;">
            Evidence Gap Analysis &nbsp;·&nbsp; {type_label}
        </div>
        <span style="color:#ffffff;display:block;">{para_white}</span>
        {badges}
    </div>'''


def make_gap_tab_content(analysis_type, method_filter="All"):
    if method_filter == "IE Quantitative":
        dff = df[df["method_eval"] == "IE Quantitative"]
    elif method_filter == "IE Qualitative":
        dff = df[df["method_eval"] == "IE Qualitative"]
    elif method_filter == "SR (Systematic Review)":
        dff = df[df["method_eval"] == "SR"]
    else:
        dff = df

    n   = len(dff)
    sfx = f" [{method_filter}, {n} studies]" if method_filter != "All" else ""

    configs = {
        "Intervention Category × Outcome Category": dict(
            rows=INTERVENTIONS, cols=OUTCOMES,
            rc="intervention_category", cc="outcome_categories_assigned",
            rl=INT_SHORT, cl={o: o for o in OUTCOMES}, h=400,
        ),
        "Intervention Category × SDG": dict(
            rows=INTERVENTIONS, cols=SDGS,
            rc="intervention_category", cc="un_sustainable_development_goal",
            rl=INT_SHORT, cl=SDG_SHORT, h=420,
        ),
        "Outcome Category × SDG": dict(
            rows=OUTCOMES, cols=SDGS,
            rc="outcome_categories_assigned", cc="un_sustainable_development_goal",
            rl={o: o for o in OUTCOMES}, cl=SDG_SHORT, h=380,
        ),
        "PSE Intervention × PSE Outcome": dict(
            rows=INTERV_PSE, cols=OUTC_PSE,
            rc="interv_PSE", cc="outc_PSE",
            rl={i: i for i in INTERV_PSE}, cl={o: o for o in OUTC_PSE}, h=480,
        ),
        "PSE Intervention × SDG": dict(
            rows=INTERV_PSE, cols=SDGS,
            rc="interv_PSE", cc="un_sustainable_development_goal",
            rl={i: i for i in INTERV_PSE}, cl=SDG_SHORT, h=480,
        ),
        "PSE Outcome × SDG": dict(
            rows=OUTC_PSE, cols=SDGS,
            rc="outc_PSE", cc="un_sustainable_development_goal",
            rl={o: o for o in OUTC_PSE}, cl=SDG_SHORT, h=460,
        ),
    }

    cfg = configs.get(analysis_type, configs["Intervention Category × Outcome Category"])
    mat = build_matrix(cfg["rows"], cfg["cols"], cfg["rc"], cfg["cc"],
                       split_col=True, dff=dff)
    summary_html = make_gap_summary(mat, cfg["rl"], cfg["cl"], analysis_type + sfx)
    fig          = make_gap_heatmap(mat, cfg["rl"], cfg["cl"],
                                    f"Evidence Map: {analysis_type}{sfx}", height=cfg["h"])
    try:
        chart_html = fig_to_html(fig)
    except Exception as e:
        chart_html = f"<p style='color:red'>Chart error: {e}</p>"
    return summary_html + chart_html


# ---------------------------------------------------------------------------
# Custom CSS — 3ie professional dark navy theme
# ---------------------------------------------------------------------------
CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');

* { font-family: 'Inter', sans-serif; }

/* Main background */
.gradio-container { background: #f0f2f5; max-width: 1200px; margin: 0 auto; }

/* Header banner */
.egm-header {
    background: linear-gradient(135deg, #0a2342 0%, #1a3a5c 60%, #1e5280 100%);
    color: white;
    padding: 28px 32px 22px;
    border-radius: 12px;
    margin-bottom: 20px;
    box-shadow: 0 4px 20px rgba(10,35,66,0.25);
}
.egm-header * { color: #ffffff !important; }
.egm-header h1 { font-size: 24px; font-weight: 700; margin: 0 0 4px; letter-spacing: -0.3px; color: #ffffff !important; }
.egm-header p  { font-size: 13px; opacity: 0.75; margin: 0; font-weight: 400; color: #ffffff !important; }
.egm-badges { display: flex; gap: 8px; margin-top: 12px; flex-wrap: wrap; }
.egm-badge {
    background: rgba(255,255,255,0.12);
    border: 1px solid rgba(255,255,255,0.2);
    border-radius: 20px;
    padding: 3px 12px;
    font-size: 11px;
    font-weight: 500;
    color: rgba(255,255,255,0.9);
}

/* Tabs */
.tab-nav button {
    font-weight: 600;
    font-size: 13px;
    color: #4a5568;
    border-radius: 8px 8px 0 0;
}
.tab-nav button.selected {
    color: #0a2342;
    border-bottom: 3px solid #0a2342;
    background: white;
}

/* Chat bubbles — Gradio 5 selectors */
.message.user > div, .user > .bubble-wrap, [data-testid="user"] .prose,
.message.user .prose, div.user { background: #0a2342 !important; color: #ffffff !important; border-radius: 14px 14px 4px 14px !important; }
.message.user > div *, .message.user p, .message.user span { color: #ffffff !important; }
.message.bot > div, .bot > .bubble-wrap, [data-testid="bot"] .prose,
.message.bot .prose   { background: white !important; border: 1px solid #e2e8f0 !important; border-radius: 14px 14px 14px 4px !important; box-shadow: 0 1px 4px rgba(0,0,0,0.06) !important; }

/* Input box */
textarea, input[type=text] {
    border: 1.5px solid #cbd5e0;
    border-radius: 10px;
    background: white;
    font-size: 14px;
    transition: border-color 0.2s;
}
textarea:focus, input[type=text]:focus { border-color: #0a2342; box-shadow: 0 0 0 3px rgba(10,35,66,0.08); }

/* Buttons */
.gr-button-primary {
    background: #0a2342;
    border: none;
    border-radius: 10px;
    font-weight: 600;
    letter-spacing: 0.2px;
    transition: background 0.2s;
}
.gr-button-primary:hover { background: #1a3a5c; }
.gr-button-secondary { border-color: #cbd5e0; border-radius: 10px; font-weight: 500; }

/* Panels / blocks */
.gr-panel, .gr-box { border-radius: 12px; border: 1px solid #e2e8f0; box-shadow: 0 1px 6px rgba(0,0,0,0.05); }

/* Slider */
.gr-slider input[type=range]::-webkit-slider-thumb { background: #0a2342; }

/* Dropdown */
.gr-dropdown { border-radius: 10px; }

/* Footer */
.egm-footer { font-size: 11px; color: #94a3b8; text-align: center; padding: 12px 0 4px; }

/* Ensure bold text inside navy cards stays white */
.egm-card b, .egm-card strong { color: #ffffff; }

"""

with gr.Blocks(title="EGM Research Assistant", theme=gr.themes.Base(), css=CUSTOM_CSS) as demo:
    gr.HTML(f"""
    <div class="egm-header">
        <h1>📚 EGM Research Assistant</h1>
        <p>Evidence Gap Map · AI-powered research synthesis · Dataset-only responses</p>
        <div class="egm-badges">
            <span class="egm-badge">📚 {TOTAL} studies</span>
            <span class="egm-badge">🔍 Semantic RAG</span>
            <span class="egm-badge">✅ Dataset-only responses</span>
            <span class="egm-badge">📊 Evidence visualised</span>
        </div>
    </div>
    """)

    with gr.Tabs():

        # ── Tab 1: Chat ──────────────────────────────────────────────────
        with gr.TabItem("💬 Research Chat"):
            with gr.Row():
                with gr.Column(scale=4):
                    chatbot = gr.Chatbot(label="Research conversation",
                                        height=480, type="messages")
                    with gr.Row():
                        msg = gr.Textbox(
                            placeholder="e.g. What does the evidence say about cash transfers on education outcomes?",
                            label="Your question", lines=2, scale=5)
                        send_btn = gr.Button("Send ↗", variant="primary", scale=1)

                with gr.Column(scale=1, min_width=220):
                    gr.Markdown("### ⚙️ Settings")
                    top_k = gr.Slider(minimum=10, maximum=358, value=25, step=5,
                        label="Studies retrieved per query",
                        info="Charts show ONLY retrieved studies — lower = more focused on your question")
                    clear_btn = gr.Button("🗑 Clear", variant="secondary")
                    gr.Markdown("### 💡 Try asking")
                    gr.Examples(examples=[
                        ["What interventions improve employment outcomes?"],
                        ["Evidence on cash transfers for education in Sub-Saharan Africa"],
                        ["How effective is microfinance for women's empowerment?"],
                        ["Which studies focus on fragile or conflict-affected states?"],
                        ["What qualitative studies exist in the dataset?"],
                    ], inputs=msg)

            chat_chart = gr.HTML(value=EMPTY_HTML, label="Evidence snapshot")

            send_btn.click(fn=chat, inputs=[msg, chatbot, top_k],
                outputs=[chatbot, chat_chart]).then(lambda: "", outputs=msg)
            msg.submit(fn=chat, inputs=[msg, chatbot, top_k],
                outputs=[chatbot, chat_chart]).then(lambda: "", outputs=msg)
            clear_btn.click(fn=clear_chat, outputs=[chatbot, chat_chart])

        # ── Tab 2: Dataset Summary ────────────────────────────────────────
        with gr.TabItem("📊 Dataset Summary"):
            # Load Plotly JS once — shared by all charts in this tab
            gr.Markdown("""
### Explore the spread of any column in the dataset
Select a column below to instantly see how studies are distributed across categories.
No AI involved — this reads directly from the dataset.
            """)
            col_dropdown = gr.Dropdown(
                choices=list(SUMMARY_COLS.keys()),
                value="PSE Intervention Group",
                label="Select column to summarise",
                interactive=True,
            )
            summary_chart = gr.HTML(
                value=make_summary_chart("Intervention Category"),
                label="Column distribution"
            )
            col_dropdown.change(
                fn=make_summary_chart,
                inputs=col_dropdown,
                outputs=summary_chart,
            )


        # ── Tab 3: Evidence Gaps ─────────────────────────────────────────
        with gr.TabItem("🔴 Evidence Gaps"):
            gr.Markdown("""
### Evidence Gap Analysis
Heatmaps showing where evidence exists and where it is missing. Cells show study count or **GAP** where no evidence exists.
🔴 Red = no studies &nbsp;|&nbsp; 🟡 Amber = limited (&lt;5) &nbsp;|&nbsp; 🔵 Navy = strong (10+)
            """)
            with gr.Row():
                gap_type = gr.Radio(
                    choices=[
                        "Intervention Category × Outcome Category",
                        "Intervention Category × SDG",
                        "Outcome Category × SDG",
                        "PSE Intervention × PSE Outcome",
                        "PSE Intervention × SDG",
                        "PSE Outcome × SDG",
                    ],
                    value="Intervention Category × Outcome Category",
                    label="Analysis type",
                    interactive=True,
                )
                gap_method = gr.Radio(
                    choices=["All", "IE Quantitative", "IE Qualitative", "SR (Systematic Review)"],
                    value="All",
                    label="Filter by method type",
                    interactive=True,
                )
            gap_output = gr.HTML(
                value=make_gap_tab_content("Intervention Category × Outcome Category", "All")
            )
            gap_type.change(fn=make_gap_tab_content, inputs=[gap_type, gap_method], outputs=gap_output)
            gap_method.change(fn=make_gap_tab_content, inputs=[gap_type, gap_method], outputs=gap_output)


    gr.HTML('<div class="egm-footer">Responses grounded exclusively in the EGM dataset · sentence-transformers + ChromaDB + Claude </div>')

if __name__ == "__main__":
    demo.launch()