"""
QPPG Core: Reusable substrate for bifurcation-based novelty detection.

Extracted and generalized from exp5_task_demo.py and exp5b_hard_task_demo.py.
Works in arbitrary dimensions. The key mechanism is mu-controlled bifurcation:
wells in the energy landscape activate/deactivate as mu changes, creating or
destroying attractor basins. This gives principled novelty detection — a new
attractor = a genuinely new representational primitive.

Usage:
    from qppg import QPPGSubstrate, QPPGOnlineClusterer

    clusterer = QPPGOnlineClusterer(d=384, n_wells=12)
    labels = clusterer.process_batch(embeddings)
    print(f"Detected {clusterer.get_n_clusters()} categories")
    print(f"Novelty events at phases: {clusterer.get_novelty_events()}")
"""

import numpy as np
import torch
import torch.nn as nn
from sklearn.cluster import KMeans, DBSCAN
from sklearn.metrics import silhouette_score
from sklearn.neighbors import NearestNeighbors


class QPPGSubstrate(nn.Module):
    """
    D-dimensional multi-well energy landscape with mu-controlled bifurcation.

    E(x, mu) = confinement * ||x||^2 + sum_k depth_k(mu) * gaussian_well_k(x)
    dx/dt = -grad_x E(x, mu)

    As mu changes, wells activate/deactivate = bifurcation = new clusters.

    Parameters:
        d: State space dimension (e.g., 384 for sentence-transformer embeddings)
        n_wells: Maximum number of wells in the energy landscape
        seed: Random seed for reproducibility
        confinement: Confinement strength (prevents points from escaping)
        sigma_base: Base well width (lower = tighter wells = more distinct basins)
    """

    def __init__(self, d, n_wells=10, seed=42, confinement=0.0001, sigma_base=1.0):
        super().__init__()
        torch.manual_seed(seed)
        self.d = d
        self.n_wells = n_wells
        self.mu = 0.0

        self.well_centers = nn.Parameter(torch.randn(n_wells, d) * 3.0)
        self.base_depths = nn.Parameter(torch.rand(n_wells) * 2.0 + 1.5)
        self.mu_sensitivity = nn.Parameter(torch.randn(n_wells) * 1.5)
        self.well_sigmas = nn.Parameter(torch.ones(n_wells) * sigma_base)
        self.confinement = confinement

    def energy(self, x):
        """Multi-well energy landscape. Returns (B,) energy values."""
        b = x.shape[0]
        well_energy = torch.zeros(b)

        for k in range(self.n_wells):
            diff = x - self.well_centers[k].unsqueeze(0)
            dist_sq = (diff ** 2).sum(dim=-1)
            depth = self.base_depths[k] + self.mu_sensitivity[k] * self.mu
            depth = torch.relu(depth)
            sigma = torch.relu(self.well_sigmas[k]) + 0.3
            well_energy -= depth * torch.exp(-dist_sq / (2 * sigma ** 2))

        confine = self.confinement * (x ** 2).sum(dim=-1)
        return well_energy + confine

    def flow(self, x):
        """dx/dt = -grad_x E(x, mu). Returns (B, d) velocity."""
        x_grad = x.detach().requires_grad_(True)
        with torch.enable_grad():
            E = self.energy(x_grad)
            grad_E = torch.autograd.grad(E.sum(), x_grad, create_graph=False)[0]
        return -grad_E.detach()

    def integrate(self, x0, n_steps=300, dt=0.05):
        """Integrate points to their attractor basins via Euler method."""
        x = x0.detach().clone()
        for _ in range(n_steps):
            dx = self.flow(x)
            x = x + dt * dx
            x = torch.clamp(x, -30.0, 30.0)
        return x

    def find_attractors(self, points, dbscan_eps=None, merge_threshold=None):
        """
        Integrate points and cluster converged states to find attractors.

        Args:
            points: (N, d) tensor of input points
            dbscan_eps: DBSCAN eps parameter. If None, computed adaptively.
            merge_threshold: Distance below which attractor basins are merged.
                           If None, computed adaptively from data.

        Returns:
            labels: (N,) integer array of cluster assignments
            n_attractors: number of distinct attractors found
            centers: (n_attractors, d) array of attractor centers in data space
        """
        x_final = self.integrate(points)
        x_np = x_final.numpy()
        pts_np = points.detach().numpy()

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

        # Adaptive DBSCAN eps based on nearest-neighbor distances in converged space
        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)
        basin_labels = db.fit_predict(x_np)

        # Assign noise points to nearest cluster
        if -1 in basin_labels:
            cluster_mask = basin_labels >= 0
            if cluster_mask.any():
                unique_clusters = sorted(set(basin_labels) - {-1})
                centers_conv = np.array([x_np[basin_labels == c].mean(axis=0)
                                        for c in unique_clusters])
                for i in np.where(basin_labels == -1)[0]:
                    dists = np.linalg.norm(centers_conv - x_np[i], axis=1)
                    basin_labels[i] = unique_clusters[dists.argmin()]
            else:
                basin_labels = np.zeros(len(x_np), dtype=int)

        # Merge basins whose data-space centroids are close
        unique_basins = sorted(set(basin_labels))
        if len(unique_basins) > 1 and merge_threshold is not None:
            basin_centroids = np.array([pts_np[basin_labels == b].mean(axis=0)
                                       for b in unique_basins])
            merged = {b: b for b in unique_basins}

            for i, b1 in enumerate(unique_basins):
                for j, b2 in enumerate(unique_basins):
                    if j <= i:
                        continue
                    dist = np.linalg.norm(basin_centroids[i] - basin_centroids[j])
                    if dist < merge_threshold:
                        root1 = b1
                        while merged[root1] != root1:
                            root1 = merged[root1]
                        root2 = b2
                        while merged[root2] != root2:
                            root2 = merged[root2]
                        if root1 != root2:
                            merged[root2] = root1

            for b in unique_basins:
                root = b
                while merged[root] != root:
                    root = merged[root]
                merged[b] = root

            basin_labels = np.array([merged[b] for b in basin_labels])

        # Relabel to 0, 1, 2, ...
        unique_labels = sorted(set(basin_labels))
        label_map = {old: new for new, old in enumerate(unique_labels)}
        labels = np.array([label_map[l] for l in basin_labels])

        n_attractors = len(unique_labels)
        centers = np.array([pts_np[labels == i].mean(axis=0)
                           for i in range(n_attractors)])

        return labels, n_attractors, centers


