"""
Trust Engine: QPPG-based trust verification for LLM outputs.

Uses bifurcation-based topology analysis to determine whether an LLM's
answer is trustworthy. Core insight: sample N solutions to the same problem,
embed them, and analyze the distribution:
  - Tight single cluster → consistent → likely correct
  - Multiple clusters (bifurcation) → inconsistent → lower trust
  - QPPG mu-sweep detects solution modes as phase transitions, not thresholds

This is model-agnostic (works on any LLM's text outputs), requires no
logprobs or model internals, and needs zero training data.

Usage:
    from qppg import TrustScorer, SolutionGenerator

    generator = SolutionGenerator()
    scorer = TrustScorer()

    solutions = generator.generate(problem, n=10)
    trust = scorer.compute_trust(solutions)
    print(f"Trust: {trust['qppg_trust']:.3f}, Answer: {trust['majority_answer']}")
"""

import os
import re
import json
import time
import hashlib
import numpy as np
from pathlib import Path
from typing import List, Dict, Optional, Tuple
from collections import Counter


# ============================================================
# Solution Generator — query LLM N times for same problem
# ============================================================

class SolutionGenerator:
    """
    Generates N diverse solutions to a problem using an LLM.

    Parameters:
        api_key: Anthropic API key (or set ANTHROPIC_API_KEY env var)
        model: Claude model to use
        temperature: Sampling temperature (>0 for diversity)
        max_tokens: Max tokens per solution
        cache_path: Path to persist solution cache
    """

    def __init__(self, api_key=None, model="claude-haiku-4-5-20251001",
                 temperature=0.8, max_tokens=500, cache_path=None):
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.cache = {}
        self.cache_path = cache_path

        if cache_path and os.path.exists(cache_path):
            with open(cache_path, 'r') as f:
                self.cache = json.load(f)

        self.api_key = api_key or os.environ.get('ANTHROPIC_API_KEY', '')

        self._client = None
        self._init_client()

        self.total_input_tokens = 0
        self.total_output_tokens = 0
        self.total_calls = 0
        self.cache_hits = 0

    def _init_client(self):
        try:
            import anthropic
            self._client = anthropic.Anthropic(api_key=self.api_key)
        except ImportError:
            print("WARNING: anthropic package not installed. pip install anthropic")
            self._client = None

    def _call_api(self, system_prompt, user_prompt, max_retries=8):
        """Call Claude API with exponential backoff retry."""
        for attempt in range(max_retries):
            try:
                response = self._client.messages.create(
                    model=self.model,
                    max_tokens=self.max_tokens,
                    temperature=self.temperature,
                    system=system_prompt,
                    messages=[{"role": "user", "content": user_prompt}]
                )
                self.total_input_tokens += response.usage.input_tokens
                self.total_output_tokens += response.usage.output_tokens
                self.total_calls += 1
                return response.content[0].text
            except Exception as e:
                import sys
                sys.stderr.write(f"  [API retry {attempt+1}/{max_retries}: {type(e).__name__}: {str(e)[:80]}]\n")
                sys.stderr.flush()
                if attempt < max_retries - 1:
                    wait = min(2 ** attempt + np.random.random(), 15.0)
                    time.sleep(wait)
                else:
                    return f"[API Error: {str(e)[:100]}]"

    def _cache_key(self, problem, idx, seed):
        """Generate cache key for a specific solution."""
        data = f"{problem}|{idx}|{seed}|{self.model}|{self.temperature}".encode()
        return hashlib.sha256(data).hexdigest()[:16]

    def generate(self, problem, n=10, seed=42):
        """
        Generate N solutions to a problem.

        Args:
            problem: str — the problem text
            n: int — number of solutions to generate
            seed: int — for cache key differentiation across runs

        Returns:
            List[str] — N solution texts
        """
        system = ("You are a math tutor. Solve the problem step by step. "
                  "End your solution with the final numerical answer on its own line "
                  "in the format: #### [number]")

        solutions = []
        for i in range(n):
            key = self._cache_key(problem, i, seed)
            if key in self.cache:
                self.cache_hits += 1
                solutions.append(self.cache[key])
            else:
                sol = self._call_api(system, problem)
                self.cache[key] = sol
                solutions.append(sol)

        return solutions

    def generate_greedy(self, problem):
        """Generate a single greedy (temp=0) solution."""
        system = ("You are a math tutor. Solve the problem step by step. "
                  "End your solution with the final numerical answer on its own line "
                  "in the format: #### [number]")

        # Temporarily override temperature
        old_temp = self.temperature
        self.temperature = 0.0
        sol = self._call_api(system, problem)
        self.temperature = old_temp
        return sol

    def save_cache(self):
        """Persist solution cache to disk."""
        if self.cache_path:
            os.makedirs(os.path.dirname(self.cache_path) or '.', exist_ok=True)
            with open(self.cache_path, 'w') as f:
                json.dump(self.cache, f)

    def get_cost_summary(self):
        """Return cumulative API cost tracking."""
        cost_input = self.total_input_tokens * 0.80 / 1_000_000
        cost_output = self.total_output_tokens * 4.00 / 1_000_000
        return {
            'total_calls': self.total_calls,
            'cache_hits': self.cache_hits,
            'total_input_tokens': self.total_input_tokens,
            'total_output_tokens': self.total_output_tokens,
            'cost_usd': round(cost_input + cost_output, 4),
        }


