"""
Blindness Detection: QPPG Layer 3 — Metacognitive Interface.

Detects when inputs fall outside all attractor basins (the system is "blind")
using five complementary signals from the energy landscape:
  1. DBSCAN noise label (point doesn't cluster with any attractor)
  2. Residual energy after integration (high = far from wells)
  3. Distance to nearest attractor center
  4. Integration velocity at termination (high = didn't converge)
  5. Combined blindness probability

Also provides energy accounting for attractor maintenance vs exploration.

Three-way classification:
  FAMILIAR — converges to known attractor basin
  NOVEL    — creates new attractor via bifurcation
  BLIND    — doesn't converge anywhere

Usage:
    from qppg import BlindnessDetector, InputClassification

    detector = BlindnessDetector(substrate=my_substrate)
    result = detector.detect(points)
    print(f"Blind: {result['blind_fraction']:.1%}")
    for i in result['blind_indices']:
        print(f"  Point {i}: {result['signals'][i].blindness_probability:.3f}")
"""

import numpy as np
import torch
from enum import Enum
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Tuple

from sklearn.cluster import DBSCAN
from sklearn.neighbors import NearestNeighbors


# ============================================================
# Types
# ============================================================

class InputClassification(Enum):
    """Three-way classification for each input."""
    FAMILIAR = "familiar"
    NOVEL = "novel"
    BLIND = "blind"


@dataclass
class BlindnessSignals:
    """Raw diagnostic signals for a single input point."""
    basin_label: int                    # DBSCAN raw label (-1 = noise)
    residual_energy: float              # Energy after integration
    distance_to_nearest: float          # Distance to nearest attractor center
    final_velocity: float               # ||dx/dt|| after integration
    convergence_steps: int              # Steps to convergence (or max_steps)
    blindness_probability: float        # Combined score in [0, 1]
    classification: InputClassification = InputClassification.FAMILIAR


@dataclass
class EnergyBudgetState:
    """Snapshot of energy accounting state."""
    total_budget: float
    maintenance_cost: float
    exploration_cost: float
    remaining: float
    n_attractors: int
    can_explore: bool
    pressure: float                     # maintenance / total (0-1)
    consolidation_triggered: bool


# ============================================================
# BlindnessDetector
# ============================================================

