"""Symmetric round-to-nearest fake quantization for linear and embedding weights."""

from __future__ import annotations

from typing import Any
import weakref

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


def _validate_rtn_args(
    weight: torch.Tensor,
    bits: int,
    granularity: str,
    group_size: int,
) -> None:
    if weight.ndim != 2:
        raise ValueError(f"RTN only supports 2-D weights, got {weight.ndim}-D")
    if not weight.is_floating_point():
        raise ValueError("RTN fake quantization requires a floating-point weight")
    if bits not in {3, 4, 8}:
        raise ValueError("bits must be 3, 4, or 8")
    if granularity not in {"per_channel", "per_group"}:
        raise ValueError("granularity must be 'per_channel' or 'per_group'")
    if not isinstance(group_size, int) or group_size <= 0:
        raise ValueError("group_size must be a positive integer")


def _rtn_groups(
    weight: torch.Tensor,
    granularity: str,
    group_size: int,
) -> tuple[torch.Tensor, int]:
    """Return padded row-local groups and the original column count."""

    rows, columns = weight.shape
    actual_group_size = columns if granularity == "per_channel" else group_size
    number_of_groups = (columns + actual_group_size - 1) // actual_group_size
    padded_columns = number_of_groups * actual_group_size
    if padded_columns != columns:
        weight = F.pad(weight, (0, padded_columns - columns))
    return weight.reshape(rows, number_of_groups, actual_group_size), columns


