"""Numerical check on real layer-0 attention and one expert; does not save a model."""
import argparse
import json
from pathlib import Path

import torch
from safetensors import safe_open

from convert_fp8_to_int8 import quantize_weight


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.{p}_proj.weight" for p in "qkvo"]
    names += [f"model.layers.0.block_sparse_moe.experts.0.w{i}.weight" for i in (1, 2, 3)]
    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():
            q, scales = quantize_weight(weight, block_scale)
            # Independent reference via explicit block expansion.
            reference = weight.float() * block_scale.repeat_interleave(128, 0).repeat_interleave(128, 1)[:weight.shape[0], :weight.shape[1]]
            error = q.float() * scales - reference
            bound = scales / 2 + scales * 1e-4 + 1e-7
            assert torch.all(error.abs() <= bound), name
            relative_rmse = (error.square().mean() / reference.square().mean()).sqrt().item()
            report.append({"name": name, "shape": list(weight.shape),
                           "relative_rmse": relative_rmse,
                           "max_absolute_error": error.abs().max().item(),
                           "rounding_error_bound": "passed"})
            print(json.dumps(report[-1]), flush=True)
    args.report.write_text(json.dumps(report, indent=2) + "\n")


if __name__ == "__main__":
    main()