Blame for tools.py

"""
tools.py — Core functions for the AI-driven topic modeling pipeline.

This module provides all analytical functions used by the TopicAgent:
  - CSV ingestion and validation
  - Text preprocessing (lowercasing, stopword removal, cleaning)
  - Topic modeling via BERTopic (with fallback to sklearn LDA)
  - Automatic human-readable label generation
  - Cross-source theme comparison (Title vs Abstract)
  - Taxonomy mapping (MAPPED / NOVEL classification)
"""

import re
import json
import logging
from typing import Dict, List, Tuple, Optional, Any

import numpy as np
import pandas as pd
import nltk
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")

# ---------------------------------------------------------------------------
# NLTK data download (idempotent)
# ---------------------------------------------------------------------------
for _resource in ("punkt", "punkt_tab", "stopwords"):
    try:
        nltk.data.find(f"tokenizers/{_resource}" if "punkt" in _resource else f"corpora/{_resource}")
    except LookupError:
        nltk.download(_resource, quiet=True)

# ---------------------------------------------------------------------------
# Reference taxonomy of known AI / business / research themes
# Used by create_taxonomy_map() for MAPPED vs NOVEL classification
# ---------------------------------------------------------------------------
KNOWN_THEMES: List[str] = [
    # AI / ML
    "artificial intelligence", "machine learning", "deep learning", "neural network",
    "natural language processing", "computer vision", "reinforcement learning",
    "generative ai", "large language model", "transformer", "chatbot",
    "recommendation system", "knowledge graph", "robotics", "autonomous",
    "explainable ai", "federated learning", "transfer learning", "ai ethics",
    "adversarial", "gan", "diffusion model", "prompt engineering",
    # Data science
    "data mining", "big data", "analytics", "data science", "data quality",
    "feature engineering", "dimensionality reduction", "clustering", "classification",
    "regression", "time series", "anomaly detection", "sentiment analysis",
    # Business / Management
    "digital transformation", "innovation", "strategy", "supply chain",
    "customer experience", "marketing", "e-commerce", "fintech", "blockchain",
    "sustainability", "corporate social responsibility", "knowledge management",
    "decision support", "business intelligence", "enterprise", "organizational",
    "human resource", "leadership", "entrepreneurship", "business model",
    # Information systems
    "information systems", "technology adoption", "user acceptance", "privacy",
    "security", "trust", "social media", "online community", "platform",
    "crowdsourcing", "cloud computing", "iot", "internet of things",
    "software engineering", "agile", "devops", "digital platform",
    # Healthcare / Society
    "healthcare", "telemedicine", "electronic health", "public health",
    "education", "e-learning", "smart city", "government", "policy",
    "ethics", "fairness", "bias", "misinformation", "content moderation",
    # Research methods
    "survey", "experiment", "case study", "meta-analysis", "bibliometric",
    "systematic review", "structural equation", "grounded theory",
]


# ===================================================================
# 1. load_csv  — Ingest and validate the CSV dataset
# ===================================================================
def load_csv(filepath: str) -> pd.DataFrame:
    """
    Load a CSV file and ensure the required columns (Title, Abstract) exist.

    Parameters
    ----------
    filepath : str
        Path to the CSV file.

    Returns
    -------
    pd.DataFrame
        DataFrame with at least 'Title' and 'Abstract' columns.

    Raises
    ------
    FileNotFoundError
        If the specified file does not exist.
    ValueError
        If required columns are missing.
    """
    logger.info("Loading CSV from %s", filepath)
    df = pd.read_csv(filepath, encoding="utf-8", on_bad_lines="skip")
    logger.info("Loaded %d rows × %d columns", len(df), len(df.columns))

    # Validate required columns (case-insensitive match)
    col_map = {str(c).strip().lower(): c for c in df.columns}
    
    # Title and Abstract are strictly required
    if "title" not in col_map or "abstract" not in col_map:
        missing = {"title", "abstract"} - set(col_map.keys())
        raise ValueError(f"CSV is missing required columns: {missing}. Found: {list(df.columns)}")
        
    # Map Title and Abstract
    rename_dict = {
        col_map["title"]: "Title",
        col_map["abstract"]: "Abstract"
    }

    # Map Keywords
    keyword_cols = ["keywords", "author keywords", "index keywords"]
    found_keyword = next((c for c in keyword_cols if c in col_map), None)
    if found_keyword:
        rename_dict[col_map[found_keyword]] = "Keywords"
        
    # Map Full Text
    fulltext_cols = ["full text", "body"]
    found_fulltext = next((c for c in fulltext_cols if c in col_map), None)
    if found_fulltext:
        rename_dict[col_map[found_fulltext]] = "Full Text"

    # Rename to canonical form
    df = df.rename(columns=rename_dict)

    # Fill missing columns
    if "Keywords" not in df.columns:
        df["Keywords"] = ""
    if "Full Text" not in df.columns:
        df["Full Text"] = df["Abstract"]

    # Drop rows where all of these are empty
    df = df.dropna(subset=["Title", "Abstract", "Keywords", "Full Text"], how="all").reset_index(drop=True)
    df["Title"] = df["Title"].fillna("")
    df["Abstract"] = df["Abstract"].fillna("")
    df["Keywords"] = df["Keywords"].fillna("")
    df["Full Text"] = df["Full Text"].fillna("")

    logger.info("After cleaning: %d usable rows", len(df))
    return df


