"""
QPPG+LLM Novelty Engine: Online novelty detection + transform library builder.

Integrates the LLM encoder (semantic descriptions → embeddings) with the
QPPG substrate (bifurcation-based clustering) and a growing transform library.

The key insight: QPPG Layer 1 excels at clustering when given good
representations (ARI=0.997 on Exp5). LLMs provide semantic understanding
of transformations that learned encoders cannot generalize (ceiling at ~0.53).
Combining them should break the generalization barrier.

Usage:
    from qppg import NoveltyEngine

    engine = NoveltyEngine(domain="vector")
    result = engine.process_batch(inputs, outputs)
    print(f"Categories: {result['n_categories']}")
    print(f"Novel: {result['novel_indices']}")
    engine.get_library().save("library.json")
"""

import json
import numpy as np
from dataclasses import dataclass, field, asdict
from typing import List, Optional, Tuple, Dict
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
from sklearn.metrics import adjusted_rand_score, silhouette_score

from .core import QPPGOnlineClusterer, compute_jacobian, jacobian_signature
from .encoder import LLMEncoder


# ============================================================
# Transform Category & Library
# ============================================================

@dataclass
class TransformCategory:
    """A discovered transformation category in the library."""
    category_id: int
    description: str                                    # Representative LLM description
    representative_descriptions: List[str] = field(default_factory=list)
    attractor_center: Optional[List[float]] = None      # Center in embedding space
    jacobian_sig: Optional[Dict] = None                 # Topological signature
    n_examples: int = 0
    discovery_phase: int = 0                            # When first detected
    example_pairs: List[Dict] = field(default_factory=list)  # Up to 10 representative pairs


class TransformLibrary:
    """
    Growing library of discovered transformation categories.

    Each category corresponds to a QPPG attractor basin. When a bifurcation
    occurs (new attractor appears), a new category is created.
    """

    def __init__(self):
        self.categories: Dict[int, TransformCategory] = {}

    def add_category(self, category: TransformCategory):
        self.categories[category.category_id] = category

    def update_category(self, category_id, new_descriptions=None, new_pairs=None):
        if category_id not in self.categories:
            return
        cat = self.categories[category_id]
        if new_descriptions:
            cat.representative_descriptions.extend(new_descriptions)
            # Keep only last 20 descriptions
            cat.representative_descriptions = cat.representative_descriptions[-20:]
        if new_pairs:
            cat.example_pairs.extend(new_pairs)
            cat.example_pairs = cat.example_pairs[-10:]
        cat.n_examples += len(new_descriptions or [])

    def get_category(self, category_id) -> Optional[TransformCategory]:
        return self.categories.get(category_id)

    def get_all_categories(self) -> List[TransformCategory]:
        return sorted(self.categories.values(), key=lambda c: c.category_id)

    @property
    def n_categories(self):
        return len(self.categories)

    def summary(self) -> str:
        lines = [f"Transform Library: {self.n_categories} categories"]
        for cat in self.get_all_categories():
            lines.append(f"  [{cat.category_id}] {cat.description[:80]}... "
                        f"({cat.n_examples} examples, phase {cat.discovery_phase})")
        return '\n'.join(lines)

    def save(self, path):
        """Save library to JSON."""
        data = {
            'n_categories': self.n_categories,
            'categories': []
        }
        for cat in self.get_all_categories():
            cat_dict = {
                'category_id': cat.category_id,
                'description': cat.description,
                'representative_descriptions': cat.representative_descriptions[:5],
                'attractor_center': cat.attractor_center,
                'jacobian_sig': cat.jacobian_sig,
                'n_examples': cat.n_examples,
                'discovery_phase': cat.discovery_phase,
                'n_example_pairs': len(cat.example_pairs),
            }
            data['categories'].append(cat_dict)

        with open(path, 'w') as f:
            json.dump(data, f, indent=2)

    def load(self, path):
        """Load library from JSON."""
        with open(path, 'r') as f:
            data = json.load(f)
        for cat_dict in data['categories']:
            cat = TransformCategory(
                category_id=cat_dict['category_id'],
                description=cat_dict['description'],
                representative_descriptions=cat_dict.get('representative_descriptions', []),
                attractor_center=cat_dict.get('attractor_center'),
                jacobian_sig=cat_dict.get('jacobian_sig'),
                n_examples=cat_dict.get('n_examples', 0),
                discovery_phase=cat_dict.get('discovery_phase', 0),
            )
            self.categories[cat.category_id] = cat


