"""
computation_clusters.py — SPECTER2 semantic clustering for computation.csv

Self-contained module that:
  1. Loads computation.csv
  2. Combines text fields → single string per paper
  3. Generates SPECTER2 embeddings (cached as .npy)
  4. UMAP 2D reduction
  5. HDBSCAN clustering
  6. Auto-labels clusters via TF-IDF keywords
  7. Returns everything needed for the Streamlit Computational Strategy Map tab
"""

import os
import hashlib
import logging
from typing import Dict, List, Tuple, Any, Optional

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

# Cache directory for embeddings
_CACHE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".comp_cache")
os.makedirs(_CACHE_DIR, exist_ok=True)


def _cache_key(texts: List[str]) -> str:
    """Stable hash of input texts for caching."""
    h = hashlib.sha256()
    for t in texts:
        h.update(t.encode("utf-8", errors="replace"))
    return h.hexdigest()[:16]


def _generate_specter2_embeddings(texts: List[str], batch_size: int = 8) -> np.ndarray:
    """
    Generate SPECTER2 embeddings via HuggingFace adapters, with .npy caching.
    """
    key = _cache_key(texts)
    cache_path = os.path.join(_CACHE_DIR, f"comp_emb_{key}.npy")

    if os.path.exists(cache_path):
        logger.info("✓ Loaded cached embeddings from %s", cache_path)
        return np.load(cache_path)

    import torch
    from transformers import AutoTokenizer
    from adapters import AutoAdapterModel

    logger.info("Loading SPECTER2 model …")
    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 = []
    total_batches = (len(texts) + batch_size - 1) // batch_size
    logger.info("Generating embeddings for %d texts (%d batches) on %s …", len(texts), total_batches, device)

    for i in range(0, len(texts), batch_size):
        batch = texts[i : i + batch_size]
        inputs = tokenizer(
            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_emb = output.last_hidden_state[:, 0, :].cpu().numpy()
            embeddings.append(batch_emb)

    result = np.vstack(embeddings)
    np.save(cache_path, result)
    logger.info("✓ Embeddings cached → %s", cache_path)
    return result


def _umap_reduce(embeddings: np.ndarray, n_components: int = 2) -> np.ndarray:
    """UMAP dimensionality reduction."""
    from umap import UMAP

    n_neighbors = min(15, len(embeddings) - 1)
    reducer = UMAP(
        n_neighbors=n_neighbors,
        n_components=n_components,
        min_dist=0.1,
        metric="cosine",
        random_state=42,
    )
    return reducer.fit_transform(embeddings)


def _hdbscan_cluster(embeddings: np.ndarray) -> np.ndarray:
    """
    HDBSCAN clustering via UMAP 5D intermediate reduction.
    Tries multiple parameter combos to find the best clustering.
    Falls back gracefully for small datasets.
    """
    from umap import UMAP
    from hdbscan import HDBSCAN

    n_samples = len(embeddings)

    # For very small datasets, use 3D intermediate
    n_intermediate = min(5, n_samples - 2)
    if n_intermediate < 2:
        n_intermediate = 2

    n_neighbors = min(15, n_samples - 1)
    umap_inter = UMAP(
        n_neighbors=n_neighbors,
        n_components=n_intermediate,
        min_dist=0.0,
        metric="cosine",
        random_state=42,
    )
    reduced = umap_inter.fit_transform(embeddings)

    best_labels = None
    best_n = 0

    # Try parameter sweep
    for mcs in [3, 4, 5, 7]:
        for ms in [2, 3]:
            if mcs >= n_samples:
                continue
            clusterer = HDBSCAN(
                min_cluster_size=mcs,
                min_samples=min(ms, mcs - 1) if mcs > 1 else 1,
                metric="euclidean",
                cluster_selection_method="eom",
            )
            labels = clusterer.fit_predict(reduced)
            n_clusters = len(set(labels)) - (1 if -1 in labels else 0)
            if n_clusters > best_n:
                best_labels = labels
                best_n = n_clusters

    if best_labels is None or best_n < 2:
        logger.warning("HDBSCAN produced < 2 clusters, using single cluster")
        best_labels = np.zeros(n_samples, dtype=int)

    # Reassign noise (-1) to nearest cluster centroid
    noise_mask = best_labels == -1
    if noise_mask.any():
        from sklearn.metrics.pairwise import cosine_similarity

        valid_clusters = sorted(set(best_labels) - {-1})
        if valid_clusters:
            centroids = np.array([embeddings[best_labels == c].mean(axis=0) for c in valid_clusters])
            sims = cosine_similarity(embeddings[noise_mask], centroids)
            nearest = sims.argmax(axis=1)
            best_labels[noise_mask] = [valid_clusters[n] for n in nearest]

    logger.info("HDBSCAN: %d clusters from %d papers", len(set(best_labels)), n_samples)
    return best_labels


def _label_clusters(texts: List[str], labels: np.ndarray) -> Dict[int, str]:
    """Generate cluster labels using TF-IDF top keywords."""
    from sklearn.feature_extraction.text import TfidfVectorizer

    cluster_ids = sorted(set(labels))

    # Merge all texts per cluster
    cluster_texts = {}
    for cid in cluster_ids:
        mask = labels == cid
        cluster_texts[cid] = " ".join([texts[i] for i in range(len(texts)) if mask[i]])

    all_docs = [cluster_texts[c] for c in cluster_ids]

    vectorizer = TfidfVectorizer(
        stop_words="english", max_df=0.95, min_df=1,
        ngram_range=(1, 2), max_features=3000,
    )
    tfidf = vectorizer.fit_transform(all_docs)
    features = vectorizer.get_feature_names_out()

    result = {}
    for i, cid in enumerate(cluster_ids):
        scores = tfidf[i].toarray().flatten()
        top_idx = scores.argsort()[-4:][::-1]
        keywords = [features[j] for j in top_idx if scores[j] > 0]
        label = " / ".join(w.title() for w in keywords[:3]) if keywords else f"Cluster {cid}"
        result[cid] = label

    return result


def _build_cluster_keywords(texts: List[str], labels: np.ndarray) -> Dict[int, List[str]]:
    """Get top-5 keywords per cluster for the stats panel."""
    from sklearn.feature_extraction.text import TfidfVectorizer

    cluster_ids = sorted(set(labels))
    cluster_texts = {}
    for cid in cluster_ids:
        mask = labels == cid
        cluster_texts[cid] = " ".join([texts[i] for i in range(len(texts)) if mask[i]])

    all_docs = [cluster_texts[c] for c in cluster_ids]

    vectorizer = TfidfVectorizer(
        stop_words="english", max_df=0.95, min_df=1,
        ngram_range=(1, 2), max_features=3000,
    )
    tfidf = vectorizer.fit_transform(all_docs)
    features = vectorizer.get_feature_names_out()

    result = {}
    for i, cid in enumerate(cluster_ids):
        scores = tfidf[i].toarray().flatten()
        top_idx = scores.argsort()[-5:][::-1]
        result[cid] = [features[j] for j in top_idx if scores[j] > 0]

    return result


def run_computation_clustering(csv_path: str = "computation.csv") -> Dict[str, Any]:
    """
    Full pipeline: Load CSV → Embed → Reduce → Cluster → Label.

    Returns a dict with all data needed by the Streamlit tab:
      - df: DataFrame with original columns + cluster_id, cluster_name, x, y
      - cluster_names: {cluster_id: label}
      - cluster_keywords: {cluster_id: [kw1, kw2, ...]}
      - strategy_types: list of unique strategy types
      - n_clusters: int
      - n_papers: int
    """
    logger.info("=" * 50)
    logger.info("Computation Strategy Clustering Pipeline")
    logger.info("=" * 50)

    # 1. Load CSV
    df = pd.read_csv(csv_path, encoding="utf-8", on_bad_lines="skip")
    df.columns = [c.strip() for c in df.columns]
    for col in df.columns:
        df[col] = df[col].fillna("")
    logger.info("Loaded %d papers from %s", len(df), csv_path)

    # 2. Combine text fields
    combined = []
    for _, row in df.iterrows():
        parts = [
            str(row.get("Paper Title", "")),
            str(row.get("Computational Strategies Used", "")),
            str(row.get("Strategy Type", "")),
            str(row.get("Supporting Keywords / Evidence", "")),
        ]
        combined.append(" ".join(p for p in parts if p.strip()))

    # 3. SPECTER2 embeddings (cached)
    embeddings = _generate_specter2_embeddings(combined)

    # 4. UMAP 2D
    coords_2d = _umap_reduce(embeddings)

    # 5. HDBSCAN
    cluster_labels = _hdbscan_cluster(embeddings)

    # 6. TF-IDF labeling
    cluster_names = _label_clusters(combined, cluster_labels)
    cluster_keywords = _build_cluster_keywords(combined, cluster_labels)

    # 7. Attach to DataFrame
    df["cluster_id"] = cluster_labels.tolist()
    df["cluster_name"] = df["cluster_id"].map(cluster_names)
    df["x"] = coords_2d[:, 0].tolist()
    df["y"] = coords_2d[:, 1].tolist()

    strategy_types = sorted(set(s.strip() for s in df["Strategy Type"].unique() if s.strip()))

    logger.info("Pipeline complete: %d papers → %d clusters", len(df), len(cluster_names))

    return {
        "df": df,
        "cluster_names": cluster_names,
        "cluster_keywords": cluster_keywords,
        "strategy_types": strategy_types,
        "n_clusters": len(cluster_names),
        "n_papers": len(df),
    }