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

class FP8RowScaledEmbedding(nn.Module):
    def __init__(self, num_embeddings, embedding_dim):
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.register_buffer("weight", torch.zeros((num_embeddings, embedding_dim), dtype=torch.float8_e4m3fn))
        self.register_buffer("weight_scale", torch.zeros((num_embeddings, 1), dtype=torch.bfloat16))

    def forward(self, input_ids):
        w_sub = self.weight[input_ids]
        s_sub = self.weight_scale[input_ids].to(torch.float32)
        deq = w_sub.to(torch.float32) * s_sub
        return deq.to(torch.bfloat16)

class FlatQuantFP8Linear(nn.Module):
    def __init__(self, in_features, out_features, bias=False):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.register_buffer("weight", torch.zeros((out_features, in_features), dtype=torch.float8_e4m3fn))
        self.register_buffer("weight_scale", torch.zeros((), 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):
        w_deq = (self.weight.to(torch.float32) * self.weight_scale.to(torch.float32)).to(x.dtype)
        return F.linear(x, w_deq, self.bias.to(x.dtype) if self.bias is not None else None)

class FlatQuantW4A8Linear(nn.Module):
    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)
        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
                child = name.rsplit(".", 1)[-1]
                setattr(parent, child, FP8RowScaledEmbedding(mod.num_embeddings, mod.embedding_dim))
            elif isinstance(mod, nn.Linear):
                parent = self.model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else self.model
                child = name.rsplit(".", 1)[-1]
                is_attn = ("self_attn" in name) or ("linear_attn" in name)
                is_down = "down_proj" in name
                if is_attn or is_down:
                    setattr(parent, child, FlatQuantFP8Linear(mod.in_features, mod.out_features, bias=mod.bias is not None))
                else:
                    setattr(parent, child, FlatQuantW4A8Linear(mod.in_features, mod.out_features, bias=mod.bias is not None, group_size=16))

    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]
        embeddings = F.normalize(embeddings, dim=-1)
        return embeddings