"""
helix_cdc/regrow/cdna_encoder.py
================================

Simple CDNA encoder for kernel syscall usage.

Encodes a tensor to CDNA format using k-means quantization.

CDNA Format:
  - Magic: b"HXZC" (4 bytes)
  - Version: 1 (4 bytes, uint32 LE)
  - Rows: M (4 bytes, uint32 LE)
  - Cols: K (4 bytes, uint32 LE)
  - Codebook: [256] float32 (1024 bytes)
  - Indices: [M, K] uint8 (M*K bytes, zlib compressed)
"""

from __future__ import annotations

import struct
import zlib
from pathlib import Path
from typing import Optional

import numpy as np

# CDNA format constants
CDNA_MAGIC = b"HXZC"
CDNA_VERSION = 1
CODEBOOK_SIZE = 256


def encode_tensor_to_cdna(
    tensor: np.ndarray,
    output_path: Path,
    tensor_name: str = "tensor",
    n_clusters: int = 256,
    max_iters: int = 10,
) -> dict:
    """
    Encode a tensor to CDNA format using k-means quantization.

    Args:
        tensor: Input tensor (2D float32)
        output_path: Path to write CDNA shard
        tensor_name: Name for logging
        n_clusters: Number of codebook entries (default: 256)
        max_iters: Max k-means iterations (default: 10)

    Returns:
        dict with encoding stats
    """
    # Ensure 2D
    if tensor.ndim == 1:
        tensor = tensor.reshape(1, -1)
    elif tensor.ndim > 2:
        tensor = tensor.reshape(-1, tensor.shape[-1])

    # Ensure float32
    tensor = tensor.astype(np.float32)
    rows, cols = tensor.shape

    # Flatten for k-means
    flat = tensor.flatten()

    # Simple k-means (sklearn-free)
    codebook, indices_flat = _simple_kmeans(flat, n_clusters, max_iters)

    # Reshape indices back to 2D
    indices = indices_flat.reshape(rows, cols).astype(np.uint8)

    # Compress indices
    indices_bytes = indices.tobytes()
    indices_compressed = zlib.compress(indices_bytes, level=6)

    # Write CDNA file
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "wb") as f:
        # Header
        f.write(CDNA_MAGIC)
        f.write(struct.pack("<I", CDNA_VERSION))
        f.write(struct.pack("<I", rows))
        f.write(struct.pack("<I", cols))

        # Codebook (always 256 entries, padded if needed)
        codebook_full = np.zeros(256, dtype=np.float32)
        codebook_full[:len(codebook)] = codebook
        f.write(codebook_full.tobytes())

        # Compressed indices
        f.write(struct.pack("<I", len(indices_bytes)))  # Uncompressed size
        f.write(struct.pack("<I", len(indices_compressed)))  # Compressed size
        f.write(indices_compressed)

    return {
        "rows": rows,
        "cols": cols,
        "codebook_size": len(codebook),
        "indices_bytes": len(indices_bytes),
        "compressed_bytes": len(indices_compressed),
        "compression_ratio": len(indices_bytes) / max(1, len(indices_compressed)),
    }


