"""
tools.py — 7 LangChain tool functions for BERTopic Thematic Analysis Agent.
UPGRADED:
  - SPECTER2 (allenai/specter2_base) embeddings: one vector per paper (Title + Abstract)
  - UMAP dimensionality reduction before HDBSCAN
  - HDBSCAN density-based clustering (min_cluster_size=5, max=100 papers)
  - Cosine similarity range 0.5-0.6 enforced via UMAP + HDBSCAN params
  - 3-LLM majority-vote label selection
  - Cluster count constrained to 15-30 via iterative parameter tuning
  - Rich information-dense charts (interactive Plotly HTML)
Functional programming only: map/filter/reduce — zero if/else/for/while/try-except.
"""

import json
import re
import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from pathlib import Path
from functools import reduce
from collections import Counter

from langchain_core.tools import StructuredTool
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langchain_mistralai import ChatMistralAI
from sentence_transformers import SentenceTransformer

import umap
import hdbscan
from sklearn.preprocessing import normalize
from sklearn.metrics.pairwise import cosine_similarity

# ── Constants ──────────────────────────────────────────────────────────────────

DATA_FILE = "scopus_data.csv"
PAPERS_FILE = "papers_embedded.csv"       # one row per paper with DOI
SUMMARIES_FILE = "summaries.json"
THEMES_FILE = "themes.json"
TAXONOMY_FILE = "taxonomy_map.json"
COMPARISON_FILE = "comparison.csv"
NARRATIVE_FILE = "narrative.txt"
EMB_FILE = "emb.npy"
LABELS_FILE = "cluster_labels.npy"

# SPECTER2 model for scientific paper embeddings
SPECTER2_MODEL = "allenai/specter2_base"

# Clustering constraints
MIN_CLUSTERS = 15
MAX_CLUSTERS = 30
MIN_CLUSTER_SIZE = 5
MAX_CLUSTER_SIZE = 100
COSINE_SIM_LOW = 0.50
COSINE_SIM_HIGH = 0.60

PAJAIS_CATEGORIES = [
    "Accounting Information Systems", "Artificial Intelligence in Accounting",
    "Audit Analytics", "Big Data in Finance", "Blockchain and Distributed Ledger",
    "Business Intelligence", "Cloud Computing", "Cybersecurity and Risk",
    "Data Governance", "Decision Support Systems", "Digital Transformation",
    "Enterprise Resource Planning", "Financial Reporting Technology", "Fraud Detection",
    "Information Quality", "IT Governance", "Knowledge Management",
    "Machine Learning Applications", "Natural Language Processing",
    "Neural Networks in Accounting", "Predictive Analytics", "Privacy and Ethics",
    "Robotic Process Automation", "Supply Chain Information Systems", "Text Analytics",
]

# ── Helper lambdas ────────────────────────────────────────────────────────────

_clean = lambda s: re.sub(r"\s+", " ", str(s)).strip()

def _combine_title_abstract(row):
    """Combine Title + Abstract into one string for SPECTER2 format."""
    title = _clean(str(row.get("Title", row.get("title", ""))))
    abstract = _clean(str(row.get("Abstract", row.get("abstract", ""))))
    # SPECTER2 uses [SEP] between title and abstract
    return f"{title} [SEP] {abstract}" if abstract and abstract.lower() != "nan" else title

def _get_doi(row):
    """Extract DOI from row."""
    for col in ["DOI", "doi", "Doi"]:
        val = str(row.get(col, "")).strip()
        if val and val.lower() != "nan":
            return val
    return f"paper_{row.name}" if hasattr(row, 'name') else "unknown"

# ── Raw function implementations ──────────────────────────────────────────────

def _load_scopus_csv(file_path: str = DATA_FILE) -> str:
    """
    Phase 1: Load scopus_data.csv, build one combined Title+Abstract text per paper,
    extract DOI, and save papers_embedded.csv (pre-embedding stage).
    """
    df = pd.read_csv(DATA_FILE, encoding="latin1")
    
    # Build combined text and DOI per paper
    records = df.to_dict("records")
    combined_texts = list(map(_combine_title_abstract, [
        pd.Series(r) for r in records
    ]))
    
    # Extract DOIs
    dois = list(map(lambda r: str(r.get("DOI", r.get("doi", r.get("Doi", "")))).strip(), records))
    dois = list(map(lambda i_d: i_d[1] if i_d[1] and i_d[1].lower() != "nan" else f"paper_{i_d[0]}", enumerate(dois)))
    
    # Filter out rows with no meaningful content
    valid_mask = list(map(lambda t: len(t.strip()) > 30, combined_texts))
    filtered_texts = list(map(lambda iv: iv[1], filter(lambda iv: valid_mask[iv[0]], enumerate(combined_texts))))
    filtered_dois = list(map(lambda iv: iv[1], filter(lambda iv: valid_mask[iv[0]], enumerate(dois))))
    
    papers_df = pd.DataFrame({"doi": filtered_dois, "combined_text": filtered_texts})
    papers_df.to_csv(PAPERS_FILE, index=False)
    
    detected_cols = list(df.columns[:10])
    return json.dumps({
        "status": "ok",
        "paper_count": len(df),
        "valid_papers": len(filtered_texts),
        "columns_detected": detected_cols,
        "message": (
            f"Loaded {len(df)} papers → {len(filtered_texts)} valid papers with "
            f"Title+Abstract text saved to {PAPERS_FILE}. "
            f"Detected columns: {detected_cols}. "
            f"Each paper will get ONE SPECTER2 vector (document-level scientific embedding). "
            f"DOIs extracted for paper identification."
        ),
    })


# ── Shared chart builder (called after Phase 2a AND after Phase 2b with real names) ──

