gpu_compute.py

#!/usr/bin/env python3
"""
GPU COMPUTE ORCHESTRA - Full CUDA Utilization for SGM-Substrate
================================================================
ALL heavy computation on GPU. ALL VRAM. ALL CUDA cores.

Components:
1. GPUGeometry   - Location vectors resident in VRAM (cupy arrays)
2. gpu_kmeans    - K-means clustering fully on GPU
3. gpu_pca       - PCA eigendecomposition on GPU
4. gpu_distances - Pairwise distance matrix on GPU
5. gpu_nearest   - Batch nearest-neighbor on GPU
6. gpu_refine    - Continuous geometry refinement on GPU
7. UltraBlaster  - STEM3 ultra kernel (8 stems, multi-stream)

Information = Geometry. The GPU IS the geometry engine.
"""

import numpy as np
import time
import re
from typing import Dict, List, Tuple, Optional, Set

try:
    import cupy as cp
    from cupy import RawKernel
    HAS_GPU = True
except ImportError:
    HAS_GPU = False


# =========================================================================
# STEM3 ULTRA KERNEL - 8 stems, fully unrolled, max throughput
# =========================================================================

ULTRA_KERNEL_SRC = r'''
extern "C" __global__ void __launch_bounds__(512, 2) stem3_ultra(
    const int* __restrict__ stems1,
    const int* __restrict__ stems2,
    unsigned char* __restrict__ scores,
    int n_pairs
) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n_pairs) return;

    const int base1 = idx * 8;
    const int base2 = idx * 8;

    const int s1_0 = stems1[base1], s1_1 = stems1[base1+1];
    const int s1_2 = stems1[base1+2], s1_3 = stems1[base1+3];
    const int s1_4 = stems1[base1+4], s1_5 = stems1[base1+5];
    const int s1_6 = stems1[base1+6], s1_7 = stems1[base1+7];

    const int s2_0 = stems2[base2], s2_1 = stems2[base2+1];
    const int s2_2 = stems2[base2+2], s2_3 = stems2[base2+3];
    const int s2_4 = stems2[base2+4], s2_5 = stems2[base2+5];
    const int s2_6 = stems2[base2+6], s2_7 = stems2[base2+7];

    int c1 = (s1_0>=0) + (s1_1>=0) + (s1_2>=0) + (s1_3>=0) +
             (s1_4>=0) + (s1_5>=0) + (s1_6>=0) + (s1_7>=0);
    int c2 = (s2_0>=0) + (s2_1>=0) + (s2_2>=0) + (s2_3>=0) +
             (s2_4>=0) + (s2_5>=0) + (s2_6>=0) + (s2_7>=0);

    int inter = 0;
    if (s1_0 >= 0 && (s1_0==s2_0||s1_0==s2_1||s1_0==s2_2||s1_0==s2_3||s1_0==s2_4||s1_0==s2_5||s1_0==s2_6||s1_0==s2_7)) inter++;
    if (s1_1 >= 0 && (s1_1==s2_0||s1_1==s2_1||s1_1==s2_2||s1_1==s2_3||s1_1==s2_4||s1_1==s2_5||s1_1==s2_6||s1_1==s2_7)) inter++;
    if (s1_2 >= 0 && (s1_2==s2_0||s1_2==s2_1||s1_2==s2_2||s1_2==s2_3||s1_2==s2_4||s1_2==s2_5||s1_2==s2_6||s1_2==s2_7)) inter++;
    if (s1_3 >= 0 && (s1_3==s2_0||s1_3==s2_1||s1_3==s2_2||s1_3==s2_3||s1_3==s2_4||s1_3==s2_5||s1_3==s2_6||s1_3==s2_7)) inter++;
    if (s1_4 >= 0 && (s1_4==s2_0||s1_4==s2_1||s1_4==s2_2||s1_4==s2_3||s1_4==s2_4||s1_4==s2_5||s1_4==s2_6||s1_4==s2_7)) inter++;
    if (s1_5 >= 0 && (s1_5==s2_0||s1_5==s2_1||s1_5==s2_2||s1_5==s2_3||s1_5==s2_4||s1_5==s2_5||s1_5==s2_6||s1_5==s2_7)) inter++;
    if (s1_6 >= 0 && (s1_6==s2_0||s1_6==s2_1||s1_6==s2_2||s1_6==s2_3||s1_6==s2_4||s1_6==s2_5||s1_6==s2_6||s1_6==s2_7)) inter++;
    if (s1_7 >= 0 && (s1_7==s2_0||s1_7==s2_1||s1_7==s2_2||s1_7==s2_3||s1_7==s2_4||s1_7==s2_5||s1_7==s2_6||s1_7==s2_7)) inter++;

    int uni = c1 + c2 - inter;
    int jac = (uni > 0) ? ((inter * 255) / uni) : 0;
    scores[idx] = (unsigned char)((jac * 192 + 128 * 64) >> 8);
}

// Vector operations kernel: batch pull-together on GPU
extern "C" __global__ void batch_pull(
    float* __restrict__ locations,  // (n_concepts, dim) flattened
    const int* __restrict__ idx_a,  // indices of concept A
    const int* __restrict__ idx_b,  // indices of concept B
    const float* __restrict__ strengths,  // pull strength per pair
    int n_pairs,
    int dim
) {
    int pair = blockIdx.x * blockDim.x + threadIdx.x;
    if (pair >= n_pairs) return;

    int a = idx_a[pair];
    int b = idx_b[pair];
    float s = strengths[pair];

    for (int d = 0; d < dim; d++) {
        float va = locations[a * dim + d];
        float vb = locations[b * dim + d];
        float mid = (va + vb) * 0.5f;
        locations[a * dim + d] = va + s * (mid - va);
        locations[b * dim + d] = vb + s * (mid - vb);
    }
}

// Pairwise cosine similarity kernel
extern "C" __global__ void pairwise_cosine(
    const float* __restrict__ X,  // (n, dim) flattened
    float* __restrict__ sims,     // (n, n) output
    int n,
    int dim
) {
    int i = blockIdx.y * blockDim.y + threadIdx.y;
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n || j >= n || j <= i) return;

    float dot = 0.0f, norm_i = 0.0f, norm_j = 0.0f;
    for (int d = 0; d < dim; d++) {
        float vi = X[i * dim + d];
        float vj = X[j * dim + d];
        dot += vi * vj;
        norm_i += vi * vi;
        norm_j += vj * vj;
    }

    float denom = sqrtf(norm_i) * sqrtf(norm_j);
    float sim = (denom > 1e-8f) ? (dot / denom) : 0.0f;
    sims[i * n + j] = sim;
    sims[j * n + i] = sim;
}
'''