def _simple_kmeans(
    data: np.ndarray,
    n_clusters: int,
    max_iters: int = 10,
    rtol: float = 0.001,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Simple k-means without sklearn dependency.

    Args:
        data: 1D array of values
        n_clusters: Number of clusters
        max_iters: Maximum iterations
        rtol: Relative convergence tolerance. Stop when max centroid
              movement drops below rtol * codebook_range. Default 0.001
              (0.1% of range). Proven: convergence profiling shows
              np.allclose never triggers within 30 iters; 0.1% relative
              threshold gives identical cosine with 3-5x fewer iterations.

    Returns:
        (centroids, assignments)
    """
    n_clusters = min(n_clusters, len(np.unique(data)))

    # Initialize centroids using percentiles
    percentiles = np.linspace(0, 100, n_clusters)
    centroids = np.percentile(data, percentiles).astype(np.float32)

    # Codebook range for relative convergence check
    cb_range = float(centroids[-1] - centroids[0])
    if cb_range < 1e-30:
        cb_range = 1.0  # degenerate: all same value
    abs_tol = rtol * cb_range

    # GPU k-means: produces identical results to CPU (A/B tested 2026-03-26).
    # Original corruption was from GPU contention (two compressions on one GPU).
    # Safe when running a single compression process.
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            try:
                return _gpu_kmeans(data, centroids, n_clusters, max_iters, abs_tol)
            except torch.cuda.OutOfMemoryError:
                pass  # Fall through to CPU
    except ImportError:
        pass

    # K-means iterations (CPU fallback)
    for _ in range(max_iters):
        # Assign points to nearest centroid
        dists = np.abs(data[:, np.newaxis] - centroids)
        assignments = np.argmin(dists, axis=1).astype(np.uint8)

        # Update centroids
        new_centroids = np.zeros_like(centroids)
        for i in range(n_clusters):
            mask = assignments == i
            if np.any(mask):
                new_centroids[i] = np.mean(data[mask])
            else:
                new_centroids[i] = centroids[i]

        # Check convergence: max centroid movement < 0.1% of codebook range
        max_delta = float(np.max(np.abs(new_centroids - centroids)))
        if max_delta < abs_tol:
            break

        centroids = new_centroids

    return centroids, assignments


def _gpu_kmeans(data_np, centroids_np, n_clusters, max_iters, abs_tol):
    """GPU-accelerated 1D k-means using PyTorch."""
    import torch
    data = torch.from_numpy(data_np).cuda()
    centroids = torch.from_numpy(centroids_np).cuda()

    for _ in range(max_iters):
        # Assign: [N, K] distances
        dists = torch.abs(data.unsqueeze(1) - centroids.unsqueeze(0))
        assignments = torch.argmin(dists, dim=1)

        # Update centroids
        new_centroids = torch.zeros_like(centroids)
        for i in range(n_clusters):
            mask = assignments == i
            if mask.any():
                new_centroids[i] = data[mask].mean()
            else:
                new_centroids[i] = centroids[i]

        max_delta = float((new_centroids - centroids).abs().max())
        if max_delta < abs_tol:
            break
        centroids = new_centroids

    assignments_np = assignments.cpu().numpy().astype(np.uint8)
    centroids_np = centroids.cpu().numpy()
    return centroids_np, assignments_np


def _vector_kmeans(
    data: np.ndarray,
    n_clusters: int,
    max_iters: int = 10,
    rtol: float = 0.001,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Multi-dimensional k-means for vector quantization.

    Args:
        data: [N_vectors, d] array of d-dimensional vectors
        n_clusters: Number of clusters (codebook entries)
        max_iters: Maximum iterations
        rtol: Relative convergence tolerance

    Returns:
        (centroids [k, d], assignments [N_vectors] uint8/uint16)
    """
    n, d = data.shape
    n_clusters = min(n_clusters, n)
    index_dtype = np.uint8 if n_clusters <= 256 else np.uint16

    # K-means++ initialization: random first centroid, then distance-weighted
    rng = np.random.RandomState(42)
    centroids = np.zeros((n_clusters, d), dtype=np.float32)
    centroids[0] = data[rng.randint(n)]

    for i in range(1, min(n_clusters, 64)):
        # Matmul trick for init distances: ||x-c||^2 = ||x||^2 + ||c||^2 - 2*x@c.T
        C_init = centroids[:i]
        C_sq = np.sum(C_init ** 2, axis=1, keepdims=True).T  # [1, i]
        # Subsample for init if data is very large
        if n > 100000:
            sub_idx = rng.choice(n, 100000, replace=False)
            sub = data[sub_idx]
        else:
            sub = data
            sub_idx = None
        X_sq = np.sum(sub ** 2, axis=1, keepdims=True)  # [sub_n, 1]
        dists = X_sq + C_sq - 2.0 * (sub @ C_init.T)  # [sub_n, i]
        np.maximum(dists, 0, out=dists)  # clip negative from float precision
        min_dists = dists.min(axis=1)
        probs = min_dists / max(min_dists.sum(), 1e-30)
        chosen = rng.choice(len(sub), p=probs)
        centroids[i] = sub[chosen]
    if n_clusters > 64:
        centroids[64:] = data[rng.choice(n, n_clusters - 64, replace=False)]

    # Convergence threshold
    cb_range = float(centroids.max() - centroids.min())
    if cb_range < 1e-30:
        cb_range = 1.0
    abs_tol = rtol * cb_range

    # GPU path
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()  # Release cached memory from prior tensors
            try:
                return _gpu_vector_kmeans(data, centroids, n_clusters, d, max_iters, abs_tol, index_dtype)
            except torch.cuda.OutOfMemoryError:
                pass  # Fall through to CPU
    except ImportError:
        pass

    # CPU iterations — use matmul trick for squared distances to avoid [N, k, d] broadcast
    # ||x - c||^2 = ||x||^2 + ||c||^2 - 2 * x @ c.T
    chunk_size = max(10000, n // 10)  # Chunk assignment to limit memory

    for _ in range(max_iters):
        # Chunked assignment to avoid [N, k] memory explosion for large N and k
        C_sq = np.sum(centroids ** 2, axis=1, keepdims=True).T  # [1, k]
        assignments = np.empty(n, dtype=np.int32)
        for start in range(0, n, chunk_size):
            end = min(start + chunk_size, n)
            X_chunk = data[start:end]
            X_sq = np.sum(X_chunk ** 2, axis=1, keepdims=True)  # [chunk, 1]
            dists = X_sq + C_sq - 2.0 * (X_chunk @ centroids.T)  # [chunk, k]
            assignments[start:end] = np.argmin(dists, axis=1)

        # Vectorized centroid update via bincount (avoids Python loop over k)
        counts = np.bincount(assignments, minlength=n_clusters).astype(np.float32)
        new_centroids = np.zeros_like(centroids)
        for dim in range(d):
            sums = np.bincount(assignments, weights=data[:, dim], minlength=n_clusters)
            new_centroids[:, dim] = np.where(
                counts > 0, sums / np.maximum(counts, 1), centroids[:, dim]
            )

        max_delta = float(np.max(np.abs(new_centroids - centroids)))
        if max_delta < abs_tol:
            break
        centroids = new_centroids

    # Final assignment (chunked)
    C_sq = np.sum(centroids ** 2, axis=1, keepdims=True).T
    assignments = np.empty(n, dtype=np.int32)
    for start in range(0, n, chunk_size):
        end = min(start + chunk_size, n)
        X_chunk = data[start:end]
        X_sq = np.sum(X_chunk ** 2, axis=1, keepdims=True)
        dists = X_sq + C_sq - 2.0 * (X_chunk @ centroids.T)
        assignments[start:end] = np.argmin(dists, axis=1)

    return centroids.astype(np.float32), assignments.astype(index_dtype)


def _gpu_vector_kmeans(data_np, centroids_np, n_clusters, d, max_iters, abs_tol, index_dtype):
    """GPU-accelerated multi-dimensional k-means using torch.cdist."""
    import torch
    data = torch.from_numpy(data_np).cuda()
    centroids = torch.from_numpy(centroids_np).cuda()

    for _ in range(max_iters):
        dists = torch.cdist(data, centroids)  # [N, k]
        assignments = torch.argmin(dists, dim=1)  # [N]

        # Vectorized centroid update via scatter_add (no Python loop over k)
        counts = torch.zeros(n_clusters, device=data.device)
        counts.scatter_add_(0, assignments, torch.ones(len(data), device=data.device))
        sums = torch.zeros(n_clusters, d, device=data.device)
        sums.scatter_add_(0, assignments.unsqueeze(1).expand(-1, d), data)
        alive = counts > 0
        new_centroids = centroids.clone()
        new_centroids[alive] = sums[alive] / counts[alive].unsqueeze(1)

        max_delta = float((new_centroids - centroids).abs().max())
        if max_delta < abs_tol:
            break
        centroids = new_centroids

    centroids_np = centroids.cpu().numpy().astype(np.float32)
    assignments_np = assignments.cpu().numpy().astype(index_dtype)
    return centroids_np, assignments_np


def decode_cdna_to_tensor(cdna_path: Path) -> np.ndarray:
    """
    Decode a CDNA shard back to tensor.

    Args:
        cdna_path: Path to CDNA shard

    Returns:
        Decoded tensor (2D float32)
    """
    with open(cdna_path, "rb") as f:
        # Read header
        magic = f.read(4)
        if magic != CDNA_MAGIC:
            raise ValueError(f"Invalid CDNA magic: {magic}")

        version = struct.unpack("<I", f.read(4))[0]
        if version != CDNA_VERSION:
            raise ValueError(f"Unsupported CDNA version: {version}")

        rows = struct.unpack("<I", f.read(4))[0]
        cols = struct.unpack("<I", f.read(4))[0]

        # Read codebook
        codebook = np.frombuffer(f.read(256 * 4), dtype=np.float32)

        # Read compressed indices
        uncompressed_size = struct.unpack("<I", f.read(4))[0]
        compressed_size = struct.unpack("<I", f.read(4))[0]
        indices_compressed = f.read(compressed_size)

        # Decompress
        indices_bytes = zlib.decompress(indices_compressed)
        indices = np.frombuffer(indices_bytes, dtype=np.uint8).reshape(rows, cols)

        # Dequantize
        tensor = codebook[indices]

    return tensor
