"""
HelixLinearSTE — Trainable compressed linear layer with straight-through estimator.

Stores weights in CDNA v3 format (codebook + uint8 indices + sidecar + SVD)
but makes codebook entries, sidecar values, and SVD factors into nn.Parameters
so gradients flow through during training.

The non-differentiable step (argmin index assignment) uses a straight-through
estimator: forward quantizes, backward pretends it didn't happen.

Periodically, indices are re-assigned via k-means to track the evolving codebook.

This is a RESEARCH module for compressed-domain training experiments.
For inference, use HelixLinear (helix_linear.py) which is production-hardened.

Work Order: WO-COMPRESSED-NATIVE-TRAINING-01
"""

from __future__ import annotations

import torch
import torch.nn as nn
from typing import Optional


class VQStraightThrough(torch.autograd.Function):
    """Straight-through estimator for VQ codebook gather.

    Forward: W_vq = codebook[indices]  (standard gather, creates [out, in] tensor)
    Backward: gradient flows to codebook entries via scatter_add.
              Each codebook entry receives the sum of gradients from all
              positions assigned to it. Adam normalizes the variance.
    """

    @staticmethod
    def forward(ctx, codebook: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
        ctx.save_for_backward(indices)
        ctx.codebook_size = codebook.shape[0]
        return codebook[indices.long()]

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        indices, = ctx.saved_tensors
        # Accumulate gradients into codebook entries
        grad_codebook = torch.zeros(ctx.codebook_size, device=grad_output.device,
                                    dtype=grad_output.dtype)
        grad_codebook.scatter_add_(0, indices.long().reshape(-1),
                                   grad_output.reshape(-1))
        return grad_codebook, None  # no gradient for indices


def _torch_kmeans(
    data: torch.Tensor,
    n_clusters: int = 256,
    max_iters: int = 10,
    rtol: float = 0.001,
) -> tuple[torch.Tensor, torch.Tensor]:
    """GPU-friendly scalar k-means. Mirrors cdna_encoder._simple_kmeans.

    Args:
        data: 1D tensor of values
        n_clusters: Number of clusters (max 256 for uint8 indices)
        max_iters: Maximum iterations
        rtol: Relative convergence tolerance

    Returns:
        (centroids [n_clusters], assignments [N] uint8)
    """
    n_unique = len(torch.unique(data))
    n_clusters = min(n_clusters, n_unique)

    # Percentile initialization
    percentiles = torch.linspace(0, 1, n_clusters, device=data.device)
    sorted_data = torch.sort(data)[0]
    idx = (percentiles * (len(sorted_data) - 1)).long()
    centroids = sorted_data[idx].float()

    cb_range = (centroids[-1] - centroids[0]).item()
    if cb_range < 1e-30:
        cb_range = 1.0
    abs_tol = rtol * cb_range

    CHUNK = 1 << 20  # 1M elements per chunk to avoid OOM on [N, 256]

    for _ in range(max_iters):
        # Chunked assignment
        assignments = torch.empty(data.shape[0], device=data.device, dtype=torch.uint8)
        for start in range(0, data.shape[0], CHUNK):
            end = min(start + CHUNK, data.shape[0])
            dists = (data[start:end, None] - centroids[None, :]).abs()
            assignments[start:end] = dists.argmin(dim=1).to(torch.uint8)

        # 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 = (new_centroids - centroids).abs().max().item()
        centroids = new_centroids
        if max_delta < abs_tol:
            break

    return centroids, assignments


class HelixLinearSTE(nn.Module):
    """Trainable compressed linear layer with VQ straight-through estimator.

    Weights are stored as codebook[indices] + sidecar + SVD, same as HelixLinear.
    But codebook, sidecar_values, and SVD factors are nn.Parameters.

    Forward materializes the full W = codebook[indices] + corrections, then
    computes x @ W.T. STE lets gradients flow through the VQ gather.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        codebook: torch.Tensor,
        indices: torch.Tensor,
        svd_rank: int = 0,
        svd_U: Optional[torch.Tensor] = None,
        svd_s: Optional[torch.Tensor] = None,
        svd_Vt: Optional[torch.Tensor] = None,
        bias: Optional[torch.Tensor] = None,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features

        # Learnable VQ codebook
        self.codebook = nn.Parameter(codebook.float().contiguous())
        # Discrete indices — NOT learnable, updated by reassign_indices()
        self.register_buffer("indices", indices.to(torch.uint8).contiguous())

        # Learnable SVD residual
        self.has_svd = svd_U is not None
        if self.has_svd:
            self.svd_U = nn.Parameter(svd_U.float().contiguous())
            self.svd_s = nn.Parameter(svd_s.float().contiguous())
            self.svd_Vt = nn.Parameter(svd_Vt.float().contiguous())
        else:
            self.svd_U = None
            self.svd_s = None
            self.svd_Vt = None

        # Optional bias
        if bias is not None:
            self.bias = nn.Parameter(bias.float().contiguous())
        else:
            self.bias = None

        # Monitoring
        self._reassign_count = 0
        self._last_indices_changed = 0

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. VQ reconstruction with STE
        W = VQStraightThrough.apply(self.codebook, self.indices)  # [out, in]

        # 2. SVD residual correction (direct gradient)
        if self.has_svd:
            W = W + (self.svd_U * self.svd_s.unsqueeze(0)) @ self.svd_Vt

        # 3. Linear operation
        out = x @ W.t()
        if self.bias is not None:
            out = out + self.bias
        return out

    @torch.no_grad()
    def reassign_indices(self):
        """Re-run k-means assignment using current codebook.

        Call every N training steps to keep indices aligned with codebook.
        """
        W = self._reconstruct_full()
        flat = W.reshape(-1)

        # Chunked nearest-centroid assignment (avoid [N, 256] OOM)
        CHUNK = 1 << 20
        new_indices = torch.empty_like(self.indices.reshape(-1))
        for start in range(0, flat.shape[0], CHUNK):
            end = min(start + CHUNK, flat.shape[0])
            dists = (flat[start:end, None] - self.codebook[None, :]).abs()
            new_indices[start:end] = dists.argmin(dim=1).to(torch.uint8)

        new_indices = new_indices.reshape(self.out_features, self.in_features)

        # Track how many changed
        changed = (new_indices != self.indices).sum().item()
        total = self.indices.numel()
        self._last_indices_changed = changed
        self._reassign_count += 1

        self.indices.copy_(new_indices)
        return changed, total

    @torch.no_grad()
    def _reconstruct_full(self) -> torch.Tensor:
        """Reconstruct full W matrix (for reassignment and diagnostics)."""
        W = self.codebook[self.indices.long()]
        if self.has_svd:
            W = W + (self.svd_U * self.svd_s.unsqueeze(0)) @ self.svd_Vt
        return W

    @torch.no_grad()
    def weight_kurtosis(self) -> float:
        """Compute excess kurtosis of the effective weight distribution."""
        W = self._reconstruct_full()
        flat = W.reshape(-1).float()
        mean = flat.mean()
        var = ((flat - mean) ** 2).mean()
        if var < 1e-12:
            return 0.0
        kurt = ((flat - mean) ** 4).mean() / (var ** 2) - 3.0
        return kurt.item()

    @torch.no_grad()
    def codebook_utilization(self) -> float:
        """Fraction of codebook entries that are actually used."""
        used = len(torch.unique(self.indices))
        return used / self.codebook.shape[0]

    @classmethod
    def from_scratch(
        cls,
        in_features: int,
        out_features: int,
        bias: bool = False,
        svd_rank: int = 0,
        n_clusters: int = 256,
        device: Optional[torch.device] = None,
    ) -> "HelixLinearSTE":
        """Initialize with Kaiming-normal weights, quantized to CDNA v3 format.

        Uses fast percentile initialization (no iterative k-means) since gradient
        descent will refine the codebook during training.
        """
        W = torch.empty(out_features, in_features, device=device or "cpu")
        nn.init.kaiming_normal_(W, mode="fan_out", nonlinearity="linear")

        flat = W.reshape(-1)

        # Fast percentile codebook init (skip iterative k-means — training refines it)
        sorted_vals = torch.sort(flat)[0]
        pct_idx = torch.linspace(0, len(sorted_vals) - 1, n_clusters).long()
        codebook = sorted_vals[pct_idx].float()

        # Fast assignment via searchsorted — O(N log K) since codebook is sorted
        # Find the bin each value falls into, then pick the nearer of the two neighbors
        bucket = torch.searchsorted(codebook, flat).clamp(0, n_clusters - 1)
        # Check if the left neighbor is closer
        left = (bucket - 1).clamp(0)
        dist_right = (flat - codebook[bucket]).abs()
        dist_left = (flat - codebook[left]).abs()
        indices = torch.where(dist_left < dist_right, left, bucket).to(torch.uint8)
        indices = indices.reshape(out_features, in_features)

        # Optional SVD of quantization residual
        svd_U = svd_s = svd_Vt = None
        if svd_rank > 0:
            W_vq = codebook[indices.long()]
            residual = W - W_vq
            U, s, Vt = torch.linalg.svd(residual, full_matrices=False)
            svd_U = U[:, :svd_rank]
            svd_s = s[:svd_rank]
            svd_Vt = Vt[:svd_rank, :]

        bias_tensor = torch.zeros(out_features, device=device or "cpu") if bias else None

        return cls(
            in_features=in_features,
            out_features=out_features,
            codebook=codebook,
            indices=indices,
            svd_rank=svd_rank,
            svd_U=svd_U,
            svd_s=svd_s,
            svd_Vt=svd_Vt,
            bias=bias_tensor,
        )

    def extra_repr(self) -> str:
        parts = [f"in_features={self.in_features}", f"out_features={self.out_features}"]
        parts.append(f"codebook={self.codebook.shape[0]}")
        if self.has_svd:
            parts.append(f"svd_rank={self.svd_U.shape[1]}")
        if self.bias is not None:
            parts.append("bias=True")
        return ", ".join(parts)
