"""WeMM-Embedding INT8: asymmetric int8 with a zero point and one scale per 64-weight block.

Two builds are served by this one file, selected by config.json:

  quant_upgraded_int4 = false   attention and down_proj are int8; the MLP gate/up, ViT and
                                merger keep their original uint8 group-16 int4
  quant_upgraded_int4 = true    every quantized matrix in the model is int8, so the int4
                                unpacking path is unused

Layout comes from `quant_scheme` and `quant_group_size` in config.json, where a group size
of 0 means one scale per output channel.

Why a zero point and why 64-weight blocks: SLQ (arXiv 2605.02404) proves symmetric grids
inflate output variance by gamma^2 with gamma = 2M/R over the dynamic range, and its
allocation counts zero-point overhead against the budget. Measured on the shapes in this
checkpoint, one scalar per matrix becomes 64-weight blocks with a zero point and output
error drops from 3.5e-04 to 1.5e-05 in 1-cosine terms.

    model = AutoModel.from_pretrained(repo, trust_remote_code=True)
    emb = model.embedding(input_ids=..., attention_mask=...)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import Qwen3_5ForConditionalGeneration

DEFAULT_GROUP = 64


def _dequant_int8(q, s, z, group):
    """(q - z) * s, with s and z shared across `group` input features."""
    shape = q.shape
    in_features = shape[-1]
    g = in_features if not group else group
    v = q.to(torch.float32).view(*shape[:-1], in_features // g, g)
    sc = s.to(torch.float32).unsqueeze(-1)
    zz = z.to(torch.float32).unsqueeze(-1)
    return ((v - zz) * sc).view(shape)


class Int8AsymLinear(nn.Module):
    def __init__(self, in_features, out_features, bias=False, group=DEFAULT_GROUP):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.group = in_features if not group else group
        self.register_buffer("weight", torch.zeros((out_features, in_features), dtype=torch.int8))
        self.register_buffer("weight_scale", torch.zeros((out_features, in_features // self.group), dtype=torch.bfloat16))
        self.register_buffer("weight_zero", torch.zeros((out_features, in_features // self.group), dtype=torch.int8))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=torch.bfloat16))
        else:
            self.register_parameter("bias", None)

    def forward(self, x):
        w = _dequant_int8(self.weight, self.weight_scale, self.weight_zero, self.group).to(x.dtype)
        return F.linear(x, w, self.bias.to(x.dtype) if self.bias is not None else None)


class Int8AsymRowEmbedding(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, group=DEFAULT_GROUP):
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.group = embedding_dim if not group else group
        self.register_buffer("weight", torch.zeros((num_embeddings, embedding_dim), dtype=torch.int8))
        self.register_buffer("weight_scale", torch.zeros((num_embeddings, embedding_dim // self.group), dtype=torch.bfloat16))
        self.register_buffer("weight_zero", torch.zeros((num_embeddings, embedding_dim // self.group), dtype=torch.int8))

    def forward(self, input_ids):
        # index_select takes integer dtypes on CPU, CUDA and MPS. F.embedding does not,
        # which is why the FP8 checkpoint needed a third party kernel shim on Apple Silicon.
        flat = input_ids.reshape(-1)
        w = torch.index_select(self.weight, 0, flat).view(*input_ids.shape, self.embedding_dim)
        s = torch.index_select(self.weight_scale, 0, flat).view(*input_ids.shape, -1)
        z = torch.index_select(self.weight_zero, 0, flat).view(*input_ids.shape, -1)
        return _dequant_int8(w, s, z, self.group).to(torch.bfloat16)


class FlatQuantW4A8Linear(nn.Module):
    """uint8 packed int4, group 16, two nibbles per byte in even/odd column order.

    Copied from the FP8 checkpoint so those tensors load byte for byte, and kept only for
    builds where the int4 modules were not upgraded to int8.
    """

    def __init__(self, in_features, out_features, bias=False, group_size=16):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.group_size = group_size
        self.register_buffer("weight", torch.zeros((out_features, in_features // 2), dtype=torch.uint8))
        self.register_buffer("weight_scale", torch.zeros((out_features, in_features // group_size), dtype=torch.bfloat16))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=torch.bfloat16))
        else:
            self.register_parameter("bias", None)

    def forward(self, x):
        low = (self.weight & 0x0F).to(torch.int8)
        high = (self.weight >> 4).to(torch.int8)
        low = torch.where(low >= 8, low - 16, low)
        high = torch.where(high >= 8, high - 16, high)
        unpacked = torch.empty(self.out_features, self.in_features, device=x.device, dtype=x.dtype)
        unpacked[:, 0::2] = low.to(x.dtype)
        unpacked[:, 1::2] = high.to(x.dtype)
        scales = self.weight_scale.to(x.dtype).repeat_interleave(self.group_size, dim=1)
        return F.linear(x, unpacked * scales, self.bias.to(x.dtype) if self.bias is not None else None)


class WeMMEmbedding(Qwen3_5ForConditionalGeneration):
    def __init__(self, config):
        super().__init__(config)
        group = int(getattr(config, "quant_group_size", DEFAULT_GROUP))
        upgraded = bool(getattr(config, "quant_upgraded_int4", False))
        for name, mod in list(self.model.named_modules()):
            if name.endswith("embed_tokens") and isinstance(mod, nn.Embedding):
                parent = self.model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else self.model
                setattr(parent, name.rsplit(".", 1)[-1],
                        Int8AsymRowEmbedding(mod.num_embeddings, mod.embedding_dim, group=group))
            elif isinstance(mod, nn.Linear):
                parent = self.model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else self.model
                is_attn = ("self_attn" in name) or ("linear_attn" in name)
                if is_attn or "down_proj" in name or upgraded:
                    new = Int8AsymLinear(mod.in_features, mod.out_features,
                                         bias=mod.bias is not None, group=group)
                else:
                    new = FlatQuantW4A8Linear(mod.in_features, mod.out_features,
                                              bias=mod.bias is not None, group_size=16)
                setattr(parent, name.rsplit(".", 1)[-1], new)

    def embedding(self, input_ids=None, attention_mask=None, **kwargs):
        self.model.rope_deltas = None
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
        last_hidden_state = outputs.last_hidden_state
        if attention_mask is not None:
            eos_positions = attention_mask.sum(dim=1) - 1
        else:
            eos_positions = torch.full((last_hidden_state.shape[0],), last_hidden_state.shape[1] - 1, device=last_hidden_state.device)
        eos_positions = eos_positions.clamp(min=0)
        batch_indices = torch.arange(last_hidden_state.size(0), device=last_hidden_state.device)
        embeddings = last_hidden_state[batch_indices, eos_positions]
        return F.normalize(embeddings, dim=-1)

    def tie_weights(self, *args, **kwargs):
        pass