"""Fixed BF16 label-row projection for Flash Next; no vocabulary-sized GEMM."""
import os


def register():
    from vllm import ModelRegistry
    ModelRegistry.register_model(
        'Qwen4ExpForConditionalGeneration',
        'jev_flash_linear:Qwen4ExpForConditionalGeneration',
    )


# Keep registration lightweight: vLLM imports the model class only when needed.
def __getattr__(name):
    if name != 'Qwen4ExpForConditionalGeneration':
        raise AttributeError(name)
    import torch
    from safetensors.torch import load_file
    from torch import nn
    from vllm.logger import init_logger
    from vllm.models.qwen3_8_flash_next.nvidia.model import (
        Qwen3_8FlashNextForConditionalGeneration as BaseModel,
    )
    logger = init_logger('vllm.jev_flash_linear')

    class Qwen4ExpForConditionalGeneration(BaseModel):
        def __init__(self, *, vllm_config, prefix=''):
            super().__init__(vllm_config=vllm_config, prefix=prefix)
            tensors = load_file(os.environ['JEV_LINEAR_HEAD'], device='cpu')
            weight = tensors['weight']
            ids = tensors['token_ids']
            self.choice_head = nn.Linear(weight.shape[1], weight.shape[0], bias=False,
                                         dtype=torch.bfloat16)
            self.choice_head.weight.data.copy_(weight)
            self.register_buffer('choice_token_ids', ids.to(self.choice_head.weight.device),
                                 persistent=False)
            self.choice_vocab_size = vllm_config.model_config.hf_text_config.vocab_size
            del self.language_model.lm_head
            del self.language_model.logits_processor
            if any(key.endswith('lm_head') for key, _ in self.named_modules()):
                raise RuntimeError('A full vocabulary output head remains')
            if self.choice_head.weight.device.type != 'meta':
                torch.testing.assert_close(self.choice_head.weight.cpu(), weight, rtol=0, atol=0)
            self._choice_head_checked = False
            logger.info('JEV compact linear head shape=%s dtype=%s device=%s; original lm_head removed',
                        tuple(weight.shape), self.choice_head.weight.dtype,
                        self.choice_head.weight.device)

        def load_weights(self, weights):
            loaded = super().load_weights(
                (key, value) for key, value in weights if key != 'lm_head.weight'
            )
            if loaded is not None:
                loaded.add('choice_head.weight')
            return loaded

        def compute_logits(self, hidden_states):
            scores = self.choice_head(hidden_states)
            # Keep the existing token-ID API. This scatter costs no full-vocab
            # projection; the only output matrix multiplication is H x 86.
            logits = scores.new_full((*scores.shape[:-1], self.choice_vocab_size), float('-inf'))
            logits.index_copy_(-1, self.choice_token_ids, scores)
            if not self._choice_head_checked:
                logger.info('JEV linear forward hidden=%s scores=%s shape=%s; finite label logits=%s',
                            hidden_states.dtype, scores.dtype, tuple(scores.shape),
                            bool(torch.isfinite(scores).all()))
                self._choice_head_checked = True
            return logits

    globals()[name] = Qwen4ExpForConditionalGeneration
    return Qwen4ExpForConditionalGeneration