"""Validate complete native parameter headers and real Q4 samples on Metal.

This is not a full 27B Mac generation test. Header fixtures are never saved.
"""
import importlib.metadata
import argparse
import json
import platform
import resource
from pathlib import Path
import mlx.core as mx
import mlx.nn as nn
import numpy as np
from mlx.utils import tree_flatten
from mlx_lm.utils import _get_classes

parser=argparse.ArgumentParser(description=__doc__)
parser.add_argument('--folder',type=Path,default=Path(__file__).resolve().parent/'mac-check')
parser.add_argument('--output',type=Path,default=Path('mac-compatibility-results.json'))
options=parser.parse_args()
folder=options.folder
assert platform.system()=='Darwin' and platform.machine()=='arm64'
assert mx.metal.is_available()
mx.set_default_device(mx.gpu)
config=json.loads((folder/'config.json').read_text())
assert config['quantization']=={'group_size':64,'bits':4,'mode':'affine'}
headers=json.loads((folder/'checkpoint-headers.json').read_text())
index=json.loads((folder/'model.safetensors.index.json').read_text())
assert set(headers)==set(index['weight_map'])
cls,args=_get_classes(config)
model=cls(args.from_dict(config))
dtype={'BF16':mx.bfloat16,'F16':mx.float16,'F32':mx.float32,'U32':mx.uint32}
fixtures={k:mx.broadcast_to(mx.array(0,dtype=dtype[v['dtype']]),v['shape']) for k,v in headers.items()}
fixtures=model.sanitize(fixtures)
nn.quantize(model,group_size=64,bits=4,mode='affine',class_predicate=lambda path,module: path+'.scales' in fixtures)
model.load_weights(list(fixtures.items()),strict=True)
assert set(dict(tree_flatten(model.parameters())))==set(headers)
assert len(model.weight_mapping())==1199
del fixtures,model

samples=mx.load(str(folder/'real-checkpoint-samples.safetensors'))
description=json.loads((folder/'samples.json').read_text())
results=[]
for item in description['samples']:
    prefix=f"sample_{item['sample']}."
    x,w,s,b,ref=[samples[prefix+k] for k in ('input','weight','scales','biases','fp32_reference')]
    assert w.dtype==mx.uint32 and s.dtype==mx.bfloat16 and b.dtype==mx.bfloat16
    native=mx.quantized_matmul(x,w,s,b,transpose=True,group_size=64,bits=4,mode='affine')
    fp32=mx.quantized_matmul(x.astype(mx.float32),w,s.astype(mx.float32),b.astype(mx.float32),transpose=True,group_size=64,bits=4,mode='affine')
    mx.eval(native,fp32,ref)
    a,r,f=[np.array(value.astype(mx.float32)) for value in (native,ref,fp32)]
    np.testing.assert_allclose(a,r,atol=0.01,rtol=0.02)
    np.testing.assert_allclose(f,r,atol=1e-4,rtol=5e-4)
    results.append(dict(item,native_metal_bf16='PASS',metal_fp32='PASS',bf16_max_absolute_error=float(np.max(np.abs(a-r))),fp32_max_absolute_error=float(np.max(np.abs(f-r)))))
result={'status':'PASS_MAC_NATIVE_MLX_FORMAT_AND_REAL_METAL_SAMPLES','platform':platform.platform(),'machine':platform.machine(),'mlx':importlib.metadata.version('mlx'),'mlx_lm':importlib.metadata.version('mlx-lm'),'device':str(mx.default_device()),'quantization':config['quantization'],'all_saved_tensor_headers_validated':len(headers),'all_source_parameters_accounted_for':1199,'strict_complete_parameter_tree':'PASS using unevaluated header fixtures; no fabricated weights saved','actual_checkpoint_samples':results,'full_model_mac_generation':'NOT_RUN: targeted native Metal format and real-weight component checks only','peak_mlx_bytes':mx.get_peak_memory(),'peak_process_rss_bytes':resource.getrusage(resource.RUSAGE_SELF).ru_maxrss}
with options.output.open('x') as stream:json.dump(result,stream,indent=2)
print(json.dumps(result,indent=2))