# ===================================================================
# 2. preprocess_text  — Clean and normalise a list of text documents
# ===================================================================
def preprocess_text(documents: List[str]) -> List[str]:
    """
    Apply professional-grade text preprocessing:
      1. Lowercase
      2. Remove URLs, emails, special characters, digits
      3. Tokenize
      4. Remove stopwords (NLTK English)
      5. Remove very short tokens (length ≤ 2)
      6. Rejoin into cleaned strings

    Parameters
    ----------
    documents : list of str
        Raw text documents.

    Returns
    -------
    list of str
        Cleaned text documents.
    """
    stop_words = set(stopwords.words("english"))
    # Extended stopwords common in academic abstracts
    stop_words.update([
        "©", "elsevier", "rights", "reserved", "doi", "http", "https",
        "vol", "pp", "fig", "table", "journal", "author", "authors",
        "study", "paper", "research", "results", "findings", "however",
        "propose", "proposed", "approach", "using", "based", "also",
        "show", "shows", "shown", "may", "used", "use", "one", "two",
        "three", "new", "well", "within", "among", "across", "toward",
        "towards", "et", "al", "ie", "eg", "cf", "thus", "therefore",
        "moreover", "furthermore", "addition", "conclusion", "conclusions",
    ])

    cleaned: List[str] = []
    for doc in documents:
        if not isinstance(doc, str) or not doc.strip():
            cleaned.append("")
            continue

        text = doc.lower()
        # Remove URLs
        text = re.sub(r"https?://\S+|www\.\S+", " ", text)
        # Remove emails
        text = re.sub(r"\S+@\S+", " ", text)
        # Remove digits and special characters but keep spaces
        text = re.sub(r"[^a-z\s]", " ", text)
        # Collapse whitespace
        text = re.sub(r"\s+", " ", text).strip()

        # Tokenize and filter
        tokens = word_tokenize(text)
        tokens = [t for t in tokens if t not in stop_words and len(t) > 2]

        cleaned.append(" ".join(tokens))

    logger.info("Preprocessed %d documents", len(cleaned))
    return cleaned