class BlindnessDetector:
    """
    Runtime failure classifier for QPPG-based systems.

    Detects when inputs fall outside all attractor basins using
    five complementary signals from the energy landscape.

    Parameters:
        substrate: QPPGSubstrate instance (if None, created on first use)
        embedding_model: sentence-transformers model for text mode
        pca_dim: PCA dimension for high-d inputs
        n_integration_steps: Steps for gradient flow integration
        dt: Integration time step
        energy_threshold: Residual energy above which point is suspect (None=adaptive)
        distance_threshold: Distance above which point is suspect (None=adaptive)
        velocity_threshold: Velocity above which point hasn't converged
        blindness_threshold: Combined probability above which to classify as BLIND
    """

    def __init__(self, substrate=None, embedding_model="all-MiniLM-L6-v2",
                 pca_dim=32, n_integration_steps=300, dt=0.05,
                 energy_threshold=None, distance_threshold=None,
                 velocity_threshold=1e-3, blindness_threshold=0.6):
        self.substrate = substrate
        self.pca_dim = pca_dim
        self.n_integration_steps = n_integration_steps
        self.dt = dt
        self.energy_threshold = energy_threshold
        self.distance_threshold = distance_threshold
        self.velocity_threshold = velocity_threshold
        self.blindness_threshold = blindness_threshold

        self._embedder = None
        self._embedding_model_name = embedding_model
        self._pca = None

        # Cumulative diagnostics
        self.total_points = 0
        self.total_blind = 0
        self.total_familiar = 0
        self.total_novel = 0

    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 _pca_reduce(self, points):
        """PCA reduction for high-d inputs."""
        from sklearn.decomposition import PCA

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

        if self._pca is None:
            self._pca = PCA(n_components=actual_dim)
            self._pca.fit(points)
        return self._pca.transform(points)

    def _ensure_substrate(self, d):
        """Create substrate if not provided."""
        if self.substrate is None:
            from .core import QPPGSubstrate
            self.substrate = QPPGSubstrate(
                d=d, n_wells=12, seed=42,
                confinement=0.001, sigma_base=1.5
            )

    def _integrate_with_diagnostics(self, x0):
        """
        Integrate points to attractor basins with diagnostic tracking.

        Returns:
            x_final: (N, d) final positions
            energies: (N,) final energy values
            velocities: (N,) final velocity magnitudes
            convergence_steps: (N,) steps to convergence per point
        """
        n_steps = self.n_integration_steps
        dt = self.dt

        x = x0.detach().clone()
        converged = torch.zeros(len(x), dtype=torch.bool)
        convergence_steps = torch.full((len(x),), n_steps, dtype=torch.int)

        for step in range(n_steps):
            dx = self.substrate.flow(x)
            velocity = torch.norm(dx, dim=-1)

            newly_converged = (~converged) & (velocity < self.velocity_threshold)
            convergence_steps[newly_converged] = step
            converged = converged | newly_converged

            x = x + dt * dx
            x = torch.clamp(x, -30.0, 30.0)

        final_energies = self.substrate.energy(x).detach()
        final_dx = self.substrate.flow(x)
        final_velocities = torch.norm(final_dx, dim=-1).detach()

        return x.detach(), final_energies, final_velocities, convergence_steps

    def _find_attractors_raw(self, integrated_points, dbscan_eps=None):
        """
        Cluster integrated points via DBSCAN, preserving noise labels (-1).

        Unlike core.py find_attractors(), does NOT reassign noise points.
        Noise = potential blind inputs.

        Returns:
            raw_labels: (N,) with -1 for noise points
            n_attractors: int (excluding noise)
            centers: (n_attractors, d) attractor centers from non-noise points
        """
        x_np = integrated_points.detach().numpy() if torch.is_tensor(integrated_points) \
            else integrated_points

        if len(x_np) < 3:
            return np.zeros(len(x_np), dtype=int), 1, x_np.mean(axis=0, keepdims=True)

        if dbscan_eps is None:
            nn = NearestNeighbors(n_neighbors=min(5, len(x_np) - 1))
            nn.fit(x_np)
            distances, _ = nn.kneighbors(x_np)
            dbscan_eps = float(np.median(distances[:, -1]) * 2.0)
            dbscan_eps = max(dbscan_eps, 0.1)

        db = DBSCAN(eps=dbscan_eps, min_samples=3)
        raw_labels = db.fit_predict(x_np)

        # Compute centers from non-noise points only
        unique_clusters = sorted(set(raw_labels) - {-1})
        n_attractors = len(unique_clusters)

        if n_attractors == 0:
            # All points are noise — extreme blindness
            return raw_labels, 0, np.empty((0, x_np.shape[1]))

        centers = np.array([x_np[raw_labels == c].mean(axis=0) for c in unique_clusters])

        # Relabel non-noise to 0, 1, 2, ... (keep -1 as noise)
        label_map = {old: new for new, old in enumerate(unique_clusters)}
        label_map[-1] = -1
        raw_labels = np.array([label_map[l] for l in raw_labels])

        return raw_labels, n_attractors, centers

    def _compute_adaptive_thresholds(self, energies, distances):
        """
        Compute adaptive thresholds from in-distribution data.

        Uses median + 2 * IQR for both energy and distance.
        Only calibrates from non-extreme values.
        """
        if self.energy_threshold is None and len(energies) > 0:
            q25, q50, q75 = np.percentile(energies, [25, 50, 75])
            iqr = q75 - q25
            self.energy_threshold = q50 + 2.0 * max(iqr, 0.01)

        if self.distance_threshold is None and len(distances) > 0:
            q25, q50, q75 = np.percentile(distances, [25, 50, 75])
            iqr = q75 - q25
            self.distance_threshold = q50 + 2.0 * max(iqr, 0.01)

    @staticmethod
    def _sigmoid(x, center, scale=1.0):
        """Sigmoid activation centered at threshold."""
        z = scale * (x - center)
        z = np.clip(z, -20.0, 20.0)
        return 1.0 / (1.0 + np.exp(-z))

    def _compute_blindness_probability(self, is_noise, energy, distance, velocity):
        """
        Combine 4 signals into single blindness probability [0, 1].

        Weights: noise=0.35, energy=0.25, distance=0.25, velocity=0.15
        """
        s1 = 1.0 if is_noise else 0.0
        s2 = float(self._sigmoid(energy, self.energy_threshold, scale=2.0))
        s3 = float(self._sigmoid(distance, self.distance_threshold, scale=1.0))
        s4 = float(self._sigmoid(velocity, self.velocity_threshold, scale=100.0))

        weights = [0.35, 0.25, 0.25, 0.15]
        return sum(w * s for w, s in zip(weights, [s1, s2, s3, s4]))

    def detect(self, points, known_centers=None):
        """
        Detect blindness in a batch of points.

        Args:
            points: (N, d) numpy array or torch tensor
            known_centers: (K, d) optional pre-computed attractor centers.
                          If None, discovered from the data.

        Returns:
            dict with:
                'signals': List[BlindnessSignals]
                'classifications': List[InputClassification]
                'blind_indices': List[int]
                'familiar_indices': List[int]
                'blind_fraction': float
                'mean_blindness': float
                'attractor_centers': ndarray
                'n_attractors': int
        """
        if isinstance(points, np.ndarray):
            pts_np = points
            pts_t = torch.tensor(points, dtype=torch.float32)
        else:
            pts_t = points.float()
            pts_np = points.detach().numpy()

        n, d = pts_np.shape
        self._ensure_substrate(d)

        # Place wells at data or known centers
        with torch.no_grad():
            if known_centers is not None:
                kc = torch.tensor(known_centers, dtype=torch.float32)
                n_wells = min(len(kc), self.substrate.n_wells)
                self.substrate.well_centers.data[:n_wells] = kc[:n_wells]
                self.substrate.well_centers.data[n_wells:] = 100.0
                self.substrate.base_depths.data[:n_wells] = 2.5
                self.substrate.base_depths.data[n_wells:] = 0.0
            else:
                # Place wells at k-means centers from data
                from sklearn.cluster import KMeans
                n_wells = min(self.substrate.n_wells, max(2, n // 5))
                km = KMeans(n_clusters=n_wells, n_init=5, random_state=42)
                km.fit(pts_np)
                self.substrate.well_centers.data[:n_wells] = torch.tensor(
                    km.cluster_centers_, dtype=torch.float32)
                self.substrate.well_centers.data[n_wells:] = 100.0
                counts = np.bincount(km.labels_, minlength=n_wells)
                depths = counts / max(counts.max(), 1) * 3.0 + 1.0
                self.substrate.base_depths.data[:n_wells] = torch.tensor(
                    depths, dtype=torch.float32)
                self.substrate.base_depths.data[n_wells:] = 0.0

            # Spread mu sensitivity
            sens = torch.linspace(-1.5, 1.5, self.substrate.n_wells)
            self.substrate.mu_sensitivity.data = sens

        # Set mu to a moderate value for detection
        self.substrate.mu = 0.5

        # Integrate with diagnostics
        x_final, energies, velocities, conv_steps = self._integrate_with_diagnostics(pts_t)

        # Cluster integrated points (raw — preserve noise labels)
        raw_labels, n_attractors, centers = self._find_attractors_raw(x_final)

        # Compute distances to nearest center
        energies_np = energies.numpy()
        velocities_np = velocities.numpy()

        if n_attractors > 0 and len(centers) > 0:
            distances = np.array([
                np.min(np.linalg.norm(centers - pts_np[i], axis=1))
                for i in range(n)
            ])
        else:
            distances = np.full(n, 999.0)

        # Calibrate adaptive thresholds from clustered points
        clustered_mask = raw_labels >= 0
        if clustered_mask.any():
            self._compute_adaptive_thresholds(
                energies_np[clustered_mask], distances[clustered_mask]
            )
        else:
            # Fallback thresholds
            if self.energy_threshold is None:
                self.energy_threshold = float(np.median(energies_np))
            if self.distance_threshold is None:
                self.distance_threshold = float(np.median(distances))

        # Compute per-point signals
        signals = []
        classifications = []
        blind_indices = []
        familiar_indices = []

        for i in range(n):
            is_noise = (raw_labels[i] == -1)
            bp = self._compute_blindness_probability(
                is_noise, energies_np[i], distances[i], velocities_np[i]
            )

            if bp >= self.blindness_threshold:
                cls = InputClassification.BLIND
                blind_indices.append(i)
            else:
                cls = InputClassification.FAMILIAR
                familiar_indices.append(i)

            sig = BlindnessSignals(
                basin_label=int(raw_labels[i]),
                residual_energy=float(energies_np[i]),
                distance_to_nearest=float(distances[i]),
                final_velocity=float(velocities_np[i]),
                convergence_steps=int(conv_steps[i]),
                blindness_probability=bp,
                classification=cls,
            )
            signals.append(sig)
            classifications.append(cls)

        # Update cumulative stats
        self.total_points += n
        self.total_blind += len(blind_indices)
        self.total_familiar += len(familiar_indices)

        blind_probs = [s.blindness_probability for s in signals]

        return {
            'signals': signals,
            'classifications': classifications,
            'blind_indices': blind_indices,
            'familiar_indices': familiar_indices,
            'blind_fraction': len(blind_indices) / n if n > 0 else 0.0,
            'mean_blindness': float(np.mean(blind_probs)) if blind_probs else 0.0,
            'attractor_centers': centers if n_attractors > 0 else np.empty((0, d)),
            'n_attractors': n_attractors,
        }

    def detect_text(self, texts, known_centers=None):
        """
        Detect blindness on text inputs.

        Embeds texts via sentence-transformers, then runs detect().

        Args:
            texts: str or List[str]
            known_centers: optional pre-computed centers in embedding space

        Returns:
            Same dict as detect(), plus 'embeddings': (N, d) array
        """
        if isinstance(texts, str):
            texts = [texts]

        embedder = self._get_embedder()
        embeddings = embedder.embed_and_normalize(texts)

        # PCA reduce if high-dimensional
        reduced = self._pca_reduce(embeddings)

        result = self.detect(reduced, known_centers=known_centers)
        result['embeddings'] = embeddings
        return result

    def classify_with_novelty(self, points, known_centers, prev_n_categories=0):
        """
        Three-way classification: FAMILIAR / NOVEL / BLIND.

        Two-pass approach:
          1. detect() to find blind points
          2. For non-blind points, check if they converged to a new attractor

        Args:
            points: (N, d) numpy array
            known_centers: (K, d) current attractor centers
            prev_n_categories: previous category count

        Returns:
            dict with:
                'classifications': List[InputClassification]
                'signals': List[BlindnessSignals]
                'n_familiar': int
                'n_novel': int
                'n_blind': int
        """
        result = self.detect(points, known_centers=known_centers)

        n_novel = 0
        n_familiar = 0
        n_blind = len(result['blind_indices'])

        # Check if new attractors appeared (novelty)
        new_n = result['n_attractors']
        has_novelty = new_n > prev_n_categories

        classifications = list(result['classifications'])
        signals = list(result['signals'])

        if has_novelty:
            # Points that converged to new attractor basins are NOVEL, not FAMILIAR
            for i in result['familiar_indices']:
                label = signals[i].basin_label
                if label >= prev_n_categories:
                    classifications[i] = InputClassification.NOVEL
                    signals[i] = BlindnessSignals(
                        basin_label=signals[i].basin_label,
                        residual_energy=signals[i].residual_energy,
                        distance_to_nearest=signals[i].distance_to_nearest,
                        final_velocity=signals[i].final_velocity,
                        convergence_steps=signals[i].convergence_steps,
                        blindness_probability=signals[i].blindness_probability,
                        classification=InputClassification.NOVEL,
                    )
                    n_novel += 1
                else:
                    n_familiar += 1
        else:
            n_familiar = len(result['familiar_indices'])

        self.total_novel += n_novel

        return {
            'classifications': classifications,
            'signals': signals,
            'n_familiar': n_familiar,
            'n_novel': n_novel,
            'n_blind': n_blind,
            'n_attractors': result['n_attractors'],
            'attractor_centers': result['attractor_centers'],
        }

    def get_diagnostics(self):
        """Return cumulative diagnostic summary."""
        total = self.total_points
        return {
            'total_points': total,
            'total_blind': self.total_blind,
            'total_familiar': self.total_familiar,
            'total_novel': self.total_novel,
            'blind_rate': self.total_blind / total if total > 0 else 0.0,
        }


# ============================================================
# EnergyAccountant
# ============================================================

class EnergyAccountant:
    """
    Tracks energy budget for attractor maintenance vs exploration.

    Implements Theorem 3.4 (Exploration-Exploitation Tradeoff):
    finite energy constrains library size. Under budget pressure,
    triggers consolidation (merge nearby attractors, shed weakest).

    Parameters:
        total_budget: Total energy budget E_max
        maintenance_cost_per_attractor: Energy per attractor per cycle
        exploration_cost: Energy cost to create a new attractor
        consolidation_pressure_threshold: pressure above this triggers consolidation
        min_attractor_depth: Minimum depth to retain an attractor
    """

    def __init__(self, total_budget=100.0,
                 maintenance_cost_per_attractor=2.0,
                 exploration_cost=10.0,
                 consolidation_pressure_threshold=0.8,
                 min_attractor_depth=0.5):
        self.total_budget = total_budget
        self.maintenance_cost_per_attractor = maintenance_cost_per_attractor
        self.exploration_cost = exploration_cost
        self.consolidation_pressure_threshold = consolidation_pressure_threshold
        self.min_attractor_depth = min_attractor_depth

        self._centers = None
        self._depths = None
        self._n_attractors = 0
        self._cumulative_exploration = 0.0
        self._cumulative_maintenance = 0.0
        self._history = []

    def register_attractors(self, centers, depths=None):
        """
        Register current attractor set.

        Args:
            centers: (K, d) attractor center positions
            depths: (K,) optional depth values (default: uniform 1.0)
        """
        self._centers = np.array(centers)
        self._n_attractors = len(centers)
        if depths is not None:
            self._depths = np.array(depths)
        else:
            self._depths = np.ones(self._n_attractors)

    def can_explore(self):
        """Whether budget allows creating a new attractor."""
        remaining = self.total_budget - self._cumulative_maintenance - self._cumulative_exploration
        return remaining >= self.exploration_cost

    def spend_exploration(self):
        """
        Deduct exploration cost from budget.
        Returns True if successful, False if insufficient budget.
        """
        if not self.can_explore():
            return False
        self._cumulative_exploration += self.exploration_cost
        return True

    def compute_maintenance_cost(self):
        """Compute total maintenance cost for current attractors."""
        if self._depths is None:
            return self._n_attractors * self.maintenance_cost_per_attractor
        return float(np.sum(self._depths) * self.maintenance_cost_per_attractor)

    def tick(self):
        """
        Advance one energy accounting cycle.

        Deducts maintenance costs, checks budget pressure,
        returns snapshot.

        Returns:
            EnergyBudgetState
        """
        maint = self.compute_maintenance_cost()
        self._cumulative_maintenance += maint
        remaining = self.total_budget - self._cumulative_maintenance - self._cumulative_exploration
        remaining = max(remaining, 0.0)
        pressure = (self._cumulative_maintenance + self._cumulative_exploration) / self.total_budget
        pressure = min(pressure, 1.0)
        consolidation = pressure >= self.consolidation_pressure_threshold

        state = EnergyBudgetState(
            total_budget=self.total_budget,
            maintenance_cost=self._cumulative_maintenance,
            exploration_cost=self._cumulative_exploration,
            remaining=remaining,
            n_attractors=self._n_attractors,
            can_explore=remaining >= self.exploration_cost,
            pressure=pressure,
            consolidation_triggered=consolidation,
        )
        self._history.append(state)
        return state

    def consolidate(self, centers, depths=None, merge_threshold=None):
        """
        Reduce attractor count to fit within budget.

        Strategy:
          1. Merge attractors closer than merge_threshold
          2. Shed attractors weaker than min_attractor_depth

        Args:
            centers: (K, d) current attractor centers
            depths: (K,) attractor depths (if None, uniform)
            merge_threshold: distance below which to merge (None=adaptive)

        Returns:
            new_centers: (K', d) consolidated centers
            merged_map: dict mapping old index → new index
            shed_indices: list of removed attractor indices
        """
        centers = np.array(centers)
        n = len(centers)
        if depths is None:
            depths = np.ones(n)
        else:
            depths = np.array(depths)

        if merge_threshold is None:
            if n >= 2:
                all_dists = []
                for i in range(n):
                    for j in range(i + 1, n):
                        all_dists.append(np.linalg.norm(centers[i] - centers[j]))
                merge_threshold = np.median(all_dists) * 0.3
            else:
                merge_threshold = 1.0

        # Step 1: Merge nearby attractors (union-find)
        parent = list(range(n))

        def find(x):
            while parent[x] != x:
                parent[x] = parent[parent[x]]
                x = parent[x]
            return x

        for i in range(n):
            for j in range(i + 1, n):
                if np.linalg.norm(centers[i] - centers[j]) < merge_threshold:
                    ri, rj = find(i), find(j)
                    if ri != rj:
                        # Keep the deeper one as root
                        if depths[ri] >= depths[rj]:
                            parent[rj] = ri
                        else:
                            parent[ri] = rj

        # Build merged groups
        groups = {}
        for i in range(n):
            root = find(i)
            if root not in groups:
                groups[root] = []
            groups[root].append(i)

        # Step 2: Shed weak attractors
        shed_indices = []
        new_centers = []
        new_depths = []
        merged_map = {}

        new_idx = 0
        for root, members in sorted(groups.items()):
            group_depth = max(depths[m] for m in members)
            if group_depth < self.min_attractor_depth:
                for m in members:
                    shed_indices.append(m)
                    merged_map[m] = -1
                continue

            # Weighted centroid
            total_d = sum(depths[m] for m in members)
            centroid = sum(depths[m] * centers[m] for m in members) / total_d
            new_centers.append(centroid)
            new_depths.append(group_depth)
            for m in members:
                merged_map[m] = new_idx
            new_idx += 1

        if len(new_centers) > 0:
            new_centers = np.array(new_centers)
        else:
            new_centers = np.empty((0, centers.shape[1]))

        # Update internal state
        self._centers = new_centers
        self._depths = np.array(new_depths) if new_depths else np.array([])
        self._n_attractors = len(new_centers)

        return new_centers, merged_map, shed_indices

    def get_state(self):
        """Return current budget state without advancing a cycle."""
        remaining = self.total_budget - self._cumulative_maintenance - self._cumulative_exploration
        remaining = max(remaining, 0.0)
        pressure = (self._cumulative_maintenance + self._cumulative_exploration) / self.total_budget
        pressure = min(pressure, 1.0)

        return EnergyBudgetState(
            total_budget=self.total_budget,
            maintenance_cost=self._cumulative_maintenance,
            exploration_cost=self._cumulative_exploration,
            remaining=remaining,
            n_attractors=self._n_attractors,
            can_explore=remaining >= self.exploration_cost,
            pressure=pressure,
            consolidation_triggered=pressure >= self.consolidation_pressure_threshold,
        )

    def get_history(self):
        """Return list of EnergyBudgetState snapshots from all ticks."""
        return list(self._history)
