"""
ActiveSamplingStrategy — Informative chain selection for cold-start calibration.
=================================================================================

Reduces labeling burden from 50+ → 5-10 chains by selecting the most
informative chains from an unlabeled pool.

Three strategies:
  diversity   — greedy max-min distance (maximize coverage of chain space)
  uncertainty — select chains closest to decision boundary (|score - 0.5| smallest)
  hybrid      — alternate: even picks = diversity, odd picks = uncertainty

Usage
-----
    from llm_guard import ActiveSamplingStrategy

    strategy = ActiveSamplingStrategy(strategy="hybrid")
    to_label = strategy.select(pool=unlabeled_chains, k=5, scores=current_scores)
    # Label these 5 chains, then call guard.calibrate_isotonic(scores, labels)
"""
from __future__ import annotations

from typing import List, Literal, Optional

import numpy as np


class ActiveSamplingStrategy:
    """
    Select the k most informative chains from an unlabeled pool.

    Parameters
    ----------
    strategy : "diversity" | "uncertainty" | "hybrid"
    embedder : "minilm" | "tfidf"
        MiniLM requires sentence_transformers. Falls back to tfidf if unavailable.
    """

    def __init__(
        self,
        strategy: Literal["diversity", "uncertainty", "hybrid"] = "hybrid",
        embedder: Literal["minilm", "tfidf"] = "minilm",
    ) -> None:
        if strategy not in ("diversity", "uncertainty", "hybrid"):
            raise ValueError(f"strategy must be diversity|uncertainty|hybrid, got {strategy!r}")
        self._strategy = strategy
        self._embedder = embedder

    def select(
        self,
        pool: List[dict],
        k: int,
        scores: Optional[List[float]] = None,
    ) -> List[dict]:
        """Return k chains from pool that are most informative for labeling."""
        if not pool:
            return []
        k = min(k, len(pool))

        if self._strategy == "diversity":
            return self._diversity(pool, k)
        if self._strategy == "uncertainty":
            if scores is None:
                raise ValueError(
                    "ActiveSamplingStrategy(strategy='uncertainty') requires scores. "
                    "Pass scores=current_risk_scores to select()."
                )
            return self._uncertainty(pool, k, scores)
        # hybrid
        if scores is None:
            return self._diversity(pool, k)
        return self._hybrid(pool, k, scores)

    def _embed(self, pool: List[dict]) -> np.ndarray:
        """Embed chains to vectors. Falls back to TF-IDF if MiniLM unavailable."""
        texts = [self._chain_text(c) for c in pool]
        if self._embedder == "minilm":
            try:
                from sentence_transformers import SentenceTransformer
                model = SentenceTransformer("all-MiniLM-L6-v2")
                return model.encode(texts, show_progress_bar=False)
            except ImportError:
                pass
        from sklearn.feature_extraction.text import TfidfVectorizer
        vec = TfidfVectorizer(max_features=200)
        return vec.fit_transform(texts).toarray().astype(np.float32)

    @staticmethod
    def _chain_text(chain: dict) -> str:
        q = chain.get("question", "")
        steps = chain.get("steps", [])
        thoughts = " ".join(s.get("thought", "") for s in steps)
        return f"{q} {thoughts}".strip()

    def _diversity(self, pool: List[dict], k: int) -> List[dict]:
        """Greedy max-min distance selection."""
        embeddings = self._embed(pool)
        n = len(pool)
        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
        norms = np.where(norms == 0, 1, norms)
        emb = embeddings / norms

        selected_idx = [0]
        min_dists = np.full(n, np.inf)

        for _ in range(k - 1):
            last = emb[selected_idx[-1]]
            dists = 1.0 - emb @ last  # cosine distance
            min_dists = np.minimum(min_dists, dists)
            min_dists[selected_idx] = -np.inf
            selected_idx.append(int(np.argmax(min_dists)))

        return [pool[i] for i in selected_idx]

    @staticmethod
    def _uncertainty(pool: List[dict], k: int, scores: List[float]) -> List[dict]:
        """Select chains with scores closest to 0.5 (highest uncertainty)."""
        distances = [abs(s - 0.5) for s in scores]
        sorted_idx = sorted(range(len(pool)), key=lambda i: distances[i])
        return [pool[i] for i in sorted_idx[:k]]

    def _hybrid(self, pool: List[dict], k: int, scores: List[float]) -> List[dict]:
        """Alternate: even picks = diversity, odd picks = uncertainty."""
        selected = []
        selected_idx = set()

        embeddings = self._embed(pool)
        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
        norms = np.where(norms == 0, 1, norms)
        emb = embeddings / norms

        diverse_selected = []

        for pick in range(k):
            avail = [i for i in range(len(pool)) if i not in selected_idx]
            if not avail:
                break

            if pick % 2 == 0:
                if not diverse_selected:
                    chosen = avail[0]
                else:
                    last_emb = emb[diverse_selected[-1]]
                    dists = [1.0 - float(emb[i] @ last_emb) for i in avail]
                    chosen = avail[int(np.argmax(dists))]
                diverse_selected.append(chosen)
            else:
                unc = [abs(scores[i] - 0.5) for i in avail]
                chosen = avail[int(np.argmin(unc))]

            selected.append(pool[chosen])
            selected_idx.add(chosen)

        return selected