# ===================================================================
# 3. run_topic_modeling  — Discover clusters + sub-topics
# ===================================================================
def run_topic_modeling(
    documents: List[str],
    source_label: str = "documents",
    min_topics: int = 100,
    max_clusters: int = 25,
    use_bertopic: bool = True,
) -> Tuple[pd.DataFrame, Any]:
    """
    Perform two-level topic modeling on a corpus of preprocessed documents.

    Strategy:
      1. First pass: Discover clusters (target 10–max_clusters range)
         using BERTopic with UMAP + HDBSCAN, or LDA fallback.
      2. Second pass: Within each cluster, extract fine-grained sub-topics
         using NMF to guarantee ≥ min_topics total topics.

    Parameters
    ----------
    documents : list of str
        Preprocessed text documents.
    source_label : str
        Label for logging (e.g. "Titles" or "Abstracts").
    min_topics : int
        Minimum total number of topics (sub-topics) required (default 100).
    max_clusters : int
        Target maximum number of high-level clusters (default 25).
    use_bertopic : bool
        Whether to attempt BERTopic first.

    Returns
    -------
    topics_df : pd.DataFrame
        Columns: cluster_id, topic_id, sub_topic_id, keywords, label, source, n_documents
    model : object
        The fitted cluster-level model for downstream inspection.
    """
    from sklearn.decomposition import NMF
    from sklearn.feature_extraction.text import TfidfVectorizer

    # Filter out empty documents
    valid_docs = [d for d in documents if d.strip()]
    if len(valid_docs) < 20:
        raise ValueError(f"Not enough valid documents ({len(valid_docs)}) for topic modeling.")

    logger.info(
        "Running topic modeling on %d %s (target: %d–%d clusters, ≥ %d topics)",
        len(valid_docs), source_label, 10, max_clusters, min_topics,
    )

    # ---- FIRST PASS: Discover clusters ----
    cluster_df = None
    model = None

    if use_bertopic:
        try:
            cluster_df, model = _run_bertopic(valid_docs, source_label, max_clusters)
        except Exception as exc:
            logger.warning("BERTopic failed (%s). Falling back to LDA.", exc)
            cluster_df = None

    if cluster_df is None or len(cluster_df) < 3:
        logger.info("Using LDA for clustering %s (target %d clusters)", source_label, max_clusters)
        cluster_df, model = _run_lda(valid_docs, source_label, max_clusters)

    n_clusters = len(cluster_df)
    logger.info("First pass: %d clusters discovered for %s", n_clusters, source_label)

    # ---- SECOND PASS: Extract sub-topics within each cluster ----
    # Assign each document to its nearest cluster based on keyword overlap
    cluster_assignments = _assign_docs_to_clusters(valid_docs, cluster_df)

    # Calculate sub-topic counts per cluster
    cluster_sizes = {}
    for cid in cluster_df["topic_id"].tolist():
        cluster_sizes[cid] = int((np.array(cluster_assignments) == cid).sum())

    total_assigned = sum(cluster_sizes.values())
    subtopic_counts = {}
    for cid in cluster_df["topic_id"].tolist():
        if cluster_sizes[cid] == 0:
            subtopic_counts[cid] = 1
            continue
        raw_share = (cluster_sizes[cid] / max(total_assigned, 1)) * min_topics
        subtopic_counts[cid] = max(2, round(raw_share))

    # Ensure total >= min_topics
    current_total = sum(subtopic_counts.values())
    if current_total < min_topics:
        deficit = min_topics - current_total
        sorted_clusters = sorted(
            cluster_df["topic_id"].tolist(),
            key=lambda c: cluster_sizes.get(c, 0), reverse=True,
        )
        i = 0
        while deficit > 0 and i < len(sorted_clusters) * 10:
            cid = sorted_clusters[i % len(sorted_clusters)]
            max_possible = max(cluster_sizes.get(cid, 0) // 2, 2)
            if subtopic_counts[cid] < max_possible:
                subtopic_counts[cid] += 1
                deficit -= 1
            i += 1

    # Extract sub-topics per cluster
    all_topics = []
    global_topic_id = 0

    for _, cluster_row in cluster_df.iterrows():
        cid = int(cluster_row["topic_id"])
        # Get documents in this cluster
        doc_indices = [i for i, a in enumerate(cluster_assignments) if a == cid]
        cluster_docs = [valid_docs[i] for i in doc_indices]
        n_subtopics = subtopic_counts.get(cid, 2)

        if len(cluster_docs) < 2:
            # Single document cluster — single topic
            all_topics.append({
                "cluster_id": cid,
                "topic_id": global_topic_id,
                "sub_topic_id": 0,
                "keywords": cluster_row["keywords"],
                "source": source_label,
                "n_documents": len(cluster_docs),
            })
            global_topic_id += 1
            continue

        n_subtopics = min(n_subtopics, len(cluster_docs))
        if n_subtopics < 2:
            n_subtopics = 1

        try:
            vectorizer = TfidfVectorizer(
                stop_words="english", max_df=0.95, min_df=1,
                ngram_range=(1, 2), max_features=5000,
            )
            tfidf_matrix = vectorizer.fit_transform(cluster_docs)
            feature_names = vectorizer.get_feature_names_out()

            if n_subtopics == 1 or tfidf_matrix.shape[0] < 2:
                scores = tfidf_matrix.sum(axis=0).A1
                top_indices = scores.argsort()[-10:][::-1]
                kw = ", ".join([feature_names[i] for i in top_indices])
                all_topics.append({
                    "cluster_id": cid,
                    "topic_id": global_topic_id,
                    "sub_topic_id": 0,
                    "keywords": kw,
                    "source": source_label,
                    "n_documents": len(cluster_docs),
                })
                global_topic_id += 1
            else:
                nmf = NMF(n_components=n_subtopics, random_state=42, max_iter=300, init="nndsvda")
                W = nmf.fit_transform(tfidf_matrix)
                H = nmf.components_
                doc_assignments = W.argmax(axis=1)

                for sub_id in range(n_subtopics):
                    top_indices = H[sub_id].argsort()[-10:][::-1]
                    kw = ", ".join([feature_names[i] for i in top_indices])
                    n_docs = int((doc_assignments == sub_id).sum())
                    if n_docs == 0:
                        n_docs = 1

                    all_topics.append({
                        "cluster_id": cid,
                        "topic_id": global_topic_id,
                        "sub_topic_id": sub_id,
                        "keywords": kw,
                        "source": source_label,
                        "n_documents": n_docs,
                    })
                    global_topic_id += 1

        except Exception as e:
            logger.warning("NMF failed for cluster %d: %s. Using cluster keywords.", cid, e)
            all_topics.append({
                "cluster_id": cid,
                "topic_id": global_topic_id,
                "sub_topic_id": 0,
                "keywords": cluster_row["keywords"],
                "source": source_label,
                "n_documents": len(cluster_docs),
            })
            global_topic_id += 1

    topics_df = pd.DataFrame(all_topics)

    # Add cluster_label from cluster keywords
    cluster_label_map = {}
    for _, row in cluster_df.iterrows():
        cid = int(row["topic_id"])
        kws = [kw.strip().title() for kw in row["keywords"].split(",")][:3]
        cluster_label_map[cid] = " / ".join(kws)
    topics_df["cluster_label"] = topics_df["cluster_id"].map(cluster_label_map).fillna("")

    logger.info(
        "Topic modeling complete for %s: %d topics in %d clusters",
        source_label, len(topics_df), n_clusters,
    )
    return topics_df, model


def _assign_docs_to_clusters(docs: List[str], cluster_df: pd.DataFrame) -> List[int]:
    """Assign each document to its most similar cluster based on keyword overlap."""
    cluster_ids = cluster_df["topic_id"].tolist()
    cluster_keywords = []
    for _, row in cluster_df.iterrows():
        kws = set(kw.strip().lower() for kw in row["keywords"].split(",") if kw.strip())
        # Also add individual words from multi-word keywords
        expanded = set()
        for kw in kws:
            expanded.update(kw.split())
        expanded.update(kws)
        cluster_keywords.append(expanded)

    assignments = []
    for doc in docs:
        doc_tokens = set(doc.lower().split())
        best_score = -1
        best_cid = cluster_ids[0]  # default to first cluster
        for cid, kws in zip(cluster_ids, cluster_keywords):
            overlap = len(doc_tokens & kws)
            if overlap > best_score:
                best_score = overlap
                best_cid = cid
        assignments.append(best_cid)

    return assignments


def _run_bertopic(docs: List[str], source_label: str, max_clusters: int):
    """Run BERTopic to discover clusters (not topics)."""
    from bertopic import BERTopic
    from umap import UMAP
    from hdbscan import HDBSCAN
    from sklearn.feature_extraction.text import CountVectorizer

    umap_model = UMAP(
        n_neighbors=10,
        n_components=5,
        min_dist=0.0,
        metric="cosine",
        random_state=42,
    )
    hdbscan_model = HDBSCAN(
        min_cluster_size=max(5, len(docs) // max_clusters),
        min_samples=2,
        prediction_data=True,
    )
    vectorizer = CountVectorizer(
        stop_words="english",
        ngram_range=(1, 2),
        max_df=0.90,
        min_df=2,
    )

    topic_model = BERTopic(
        umap_model=umap_model,
        hdbscan_model=hdbscan_model,
        vectorizer_model=vectorizer,
        nr_topics=max_clusters,
        top_n_words=10,
        verbose=False,
    )

    topics, _probs = topic_model.fit_transform(docs)
    info = topic_model.get_topic_info()
    # Exclude outlier topic (-1)
    info = info[info["Topic"] != -1].reset_index(drop=True)

    rows = []
    for _, row in info.iterrows():
        tid = int(row["Topic"])
        topic_words = topic_model.get_topic(tid)
        kw = ", ".join([w for w, _ in topic_words[:10]])
        rows.append({"topic_id": tid, "keywords": kw, "source": source_label})

    df = pd.DataFrame(rows)
    return df, topic_model


def _run_lda(docs: List[str], source_label: str, n_clusters: int):
    """Run sklearn LDA to discover clusters."""
    from sklearn.decomposition import LatentDirichletAllocation
    from sklearn.feature_extraction.text import CountVectorizer

    vectorizer = CountVectorizer(
        stop_words="english",
        max_df=0.90,
        min_df=2,
        ngram_range=(1, 2),
        max_features=10000,
    )
    dtm = vectorizer.fit_transform(docs)
    feature_names = vectorizer.get_feature_names_out()

    lda = LatentDirichletAllocation(
        n_components=n_clusters,
        max_iter=25,
        learning_method="online",
        random_state=42,
        n_jobs=-1,
    )
    lda.fit(dtm)

    rows = []
    for idx, component in enumerate(lda.components_):
        top_indices = component.argsort()[-10:][::-1]
        kw = ", ".join([feature_names[i] for i in top_indices])
        rows.append({"topic_id": idx, "keywords": kw, "source": source_label})

    df = pd.DataFrame(rows)
    return df, lda


# ===================================================================
# 4. generate_labels  — Create human-readable labels for each topic
# ===================================================================
def generate_labels(
    topics_df: pd.DataFrame,
    use_llm: bool = False,
    groq_api_key: Optional[str] = None,
) -> pd.DataFrame:
    """
    Generate a short human-readable label for every topic.

    Strategy:
      - If use_llm=True and a Groq API key is provided, use the Groq LLM
        (llama-3.3-70b-versatile, free tier) to produce contextual labels.
      - Otherwise, apply a heuristic: capitalise the first 3–4 keywords.

    Parameters
    ----------
    topics_df : pd.DataFrame
        Must contain columns 'topic_id' and 'keywords'.
    use_llm : bool
        Whether to use the Groq LLM for label generation.
    groq_api_key : str, optional
        Groq API key, required if use_llm is True.

    Returns
    -------
    pd.DataFrame
        Same DataFrame with an additional 'label' column.
    """
    if use_llm and groq_api_key:
        logger.info("Generating labels using Groq LLM …")
        topics_df = _generate_labels_llm(topics_df, groq_api_key)
    else:
        logger.info("Generating labels using keyword heuristic …")
        topics_df = _generate_labels_heuristic(topics_df)

    return topics_df


def _generate_labels_heuristic(df: pd.DataFrame) -> pd.DataFrame:
    """Create labels from the top keywords of each topic."""
    labels = []
    for _, row in df.iterrows():
        kws = [kw.strip() for kw in row["keywords"].split(",")]
        # Take the first 3-4 non-trivial keywords and title-case them
        candidates = [kw.title() for kw in kws if len(kw) > 2][:4]
        label = " / ".join(candidates) if candidates else f"Topic {row['topic_id']}"
        labels.append(label)
    df = df.copy()
    df["label"] = labels
    return df


def _generate_labels_llm(df: pd.DataFrame, api_key: str) -> pd.DataFrame:
    """Use Groq API to generate contextual labels for topics (batched)."""
    import time
    try:
        from groq import Groq
    except ImportError:
        logger.warning("groq package not installed. Falling back to heuristic labels.")
        return _generate_labels_heuristic(df)

    client = Groq(api_key=api_key)
    labels = []

    # Process in batches to avoid rate limits
    batch_size = 10
    for batch_start in range(0, len(df), batch_size):
        batch = df.iloc[batch_start:batch_start + batch_size]
        prompt_lines = []
        for _, row in batch.iterrows():
            prompt_lines.append(f"Topic {row['topic_id']}: keywords = [{row['keywords']}]")

        prompt = (
            "You are a research taxonomy expert. For each topic below, "
            "generate a concise, descriptive label (3-6 words) that captures "
            "the theme of the keywords. Return ONLY a JSON list of objects "
            'with keys "topic_id" and "label". No extra text.\n\n'
            + "\n".join(prompt_lines)
        )

        try:
            chat = client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.3,
                max_tokens=1024,
            )
            resp = chat.choices[0].message.content.strip()
            # Parse JSON from the response
            # Find JSON array in response
            json_match = re.search(r"\[.*\]", resp, re.DOTALL)
            if json_match:
                batch_labels = json.loads(json_match.group())
                label_map = {item["topic_id"]: item["label"] for item in batch_labels}
                for _, row in batch.iterrows():
                    labels.append(label_map.get(row["topic_id"], f"Topic {row['topic_id']}"))
            else:
                # Fallback for this batch
                for _, row in batch.iterrows():
                    kws = [kw.strip().title() for kw in row["keywords"].split(",")][:4]
                    labels.append(" / ".join(kws))
        except Exception as exc:
            logger.warning("Groq API error for batch starting at %d: %s", batch_start, exc)
            for _, row in batch.iterrows():
                kws = [kw.strip().title() for kw in row["keywords"].split(",")][:4]
                labels.append(" / ".join(kws))

        # Rate-limit courtesy delay
        time.sleep(0.5)

    df = df.copy()
    df["label"] = labels
    return df


# ===================================================================
# 5. compare_themes  — Cross-compare title vs abstract topics
# ===================================================================
def compare_themes(
    title_topics: pd.DataFrame,
    abstract_topics: pd.DataFrame,
) -> pd.DataFrame:
    """
    Build a comparison table showing dominant themes from titles and
    abstracts side-by-side.

    Matching strategy:
      - Compute keyword overlap (Jaccard similarity) between every
        title-topic and abstract-topic pair.
      - For each title-topic, find the best matching abstract-topic.
      - Report similarity score and alignment status.

    Parameters
    ----------
    title_topics : pd.DataFrame
        Topics extracted from titles (with 'topic_id', 'keywords', 'label').
    abstract_topics : pd.DataFrame
        Topics extracted from abstracts (with 'topic_id', 'keywords', 'label').

    Returns
    -------
    pd.DataFrame
        Comparison table with columns:
          title_topic_id, title_label, title_keywords,
          abstract_topic_id, abstract_label, abstract_keywords,
          similarity, alignment
    """
    logger.info("Comparing themes: %d title topics × %d abstract topics",
                len(title_topics), len(abstract_topics))

    def _keywords_set(kw_str: str) -> set:
        return set(kw.strip().lower() for kw in kw_str.split(",") if kw.strip())

    rows = []
    for _, t_row in title_topics.iterrows():
        t_kws = _keywords_set(t_row["keywords"])
        best_sim = 0.0
        best_match = None

        for _, a_row in abstract_topics.iterrows():
            a_kws = _keywords_set(a_row["keywords"])
            if not t_kws or not a_kws:
                continue
            # Jaccard similarity
            intersection = len(t_kws & a_kws)
            union = len(t_kws | a_kws)
            sim = intersection / union if union else 0.0
            if sim > best_sim:
                best_sim = sim
                best_match = a_row

        alignment = (
            "Strong" if best_sim >= 0.4
            else "Moderate" if best_sim >= 0.2
            else "Weak" if best_sim > 0
            else "No Match"
        )

        rows.append({
            "title_topic_id": t_row["topic_id"],
            "title_label": t_row.get("label", ""),
            "title_keywords": t_row["keywords"],
            "abstract_topic_id": best_match["topic_id"] if best_match is not None else None,
            "abstract_label": best_match.get("label", "") if best_match is not None else "",
            "abstract_keywords": best_match["keywords"] if best_match is not None else "",
            "similarity": round(best_sim, 4),
            "alignment": alignment,
        })

    comparison_df = pd.DataFrame(rows)
    logger.info("Theme comparison complete: %d rows", len(comparison_df))
    return comparison_df


# ===================================================================
# 6. create_taxonomy_map  — Classify themes as MAPPED or NOVEL
# ===================================================================
def create_taxonomy_map(
    topics_df: pd.DataFrame,
    known_themes: Optional[List[str]] = None,
    threshold: float = 0.15,
) -> Dict[str, Any]:
    """
    Classify each topic as either MAPPED (similar to a well-known
    AI / business / IS research theme) or NOVEL (previously unseen).

    Heuristic:
      For each topic's keyword set, compute its best token-overlap
      ratio against the known themes list.  If the ratio exceeds the
      threshold, label it as MAPPED; otherwise NOVEL.

    Parameters
    ----------
    topics_df : pd.DataFrame
        Must contain 'topic_id', 'keywords', and 'label' columns.
    known_themes : list of str, optional
        Reference themes (defaults to the built-in KNOWN_THEMES).
    threshold : float
        Minimum overlap ratio to classify as MAPPED.

    Returns
    -------
    dict
        JSON-serialisable taxonomy map:
        {
          "metadata": { ... },
          "mapped": [ {topic_id, label, keywords, matched_theme, score}, ... ],
          "novel":  [ {topic_id, label, keywords, score}, ... ],
        }
    """
    if known_themes is None:
        known_themes = KNOWN_THEMES

    logger.info("Building taxonomy map for %d topics (threshold=%.2f)", len(topics_df), threshold)

    mapped: List[Dict] = []
    novel: List[Dict] = []

    known_tokens_list = [set(theme.lower().split()) for theme in known_themes]

    for _, row in topics_df.iterrows():
        topic_tokens = set(
            kw.strip().lower()
            for kw in row["keywords"].split(",")
            if kw.strip()
        )
        # Also include individual words from multi-word keywords
        expanded_tokens = set()
        for token in topic_tokens:
            expanded_tokens.update(token.split())
        expanded_tokens.update(topic_tokens)

        best_score = 0.0
        best_theme = ""
        for theme_str, theme_tokens in zip(known_themes, known_tokens_list):
            if not expanded_tokens or not theme_tokens:
                continue
            intersection = len(expanded_tokens & theme_tokens)
            union_size = len(expanded_tokens | theme_tokens)
            score = intersection / union_size if union_size else 0.0
            if score > best_score:
                best_score = score
                best_theme = theme_str

        entry = {
            "topic_id": int(row["topic_id"]),
            "label": row.get("label", ""),
            "keywords": row["keywords"],
            "score": round(best_score, 4),
        }

        if best_score >= threshold:
            entry["matched_theme"] = best_theme
            entry["classification"] = "MAPPED"
            mapped.append(entry)
        else:
            entry["classification"] = "NOVEL"
            novel.append(entry)

    taxonomy = {
        "metadata": {
            "total_topics": len(topics_df),
            "mapped_count": len(mapped),
            "novel_count": len(novel),
            "threshold": threshold,
        },
        "mapped": mapped,
        "novel": novel,
    }

    logger.info("Taxonomy: %d MAPPED, %d NOVEL", len(mapped), len(novel))
    return taxonomy

# ===================================================================
# TCCM Extraction Tool
# ===================================================================
def extract_tccm(documents: List[str], api_key: str) -> pd.DataFrame:
    # Basic full-text TCCM extraction stub
    logger.info("Extracting TCCM from full text...")
    return pd.DataFrame([{"TCCM": "Sample"} for _ in documents])


# ===================================================================
# SPECTER2 Pipeline Functions
# ===================================================================

def generate_specter2_embeddings(papers: List[Dict], batch_size: int = 16) -> np.ndarray:
    """
    Generate embeddings using SPECTER2 via adapters.
    papers: list of dicts with 'title' and 'abstract' keys.
    """
    import torch
    from transformers import AutoTokenizer
    from adapters import AutoAdapterModel

    logger.info("Loading SPECTER2 model and tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained('allenai/specter2_base')
    model = AutoAdapterModel.from_pretrained('allenai/specter2_base')
    model.load_adapter("allenai/specter2", source="hf", load_as="proximity", set_active=True)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    model.eval()

    embeddings = []
    logger.info("Generating embeddings for %d papers using %s...", len(papers), device)

    for i in range(0, len(papers), batch_size):
        batch = papers[i:i+batch_size]
        text_batch = [d['title'] + tokenizer.sep_token + (d.get('abstract') or '') for d in batch]

        inputs = tokenizer(
            text_batch, padding=True, truncation=True,
            return_tensors="pt", return_token_type_ids=False, max_length=512,
        )
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            output = model(**inputs)
            batch_embeddings = output.last_hidden_state[:, 0, :].cpu().numpy()
            embeddings.append(batch_embeddings)

    return np.vstack(embeddings)


def store_in_chroma(embeddings, documents, metadatas, collection_name="specter_vectors"):
    """Store embeddings in ChromaDB for persistence."""
    import chromadb

    client = chromadb.PersistentClient(path="./chroma_db")
    # Delete existing collection to avoid ID conflicts
    try:
        client.delete_collection(name=collection_name)
    except Exception:
        pass
    collection = client.get_or_create_collection(name=collection_name)

    ids = [str(i) for i in range(len(documents))]
    collection.add(
        embeddings=embeddings.tolist(),
        documents=documents,
        metadatas=metadatas,
        ids=ids,
    )
    return collection


def load_embeddings_from_chroma(collection_name="specter_vectors"):
    """Load cached embeddings from ChromaDB."""
    import chromadb

    client = chromadb.PersistentClient(path="./chroma_db")
    collection = client.get_collection(name=collection_name)
    result = collection.get(include=["embeddings", "documents", "metadatas"])

    if not result["embeddings"]:
        raise ValueError("No embeddings found in ChromaDB.")

    embeddings = np.array(result["embeddings"])
    documents = result["documents"]
    metadatas = result["metadatas"]
    return embeddings, documents, metadatas


def adaptive_hdbscan_clustering(
    embeddings: np.ndarray,
    min_clusters: int = 10,
    max_clusters: int = 40,
) -> tuple:
    """
    Adaptive HDBSCAN clustering with UMAP dimensionality reduction.
    Tries different min_cluster_size values to land within the target cluster range.

    Returns: (labels, reduced_2d, n_clusters, params_used)
    """
    from umap import UMAP
    from hdbscan import HDBSCAN

    logger.info(
        "Adaptive HDBSCAN: targeting %d–%d clusters for %d embeddings",
        min_clusters, max_clusters, len(embeddings),
    )

    # UMAP reduction to 5D for clustering
    umap_5d = UMAP(
        n_neighbors=15, n_components=5, min_dist=0.0,
        metric="cosine", random_state=42,
    )
    reduced_5d = umap_5d.fit_transform(embeddings)

    # UMAP reduction to 2D for visualization
    umap_2d = UMAP(
        n_neighbors=15, n_components=2, min_dist=0.1,
        metric="cosine", random_state=42,
    )
    reduced_2d = umap_2d.fit_transform(embeddings)

    # Try different min_cluster_size to find good cluster count
    best_labels = None
    best_n_clusters = 0
    best_params = {}
    best_diff = float("inf")

    target_mid = (min_clusters + max_clusters) / 2

    for mcs in [5, 8, 10, 15, 20, 25, 30, 40, 50, 60, 80, 100, 120, 150]:
        for ms in [2, 3, 5]:
            clusterer = HDBSCAN(
                min_cluster_size=mcs,
                min_samples=ms,
                metric="euclidean",
                cluster_selection_method="eom",
            )
            labels = clusterer.fit_predict(reduced_5d)
            n_clusters = len(set(labels)) - (1 if -1 in labels else 0)

            diff = abs(n_clusters - target_mid)
            in_range = min_clusters <= n_clusters <= max_clusters

            if in_range and diff < best_diff:
                best_labels = labels
                best_n_clusters = n_clusters
                best_params = {"min_cluster_size": mcs, "min_samples": ms}
                best_diff = diff

            if n_clusters == round(target_mid):
                break

    # Fallback: if nothing landed in range, pick closest
    if best_labels is None:
        logger.warning("Could not find params for target range, using best effort")
        for mcs in [15, 20, 25, 30, 40, 50]:
            clusterer = HDBSCAN(
                min_cluster_size=mcs, min_samples=3,
                metric="euclidean", cluster_selection_method="eom",
            )
            labels = clusterer.fit_predict(reduced_5d)
            n_clusters = len(set(labels)) - (1 if -1 in labels else 0)
            if n_clusters >= 5:
                best_labels = labels
                best_n_clusters = n_clusters
                best_params = {"min_cluster_size": mcs, "min_samples": 3}
                break

    if best_labels is None:
        raise RuntimeError("HDBSCAN could not produce any valid clusters.")

    logger.info(
        "HDBSCAN selected: %d clusters (params: %s)",
        best_n_clusters, best_params,
    )
    return best_labels, reduced_2d, best_n_clusters, best_params


def c_tf_idf(documents, m, ngram_range=(1, 1)):
    """Compute class-based TF-IDF for topic keyword extraction."""
    from sklearn.feature_extraction.text import CountVectorizer

    count = CountVectorizer(ngram_range=ngram_range, stop_words="english").fit(documents)
    t = count.transform(documents).toarray()
    w = t.sum(axis=1)
    tf = np.divide(t.T, w)
    sum_t = t.sum(axis=0)
    idf = np.log(np.divide(m, sum_t)).reshape(-1, 1)
    tf_idf = np.multiply(tf, idf)
    return tf_idf, count


def extract_top_n_words_per_topic(tf_idf, count, docs_per_topic, n=10):
    """Extract top N keywords per topic from c-TF-IDF matrix."""
    words = count.get_feature_names_out()
    labels = list(docs_per_topic.topic_id)
    tf_idf_transposed = tf_idf.T
    indices = tf_idf_transposed.argsort()[:, -n:]
    top_n_words = {
        label: [(words[j], tf_idf_transposed[i][j]) for j in indices[i]][::-1]
        for i, label in enumerate(labels)
    }
    return top_n_words


def extract_representative_sentences(
    df: pd.DataFrame,
    embeddings: np.ndarray,
    cluster_labels: np.ndarray,
    top_n: int = 5,
) -> Dict[int, List[str]]:
    """
    For each cluster, find the top_n documents closest to the cluster centroid.
    Returns {cluster_id: [sentence1, sentence2, ...]}.
    """
    from sklearn.metrics.pairwise import cosine_similarity

    result = {}
    unique_labels = set(cluster_labels)

    for cid in unique_labels:
        if cid == -1:
            continue
        mask = cluster_labels == cid
        cluster_embs = embeddings[mask]
        cluster_indices = np.where(mask)[0]

        if len(cluster_embs) == 0:
            continue

        centroid = cluster_embs.mean(axis=0).reshape(1, -1)
        sims = cosine_similarity(centroid, cluster_embs)[0]
        top_indices = sims.argsort()[-top_n:][::-1]

        sentences = []
        for idx in top_indices:
            doc_idx = cluster_indices[idx]
            title = str(df.iloc[doc_idx].get("Title", ""))
            abstract = str(df.iloc[doc_idx].get("Abstract", ""))
            # Use title + first sentence of abstract
            first_sent = abstract.split(".")[0] if abstract else ""
            rep = f"{title}. {first_sent}".strip()
            if rep:
                sentences.append(rep)
        result[int(cid)] = sentences

    return result


def get_council_debate(keywords: str, groq_api_key: str) -> Dict[str, str]:
    """
    Run a 3-agent AI Council debate to generate a topic label.
    Returns dict with: expert, analyst, anti_hallucination, final_label.
    """
    import time

    try:
        from groq import Groq
    except ImportError:
        kws = [kw.strip().title() for kw in keywords.split(",")][:3]
        return {
            "expert": ", ".join(kws),
            "analyst": ", ".join(kws),
            "anti_hallucination": "N/A",
            "final_label": ", ".join(kws),
        }

    client = Groq(api_key=groq_api_key)

    # Persona 1: Domain Expert
    prompt_expert = (
        f"You are a scientific Domain Expert. Keywords: [{keywords}]. "
        "Propose a 3-5 word descriptive label for this academic topic. "
        "Return ONLY the label."
    )

    # Persona 2: Data Analyst
    prompt_analyst = (
        f"You are a Data Analyst categorizing research clusters. Keywords: [{keywords}]. "
        "Give a short, precise categorical name (3-5 words). Return ONLY the name."
    )

    try:
        res_expert = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[{"role": "user", "content": prompt_expert}],
            temperature=0.4, max_tokens=20,
        ).choices[0].message.content.strip()

        time.sleep(0.3)

        res_analyst = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[{"role": "user", "content": prompt_analyst}],
            temperature=0.4, max_tokens=20,
        ).choices[0].message.content.strip()

        time.sleep(0.3)

        # Persona 3: Anti-hallucination checker
        prompt_check = (
            f"Keywords: [{keywords}].\n"
            f"Expert label: '{res_expert}'.\n"
            f"Analyst label: '{res_analyst}'.\n"
            "Check: are these labels grounded in the keywords? "
            "If yes, synthesize into a final 3-5 word label. "
            "If not, create a grounded label from the keywords only. "
            "Return ONLY the final label."
        )

        res_judge = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[{"role": "user", "content": prompt_check}],
            temperature=0.2, max_tokens=20,
        ).choices[0].message.content.strip()
        res_judge = res_judge.replace('"', '').replace("'", "").strip()

        time.sleep(0.3)

        return {
            "expert": res_expert,
            "analyst": res_analyst,
            "anti_hallucination": res_judge,
            "final_label": res_judge,
        }

    except Exception as e:
        logger.warning("Council debate failed: %s", e)
        kws = [kw.strip().title() for kw in keywords.split(",")][:3]
        return {
            "expert": ", ".join(kws),
            "analyst": ", ".join(kws),
            "anti_hallucination": "fallback",
            "final_label": ", ".join(kws),
        }


