"""CPU checks of installed Kunlun INT8 loader and BF16 MLA weight preparation.

Uses synthetic layers and mocked TP rank/size, not a full engine startup.
"""
import copy
import os
from types import SimpleNamespace
from unittest.mock import patch

os.environ["VLLM_KUNLUN_ENABLE_INT8_BMM"] = "0"

import torch
from convert_deepseek_fp8_to_int8 import quantization_config


def main():
    from vllm.platforms import current_platform
    current_platform.pre_register_and_update()
    from vllm_kunlun.quantization.compressed_tensors.compressed_tensors import KunlunCompressedTensorsConfig
    from vllm_kunlun.quantization.compressed_tensors.compressed_tensors_moe import KunlunCompressedTensorsMoEMethod
    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
    cfg = KunlunCompressedTensorsConfig.from_config(copy.deepcopy(quantization_config()))
    cfg.packed_modules_mapping = {
        "gate_up_proj": ["gate_proj", "up_proj"],
        "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
    }
    for name, widths in [
        ("model.layers.0.self_attn.fused_qkv_a_proj", [128, 64]),
        ("model.layers.0.self_attn.q_b_proj", [128]),
        ("model.layers.0.self_attn.o_proj", [128]),
    ]:
        layer = torch.nn.Linear(128, sum(widths), bias=False)
        scheme = cfg.get_scheme(layer, name)
        del layer._parameters['weight']
        with patch('vllm.model_executor.parameter.get_tensor_model_parallel_rank', return_value=0), \
                patch('vllm.model_executor.parameter.get_tensor_model_parallel_world_size', return_value=1):
            scheme.create_weights(layer, widths, 128, torch.bfloat16, lambda *a, **k: None)
        assert type(scheme.kernel).__name__ == 'KunlunScaledMMLinearKernel'
        assert layer.weight.dtype == torch.int8
        assert tuple(layer.weight_scale.shape) == (sum(widths), 1)
        print('LINEAR_PASS', name, type(scheme.kernel).__name__)

    class RoutedExperts(torch.nn.Module):
        pass

    experts = RoutedExperts()
    experts.moe_config = SimpleNamespace(is_act_and_mul=True)
    method = KunlunCompressedTensorsMoEMethod.get_moe_method(cfg, experts, 'model.layers.3.mlp.experts')
    assert type(method).__name__ == 'KunlunCompressedTensorsW8A8Int8MoEMethod'
    method.create_weights(experts, 2, 128, 128, torch.bfloat16)
    assert tuple(experts.w13_weight_scale.shape) == (2, 256, 1)
    with torch.no_grad():
        experts.w13_weight_scale.fill_(0.25)
        experts.w2_weight_scale.fill_(0.25)
    method.process_weights_after_loading(experts)
    assert experts.w13_weight_scale.flatten()[0].item() == 31.75
    print('MOE_SCALE_PASS', type(method).__name__)
    for name in ['lm_head', 'model.layers.3.mlp.gate', 'model.layers.0.self_attn.kv_b_proj',
                 'model.layers.0.mlp.gate_up_proj', 'model.layers.0.mlp.down_proj',
                 'model.layers.3.mlp.shared_experts.gate_up_proj',
                 'model.layers.3.mlp.shared_experts.down_proj']:
        assert cfg.get_scheme(torch.nn.Linear(128, 128), name) is None
    print('IGNORE_PASS: head, router, MLA kv_b_proj, dense MLPs, shared experts')

    from vllm_kunlun.v1.attention.backends.mla.common import MLACommonImpl
    projection = torch.nn.Linear(128, 256, bias=False, dtype=torch.bfloat16)
    projection.quant_method = UnquantizedLinearMethod()
    impl = SimpleNamespace(kv_b_proj=projection, kv_lora_rank=128,
                           num_heads=2, qk_nope_head_dim=64, v_head_dim=64)
    with torch.inference_mode():
        MLACommonImpl.process_weights_after_loading(impl, torch.bfloat16)
    assert tuple(impl.W_UK_T.shape) == (2, 64, 128)
    assert tuple(impl.W_UV.shape) == (2, 128, 64)
    assert impl.W_UK_T.dtype == impl.W_UV.dtype == torch.bfloat16
    expected = projection.weight.T.reshape(128, 2, 128)
    assert torch.equal(impl.W_UK_T, expected[:, :, :64].permute(1, 2, 0))
    assert torch.equal(impl.W_UV, expected[:, :, 64:].transpose(0, 1))
    print('MLA_BF16_PREPARATION_PASS')


if __name__ == '__main__':
    main()