"""Core Sherry 3:4 ternary fake quantization and Arenas scheduling."""

from __future__ import annotations

import math
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F


def _validate_quant_shape(weight: torch.Tensor, group_size: int) -> None:
    """Validate constraints that the original implementation left to reshapes."""

    if weight.ndim != 2:
        raise ValueError(f"Sherry only supports 2-D linear weights, got {weight.ndim}-D")
    if not weight.is_floating_point():
        raise ValueError("Sherry fake quantization requires a floating-point weight")
    if not isinstance(group_size, int) or group_size <= 0:
        raise ValueError("group_size must be a positive integer")
    if group_size % 4:
        raise ValueError("group_size must be divisible by the 3:4 block width (4)")
    if weight.shape[1] % group_size:
        raise ValueError(
            f"in_features ({weight.shape[1]}) must be divisible by group_size ({group_size}); "
            "groups may not cross output-channel rows"
        )


class _NMQuantizeSTE(torch.autograd.Function):
    """Autograd implementation whose backward is the unmodified identity STE."""

    @staticmethod
    def forward(ctx: Any, weight: torch.Tensor, group_size: int) -> torch.Tensor:
        del ctx
        _validate_quant_shape(weight, group_size)

        original_shape = weight.shape
        blocks = weight.reshape(original_shape[0], original_shape[1] // 4, 4)
        survivor_indices = blocks.abs().topk(k=3, dim=-1).indices
        mask = torch.zeros_like(blocks, dtype=torch.bool)
        mask.scatter_(-1, survivor_indices, True)
        sparse = blocks.masked_fill(~mask, 0).reshape(-1, group_size)

        # The original accumulates in the latent dtype (often bf16).  Sherry notes
        # identify that as a precision bug, so accumulation is deliberately fp32.
        scale = sparse.float().abs().sum(dim=-1, keepdim=True) / (group_size * 3 / 4)
        signs = sparse.sign().float()
        quantized = (signs * scale).reshape(original_shape)
        return quantized.to(dtype=weight.dtype)

    @staticmethod
    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]:
        del ctx
        # Unlike mask-aware pruning estimators, the original Sherry rule passes
        # gradients through pruned lanes too.  Preserve that pure STE exactly.
        return grad_output, None


def nm_quantize(weight: torch.Tensor, group_size: int = 128) -> torch.Tensor:
    """Fake-quantize a linear weight with dynamic 3:4 sparsity and ternary scales.

    Every consecutive four input lanes retains the three largest magnitudes.  A
    group shares ``delta = abs(nonzero).mean()`` and survivors become ``+/-delta``.

    Differences from AngelSlim Sherry: the public API contains only live
    arguments, validates row-local grouping, and computes scale reductions in
    fp32.  Its backward remains the original, unclipped, mask-free STE.
    """

    return _NMQuantizeSTE.apply(weight, group_size)