# ============================================================
# Answer Extraction
# ============================================================

def extract_numerical_answer(text):
    """
    Extract the final numerical answer from a solution text.

    Looks for "#### [number]" format first (GSM8K style),
    then falls back to last number in the text.

    Returns:
        str or None — the extracted answer as a string, or None if not found
    """
    # Try GSM8K-style "#### number" format
    match = re.findall(r'####\s*(-?[\d,]+\.?\d*)', text)
    if match:
        return match[-1].replace(',', '')

    # Fallback: last number in text
    numbers = re.findall(r'-?[\d,]+\.?\d*', text)
    if numbers:
        return numbers[-1].replace(',', '')

    return None


# ============================================================
# Trust Scorer — QPPG bifurcation-based trust quantification
# ============================================================

class TrustScorer:
    """
    Computes trust metrics for a set of LLM solutions using:
    1. Majority vote fraction (self-consistency baseline)
    2. Embedding tightness (mean pairwise cosine similarity)
    3. QPPG bifurcation count (number of solution modes via mu-sweep)
    4. Combined trust score

    Parameters:
        embedding_model: sentence-transformers model name
        pca_dim: PCA dimension for QPPG (reduce from 384d)
        n_wells: Number of wells in QPPG substrate
        mu_range: (min, max) for mu sweep
        mu_steps: Number of mu values to test
    """

    def __init__(self, embedding_model="all-MiniLM-L6-v2",
                 pca_dim=32, n_wells=8,
                 mu_range=(-2.0, 3.0), mu_steps=20):
        self.pca_dim = pca_dim
        self.n_wells = n_wells
        self.mu_range = mu_range
        self.mu_steps = mu_steps

        # Lazy-loaded components
        self._embedder = None
        self._embedding_model_name = embedding_model
        self._pca = None

    def _get_embedder(self):
        """Lazy-load the text embedder."""
        if self._embedder is None:
            from .encoder import TextEmbedder
            self._embedder = TextEmbedder(model_name=self._embedding_model_name)
        return self._embedder

    def _embed_solutions(self, solutions):
        """Embed solution texts. Returns (N, embed_dim) normalized embeddings."""
        embedder = self._get_embedder()
        return embedder.embed_and_normalize(solutions)

    def _compute_tightness(self, embeddings):
        """
        Compute embedding cluster tightness as mean pairwise cosine similarity.
        Range: -1 to 1 (higher = tighter = more consistent).
        For L2-normalized embeddings, cosine sim = dot product.
        """
        n = len(embeddings)
        if n < 2:
            return 1.0

        sim_matrix = embeddings @ embeddings.T
        # Extract upper triangle (excluding diagonal)
        mask = np.triu(np.ones((n, n), dtype=bool), k=1)
        pairwise_sims = sim_matrix[mask]
        return float(np.mean(pairwise_sims))

    def _compute_spread(self, embeddings):
        """Compute embedding spread as std of pairwise distances."""
        n = len(embeddings)
        if n < 2:
            return 0.0

        dists = []
        for i in range(n):
            for j in range(i + 1, n):
                dists.append(np.linalg.norm(embeddings[i] - embeddings[j]))
        return float(np.std(dists))

    def _pca_reduce(self, embeddings):
        """Reduce dimensionality via PCA for QPPG."""
        from sklearn.decomposition import PCA

        n, d = embeddings.shape
        actual_dim = min(self.pca_dim, n - 1, d)
        if actual_dim < 2:
            return embeddings[:, :2] if d >= 2 else embeddings

        pca = PCA(n_components=actual_dim)
        return pca.fit_transform(embeddings)

    def _qppg_bifurcation(self, embeddings):
        """
        Detect number of solution modes via QPPG mu-sweep.

        Uses QPPGSubstrate directly (not QPPGOnlineClusterer) because
        we have small N (~10 points) and want fine-grained control.

        Returns:
            n_clusters: int — number of distinct solution modes
            best_mu: float — optimal mu value
            best_sil: float — silhouette score at optimal mu
        """
        import torch
        from sklearn.metrics import silhouette_score
        from .core import QPPGSubstrate

        pca_emb = self._pca_reduce(embeddings)
        n, d = pca_emb.shape

        if n < 3:
            return 1, 0.0, -1.0

        # Create substrate with wells at data point locations
        n_wells_actual = min(self.n_wells, n)
        substrate = QPPGSubstrate(
            d=d, n_wells=n_wells_actual, seed=42,
            confinement=0.001, sigma_base=1.5
        )

        # Place wells at data points (or k-means centers if n > n_wells)
        with torch.no_grad():
            pts = torch.tensor(pca_emb, dtype=torch.float32)
            if n <= n_wells_actual:
                substrate.well_centers.data[:n] = pts
                substrate.well_centers.data[n:] = 100.0  # push unused wells far away
                substrate.base_depths.data[:n] = 2.5
                substrate.base_depths.data[n:] = 0.0
            else:
                from sklearn.cluster import KMeans
                km = KMeans(n_clusters=n_wells_actual, n_init=5, random_state=42)
                km.fit(pca_emb)
                substrate.well_centers.data[:n_wells_actual] = torch.tensor(
                    km.cluster_centers_, dtype=torch.float32)
                counts = np.bincount(km.labels_, minlength=n_wells_actual)
                depths = counts / max(counts.max(), 1) * 3.0 + 1.0
                substrate.base_depths.data[:n_wells_actual] = torch.tensor(
                    depths, dtype=torch.float32)

            # Set mu sensitivity: spread wells across mu range
            sens = torch.linspace(-1.5, 1.5, n_wells_actual)
            substrate.mu_sensitivity.data[:n_wells_actual] = sens

        # Mu sweep
        best_mu = self.mu_range[0]
        best_sil = -1.0
        best_n = 1

        for mu_val in np.linspace(self.mu_range[0], self.mu_range[1], self.mu_steps):
            substrate.mu = mu_val
            labels, n_clusters, centers = substrate.find_attractors(pts)

            if n_clusters >= 2:
                try:
                    sil = silhouette_score(pca_emb, labels)
                except ValueError:
                    sil = -1.0

                if sil > best_sil:
                    best_sil = sil
                    best_mu = mu_val
                    best_n = n_clusters

        return best_n, best_mu, best_sil

    def compute_trust(self, solutions, answers=None):
        """
        Compute trust metrics for a set of solutions to the same problem.

        Args:
            solutions: List[str] — N solution texts
            answers: List[str] or None — extracted final answers.
                     If None, extracts from solution texts automatically.

        Returns:
            dict with trust metrics:
                majority_answer: str — most common answer
                majority_fraction: float — fraction agreeing with majority (0-1)
                embedding_tightness: float — mean pairwise cosine sim (0-1)
                embedding_spread: float — std of pairwise distances
                n_clusters_qppg: int — solution modes from QPPG
                bifurcation_score: float — 1/n_clusters (1.0=single mode)
                qppg_trust: float — combined trust score (0-1)
                answer_entropy: float — entropy of answer distribution
                embeddings: ndarray — (N, embed_dim) raw embeddings
        """
        n = len(solutions)

        # Extract answers if not provided
        if answers is None:
            answers = [extract_numerical_answer(s) for s in solutions]

        # --- Majority vote ---
        valid_answers = [a for a in answers if a is not None]
        if valid_answers:
            counter = Counter(valid_answers)
            majority_answer = counter.most_common(1)[0][0]
            majority_fraction = counter[majority_answer] / n
        else:
            majority_answer = None
            majority_fraction = 0.0

        # Answer entropy
        if valid_answers:
            probs = np.array([c / len(valid_answers) for c in counter.values()])
            answer_entropy = float(-np.sum(probs * np.log2(probs + 1e-10)))
        else:
            answer_entropy = 0.0

        # --- Embedding metrics ---
        embeddings = self._embed_solutions(solutions)
        tightness = self._compute_tightness(embeddings)
        spread = self._compute_spread(embeddings)

        # --- QPPG bifurcation ---
        n_clusters, best_mu, best_sil = self._qppg_bifurcation(embeddings)
        bifurcation_score = 1.0 / n_clusters

        # --- Combined trust score ---
        # Normalize tightness from [-1, 1] to [0, 1]
        tightness_norm = (tightness + 1.0) / 2.0
        # Combined: weighted sum
        qppg_trust = (0.4 * tightness_norm +
                      0.3 * bifurcation_score +
                      0.3 * majority_fraction)

        return {
            'majority_answer': majority_answer,
            'majority_fraction': majority_fraction,
            'embedding_tightness': tightness,
            'embedding_spread': spread,
            'n_clusters_qppg': n_clusters,
            'bifurcation_score': bifurcation_score,
            'qppg_trust': qppg_trust,
            'answer_entropy': answer_entropy,
            'best_mu': best_mu,
            'best_silhouette': best_sil,
            'embeddings': embeddings,
        }

    def compute_trust_batch(self, problems_solutions):
        """
        Compute trust for a batch of problems.

        Args:
            problems_solutions: List[List[str]] — for each problem, a list of N solutions

        Returns:
            List[dict] — trust metrics for each problem
        """
        results = []
        for i, solutions in enumerate(problems_solutions):
            trust = self.compute_trust(solutions)
            results.append(trust)
            if (i + 1) % 50 == 0:
                print(f"  Scored {i+1}/{len(problems_solutions)} problems")
        return results
