import collections,json,pathlib
import torch


def audit_loaded_parameters(headers, params, diagnostics_path):
    counts=collections.Counter();validated_scales=0;native_scales=0;bitcast_scales=0;checked=0
    result={'passed':False}
    try:
        assert set(params)==set(headers), (set(params)-set(headers),set(headers)-set(params))
        dtype_map={'BF16':torch.bfloat16,'F32':torch.float32,'F8_E4M3':torch.float8_e4m3fn,'U8':torch.uint8}
        for name,h in headers.items():
            t=params[name]
            assert list(t.shape)==h['shape'],(name,list(t.shape),h['shape'])
            is_scale=h['dtype']=='U8' and name.endswith('.weight_scale_inv')
            if is_scale:
                assert t.dtype in (torch.uint8,torch.float8_e8m0fnu),(name,t.dtype)
                validated_scales+=1
                native_scales+=int(t.dtype==torch.uint8)
                bitcast_scales+=int(t.dtype==torch.float8_e8m0fnu)
            else:
                assert t.dtype==dtype_map[h['dtype']],(name,t.dtype,dtype_map[h['dtype']])
            counts[str(t.dtype)]+=1;checked+=1
        assert len(headers)==19434,len(headers)
        assert counts['torch.float8_e4m3fn']==9216,dict(counts)
        assert validated_scales==9216,(validated_scales,native_scales,bitcast_scales)
        assert native_scales+bitcast_scales==validated_scales
        result['passed']=True
    except BaseException as exc:
        result['error']=repr(exc)
        raise
    finally:
        result.update(checked_tensors=checked,dtype_counts=dict(counts),validated_scale_tensors=validated_scales,native_U8_scale_tensors=native_scales,logical_U8_to_E8M0_scales=bitcast_scales)
        pathlib.Path(diagnostics_path).write_text(json.dumps(result,indent=2)+chr(10))
    return result


def normalize_loading_info(info):
    # Preserve loader sets as sorted lists before report/progress serialization.
    def convert(value):
        if isinstance(value, dict):
            return {key:convert(item) for key,item in value.items()}
        if isinstance(value, (set,frozenset)):
            return sorted(convert(item) for item in value)
        if isinstance(value, (list,tuple)):
            return [convert(item) for item in value]
        if value is None or isinstance(value,(str,int,float,bool)):
            return value
        raise TypeError(f'Unsupported loading metadata type: {type(value).__name__}')
    return convert(info)