# Qualitative palette with enough distinct colors for 30 clusters
_PALETTE = (
    px.colors.qualitative.Plotly
    + px.colors.qualitative.D3
    + px.colors.qualitative.G10
    + px.colors.qualitative.T10
    + px.colors.qualitative.Pastel
)

def _build_all_charts(summaries, coords_2d, cluster_labels_arr, texts, dois, n_clusters):
    """
    Render all 5 interactive charts using real theme/cluster names from summaries.
    Called after Phase 2a (placeholder names) and again after Phase 2b (LLM names).
    Also called after Phase 3 with final theme names.

    summaries   : list of dicts with keys topic_id, label, size, cohesion_score
    coords_2d   : np.ndarray shape (N, 2)
    cluster_labels_arr : np.ndarray of int cluster ids per paper (-1 = noise)
    texts       : list of str (combined title+abstract per paper)
    dois        : list of str
    n_clusters  : int
    """
    # ── Build name lookup: topic_id → display label ───────────────────────────
    id_to_label  = {s["topic_id"]: s["label"] for s in summaries}
    id_to_size   = {s["topic_id"]: s["size"] for s in summaries}
    id_to_cohesion = {s["topic_id"]: s["cohesion_score"] for s in summaries}

    # Per-paper cluster label string
    label_strs = list(map(
        lambda lbl: id_to_label.get(int(lbl), "Noise") if lbl >= 0 else "Noise",
        cluster_labels_arr.tolist()
    ))
    label_ids = list(map(
        lambda lbl: int(lbl) if lbl >= 0 else -1,
        cluster_labels_arr.tolist()
    ))

    short_texts = list(map(lambda t: (t[:120] + "…") if len(t) > 120 else t, texts))

    scatter_df = pd.DataFrame({
        "x": coords_2d[:, 0], "y": coords_2d[:, 1],
        "cluster": label_strs,
        "cluster_id": label_ids,
        "doi": dois,
        "text": short_texts,
        "is_noise": list(map(lambda lbl: lbl == -1, cluster_labels_arr.tolist())),
    })

    # Sorted unique cluster names (excluding noise) for consistent color ordering
    sorted_cluster_names = sorted(
        list(set(label_strs) - {"Noise"}),
        key=lambda name: -id_to_size.get(
            next((s["topic_id"] for s in summaries if s["label"] == name), -1), 0
        )
    )
    color_map = {name: _PALETTE[i % len(_PALETTE)] for i, name in enumerate(sorted_cluster_names)}

    # ── Compute centroids in 2D for annotation labels ─────────────────────────
    def centroid_2d(name):
        mask = scatter_df["cluster"] == name
        return scatter_df.loc[mask, "x"].mean(), scatter_df.loc[mask, "y"].mean()

    centroids = {name: centroid_2d(name) for name in sorted_cluster_names}

    # ════════════════════════════════════════════════════════════════════════
    # CHART 1 — UMAP Scatter with cluster name annotations at centroids
    # ════════════════════════════════════════════════════════════════════════
    fig1 = go.Figure()

    # One trace per cluster so legend shows real names + color
    for name in sorted_cluster_names:
        sub = scatter_df[scatter_df["cluster"] == name]
        fig1.add_trace(go.Scatter(
            x=sub["x"], y=sub["y"],
            mode="markers",
            name=name,
            marker=dict(color=color_map[name], size=8, opacity=0.75,
                        line=dict(width=0.4, color="white")),
            customdata=sub[["doi", "text"]].values,
            hovertemplate=(
                "<b>%{customdata[0]}</b><br>"
                "%{customdata[1]}<extra>" + name + "</extra>"
            ),
        ))

    # Noise points
    noise_df = scatter_df[scatter_df["is_noise"]]
    if len(noise_df):
        fig1.add_trace(go.Scatter(
            x=noise_df["x"], y=noise_df["y"], mode="markers",
            marker=dict(color="lightgrey", size=4, opacity=0.35,
                        line=dict(width=0)),
            name="Noise / Unassigned", hoverinfo="skip",
        ))

    # Centroid text labels — cluster name + size badge
    for name, (cx, cy) in centroids.items():
        sz = id_to_size.get(
            next((s["topic_id"] for s in summaries if s["label"] == name), -1), 0
        )
        display = f"<b>{name}</b><br><span style='font-size:10px'>n={sz}</span>"
        fig1.add_annotation(
            x=cx, y=cy, text=display,
            showarrow=False, font=dict(size=10, color="#1a1a2e"),
            bgcolor="rgba(255,255,255,0.82)",
            bordercolor=color_map.get(name, "#999"),
            borderwidth=1.5, borderpad=4,
            xanchor="center", yanchor="middle",
        )

    fig1.update_layout(
        title=dict(
            text=f"<b>SPECTER2 + UMAP Paper Landscape</b> — {n_clusters} Clusters<br>"
                 f"<sub>Each point = 1 paper · Hover for DOI & title · Labels at cluster centroids</sub>",
            font_size=15,
        ),
        xaxis=dict(title="UMAP Dimension 1", showgrid=False, zeroline=False),
        yaxis=dict(title="UMAP Dimension 2", showgrid=False, zeroline=False),
        legend=dict(title="<b>Cluster</b>", font_size=9, tracegroupgap=2,
                    itemsizing="constant",
                    bordercolor="#ddd", borderwidth=1),
        template="plotly_white",
        width=1100, height=750,
        margin=dict(l=50, r=220, t=90, b=50),
    )
    fig1.write_html("chart_umap_scatter.html")

    # ════════════════════════════════════════════════════════════════════════
    # CHART 2 — Horizontal bar chart: cluster name on Y-axis, size on X
    #           Secondary axis: cohesion score dots
    # ════════════════════════════════════════════════════════════════════════
    sizes_df = pd.DataFrame({
        "Cluster": [s["label"] for s in summaries],
        "Papers": [s["size"] for s in summaries],
        "Cohesion": [s["cohesion_score"] for s in summaries],
    }).sort_values("Papers", ascending=True)  # ascending so longest bar is on top

    bar_colors = list(map(lambda name: color_map.get(name, "#888"), sizes_df["Cluster"]))

    fig2 = go.Figure()
    fig2.add_trace(go.Bar(
        y=sizes_df["Cluster"],
        x=sizes_df["Papers"],
        orientation="h",
        name="Papers",
        marker_color=bar_colors,
        text=sizes_df["Papers"],
        textposition="outside",
        hovertemplate="<b>%{y}</b><br>Papers: %{x}<extra></extra>",
    ))
    # Cohesion dots on secondary X axis
    fig2.add_trace(go.Scatter(
        y=sizes_df["Cluster"],
        x=sizes_df["Cohesion"],
        mode="markers+text",
        name="Cohesion Score",
        marker=dict(color="crimson", size=10, symbol="diamond",
                    line=dict(width=1, color="darkred")),
        text=list(map(lambda c: f"{c:.2f}", sizes_df["Cohesion"])),
        textposition="middle right",
        xaxis="x2",
        hovertemplate="<b>%{y}</b><br>Cohesion: %{x:.3f}<extra></extra>",
    ))
    max_papers = max(sizes_df["Papers"])
    fig2.update_layout(
        title=dict(
            text="<b>Cluster Sizes & Intra-Cluster Cohesion</b><br>"
                 "<sub>Bars = paper count · Red diamonds = mean cosine similarity to centroid</sub>",
            font_size=14,
        ),
        xaxis=dict(title="Number of Papers", range=[0, max_papers * 1.18]),
        xaxis2=dict(title="Cohesion (cosine)", overlaying="x", side="top",
                    range=[0, 1.05], showgrid=False),
        yaxis=dict(title="", automargin=True, tickfont=dict(size=11)),
        legend=dict(x=1.01, y=1, bordercolor="#ddd", borderwidth=1),
        template="plotly_white",
        height=max(500, 28 * len(summaries)),
        margin=dict(l=220, r=120, t=100, b=50),
        bargap=0.28,
    )
    # Target cohesion band
    fig2.add_vrect(
        x0=COSINE_SIM_LOW, x1=COSINE_SIM_HIGH, xref="x2",
        fillcolor="rgba(0,180,0,0.07)", line_width=0,
        annotation_text="Target cohesion band", annotation_position="top left",
    )
    fig2.write_html("chart_cluster_sizes.html")

    # ════════════════════════════════════════════════════════════════════════
    # CHART 3 — Treemap: area = papers, color = cohesion, label = cluster name + n
    # ════════════════════════════════════════════════════════════════════════
    treemap_df = pd.DataFrame({
        "label": [f"{s['label']}<br>n={s['size']}" for s in summaries],
        "plain_label": [s["label"] for s in summaries],
        "size": [s["size"] for s in summaries],
        "cohesion": [s["cohesion_score"] for s in summaries],
    })
    fig3 = px.treemap(
        treemap_df, path=["label"], values="size", color="cohesion",
        color_continuous_scale="RdYlGn",
        color_continuous_midpoint=0.55,
        hover_data={"plain_label": True, "cohesion": ":.3f", "size": True, "label": False},
        title="<b>Cluster Treemap</b> — Area = Papers · Color = Cohesion Score (green=high)",
    )
    fig3.update_traces(
        textfont_size=12,
        hovertemplate="<b>%{customdata[0]}</b><br>Papers: %{value}<br>Cohesion: %{customdata[1]:.3f}<extra></extra>",
    )
    fig3.update_layout(
        coloraxis_colorbar=dict(title="Cohesion", tickformat=".2f"),
        margin=dict(t=60, l=10, r=10, b=10),
        height=580,
    )
    fig3.write_html("chart_treemap.html")

    # ════════════════════════════════════════════════════════════════════════
    # CHART 4 — Cumulative coverage with cluster name annotations at steps
    # ════════════════════════════════════════════════════════════════════════
    sorted_summaries = sorted(summaries, key=lambda s: -s["size"])
    sorted_sizes = [s["size"] for s in sorted_summaries]
    sorted_names = [s["label"] for s in sorted_summaries]
    total_clustered = sum(sorted_sizes)
    cumulative = list(map(
        lambda i: sum(sorted_sizes[:i + 1]) / total_clustered * 100,
        range(len(sorted_sizes))
    ))
    x_vals = list(range(1, len(sorted_sizes) + 1))

    fig4 = go.Figure()
    fig4.add_trace(go.Scatter(
        y=cumulative, x=x_vals,
        mode="lines+markers",
        name="Cumulative Coverage",
        fill="tozeroy", fillcolor="rgba(0,100,255,0.07)",
        line=dict(color="royalblue", width=2.5),
        marker=dict(size=8, color="royalblue"),
        customdata=list(zip(sorted_names, sorted_sizes)),
        hovertemplate="<b>%{customdata[0]}</b><br>Rank: %{x}<br>Papers: %{customdata[1]}<br>Coverage: %{y:.1f}%<extra></extra>",
    ))
    fig4.add_hline(y=80, line_dash="dash", line_color="orange",
                   annotation_text="80% coverage", annotation_font_size=11)
    fig4.add_hline(y=50, line_dash="dot", line_color="green",
                   annotation_text="50% coverage", annotation_font_size=11)
    # Annotate first 8 cluster names along the curve
    for i, (name, cum) in enumerate(zip(sorted_names[:8], cumulative[:8])):
        short = name[:22] + "…" if len(name) > 22 else name
        fig4.add_annotation(
            x=i + 1, y=cum,
            text=short, showarrow=True, arrowhead=2, arrowsize=0.8,
            ax=20, ay=-30 if i % 2 == 0 else 30,
            font=dict(size=9, color="#333"),
            bgcolor="rgba(255,255,255,0.8)", borderpad=2,
        )
    fig4.update_layout(
        title=dict(
            text="<b>Cumulative Paper Coverage by Cluster Rank</b><br>"
                 "<sub>Hover each point to see cluster name and paper count</sub>",
            font_size=14,
        ),
        xaxis=dict(title="Cluster Rank (largest first)", tickmode="linear", dtick=1),
        yaxis=dict(title="Cumulative Coverage %", range=[0, 105]),
        template="plotly_white", height=500,
        margin=dict(l=60, r=60, t=90, b=60),
    )
    fig4.write_html("chart_cumulative.html")

    # ════════════════════════════════════════════════════════════════════════
    # CHART 5 — Cohesion histogram + per-cluster dot strip
    # ════════════════════════════════════════════════════════════════════════
    cohesion_scores = [s["cohesion_score"] for s in summaries]
    cluster_names_list = [s["label"] for s in summaries]

    fig5 = go.Figure()
    fig5.add_trace(go.Histogram(
        x=cohesion_scores, nbinsx=12, name="All Clusters",
        marker_color="steelblue", opacity=0.55,
        hovertemplate="Cohesion: %{x:.2f}<br>Count: %{y}<extra></extra>",
    ))
    # Strip dots colored by cluster name
    jitter = np.random.default_rng(0).uniform(-0.4, 0.4, len(cohesion_scores))
    fig5.add_trace(go.Scatter(
        x=cohesion_scores,
        y=list(map(lambda j: j, jitter)),
        mode="markers+text",
        text=list(map(lambda n: n[:18], cluster_names_list)),
        textposition="top center",
        textfont=dict(size=8),
        marker=dict(
            color=list(map(lambda n: color_map.get(n, "#888"), cluster_names_list)),
            size=11, opacity=0.9, line=dict(width=0.8, color="white"),
            symbol="circle",
        ),
        name="Individual Clusters",
        customdata=list(zip(cluster_names_list, [s["size"] for s in summaries])),
        hovertemplate="<b>%{customdata[0]}</b><br>Cohesion: %{x:.3f}<br>Papers: %{customdata[1]}<extra></extra>",
        yaxis="y2",
    ))
    fig5.add_vrect(
        x0=COSINE_SIM_LOW, x1=COSINE_SIM_HIGH,
        fillcolor="rgba(0,180,0,0.12)", line_width=1.5,
        line_color="green",
        annotation_text=f"Target range [{COSINE_SIM_LOW}–{COSINE_SIM_HIGH}]",
        annotation_position="top left", annotation_font_size=11,
    )
    fig5.update_layout(
        title=dict(
            text="<b>Intra-Cluster Cohesion Distribution</b><br>"
                 "<sub>Green band = target cosine similarity range · Each dot = one cluster</sub>",
            font_size=14,
        ),
        xaxis=dict(title="Mean Cosine Similarity to Cluster Centroid", range=[0, 1.05]),
        yaxis=dict(title="Number of Clusters", showgrid=True),
        yaxis2=dict(title="", overlaying="y", range=[-2.5, 4], showticklabels=False, showgrid=False),
        template="plotly_white", height=480,
        legend=dict(x=0.01, y=0.99),
        margin=dict(l=60, r=40, t=90, b=60),
    )
    fig5.write_html("chart_cohesion_hist.html")