class SherryLinear(nn.Linear):
    """``nn.Linear``-compatible Sherry layer with an externally driven residual.

    Differences from AngelSlim ``Arenas``: ``eps`` is a persistent buffer, no
    forward counter exists, schedule constants are not hard-coded, the typo in
    the terminal branch is removed with that branch, and the master parameters
    remain fp32.  The bias is added once rather than once per parallel path.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        *,
        group_size: int = 128,
        device: torch.device | str | None = None,
    ) -> None:
        if in_features % group_size:
            raise ValueError(
                f"in_features ({in_features}) must be divisible by group_size ({group_size})"
            )
        super().__init__(
            in_features,
            out_features,
            bias=bias,
            device=device,
            dtype=torch.float32,
        )
        self.group_size = group_size
        self.register_buffer("eps", torch.tensor(0.0, dtype=torch.float32))
        # Cached host float so forward can skip the residual path without
        # synchronizing on ``self.eps.item()`` every call.
        self._eps_py = 0.0

    def set_eps(self, value: float | torch.Tensor) -> None:
        """Update the residual coefficient buffer and its host-side cache."""

        eps = float(value.item() if isinstance(value, torch.Tensor) else value)
        self.eps.fill_(eps)
        self._eps_py = eps

    def _load_from_state_dict(
        self,
        state_dict: dict[str, torch.Tensor],
        prefix: str,
        local_metadata: dict[str, Any],
        strict: bool,
        missing_keys: list[str],
        unexpected_keys: list[str],
        error_msgs: list[str],
    ) -> None:
        """Accept native Linear checkpoints while still persisting ``eps``.

        AngelSlim's Python ``eps`` never appeared in checkpoints.  Here it is a
        buffer, but treating an absent epsilon as zero preserves strict loading
        compatibility with existing native ``nn.Linear`` state dictionaries.
        """

        eps_key = prefix + "eps"
        if eps_key not in state_dict:
            state_dict[eps_key] = self.eps.detach().clone()
        super()._load_from_state_dict(
            state_dict,
            prefix,
            local_metadata,
            strict,
            missing_keys,
            unexpected_keys,
            error_msgs,
        )
        self._eps_py = float(self.eps.detach().item())

    @classmethod
    def from_linear(cls, linear: nn.Linear, *, group_size: int = 128) -> "SherryLinear":
        """Copy a native linear into an fp32-master Sherry layer.

        Unlike the original model-construction path, this wrapper works on any
        already-created Transformers model and upgrades bf16/fp16 weights to the
        fp32 latent master representation required for stable QAT.
        """

        converted = cls(
            linear.in_features,
            linear.out_features,
            bias=linear.bias is not None,
            group_size=group_size,
            device=linear.weight.device,
        )
        with torch.no_grad():
            converted.weight.copy_(linear.weight.detach().float())
            if linear.bias is not None and converted.bias is not None:
                converted.bias.copy_(linear.bias.detach().float())
        converted.weight.requires_grad_(linear.weight.requires_grad)
        if linear.bias is not None and converted.bias is not None:
            converted.bias.requires_grad_(linear.bias.requires_grad)
        converted.train(linear.training)
        return converted

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """Apply quantized and annealed full-precision paths.

        AngelSlim mutates its schedule on every forward (including checkpoint
        recomputation).  This implementation only reads the ``eps`` buffer; an
        :class:`ArenasScheduler` must update it once per optimizer step.
        When ``eps`` is zero the full-precision residual matmul is skipped.
        """

        quantized_weight = nm_quantize(self.weight, self.group_size).to(input.dtype)
        output = F.linear(input, quantized_weight, None)
        if self._eps_py != 0.0:
            master_for_compute = self.weight.to(input.dtype)
            output = output + self.eps.to(input.dtype) * F.linear(
                input, master_for_compute, None
            )
        if self.bias is not None:
            output = output + self.bias.to(input.dtype)
        return output


class ArenasScheduler:
    """Optimizer-step scheduler for the additive Arenas residual coefficient.

    ``eps`` rises linearly from zero to one during warmup and then decays to zero.
    The scheduler is intentionally external: unlike AngelSlim's forward-local
    counter, gradient accumulation and activation checkpointing cannot advance it
    accidentally.  Layer ``eps`` values persist in model checkpoints.
    """

    def __init__(
        self,
        total_steps: int,
        warmup_frac: float = 0.1,
        schedule: str = "cosine",
    ) -> None:
        if not isinstance(total_steps, int) or total_steps < 2:
            raise ValueError("total_steps must be an integer of at least 2")
        if not 0.0 < warmup_frac < 1.0:
            raise ValueError("warmup_frac must be strictly between 0 and 1")
        if schedule not in {"cosine", "linear"}:
            raise ValueError("schedule must be 'cosine' or 'linear'")
        self.total_steps = total_steps
        self.warmup_frac = float(warmup_frac)
        self.schedule = schedule
        self.warmup_steps = min(total_steps - 1, max(1, int(total_steps * warmup_frac)))
        self.current_step = 0

    @property
    def eps(self) -> float:
        """Return epsilon at the current optimizer step."""

        step = min(self.current_step, self.total_steps)
        if step <= self.warmup_steps:
            return step / self.warmup_steps
        progress = (step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
        if self.schedule == "linear":
            return max(0.0, 1.0 - progress)
        return 0.5 * (1.0 + math.cos(math.pi * progress))

    def step(self) -> float:
        """Advance exactly one optimizer step and return the new epsilon."""

        self.current_step = min(self.current_step + 1, self.total_steps)
        return self.eps

    @torch.no_grad()
    def set_eps(self, model: nn.Module) -> list[str]:
        """Write the current epsilon to every Sherry layer and return its names.

        This replaces AngelSlim's non-resumable Python scalar mutation.  Call it
        after ``optimizer.step()`` (and after :meth:`step`) so one scheduler tick
        always corresponds to one optimizer update.
        """

        updated: list[str] = []
        for name, module in model.named_modules():
            if isinstance(module, SherryLinear):
                module.set_eps(self.eps)
                updated.append(name)
        return updated

    def state_dict(self) -> dict[str, int | float | str]:
        """Serialize external progress; layer epsilon itself is model-persistent."""

        return {
            "total_steps": self.total_steps,
            "warmup_frac": self.warmup_frac,
            "schedule": self.schedule,
            "current_step": self.current_step,
        }

    def load_state_dict(self, state: dict[str, int | float | str]) -> None:
        """Restore progress after validating immutable schedule parameters."""

        expected = (self.total_steps, self.warmup_frac, self.schedule)
        found = (state["total_steps"], state["warmup_frac"], state["schedule"])
        if found != expected:
            raise ValueError(f"scheduler configuration mismatch: expected {expected}, got {found}")
        step = int(state["current_step"])
        if not 0 <= step <= self.total_steps:
            raise ValueError("invalid current_step in scheduler state")
        self.current_step = step