def _quantize_groups(
    groups: torch.Tensor,
    bits: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Quantize groups in fp32 with deployment-canonical fp16 scales."""

    qmax = 2 ** (bits - 1) - 1
    scales = groups.float().abs().amax(dim=-1, keepdim=True) / qmax
    # The physical RTN format stores one fp16 scale per group.  Canonicalizing
    # here makes QAT/inference consume exactly the scale precision that export
    # can preserve, including when the latent master is fp32.
    scales = scales.to(torch.float16).float()
    safe_scales = torch.where(scales == 0, torch.ones_like(scales), scales)
    integers = torch.round(groups.float() / safe_scales).clamp(-qmax, qmax)
    return integers, scales


class _RTNQuantizeSTE(torch.autograd.Function):
    """Autograd implementation with the same pure identity STE as Sherry."""

    @staticmethod
    def forward(
        ctx: Any,
        weight: torch.Tensor,
        bits: int,
        granularity: str,
        group_size: int,
    ) -> torch.Tensor:
        del ctx
        _validate_rtn_args(weight, bits, granularity, group_size)
        groups, columns = _rtn_groups(weight, granularity, group_size)
        integers, scales = _quantize_groups(groups, bits)
        quantized = (integers * scales).reshape(weight.shape[0], -1)[:, :columns]
        return quantized.to(weight.dtype)

    @staticmethod
    def backward(
        ctx: Any,
        grad_output: torch.Tensor,
    ) -> tuple[torch.Tensor, None, None, None]:
        del ctx
        return grad_output, None, None, None


def rtn_quantize(
    weight: torch.Tensor,
    bits: int = 8,
    granularity: str = "per_channel",
    group_size: int = 128,
) -> torch.Tensor:
    """Apply symmetric absmax RTN fake quantization with an identity STE.

    ``per_channel`` assigns one scale to every row (the output channel for a
    linear weight, or one embedding vector).  ``per_group`` divides each row
    along its last dimension and permits a shorter final group.  The signed
    integer grid is ``[-127, 127]`` for W8, ``[-7, 7]`` for W4, or
    ``[-3, 3]`` for W3.
    """

    return _RTNQuantizeSTE.apply(weight, bits, granularity, group_size)


class RTNLinear(nn.Linear):
    """``nn.Linear`` with an fp32 master weight and RTN fake-quant forward."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        *,
        bits: int = 8,
        granularity: str = "per_channel",
        group_size: int = 128,
        device: torch.device | str | None = None,
    ) -> None:
        # Validate configuration without imposing a divisibility requirement.
        _validate_rtn_args(
            torch.empty(out_features, in_features), bits, granularity, group_size
        )
        super().__init__(
            in_features,
            out_features,
            bias=bias,
            device=device,
            dtype=torch.float32,
        )
        self.bits = bits
        self.granularity = granularity
        self.group_size = group_size

    @classmethod
    def from_linear(
        cls,
        linear: nn.Linear,
        *,
        bits: int = 8,
        granularity: str = "per_channel",
        group_size: int = 128,
    ) -> "RTNLinear":
        converted = cls(
            linear.in_features,
            linear.out_features,
            bias=linear.bias is not None,
            bits=bits,
            granularity=granularity,
            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:
        quantized_weight = rtn_quantize(
            self.weight,
            bits=self.bits,
            granularity=self.granularity,
            group_size=self.group_size,
        ).to(input.dtype)
        bias = self.bias.to(input.dtype) if self.bias is not None else None
        return F.linear(input, quantized_weight, bias)


class RTNEmbedding(nn.Embedding):
    """Embedding with an fp32 master and configurable RTN fake quantization.

    A forward-local cache lets a tied output head consume the exact same
    fake-quant tensor used by the lookup.  This is important for a tied
    embedding/head matrix: merely tying the fp32 master parameters would still
    permit the two call sites to fake-quantize independently.
    """

    def __init__(
        self,
        num_embeddings: int,
        embedding_dim: int,
        padding_idx: int | None = None,
        max_norm: float | None = None,
        norm_type: float = 2.0,
        scale_grad_by_freq: bool = False,
        sparse: bool = False,
        *,
        bits: int = 8,
        granularity: str = "per_channel",
        group_size: int = 128,
        compute_dtype: torch.dtype = torch.float32,
        device: torch.device | str | None = None,
    ) -> None:
        _validate_rtn_args(
            torch.empty(num_embeddings, embedding_dim), bits, granularity, group_size
        )
        super().__init__(
            num_embeddings,
            embedding_dim,
            padding_idx=padding_idx,
            max_norm=max_norm,
            norm_type=norm_type,
            scale_grad_by_freq=scale_grad_by_freq,
            sparse=sparse,
            device=device,
            dtype=torch.float32,
        )
        self.bits = bits
        self.granularity = granularity
        self.group_size = group_size
        self.compute_dtype = compute_dtype
        object.__setattr__(self, "_forward_quantized_weight", None)
        object.__setattr__(self, "_eval_quantized_weight", None)

    @classmethod
    def from_embedding(
        cls,
        embedding: nn.Embedding,
        *,
        bits: int = 8,
        granularity: str = "per_channel",
        group_size: int = 128,
    ) -> "RTNEmbedding":
        converted = cls(
            embedding.num_embeddings,
            embedding.embedding_dim,
            padding_idx=embedding.padding_idx,
            max_norm=embedding.max_norm,
            norm_type=embedding.norm_type,
            scale_grad_by_freq=embedding.scale_grad_by_freq,
            sparse=embedding.sparse,
            bits=bits,
            granularity=granularity,
            group_size=group_size,
            compute_dtype=embedding.weight.dtype,
            device=embedding.weight.device,
        )
        with torch.no_grad():
            converted.weight.copy_(embedding.weight.detach().float())
        converted.weight.requires_grad_(embedding.weight.requires_grad)
        converted.train(embedding.training)
        return converted

    def _make_quantized_weight(self) -> torch.Tensor:
        return rtn_quantize(
            self.weight,
            bits=self.bits,
            granularity=self.granularity,
            group_size=self.group_size,
        ).to(self.compute_dtype)

    def cache_eval_weight(self, quantized_weight: torch.Tensor | None) -> None:
        """Install/remove a detached persistent fake-quant matrix for eval."""

        object.__setattr__(self, "_eval_quantized_weight", quantized_weight)
        object.__setattr__(self, "_forward_quantized_weight", None)

    def _lookup_quantized_weight(self) -> torch.Tensor:
        quantized_weight = self._eval_quantized_weight
        if quantized_weight is None:
            quantized_weight = self._make_quantized_weight()
        object.__setattr__(self, "_forward_quantized_weight", quantized_weight)
        return quantized_weight

    def take_tied_quantized_weight(self) -> torch.Tensor:
        """Consume the precise tensor produced by the preceding lookup."""

        quantized_weight = self._forward_quantized_weight
        if quantized_weight is None:
            raise RuntimeError(
                "tied RTN lm_head ran without a preceding embedding lookup; "
                "the shared fake-quant invariant cannot be guaranteed"
            )
        object.__setattr__(self, "_forward_quantized_weight", None)
        return quantized_weight

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        quantized_weight = self._lookup_quantized_weight()
        return F.embedding(
            input,
            quantized_weight,
            self.padding_idx,
            self.max_norm,
            self.norm_type,
            self.scale_grad_by_freq,
            self.sparse,
        )


class TiedRTNLMHead(nn.Module):
    """Bias-free projection tied to an :class:`RTNEmbedding` master and cache."""

    def __init__(self, embedding: RTNEmbedding) -> None:
        super().__init__()
        self.in_features = embedding.embedding_dim
        self.out_features = embedding.num_embeddings
        # Register the same Parameter at the conventional checkpoint key while
        # keeping the module reference weak/non-registered (no module cycle).
        self.weight = embedding.weight
        object.__setattr__(self, "_embedding_ref", weakref.ref(embedding))

    @property
    def embedding(self) -> RTNEmbedding:
        embedding = self._embedding_ref()
        if embedding is None:
            raise RuntimeError("tied RTN embedding no longer exists")
        return embedding

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        embedding = self.embedding
        if self.weight is not embedding.weight:
            raise RuntimeError("RTN embedding/lm_head fp32 master tie was broken")
        quantized_weight = embedding.take_tied_quantized_weight()
        return F.linear(input, quantized_weight, None)


def tie_rtn_lm_head(model: nn.Module) -> tuple[RTNEmbedding, TiedRTNLMHead]:
    """Reconnect MOSS input/output weights through one shared fake-quant path."""

    embedding = model.get_input_embeddings()
    if not isinstance(embedding, RTNEmbedding):
        raise TypeError(f"expected RTNEmbedding input, got {type(embedding).__name__}")
    head = TiedRTNLMHead(embedding)
    model.set_output_embeddings(head)
    if model.get_output_embeddings() is not head:
        raise RuntimeError("model rejected the tied RTN output head")
    if head.weight is not embedding.weight or head.weight.data_ptr() != embedding.weight.data_ptr():
        raise RuntimeError("embedding/lm_head fp32 master tie verification failed")
    return embedding, head