def build_cluster_graph(
    topics_df: pd.DataFrame,
    embeddings: np.ndarray,
    cluster_labels: np.ndarray,
    top_n_edges: int = 20,
) -> List[Dict]:
    """
    Build a similarity graph between clusters based on centroid cosine similarity.
    Returns list of edges: [{source, target, weight}, ...].
    """
    from sklearn.metrics.pairwise import cosine_similarity

    unique_clusters = sorted([c for c in set(cluster_labels) if c != -1])
    if len(unique_clusters) < 2:
        return []

    # Compute centroids
    centroids = []
    for cid in unique_clusters:
        mask = cluster_labels == cid
        centroids.append(embeddings[mask].mean(axis=0))
    centroids = np.array(centroids)

    # Pairwise similarity
    sim_matrix = cosine_similarity(centroids)

    edges = []
    for i in range(len(unique_clusters)):
        for j in range(i + 1, len(unique_clusters)):
            edges.append({
                "source": int(unique_clusters[i]),
                "target": int(unique_clusters[j]),
                "weight": round(float(sim_matrix[i][j]), 4),
            })

    # Keep top N edges by weight
    edges.sort(key=lambda e: e["weight"], reverse=True)
    return edges[:top_n_edges]


def merge_similar_clusters(
    topics_df: pd.DataFrame,
    cluster_labels: np.ndarray,
    df: pd.DataFrame,
    threshold: float = 0.75,
    groq_api_key: Optional[str] = None,
) -> tuple:
    """
    Merge clusters with highly similar keyword sets.
    Returns (new_topics_df, new_labels).
    """
    from sklearn.feature_extraction.text import TfidfVectorizer
    from sklearn.metrics.pairwise import cosine_similarity

    if len(topics_df) < 2:
        return topics_df, cluster_labels

    # Compute pairwise similarity of keyword strings
    vectorizer = TfidfVectorizer(stop_words="english")
    kw_matrix = vectorizer.fit_transform(topics_df["keywords"].tolist())
    sim_matrix = cosine_similarity(kw_matrix)

    # Find pairs to merge
    merge_map = {}
    topic_ids = topics_df["topic_id"].tolist()

    for i in range(len(topic_ids)):
        for j in range(i + 1, len(topic_ids)):
            if sim_matrix[i][j] >= threshold:
                src = topic_ids[j]
                tgt = topic_ids[i]
                # Follow existing merge chains
                while tgt in merge_map:
                    tgt = merge_map[tgt]
                merge_map[src] = tgt

    if not merge_map:
        return topics_df, cluster_labels

    logger.info("Merging %d cluster pairs (threshold=%.2f)", len(merge_map), threshold)

    # Remap labels
    new_labels = np.array(cluster_labels, copy=True)
    for old_id, new_id in merge_map.items():
        new_labels[new_labels == old_id] = new_id

    # Rebuild topics_df for remaining clusters
    remaining_ids = set(topic_ids) - set(merge_map.keys())
    new_topics_df = topics_df[topics_df["topic_id"].isin(remaining_ids)].copy()
    new_topics_df = new_topics_df.reset_index(drop=True)

    return new_topics_df, new_labels


def extract_subtopics_within_clusters(
    df: pd.DataFrame,
    cluster_labels: np.ndarray,
    min_topics: int = 100,
    n_keywords: int = 10,
) -> pd.DataFrame:
    """
    Extract fine-grained sub-topics within each cluster using NMF.
    Delegates to the implementation in run_specter.py.
    """
    from run_specter import extract_subtopics_within_clusters as _extract
    return _extract(df, cluster_labels, min_topics, n_keywords)