# ============================================================
# Novelty Engine
# ============================================================

class NoveltyEngine:
    """
    QPPG+LLM Novelty Detection Engine.

    Online system that:
    1. Receives (input, output) pairs from any domain
    2. Uses Claude Haiku to describe the transformation
    3. Embeds the description via sentence-transformers
    4. Feeds embeddings into QPPG substrate for clustering
    5. Detects bifurcation events = genuinely new transform types
    6. Maintains a growing TransformLibrary

    Parameters:
        api_key: Anthropic API key
        domain: "vector" (R^8) or "grid" (ARC grids)
        n_wells: QPPG wells (default 12)
        embedding_model: sentence-transformers model
        cache_path: Description cache path
        use_pca: Whether to reduce embedding dimension before QPPG
        pca_dim: Target dimension if use_pca=True
        use_qppg: If False, uses K-Means instead of QPPG (ablation)
        llm_model: Claude model to use
    """

    def __init__(self, api_key=None, domain="vector", n_wells=12,
                 embedding_model="all-MiniLM-L6-v2", cache_path=None,
                 use_pca=False, pca_dim=32, use_qppg=True,
                 llm_model="claude-haiku-4-5-20251001"):
        self.domain = domain
        self.use_pca = use_pca
        self.pca_dim = pca_dim
        self.use_qppg = use_qppg

        # LLM encoder
        self.encoder = LLMEncoder(
            api_key=api_key, domain=domain,
            embedding_model=embedding_model,
            cache_path=cache_path, llm_model=llm_model
        )

        # PCA reducer (fit on first batch)
        self.pca = None

        # QPPG clusterer — dimension depends on PCA
        effective_dim = pca_dim if use_pca else 384  # MiniLM default
        if use_qppg:
            self.clusterer = QPPGOnlineClusterer(
                d=effective_dim, n_wells=n_wells, seed=42,
                confinement=0.0001, sigma_base=1.0
            )
        else:
            self.clusterer = None

        # Transform library
        self.library = TransformLibrary()

        # State tracking
        self.phase_counter = 0
        self.all_embeddings = []
        self.all_descriptions = []
        self.prev_n_categories = 0

    def process_batch(self, inputs, outputs, phase_idx=None):
        """
        Process a batch of (input, output) pairs.

        Args:
            inputs: list of numpy arrays or (N, d_in) array
            outputs: list of numpy arrays or (N, d_out) array
            phase_idx: Optional phase identifier

        Returns:
            dict with:
                'labels': (N,) category assignments
                'novel_indices': list of indices that triggered novelty
                'n_categories': current library size
                'descriptions': list of text descriptions
                'embeddings': (N, d_embed) raw embeddings
                'cost': API cost summary
        """
        if phase_idx is None:
            phase_idx = self.phase_counter
            self.phase_counter += 1

        # Step 1: LLM encode
        embeddings, descriptions = self.encoder.encode_with_descriptions(inputs, outputs)
        self.all_embeddings.append(embeddings)
        self.all_descriptions.extend(descriptions)

        # Step 2: Optional PCA
        all_emb = np.concatenate(self.all_embeddings)
        if self.use_pca:
            if self.pca is None and len(all_emb) >= self.pca_dim:
                self.pca = PCA(n_components=self.pca_dim)
                self.pca.fit(all_emb)
            if self.pca is not None:
                all_emb_reduced = self.pca.transform(all_emb)
            else:
                all_emb_reduced = all_emb
        else:
            all_emb_reduced = all_emb

        # Step 3: Clustering
        if self.use_qppg:
            # Reset clusterer state for full re-clustering
            self.clusterer.all_data = []
            self.clusterer.n_clusters_detected = 0
            labels = self.clusterer.process_batch(all_emb_reduced, phase_idx=phase_idx)
            n_categories = self.clusterer.get_n_clusters()
        else:
            # K-Means fallback (ablation)
            max_k = min(20, len(all_emb_reduced) // 5)
            best_k, best_sil = 2, -1.0
            for k in range(2, max_k + 1):
                km = KMeans(n_clusters=k, n_init=5, random_state=42)
                km_labels = km.fit_predict(all_emb_reduced)
                try:
                    sil = silhouette_score(all_emb_reduced, km_labels)
                    if sil > best_sil:
                        best_sil = sil
                        best_k = k
                except ValueError:
                    pass
            km = KMeans(n_clusters=best_k, n_init=10, random_state=42)
            labels = km.fit_predict(all_emb_reduced)
            n_categories = best_k

        # Get labels for just the current batch
        batch_start = len(all_emb) - len(embeddings)
        batch_labels = labels[batch_start:]

        # Step 4: Detect novelty
        novel_indices = []
        if n_categories > self.prev_n_categories:
            # New categories were created — find which batch items are in new categories
            old_categories = set(range(self.prev_n_categories))
            for i, label in enumerate(batch_labels):
                if label not in old_categories:
                    novel_indices.append(i)

        # Step 5: Update library
        self._update_library(labels, all_emb_reduced, descriptions[-len(embeddings):],
                            inputs, outputs, phase_idx, batch_start)

        self.prev_n_categories = n_categories

        return {
            'labels': batch_labels,
            'novel_indices': novel_indices,
            'n_categories': n_categories,
            'descriptions': descriptions,
            'embeddings': embeddings,
            'cost': self.encoder.get_cost_summary(),
        }

    def _update_library(self, labels, embeddings, batch_descriptions,
                       inputs, outputs, phase_idx, batch_start):
        """Update the transform library with new batch results."""
        if isinstance(inputs, np.ndarray) and inputs.ndim == 2:
            inputs = [inputs[i] for i in range(len(inputs))]
            outputs = [outputs[i] for i in range(len(outputs))]

        batch_labels = labels[batch_start:]
        unique_labels = set(batch_labels)

        for cat_id in unique_labels:
            mask = batch_labels == cat_id
            cat_descriptions = [d for d, m in zip(batch_descriptions, mask) if m]

            if cat_id not in self.library.categories:
                # New category
                cat_embeddings = embeddings[labels == cat_id]
                center = cat_embeddings.mean(axis=0).tolist()

                # Pick most representative description (closest to centroid)
                if len(cat_descriptions) > 0:
                    representative = cat_descriptions[0]
                else:
                    representative = "Unknown transform"

                cat = TransformCategory(
                    category_id=int(cat_id),
                    description=representative,
                    representative_descriptions=cat_descriptions[:5],
                    attractor_center=center,
                    n_examples=int(mask.sum()),
                    discovery_phase=phase_idx,
                )
                self.library.add_category(cat)
            else:
                # Update existing category
                self.library.update_category(
                    int(cat_id),
                    new_descriptions=cat_descriptions,
                )

    def cluster_static(self, embeddings, method="kmeans", n_clusters=None):
        """
        Static clustering (non-streaming) for evaluation.
        Uses K-Means or QPPG on pre-computed embeddings.

        Args:
            embeddings: (N, d) numpy array
            method: "kmeans" or "qppg"
            n_clusters: For kmeans, fixed k. If None, auto-select via silhouette.

        Returns:
            labels: (N,) integer array
        """
        if method == "kmeans":
            if n_clusters is not None:
                km = KMeans(n_clusters=n_clusters, n_init=10, random_state=42)
                return km.fit_predict(embeddings)
            else:
                # Auto-select k via silhouette
                max_k = min(30, len(embeddings) // 5)
                best_k, best_sil = 2, -1.0
                for k in range(2, max_k + 1):
                    km = KMeans(n_clusters=k, n_init=5, random_state=42)
                    km_labels = km.fit_predict(embeddings)
                    try:
                        sil = silhouette_score(embeddings, km_labels)
                        if sil > best_sil:
                            best_sil = sil
                            best_k = k
                    except ValueError:
                        pass
                km = KMeans(n_clusters=best_k, n_init=10, random_state=42)
                return km.fit_predict(embeddings)

        elif method == "qppg":
            clusterer = QPPGOnlineClusterer(
                d=embeddings.shape[1], n_wells=12, seed=42
            )
            clusterer.all_data = []
            return clusterer.process_batch(embeddings, phase_idx=0)
        else:
            raise ValueError(f"Unknown method: {method}")

    def get_library(self) -> TransformLibrary:
        return self.library

    def get_cost_summary(self) -> dict:
        return self.encoder.get_cost_summary()

    def save_state(self, path):
        """Save engine state (library + cache)."""
        self.library.save(path)
        self.encoder.save_cache()

    def load_state(self, path):
        """Load engine state."""
        self.library.load(path)