STOP_WORDS = {'the','is','are','was','in','on','at','to','of','and','or','for','with','be','it','an','a'}

def _stem_to_int(s):
    if len(s) < 3:
        return -1
    return ord(s[0]) | (ord(s[1]) << 8) | (ord(s[2]) << 16)

def extract_stems(text, max_stems=8):
    words = re.findall(r'[a-z]+', text.lower())
    stems = []
    seen = set()
    for w in words:
        if len(w) > 2 and w not in STOP_WORDS:
            s = w[:3]
            if s not in seen:
                seen.add(s)
                stems.append(_stem_to_int(s))
                if len(stems) >= max_stems:
                    break
    while len(stems) < max_stems:
        stems.append(-1)
    return stems


class GPUCompute:
    """
    Full GPU compute orchestra. Manages VRAM, kernels, and all GPU operations.
    """

    def __init__(self):
        self.available = HAS_GPU
        self.dim = 64
        self.kernels = {}
        self.streams = []
        self.total_ops = 0
        self.total_time = 0.0
        self._gpu_locs = None  # cupy array of all location vectors
        self._gpu_keys = []    # ordered list of concept keys
        self._gpu_key_to_idx = {}  # key -> index in gpu_locs

        if self.available:
            try:
                self._init_kernels()
                self._init_streams()
                self._log_gpu_info()
            except Exception as e:
                print(f"GPU init failed: {e}")
                self.available = False

    def _init_kernels(self):
        self.kernels['stem3_ultra'] = RawKernel(ULTRA_KERNEL_SRC, 'stem3_ultra')
        self.kernels['batch_pull'] = RawKernel(ULTRA_KERNEL_SRC, 'batch_pull')
        self.kernels['pairwise_cosine'] = RawKernel(ULTRA_KERNEL_SRC, 'pairwise_cosine')

    def _init_streams(self):
        self.streams = [cp.cuda.Stream() for _ in range(4)]

    def _log_gpu_info(self):
        device = cp.cuda.Device()
        props = cp.cuda.runtime.getDeviceProperties(device.id)
        mem = cp.cuda.runtime.memGetInfo()
        self.gpu_name = props['name'].decode()
        self.sms = props['multiProcessorCount']
        self.cores = self.sms * 128
        self.vram_total = mem[1]
        self.vram_free = mem[0]
        print(f"GPU Orchestra: {self.gpu_name}")
        print(f"  CUDA Cores: {self.cores}, SMs: {self.sms}")
        print(f"  VRAM: {self.vram_free/1e9:.2f} GB free / {self.vram_total/1e9:.2f} GB total")

    # =====================================================================
    # VRAM-RESIDENT GEOMETRY
    # =====================================================================

    def sync_to_gpu(self, locations: Dict[str, np.ndarray]):
        """Upload all location vectors to VRAM. Returns GPU array."""
        if not self.available or not locations:
            return

        keys = sorted(locations.keys())
        n = len(keys)
        dim = len(next(iter(locations.values())))
        self.dim = dim

        # Build matrix and upload
        matrix = np.zeros((n, dim), dtype=np.float32)
        for i, k in enumerate(keys):
            matrix[i] = locations[k]

        self._gpu_locs = cp.asarray(matrix)
        self._gpu_keys = keys
        self._gpu_key_to_idx = {k: i for i, k in enumerate(keys)}

    def sync_to_cpu(self, locations: Dict[str, np.ndarray]):
        """Download GPU geometry back to CPU dict."""
        if self._gpu_locs is None:
            return

        cpu_locs = cp.asnumpy(self._gpu_locs)
        for i, k in enumerate(self._gpu_keys):
            if k in locations:
                locations[k] = cpu_locs[i]

    def gpu_locs_count(self) -> int:
        return len(self._gpu_keys) if self._gpu_keys else 0

    # =====================================================================
    # GPU K-MEANS
    # =====================================================================

    def gpu_kmeans(self, X_np: np.ndarray, k: int, max_iter: int = 50) -> Tuple[np.ndarray, np.ndarray]:
        """
        K-means fully on GPU. Returns (assignments, centers) as numpy.
        """
        if not self.available:
            return self._cpu_kmeans(X_np, k, max_iter)

        n = len(X_np)
        X = cp.asarray(X_np.astype(np.float32))

        # Random init
        idx = cp.random.choice(n, k, replace=False)
        centers = X[idx].copy()

        for _ in range(max_iter):
            # Pairwise distances: (n, k)
            # Use broadcasting: ||x - c||^2 = ||x||^2 + ||c||^2 - 2*x.c
            x_sq = (X * X).sum(axis=1, keepdims=True)  # (n, 1)
            c_sq = (centers * centers).sum(axis=1, keepdims=True).T  # (1, k)
            dists = x_sq + c_sq - 2.0 * X @ centers.T  # (n, k)

            assignments = cp.argmin(dists, axis=1)

            # Update centers
            new_centers = cp.zeros_like(centers)
            for j in range(k):
                mask = assignments == j
                if mask.any():
                    new_centers[j] = X[mask].mean(axis=0)
                else:
                    new_centers[j] = centers[j]

            if cp.allclose(centers, new_centers, atol=1e-6):
                break
            centers = new_centers

        self.total_ops += n * k * max_iter
        return cp.asnumpy(assignments), cp.asnumpy(centers)

    def _cpu_kmeans(self, X, k, max_iter):
        n = len(X)
        idx = np.random.choice(n, k, replace=False)
        centers = X[idx].copy()
        assignments = np.zeros(n, dtype=int)
        for _ in range(max_iter):
            x_sq = (X * X).sum(axis=1, keepdims=True)
            c_sq = (centers * centers).sum(axis=1, keepdims=True).T
            dists = x_sq + c_sq - 2.0 * X @ centers.T
            assignments = np.argmin(dists, axis=1)
            new_centers = np.zeros_like(centers)
            for j in range(k):
                mask = assignments == j
                if mask.any():
                    new_centers[j] = X[mask].mean(axis=0)
                else:
                    new_centers[j] = centers[j]
            if np.allclose(centers, new_centers):
                break
            centers = new_centers
        return assignments, centers

    # =====================================================================
    # GPU PCA
    # =====================================================================

    def gpu_pca(self, X_np: np.ndarray, n_components: int = 10) -> Tuple[np.ndarray, np.ndarray]:
        """
        PCA fully on GPU. Returns (eigenvalues, eigenvectors) as numpy.
        """
        if not self.available:
            X_centered = X_np - X_np.mean(axis=0)
            cov = np.cov(X_centered.T)
            eigenvalues, eigenvectors = np.linalg.eigh(cov)
            idx = np.argsort(eigenvalues)[::-1]
            return eigenvalues[idx[:n_components]], eigenvectors[:, idx[:n_components]]

        X = cp.asarray(X_np.astype(np.float32))
        mean = X.mean(axis=0)
        X_centered = X - mean

        cov = cp.cov(X_centered.T)
        eigenvalues, eigenvectors = cp.linalg.eigh(cov)

        idx = cp.argsort(eigenvalues)[::-1]
        result_vals = cp.asnumpy(eigenvalues[idx[:n_components]])
        result_vecs = cp.asnumpy(eigenvectors[:, idx[:n_components]])

        self.total_ops += X_np.shape[0] * X_np.shape[1]
        return result_vals, result_vecs

    # =====================================================================
    # GPU PAIRWISE DISTANCES
    # =====================================================================

    def gpu_pairwise_distances(self, X_np: np.ndarray) -> np.ndarray:
        """
        Compute pairwise Euclidean distance matrix on GPU.
        Returns (n, n) numpy array.
        """
        if not self.available:
            sq = (X_np * X_np).sum(axis=1)
            return np.sqrt(np.maximum(sq[:, None] + sq[None, :] - 2.0 * X_np @ X_np.T, 0.0))

        X = cp.asarray(X_np.astype(np.float32))
        sq = (X * X).sum(axis=1)
        dists = cp.sqrt(cp.maximum(sq[:, None] + sq[None, :] - 2.0 * X @ X.T, 0.0))

        self.total_ops += X_np.shape[0] ** 2
        return cp.asnumpy(dists)

    # =====================================================================
    # GPU NEAREST NEIGHBOR (BATCH)
    # =====================================================================

    def gpu_find_nearest(self, query_key: str, exclude: Set[str] = None, n: int = 1) -> Optional[str]:
        """Find nearest concept using GPU-resident geometry."""
        if self._gpu_locs is None or query_key not in self._gpu_key_to_idx:
            return None

        exclude = exclude or set()
        exclude.add(query_key)

        qi = self._gpu_key_to_idx[query_key]
        query = self._gpu_locs[qi:qi+1]  # (1, dim)

        # Compute distances to all concepts on GPU
        diffs = self._gpu_locs - query  # (n, dim)
        dists = (diffs * diffs).sum(axis=1)  # (n,)

        # Mask excluded
        for ex in exclude:
            if ex in self._gpu_key_to_idx:
                dists[self._gpu_key_to_idx[ex]] = cp.float32(1e30)

        best_idx = int(cp.argmin(dists))
        return self._gpu_keys[best_idx]

    def gpu_find_nearest_batch(self, query_keys: List[str], exclude_per_query: List[Set[str]] = None) -> List[Optional[str]]:
        """Batch nearest-neighbor on GPU. Much faster for multiple queries."""
        if self._gpu_locs is None:
            return [None] * len(query_keys)

        results = []
        for i, qk in enumerate(query_keys):
            excl = exclude_per_query[i] if exclude_per_query else set()
            results.append(self.gpu_find_nearest(qk, exclude=excl))
        return results

    # =====================================================================
    # GPU BATCH VECTOR PULLS
    # =====================================================================

    def gpu_batch_pull(self, pairs: List[Tuple[str, str]], strengths: List[float]):
        """
        Batch pull-together operation on GPU.
        Modifies GPU-resident location vectors directly.
        """
        if self._gpu_locs is None or not pairs:
            return 0

        idx_a = []
        idx_b = []
        str_list = []

        for (a, b), s in zip(pairs, strengths):
            if a in self._gpu_key_to_idx and b in self._gpu_key_to_idx:
                idx_a.append(self._gpu_key_to_idx[a])
                idx_b.append(self._gpu_key_to_idx[b])
                str_list.append(s)

        if not idx_a:
            return 0

        n = len(idx_a)
        d_idx_a = cp.asarray(np.array(idx_a, dtype=np.int32))
        d_idx_b = cp.asarray(np.array(idx_b, dtype=np.int32))
        d_strengths = cp.asarray(np.array(str_list, dtype=np.float32))

        blocks = (n + 255) // 256
        self.kernels['batch_pull'](
            (blocks,), (256,),
            (self._gpu_locs, d_idx_a, d_idx_b, d_strengths, np.int32(n), np.int32(self.dim))
        )
        cp.cuda.Stream.null.synchronize()

        self.total_ops += n * self.dim
        return n

    # =====================================================================
    # GPU STEM3 ULTRA - Multi-stream text similarity
    # =====================================================================

    def stem3_batch(self, texts1: List[str], texts2: List[str]) -> np.ndarray:
        """
        Compute STEM3 similarity for paired text lists using ultra kernel.
        Returns float array of scores [0, 1].
        """
        if not self.available:
            return np.zeros(len(texts1))

        n = len(texts1)
        if n == 0:
            return np.array([])

        # Extract stems on CPU (fast, string processing)
        stems1 = np.zeros(n * 8, dtype=np.int32)
        stems2 = np.zeros(n * 8, dtype=np.int32)

        for i in range(n):
            s1 = extract_stems(texts1[i])
            s2 = extract_stems(texts2[i])
            stems1[i*8:(i+1)*8] = s1
            stems2[i*8:(i+1)*8] = s2

        # Upload and compute on GPU
        d_stems1 = cp.asarray(stems1)
        d_stems2 = cp.asarray(stems2)
        d_scores = cp.zeros(n, dtype=cp.uint8)

        blocks = (n + 511) // 512
        t0 = time.perf_counter()
        self.kernels['stem3_ultra']((blocks,), (512,),
                                    (d_stems1, d_stems2, d_scores, np.int32(n)))
        cp.cuda.Stream.null.synchronize()
        elapsed = time.perf_counter() - t0

        self.total_time += elapsed
        self.total_ops += n

        return cp.asnumpy(d_scores) / 255.0

    def stem3_pairwise(self, texts: List[str]) -> np.ndarray:
        """
        Compute ALL pairwise STEM3 similarities for a list of texts.
        Returns (n, n) similarity matrix.
        Uses multi-stream for maximum throughput.
        """
        n = len(texts)
        if n < 2 or not self.available:
            return np.zeros((n, n))

        # Extract stems for all texts
        all_stems = [extract_stems(t) for t in texts]

        # Build all pairs
        n_pairs = n * (n - 1) // 2
        stems1 = np.zeros(n_pairs * 8, dtype=np.int32)
        stems2 = np.zeros(n_pairs * 8, dtype=np.int32)

        pair_idx = 0
        pair_map = []
        for i in range(n):
            for j in range(i + 1, n):
                stems1[pair_idx*8:(pair_idx+1)*8] = all_stems[i]
                stems2[pair_idx*8:(pair_idx+1)*8] = all_stems[j]
                pair_map.append((i, j))
                pair_idx += 1

        # Upload to GPU
        d_stems1 = cp.asarray(stems1)
        d_stems2 = cp.asarray(stems2)
        d_scores = cp.zeros(n_pairs, dtype=cp.uint8)

        # Multi-stream execution
        pairs_per_stream = n_pairs // len(self.streams) + 1
        blocks_per_stream = (pairs_per_stream + 511) // 512

        t0 = time.perf_counter()
        for si, stream in enumerate(self.streams):
            offset = si * pairs_per_stream
            count = min(pairs_per_stream, n_pairs - offset)
            if count <= 0:
                break
            s1_off = offset * 8
            s2_off = offset * 8
            blocks = (count + 511) // 512
            with stream:
                self.kernels['stem3_ultra'](
                    (blocks,), (512,),
                    (d_stems1[s1_off:], d_stems2[s2_off:], d_scores[offset:], np.int32(count))
                )

        for stream in self.streams:
            stream.synchronize()
        elapsed = time.perf_counter() - t0

        self.total_time += elapsed
        self.total_ops += n_pairs

        # Build similarity matrix
        scores = cp.asnumpy(d_scores) / 255.0
        sim_matrix = np.zeros((n, n))
        for k, (i, j) in enumerate(pair_map):
            sim_matrix[i, j] = scores[k]
            sim_matrix[j, i] = scores[k]

        return sim_matrix

    # =====================================================================
    # GPU GEOMETRY REFINEMENT - Continuous GPU utilization
    # =====================================================================

    def gpu_refine_geometry(self, seed_locs: Dict[str, np.ndarray],
                           mind_locs: Dict[str, np.ndarray],
                           sentences: List[str],
                           plasticity: float = 0.3) -> dict:
        """
        Full GPU geometry refinement pipeline:
        1. Compute STEM3 pairwise similarity for all sentences
        2. Extract content words per sentence
        3. Batch-pull similar concept vectors closer in BOTH Seed and Mind

        Returns stats dict.
        """
        if not self.available or len(sentences) < 4:
            return {"pairs": 0, "reinforced": 0, "rate": "N/A"}

        import random
        MAX_BATCH = 800  # More pairs = more GPU utilization
        if len(sentences) > MAX_BATCH:
            batch = random.sample(sentences, MAX_BATCH)
        else:
            batch = sentences

        t0 = time.perf_counter()

        # 1. GPU STEM3 pairwise
        sim_matrix = self.stem3_pairwise(batch)

        # 2. Extract content words per sentence
        sentence_words = []
        for s in batch:
            words = re.findall(r'[a-z]+', s.lower())
            content = [w for w in words if len(w) > 2 and w not in STOP_WORDS]
            sentence_words.append(content[:8])

        # 3. Find high-similarity pairs and batch-pull
        pull_strength = 0.02 * plasticity
        n = len(batch)
        seed_pairs = []
        seed_strengths = []
        mind_pairs = []
        mind_strengths = []

        for i in range(n):
            for j in range(i + 1, n):
                sim = sim_matrix[i, j]
                if sim < 0.25:
                    continue

                words_i = sentence_words[i]
                words_j = sentence_words[j]
                if not words_i or not words_j:
                    continue

                for wi in words_i[:5]:
                    for wj in words_j[:5]:
                        if wi == wj:
                            continue
                        s = sim * pull_strength

                        if wi in seed_locs and wj in seed_locs:
                            seed_pairs.append((wi, wj))
                            seed_strengths.append(s)

                        if wi in mind_locs and wj in mind_locs:
                            mind_pairs.append((wi, wj))
                            mind_strengths.append(s)

        # 4. Apply pulls on CPU (GPU batch_pull requires VRAM-resident locs)
        seed_reinforced = 0
        for (wi, wj), s in zip(seed_pairs, seed_strengths):
            vec_i = seed_locs[wi]
            vec_j = seed_locs[wj]
            delta = (vec_j - vec_i) * s
            seed_locs[wi] = vec_i + delta
            seed_locs[wj] = vec_j - delta
            seed_reinforced += 1

        mind_reinforced = 0
        for (wi, wj), s in zip(mind_pairs, mind_strengths):
            vec_i = mind_locs[wi]
            vec_j = mind_locs[wj]
            delta = (vec_j - vec_i) * s
            mind_locs[wi] = vec_i + delta
            mind_locs[wj] = vec_j - delta
            mind_reinforced += 1

        elapsed = time.perf_counter() - t0
        n_pairs = n * (n - 1) // 2
        rate = n_pairs / max(elapsed, 1e-9)

        return {
            "pairs": n_pairs,
            "seed_reinforced": seed_reinforced,
            "mind_reinforced": mind_reinforced,
            "elapsed_ms": elapsed * 1000,
            "rate": f"{rate/1e6:.1f}M pairs/sec",
        }

    # =====================================================================
    # GPU INTROSPECT - PCA + K-means fully on GPU
    # =====================================================================

    def gpu_introspect_mind(self, locations: Dict[str, np.ndarray], dim: int) -> dict:
        """
        Full Mind introspection on GPU: PCA + K-means.
        Returns dict with axes, clusters, cluster_centers.
        """
        if len(locations) < 5:
            return {"axes": None, "clusters": {}, "centers": None}

        concepts = list(locations.keys())
        X = np.array([locations[c] for c in concepts], dtype=np.float32)

        # GPU PCA
        eigenvalues, eigenvectors = self.gpu_pca(X, n_components=min(10, dim))

        # GPU K-means
        n_clusters = min(max(3, len(concepts) // 5), 10)
        assignments, centers = self.gpu_kmeans(X, n_clusters)

        # Build cluster dict
        clusters = {}
        for i, cid in enumerate(assignments):
            cid = int(cid)
            if cid not in clusters:
                clusters[cid] = set()
            clusters[cid].add(concepts[i])

        return {
            "axes": eigenvectors,
            "eigenvalues": eigenvalues,
            "clusters": clusters,
            "centers": centers,
        }

    def gpu_introspect_seed(self, locations: Dict[str, np.ndarray]) -> dict:
        """
        Full Seed cluster discovery on GPU.
        Uses GPU pairwise distances + threshold-based clustering.
        """
        tokens = list(locations.keys())
        n = len(tokens)
        if n < 5:
            return {"clusters": {}}

        MAX_CLUSTER = 2000
        if n > MAX_CLUSTER:
            idx = np.random.choice(n, MAX_CLUSTER, replace=False)
            work_tokens = [tokens[i] for i in idx]
            work_locs = np.array([locations[tokens[i]] for i in idx], dtype=np.float32)
        else:
            work_tokens = list(tokens)
            work_locs = np.array([locations[t] for t in tokens], dtype=np.float32)

        # GPU pairwise distances
        distances = self.gpu_pairwise_distances(work_locs)
        sn = len(work_tokens)

        threshold = np.median(distances[distances > 0])

        visited = set()
        clusters = {}
        cluster_id = 0

        for i in range(sn):
            if i in visited:
                continue
            cluster = {work_tokens[i]}
            queue = [i]
            visited.add(i)

            while queue:
                current = queue.pop(0)
                neighbors = np.where(distances[current] < threshold)[0]
                for j in neighbors:
                    if j not in visited:
                        visited.add(j)
                        queue.append(j)
                        cluster.add(work_tokens[j])

            if len(cluster) > 1:
                clusters[cluster_id] = cluster
                cluster_id += 1

        return {"clusters": clusters}

    # =====================================================================
    # STATUS
    # =====================================================================

    def status(self) -> dict:
        if not self.available:
            return {"available": False, "gpu": "none"}

        rate = self.total_ops / max(self.total_time, 1e-9) if self.total_time > 0 else 0
        mem = cp.cuda.runtime.memGetInfo()

        return {
            "available": True,
            "gpu": self.gpu_name,
            "cores": self.cores,
            "vram_used_mb": (self.vram_total - mem[0]) / 1e6,
            "vram_free_mb": mem[0] / 1e6,
            "gpu_locs": self.gpu_locs_count(),
            "total_ops": self.total_ops,
            "total_time_ms": self.total_time * 1000,
            "rate": f"{rate/1e6:.1f}M ops/sec" if self.total_ops > 0 else "N/A",
        }


# Singleton
_gpu = None

def get_gpu() -> GPUCompute:
    global _gpu
    if _gpu is None:
        _gpu = GPUCompute()
    return _gpu