from collections import Counter, defaultdict
import hashlib, json
from pathlib import Path

def verify(model, path, *, peft_prefix=False):
    import torch
    from safetensors import safe_open
    path = Path(path)
    mapping = json.loads((path / 'model.safetensors.index.json').read_text())['weight_map']
    wanted = {name for name in mapping if name.endswith(('.A_log', '.dt_bias'))}
    assert wanted
    actual = {}
    for name, p in model.named_parameters():
        if not name.endswith(('.A_log', '.dt_bias')):
            continue
        if peft_prefix:
            assert name.startswith('base_model.model.')
            name = name[len('base_model.model.'):]
        assert name not in actual
        actual[name] = p
    assert set(actual) == wanted, {'missing': sorted(wanted - set(actual)), 'extra': sorted(set(actual) - wanted)}
    groups = defaultdict(list)
    for name in wanted:
        groups[mapping[name]].append(name)
    evidence = []
    for filename, names in sorted(groups.items()):
        with safe_open(str(path / filename), framework='pt', device='cpu') as source:
            for name in sorted(names):
                saved = source.get_tensor(name)
                p = actual[name]
                assert p.dtype == torch.float32, (name, str(p.dtype), 'must preserve FP32 before loading')
                assert p.device.type == 'cuda', (name, 'must verify final device placement')
                value = p.detach().cpu()
                assert torch.equal(value, saved.float()), (name, 'checkpoint bits changed; upcasting rounded BF16 is insufficient')
                evidence.append(dict(name=name, elements=p.numel(), source_dtype=str(saved.dtype), runtime_dtype=str(p.dtype), device=str(p.device), fp32_value_sha256=hashlib.sha256(value.contiguous().view(torch.uint8).numpy().tobytes()).hexdigest(), saved_value_exact=True))
    return dict(status='PASS_NATIVE_GDN_FP32_CHECKPOINT_VALUES', tensors=len(evidence), source_dtypes=dict(Counter((r['source_dtype'] for r in evidence))), all_runtime_fp32=True, all_values_equal_saved_checkpoint=True, checked_after_final_device_placement=True, values=evidence)