"""Check installed Kunlun loader's configuration and CPU parameter contracts.

This uses synthetic layers and a mocked TP rank, not a full engine startup.
"""
import copy
from types import SimpleNamespace
from unittest.mock import patch

import torch

from convert_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

    cfg = KunlunCompressedTensorsConfig.from_config(copy.deepcopy(quantization_config()))
    cfg.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
    linear = torch.nn.Linear(128, 128, bias=False)
    scheme = cfg.get_scheme(linear, "model.layers.0.self_attn.qkv_proj")
    # CPU-only parameter allocation; no distributed processes or accelerator work.
    del linear._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(linear, [128], 128, torch.bfloat16, lambda *a, **k: None)
    assert type(scheme.kernel).__name__ == "KunlunScaledMMLinearKernel"
    assert linear.weight.dtype == torch.int8
    assert tuple(linear.weight_scale.shape) == (128, 1)
    print("LINEAR_PASS", type(scheme.kernel).__name__)

    class RoutedExperts(torch.nn.Module):
        pass

    experts = RoutedExperts()
    experts.ckpt_gate_proj_name = "w1"
    experts.ckpt_down_proj_name = "w2"
    experts.ckpt_up_proj_name = "w3"
    experts.moe_config = SimpleNamespace(is_act_and_mul=True)
    method = KunlunCompressedTensorsMoEMethod.get_moe_method(
        cfg, experts, "model.layers.0.block_sparse_moe.experts")
    assert type(method).__name__ == "KunlunCompressedTensorsW8A8Int8MoEMethod"
    method.create_weights(experts, 2, 128, 128, torch.bfloat16)
    assert tuple(experts.w13_weight.shape) == (2, 256, 128)
    assert tuple(experts.w13_weight_scale.shape) == (2, 256, 1)
    assert tuple(experts.w2_weight_scale.shape) == (2, 128, 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_PASS", type(method).__name__)
    print("SCALE_PASS: checkpoint scale * 127 = runtime absmax")
    assert cfg.get_scheme(torch.nn.Linear(128, 128), "lm_head") is None
    assert cfg.get_scheme(torch.nn.Linear(128, 128), "model.layers.0.block_sparse_moe.gate") is None
    print("IGNORE_PASS: lm_head and router gate")


if __name__ == "__main__":
    main()