"""Check real DeepSeek layer-0 MLA/Dense and layer-3 shared/routed experts."""
import argparse
import json
from pathlib import Path

import torch
from safetensors import safe_open
from convert_deepseek_fp8_to_int8 import quantize_weight, dequantize_bf16, is_bf16_projection


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--source', required=True, type=Path)
    parser.add_argument('--report', required=True, type=Path)
    args = parser.parse_args()
    torch.set_num_threads(8)
    mapping = json.loads((args.source / 'model.safetensors.index.json').read_text())['weight_map']
    names = [f'model.layers.0.self_attn.{name}.weight' for name in
             ['q_a_proj', 'q_b_proj', 'kv_a_proj_with_mqa', 'kv_b_proj', 'o_proj']]
    names += ['model.layers.0.mlp.gate_proj.weight',
              'model.layers.3.mlp.shared_experts.gate_proj.weight',
              'model.layers.3.mlp.experts.0.gate_proj.weight',
              'model.layers.3.mlp.experts.0.down_proj.weight']
    report = []
    for name in names:
        scale_name = name.removesuffix('weight') + 'weight_scale_inv'
        with safe_open(str(args.source / mapping[name]), framework='pt') as reader:
            weight = reader.get_tensor(name)
        with safe_open(str(args.source / mapping[scale_name]), framework='pt') as reader:
            block_scale = reader.get_tensor(scale_name)
        with torch.inference_mode():
            reference = weight.float() * block_scale.repeat_interleave(128, 0).repeat_interleave(128, 1)[:weight.shape[0], :weight.shape[1]]
            if is_bf16_projection(name):
                value = dequantize_bf16(weight, block_scale)
                assert torch.equal(value, reference.bfloat16())
                error = value.float() - reference
                check = 'exact BF16 rounding'
            else:
                q, scales = quantize_weight(weight, block_scale)
                error = q.float() * scales - reference
                assert torch.all(error.abs() <= scales / 2 + scales * 1e-4 + 1e-7), name
                check = 'half-scale rounding bound'
            item = {'name': name, 'shape': list(weight.shape), 'check': check,
                    'relative_rmse': (error.square().mean() / reference.square().mean()).sqrt().item()}
            report.append(item)
            print(json.dumps(item), flush=True)
    args.report.write_text(json.dumps(report, indent=2) + '\n')


if __name__ == '__main__':
    main()