def _build_themes_chart(themes):
    """
    Generate two rich charts for themes.json:
    1. chart_themes_detail.html  — sunburst showing theme → evidence sentences,
       with cohesion badge and DOI strip
    2. chart_themes_radar.html   — radar/spider chart comparing themes on size & cohesion
    """
    if not themes:
        return

    # ════════════════════════════════════════════════════════════════════════
    # CHART A — Sunburst: root → theme → top evidence sentences
    # ════════════════════════════════════════════════════════════════════════
    ids, labels, parents, values, colors, hovers = [], [], [], [], [], []

    # Root node
    total_papers = sum(t["size"] for t in themes)
    ids.append("root")
    labels.append(f"All Themes<br>n={total_papers}")
    parents.append("")
    values.append(total_papers)
    colors.append(0.5)
    hovers.append(f"Total papers across {len(themes)} themes")

    palette_sun = _PALETTE

    for i, t in enumerate(themes):
        tid = f"theme_{t['theme_id']}"
        theme_label = f"<b>{t['name']}</b><br>n={t['size']} · coh={t.get('cohesion_score',0):.2f}"
        ids.append(tid)
        labels.append(theme_label)
        parents.append("root")
        values.append(t["size"])
        colors.append(t.get("cohesion_score", 0.5))
        hovers.append(
            f"<b>{t['name']}</b><br>"
            f"Papers: {t['size']}<br>"
            f"Cohesion: {t.get('cohesion_score',0):.3f}<br>"
            f"Reasoning: {t.get('reasoning','')[:120]}"
        )
        # Evidence sentence children
        for j, ev in enumerate(t.get("evidence", [])[:3]):
            eid = f"{tid}_ev{j}"
            short_ev = (ev[:80] + "…") if len(ev) > 80 else ev
            ids.append(eid)
            labels.append(short_ev)
            parents.append(tid)
            values.append(max(1, t["size"] // 3))
            colors.append(t.get("cohesion_score", 0.5) * 0.9)
            hovers.append(ev[:300])

    fig_sun = go.Figure(go.Sunburst(
        ids=ids, labels=labels, parents=parents, values=values,
        customdata=hovers,
        hovertemplate="%{customdata}<extra></extra>",
        branchvalues="total",
        marker=dict(
            colors=colors,
            colorscale="RdYlGn",
            cmin=0.3, cmid=0.55, cmax=0.85,
            colorbar=dict(title="Cohesion", thickness=15, x=1.02),
            line=dict(width=1, color="white"),
        ),
        textfont=dict(size=11),
        insidetextorientation="radial",
        maxdepth=2,
    ))
    fig_sun.update_layout(
        title=dict(
            text="<b>Themes Structure — Sunburst View</b><br>"
                 "<sub>Inner ring = themes (area ∝ papers) · Outer ring = representative evidence · Color = cohesion</sub>",
            font_size=14,
        ),
        height=700, width=750,
        margin=dict(t=90, l=10, r=120, b=10),
    )
    fig_sun.write_html("chart_themes_detail.html")

    # ════════════════════════════════════════════════════════════════════════
    # CHART B — Horizontal grouped bar: themes compared on size + cohesion
    #           with evidence preview on hover
    # ════════════════════════════════════════════════════════════════════════
    themes_sorted = sorted(themes, key=lambda t: -t["size"])
    theme_names = [t["name"] for t in themes_sorted]
    sizes = [t["size"] for t in themes_sorted]
    cohesions = [t.get("cohesion_score", 0) for t in themes_sorted]
    reasonings = [t.get("reasoning", "")[:150] for t in themes_sorted]
    evidences = [" | ".join(t.get("evidence", [])[:2])[:200] for t in themes_sorted]

    bar_colors2 = [_PALETTE[i % len(_PALETTE)] for i in range(len(themes_sorted))]

    fig_b = go.Figure()
    fig_b.add_trace(go.Bar(
        y=theme_names, x=sizes, orientation="h",
        name="Papers",
        marker_color=bar_colors2,
        text=sizes, textposition="outside",
        customdata=list(zip(reasonings, evidences)),
        hovertemplate=(
            "<b>%{y}</b><br>"
            "Papers: %{x}<br>"
            "Reasoning: %{customdata[0]}<br>"
            "Evidence: %{customdata[1]}<extra></extra>"
        ),
    ))
    fig_b.add_trace(go.Scatter(
        y=theme_names, x=cohesions,
        mode="markers+text",
        name="Cohesion",
        marker=dict(color="crimson", size=12, symbol="diamond",
                    line=dict(width=1, color="darkred")),
        text=list(map(lambda c: f"{c:.2f}", cohesions)),
        textposition="middle right",
        xaxis="x2",
        hovertemplate="<b>%{y}</b><br>Cohesion: %{x:.3f}<extra></extra>",
    ))
    max_sz = max(sizes) if sizes else 1
    fig_b.update_layout(
        title=dict(
            text="<b>Final Themes — Papers & Cohesion</b><br>"
                 "<sub>Hover bars for reasoning & evidence · Red diamonds = cohesion score</sub>",
            font_size=14,
        ),
        xaxis=dict(title="Papers", range=[0, max_sz * 1.2]),
        xaxis2=dict(title="Cohesion Score", overlaying="x", side="top",
                    range=[0, 1.1], showgrid=False),
        yaxis=dict(automargin=True, tickfont=dict(size=11)),
        legend=dict(x=1.02, y=1, bordercolor="#ddd", borderwidth=1),
        template="plotly_white",
        height=max(450, 32 * len(themes_sorted)),
        margin=dict(l=220, r=120, t=100, b=50),
        bargap=0.3,
    )
    fig_b.write_html("chart_themes_radar.html")


def _run_bertopic_discovery(dummy: str = "") -> str:
    """
    Phase 2a: Embed papers with SPECTER2, reduce with UMAP, cluster with HDBSCAN.
    Iterates parameters to achieve 15-30 clusters. Generates rich visualisation charts.
    """
    papers_df = pd.read_csv(PAPERS_FILE)
    texts = papers_df["combined_text"].tolist()
    dois = papers_df["doi"].tolist()
    
    # ── SPECTER2 Embeddings ───────────────────────────────────────────────────
    print(f"[SPECTER2] Loading model: {SPECTER2_MODEL}")
    model = SentenceTransformer(SPECTER2_MODEL)
    print(f"[SPECTER2] Encoding {len(texts)} papers...")
    embeddings = model.encode(
        texts,
        normalize_embeddings=True,
        show_progress_bar=True,
        batch_size=16,
    )
    np.save(EMB_FILE, embeddings)
    print(f"[SPECTER2] Embeddings shape: {embeddings.shape}")
    
    # ── UMAP Reduction ────────────────────────────────────────────────────────
    # n_neighbors controls local/global structure; cosine metric for scientific text
    reducer = umap.UMAP(
        n_components=10,          # reduce to 10D before HDBSCAN
        n_neighbors=15,
        min_dist=0.0,
        metric="cosine",
        random_state=42,
    )
    reduced = reducer.fit_transform(embeddings)
    print(f"[UMAP] Reduced to shape: {reduced.shape}")
    
    # ── HDBSCAN with iterative parameter search for 15-30 clusters ───────────
    def try_hdbscan(min_size, min_samples):
        clusterer = hdbscan.HDBSCAN(
            min_cluster_size=min_size,
            min_samples=min_samples,
            metric="euclidean",
            cluster_selection_method="eom",
            prediction_data=True,
        )
        lbl = clusterer.fit_predict(reduced)
        valid_clusters = set(lbl[lbl >= 0])
        # Enforce max cluster size constraint
        cluster_sizes = Counter(lbl[lbl >= 0].tolist())
        oversized = sum(1 for sz in cluster_sizes.values() if sz > MAX_CLUSTER_SIZE)
        return lbl, len(valid_clusters), oversized, clusterer
    
    # Parameter grid search — pick first combo giving 15-30 valid clusters
    param_grid = [
        (ms, msa)
        for ms in [15, 12, 10, 8, 7, 6, 5]
        for msa in [5, 4, 3, 2]
    ]
    
    best_labels, best_n, best_clusterer = None, 0, None
    results = list(map(lambda params: (params, *try_hdbscan(*params)), param_grid))
    
    # Filter to those in [15, 30] range first
    valid_results = list(filter(lambda r: MIN_CLUSTERS <= r[2] <= MAX_CLUSTERS and r[3] == 0, results))
    
    chosen = (
        valid_results[0] if valid_results
        else sorted(results, key=lambda r: abs(r[2] - 20))[0]  # closest to 20
    )
    chosen_params, chosen_labels, chosen_n, _, best_clusterer = chosen
    best_labels = chosen_labels
    print(f"[HDBSCAN] Best params: min_cluster_size={chosen_params[0]}, min_samples={chosen_params[1]}")
    print(f"[HDBSCAN] Clusters found: {chosen_n}")
    
    np.save(LABELS_FILE, best_labels)
    
    # ── Build summaries per cluster ───────────────────────────────────────────
    valid_cluster_ids = sorted(set(best_labels[best_labels >= 0].tolist()))
    
    # Cosine similarity quality check per cluster
    def compute_cluster_cohesion(label):
        indices = [i for i, lbl in enumerate(best_labels) if lbl == label]
        if len(indices) < 2:
            return 0.0
        cluster_emb = embeddings[indices]
        centroid = cluster_emb.mean(axis=0, keepdims=True)
        sims = cosine_similarity(cluster_emb, centroid).flatten()
        return float(np.mean(sims))
    
    def build_topic(label):
        indices = [i for i, lbl in enumerate(best_labels) if lbl == label]
        cluster_emb = embeddings[indices]
        centroid = cluster_emb.mean(axis=0)
        # Top 3 papers closest to centroid (for LLM labeling)
        dists = np.linalg.norm(cluster_emb - centroid, axis=1)
        top3_idx = np.argsort(dists)[:3].tolist()
        top3_titles = list(map(lambda i: texts[indices[i]][:200], top3_idx))
        top3_dois = list(map(lambda i: dois[indices[i]], top3_idx))
        cohesion = compute_cluster_cohesion(label)
        return {
            "topic_id": int(label),
            "label": f"Cluster_{label}",
            "top_sentences": top3_titles,
            "top_dois": top3_dois,
            "size": len(indices),
            "cohesion_score": round(cohesion, 4),
            "approved": False,
            "rename_to": "",
            "reasoning": "",
        }
    
    summaries = list(map(build_topic, valid_cluster_ids))
    noise_count = int(np.sum(best_labels == -1))
    
    with open(SUMMARIES_FILE, "w") as f:
        json.dump(summaries, f, indent=2)

    # ── 2D UMAP coords (saved for re-use in later chart regeneration) ─────────
    reducer_2d = umap.UMAP(n_components=2, n_neighbors=15, min_dist=0.1, metric="cosine", random_state=42)
    coords_2d = reducer_2d.fit_transform(embeddings)
    np.save("coords_2d.npy", coords_2d)

    # ── Build initial charts with placeholder Cluster_N names ─────────────────
    _build_all_charts(summaries, coords_2d, best_labels, texts, dois, chosen_n)

    avg_cohesion = float(np.mean([s["cohesion_score"] for s in summaries]))
    in_target_range = sum(1 for s in summaries if COSINE_SIM_LOW <= s["cohesion_score"] <= COSINE_SIM_HIGH)
    
    return json.dumps({
        "status": "ok",
        "cluster_count": chosen_n,
        "noise_papers": noise_count,
        "total_papers": len(texts),
        "hdbscan_params": {"min_cluster_size": chosen_params[0], "min_samples": chosen_params[1]},
        "avg_cohesion": round(avg_cohesion, 4),
        "clusters_in_target_cosine_range": in_target_range,
        "charts": [
            "chart_umap_scatter.html",
            "chart_cluster_sizes.html",
            "chart_treemap.html",
            "chart_cumulative.html",
            "chart_cohesion_hist.html",
        ],
        "message": (
            f"SPECTER2 embedded {len(texts)} papers → UMAP (10D) → HDBSCAN → "
            f"{chosen_n} clusters (noise: {noise_count} papers). "
            f"Avg cohesion: {avg_cohesion:.3f}. "
            f"{in_target_range}/{chosen_n} clusters in target cosine range [{COSINE_SIM_LOW}-{COSINE_SIM_HIGH}]. "
            f"5 rich interactive charts generated."
        ),
    })


def _label_topics_with_llm(dummy: str = "") -> str:
    """
    Phase 2b: Label each cluster using 3-LLM majority-vote approach.
    Each cluster's top-3 paper titles are sent to Mistral 3 times;
    the most common label wins.
    """
    with open(SUMMARIES_FILE) as f:
        summaries = json.load(f)
    
    llm = ChatMistralAI(model="mistral-small-latest", temperature=0.3)
    parser = JsonOutputParser()
    prompt = PromptTemplate(
        input_variables=["titles"],
        template=(
            "You are an expert academic thematic analyst specialising in IS/technology journals.\n"
            "Given these 3 representative paper title+abstract excerpts from the same research cluster:\n\n"
            "{titles}\n\n"
            "Generate a concise, academically precise cluster label (4-7 words) that captures "
            "the unifying research theme. This will appear in a journal paper.\n\n"
            'Return ONLY valid JSON: {{"label": "...", "reasoning": "one sentence explanation"}}\n'
            "No markdown, no preamble, no extra keys."
        ),
    )
    
    def vote_label(topic):
        excerpts = "\n".join(
            list(map(lambda s: f"- {s[:200]}", topic["top_sentences"][:3]))
        )
        # 3 independent calls
        raw_results = list(map(
            lambda _: (prompt | llm | parser).invoke({"titles": excerpts}),
            range(3)
        ))
        # Normalise to dict
        parsed = list(map(
            lambda r: (r[0] if isinstance(r, list) else r),
            raw_results
        ))
        labels = list(map(lambda r: r.get("label", topic["label"]), parsed))
        # Majority vote: pick most common label
        label_counts = Counter(labels)
        winner = label_counts.most_common(1)[0][0]
        winner_reasoning = next(
            (r.get("reasoning", "") for r in parsed if r.get("label") == winner),
            parsed[0].get("reasoning", "") if parsed else ""
        )
        all_candidates = list(label_counts.keys())
        return {
            **topic,
            "label": winner,
            "label_candidates": all_candidates,
            "reasoning": winner_reasoning,
        }
    
    labeled = list(map(vote_label, summaries))

    with open(SUMMARIES_FILE, "w") as f:
        json.dump(labeled, f, indent=2)

    # ── Regenerate all charts with real LLM-voted names ───────────────────────
    coords_2d_path = Path("coords_2d.npy")
    labels_path = Path(LABELS_FILE)
    papers_path = Path(PAPERS_FILE)
    charts_rebuilt = False
    if coords_2d_path.exists() and labels_path.exists() and papers_path.exists():
        coords_2d = np.load("coords_2d.npy")
        cluster_labels_arr = np.load(LABELS_FILE)
        papers_df2 = pd.read_csv(PAPERS_FILE)
        _build_all_charts(
            labeled,
            coords_2d,
            cluster_labels_arr,
            papers_df2["combined_text"].tolist(),
            papers_df2["doi"].tolist(),
            len(set(cluster_labels_arr[cluster_labels_arr >= 0])),
        )
        charts_rebuilt = True

    return json.dumps({
        "status": "ok",
        "labeled_count": len(labeled),
        "charts_rebuilt_with_real_names": charts_rebuilt,
        "message": (
            f"Labeled {len(labeled)} clusters using 3-LLM majority vote. "
            f"{'All 5 charts regenerated with real cluster names and centroid labels.' if charts_rebuilt else 'Charts could not be rebuilt (missing cache files).'} "
            f"Review the topic table, approve/rename clusters, then submit."
        ),
    })


def _consolidate_into_themes(dummy: str = "") -> str:
    """Phase 3: Merge approved topics from summaries.json into themes.json."""
    with open(SUMMARIES_FILE) as f:
        summaries = json.load(f)
    approved = list(filter(lambda t: t.get("approved", False), summaries))
    themes = list(map(lambda t: {
        "theme_id": t["topic_id"],
        "name": t.get("rename_to") or t["label"],
        "original_label": t["label"],
        "label_candidates": t.get("label_candidates", []),
        "evidence": t["top_sentences"],
        "top_dois": t.get("top_dois", []),
        "size": t["size"],
        "cohesion_score": t.get("cohesion_score", 0.0),
        "reasoning": t.get("reasoning", ""),
    }, approved))
    with open(THEMES_FILE, "w") as f:
        json.dump(themes, f, indent=2)

    # ── Generate themes.json visual (sunburst + detail card chart) ────────────
    _build_themes_chart(themes)

    # ── Regenerate scatter/bar/treemap with rename_to labels ──────────────────
    coords_2d_path = Path("coords_2d.npy")
    labels_path = Path(LABELS_FILE)
    papers_path = Path(PAPERS_FILE)
    if coords_2d_path.exists() and labels_path.exists() and papers_path.exists():
        coords_2d = np.load("coords_2d.npy")
        cluster_labels_arr = np.load(LABELS_FILE)
        papers_df2 = pd.read_csv(PAPERS_FILE)
        # Build a patched summaries list using final theme names
        theme_name_map = {t["theme_id"]: t["name"] for t in themes}
        with open(SUMMARIES_FILE) as f:
            summaries_all = json.load(f)
        patched = list(map(
            lambda s: {**s, "label": theme_name_map.get(s["topic_id"], s["label"])},
            summaries_all
        ))
        _build_all_charts(
            patched,
            coords_2d,
            cluster_labels_arr,
            papers_df2["combined_text"].tolist(),
            papers_df2["doi"].tolist(),
            len(themes),
        )

    return json.dumps({
        "status": "ok", "theme_count": len(themes),
        "message": (
            f"Consolidated {len(themes)} approved clusters into themes.json. "
            f"Themes visualization chart generated (chart_themes_detail.html). "
            f"All scatter/bar/treemap charts updated with final theme names."
        ),
    })


def _compare_with_taxonomy(dummy: str = "") -> str:
    """Phase 5.5: Map themes to PAJAIS 25 categories via Mistral. Saves taxonomy_map.json."""
    with open(THEMES_FILE) as f:
        themes = json.load(f)
    llm = ChatMistralAI(model="mistral-small-latest", temperature=0.1)
    parser = JsonOutputParser()
    prompt = PromptTemplate(
        input_variables=["theme_name", "evidence", "categories"],
        template=(
            "Map this research theme to the most relevant PAJAIS category.\n\n"
            "Theme: {theme_name}\nEvidence: {evidence}\n\n"
            "PAJAIS Categories:\n{categories}\n\n"
            'Return ONLY JSON: {{"pajais_category": "...", "confidence": 0.0-1.0, "rationale": "..."}}'
        ),
    )
    categories_str = "\n".join(list(map(lambda c: f"- {c}", PAJAIS_CATEGORIES)))

    def map_theme(t):
        result = (prompt | llm | parser).invoke({
            "theme_name": t["name"], "evidence": " ".join(t["evidence"][:2]),
            "categories": categories_str,
        })
        result = result[0] if isinstance(result, list) else result
        return {**t, "pajais_mapping": result}

    mapped = list(map(map_theme, themes))
    with open(TAXONOMY_FILE, "w") as f:
        json.dump(mapped, f, indent=2)
    return json.dumps({
        "status": "ok", "mapped_count": len(mapped),
        "message": f"Mapped {len(mapped)} themes to PAJAIS taxonomy.",
    })


def _generate_comparison_csv(dummy: str = "") -> str:
    """Phase 6a: Build comparison.csv from taxonomy_map.json."""
    with open(TAXONOMY_FILE) as f:
        taxonomy = json.load(f)
    rows = list(map(lambda e: {
        "Theme Name": e.get("name", ""),
        "Original Label": e.get("original_label", ""),
        "Label Candidates (3-vote)": " | ".join(e.get("label_candidates", [])),
        "Topic Size (Papers)": e.get("size", 0),
        "Cohesion Score": e.get("cohesion_score", ""),
        "PAJAIS Category": e.get("pajais_mapping", {}).get("pajais_category", ""),
        "Confidence": e.get("pajais_mapping", {}).get("confidence", ""),
        "Rationale": e.get("pajais_mapping", {}).get("rationale", ""),
        "Top Evidence": " | ".join(e.get("evidence", [])[:2]),
        "Top DOIs": " | ".join(e.get("top_dois", [])[:3]),
    }, taxonomy))
    pd.DataFrame(rows).to_csv(COMPARISON_FILE, index=False)
    return json.dumps({
        "status": "ok", "rows": len(rows), "file": COMPARISON_FILE,
        "message": f"Comparison CSV with {len(rows)} rows saved (includes cohesion scores, 3-vote labels, DOIs).",
    })


def _export_narrative(dummy: str = "") -> str:
    """Phase 6b: Generate 500-word academic Section 7 and save narrative.txt."""
    with open(TAXONOMY_FILE) as f:
        taxonomy = json.load(f)
    themes_text = "\n".join(list(map(
        lambda t: (
            f"Theme '{t['name']}' (papers={t['size']}, cohesion={t.get('cohesion_score',0):.3f}) "
            f"maps to PAJAIS '{t.get('pajais_mapping', {}).get('pajais_category', 'N/A')}' — "
            f"{t.get('pajais_mapping', {}).get('rationale', '')}."
        ),
        taxonomy,
    )))
    llm = ChatMistralAI(model="mistral-small-latest", temperature=0.4)
    prompt = PromptTemplate(
        input_variables=["themes_text"],
        template=(
            "Write Section 7 'Discussion of Themes' (~500 words, formal academic tone) "
            "suitable for a journal paper in the International Journal of Electronic Commerce.\n\n"
            "Themes, cluster sizes, cohesion scores, and PAJAIS mappings:\n{themes_text}\n\n"
            "Write the discussion section:"
        ),
    )
    result = (prompt | llm).invoke({"themes_text": themes_text})
    narrative = result.content if hasattr(result, "content") else str(result)
    Path(NARRATIVE_FILE).write_text(narrative)
    return json.dumps({
        "status": "ok", "word_count": len(narrative.split()),
        "file": NARRATIVE_FILE, "message": "Narrative saved to narrative.txt.",
    })


# ── Wrap with StructuredTool ──────────────────────────────────────────────────

load_scopus_csv = StructuredTool.from_function(
    func=_load_scopus_csv, name="load_scopus_csv", handle_tool_error=True,
    description="Load scopus_data.csv, combine Title+Abstract per paper, extract DOIs, save papers_embedded.csv.",
)
run_bertopic_discovery = StructuredTool.from_function(
    func=_run_bertopic_discovery, name="run_bertopic_discovery", handle_tool_error=True,
    description=(
        "Embed papers with SPECTER2, reduce with UMAP (cosine, 10D), cluster with HDBSCAN "
        "(min_cluster_size=5, 15-30 clusters). Saves summaries.json, generates 5 rich charts."
    ),
)
label_topics_with_llm = StructuredTool.from_function(
    func=_label_topics_with_llm, name="label_topics_with_llm", handle_tool_error=True,
    description="Label clusters using 3-LLM majority-vote. Each cluster submitted 3× to Mistral; best label wins.",
)
consolidate_into_themes = StructuredTool.from_function(
    func=_consolidate_into_themes, name="consolidate_into_themes", handle_tool_error=True,
    description="Merge approved clusters into themes.json.",
)
compare_with_taxonomy = StructuredTool.from_function(
    func=_compare_with_taxonomy, name="compare_with_taxonomy", handle_tool_error=True,
    description="Map themes to PAJAIS 25 categories via Mistral. Saves taxonomy_map.json.",
)
generate_comparison_csv = StructuredTool.from_function(
    func=_generate_comparison_csv, name="generate_comparison_csv", handle_tool_error=True,
    description="Export comparison.csv (with cohesion scores, 3-vote labels, DOIs) from taxonomy_map.json.",
)
export_narrative = StructuredTool.from_function(
    func=_export_narrative, name="export_narrative", handle_tool_error=True,
    description="Generate 500-word academic Section 7 for IJEC journal and save narrative.txt.",
)