class QPPGOnlineClusterer:
    """
    Online clustering via bifurcation-driven attractor creation.

    Places wells at k-means candidates, then sweeps mu to find the optimal
    number of attractor basins via silhouette score. When n_clusters increases
    across phases, that's a bifurcation event = genuinely new category detected.

    Parameters:
        d: Input dimension
        n_wells: Maximum wells in energy landscape (default 12)
        mu_range: (min, max) for mu sweep (default (-3.0, 4.0))
        mu_steps: Number of mu values to test (default 30)
        seed: Random seed
        confinement: Substrate confinement strength
        sigma_base: Base well width
    """

    def __init__(self, d, n_wells=12, mu_range=(-3.0, 4.0), mu_steps=30,
                 seed=42, confinement=0.0001, sigma_base=1.0):
        self.d = d
        self.n_wells = n_wells
        self.mu_range = mu_range
        self.mu_steps = mu_steps
        self.seed = seed

        self.substrate = QPPGSubstrate(
            d=d, n_wells=n_wells, seed=seed,
            confinement=confinement, sigma_base=sigma_base
        )
        self.n_clusters_detected = 0
        self.all_data = []
        self.novelty_events = []  # (phase_idx, old_n, new_n)
        self.mu_history = []

    def _place_wells_at_data(self, all_pts):
        """Place wells at k-means cluster candidates (over-seeded)."""
        n_candidates = min(self.n_wells, max(2, len(all_pts) // 20))
        km = KMeans(n_clusters=n_candidates, n_init=5, random_state=self.seed)
        km.fit(all_pts)

        with torch.no_grad():
            centers = torch.tensor(km.cluster_centers_, dtype=torch.float32)
            self.substrate.well_centers.data[:n_candidates] = centers
            self.substrate.well_centers.data[n_candidates:] = 100.0

            counts = np.bincount(km.labels_, minlength=n_candidates)
            depths = counts / max(counts.max(), 1) * 3.0 + 1.0
            self.substrate.base_depths.data[:n_candidates] = torch.tensor(
                depths, dtype=torch.float32)
            self.substrate.base_depths.data[n_candidates:] = 0.0

            sorted_idx = np.argsort(-counts)
            sens = torch.zeros(self.n_wells)
            for rank, idx in enumerate(sorted_idx):
                sens[idx] = (rank - n_candidates // 2) * 0.5
            self.substrate.mu_sensitivity.data[:n_candidates] = sens[:n_candidates]

    def _find_best_mu(self, points_t):
        """Sweep mu to find the value producing highest silhouette score."""
        best_mu = self.mu_range[0]
        best_sil = -1.0
        best_labels = np.zeros(len(points_t), dtype=int)
        best_n = 1

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

            if n_clusters >= 2 and n_clusters <= self.n_wells:
                pts_np = points_t.detach().numpy()
                try:
                    sil = silhouette_score(pts_np, labels)
                except ValueError:
                    sil = -1.0

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

        self.substrate.mu = best_mu
        return best_labels, best_n, best_sil

    def process_batch(self, points, phase_idx=0):
        """
        Process a batch of points and return cluster assignments.

        Args:
            points: (N, d) numpy array of input points
            phase_idx: Phase identifier for tracking novelty events

        Returns:
            labels: (N,) integer array of cluster assignments
        """
        self.all_data.append(points)
        all_pts = np.concatenate(self.all_data)
        points_t = torch.tensor(all_pts, dtype=torch.float32)

        self._place_wells_at_data(all_pts)
        labels, n_clusters, sil = self._find_best_mu(points_t)
        self.mu_history.append(self.substrate.mu)

        if n_clusters > self.n_clusters_detected:
            self.novelty_events.append((phase_idx, self.n_clusters_detected, n_clusters))

        self.n_clusters_detected = n_clusters

        # Return labels for the most recent batch only
        start = len(all_pts) - len(points)
        return labels[start:]

    def process_all(self, points):
        """
        Process all points at once (non-streaming mode).

        Args:
            points: (N, d) numpy array

        Returns:
            labels: (N,) integer array
        """
        points_t = torch.tensor(points, dtype=torch.float32)
        self._place_wells_at_data(points)
        labels, n_clusters, sil = self._find_best_mu(points_t)
        self.n_clusters_detected = n_clusters
        return labels

    def get_n_clusters(self):
        return self.n_clusters_detected

    def get_novelty_events(self):
        return self.novelty_events

    def get_mu_history(self):
        return self.mu_history


# ============================================================
# Bifurcation Analysis Utilities (from exp2_topological_novelty.py)
# ============================================================

def compute_jacobian(f, center, d, eps=1e-3):
    """
    Compute the Jacobian of f at center via finite differences.

    Args:
        f: Function taking (B, d) tensor → (B, d) tensor (the flow field)
        center: (d,) numpy array — the point at which to compute the Jacobian
        d: Dimension of the state space
        eps: Finite difference step size

    Returns:
        J: (d, d) numpy array — the Jacobian matrix
    """
    J = np.zeros((d, d))
    x0 = torch.tensor(center, dtype=torch.float32).unsqueeze(0)
    f0 = f(x0).squeeze(0).numpy()

    for i in range(d):
        x_plus = x0.clone()
        x_plus[0, i] += eps
        f_plus = f(x_plus).squeeze(0).numpy()
        J[:, i] = (f_plus - f0) / eps

    return J


def jacobian_signature(J):
    """
    Extract topological signature from a Jacobian matrix.

    Returns:
        dict with eigenvalue spectrum, stability classification, spectral gap
    """
    eigenvalues = np.linalg.eigvals(J)
    eigenvalues_real = np.sort(np.real(eigenvalues))

    n_stable = np.sum(eigenvalues_real < -1e-6)
    n_unstable = np.sum(eigenvalues_real > 1e-6)
    n_marginal = len(eigenvalues_real) - n_stable - n_unstable

    spectral_gap = float(eigenvalues_real[-1] - eigenvalues_real[0]) if len(eigenvalues_real) > 1 else 0.0
    trace = float(np.real(np.trace(J)))

    return {
        'eigenvalues_real': eigenvalues_real.tolist(),
        'n_stable': int(n_stable),
        'n_unstable': int(n_unstable),
        'n_marginal': int(n_marginal),
        'spectral_gap': spectral_gap,
        'trace': trace,
    }


def signature_distance(sig1, sig2):
    """
    Compute distance between two Jacobian signatures.
    Uses L1 on sorted eigenvalue spectra + stability type penalty.
    """
    ev1 = np.array(sig1['eigenvalues_real'])
    ev2 = np.array(sig2['eigenvalues_real'])

    # Pad shorter to match longer
    max_len = max(len(ev1), len(ev2))
    ev1_pad = np.pad(ev1, (0, max_len - len(ev1)))
    ev2_pad = np.pad(ev2, (0, max_len - len(ev2)))

    spectral_dist = np.mean(np.abs(ev1_pad - ev2_pad))
    stability_dist = abs(sig1['n_stable'] - sig2['n_stable']) + \
                     abs(sig1['n_unstable'] - sig2['n_unstable'])

    return spectral_dist + 0.1 * stability_dist
