import argparse, hashlib, json, sys, time
from pathlib import Path
import numpy as np
import torch
from portable_model import PortableRecognizer, Compact

R = Path(__file__).resolve().parent
sys.path.insert(0, str(R/'source'))
from recognizer import Recognizer

def load(path):
    state = torch.load(path, map_location='cpu', weights_only=True)['model']
    source = Recognizer(classes=state['head.weight'].shape[0], gradient_checkpointing=False).eval()
    source.load_state_dict(state)
    return source

def main():
    p=argparse.ArgumentParser()
    p.add_argument('--checkpoint',type=Path,default=R/'source/recognizer-095000.pt')
    p.add_argument('--out',type=Path,default=R/'models')
    p.add_argument('--coreml',action='store_true')
    p.add_argument('--onnx',action='store_true')
    p.add_argument('--hqq',action='store_true')
    p.add_argument('--blocks',default='16,32,64')
    a=p.parse_args();a.out.mkdir(parents=True,exist_ok=True)
    torch.set_num_threads(2);torch.manual_seed(18404)
    source=load(a.checkpoint);model=PortableRecognizer(source).eval()
    controls=[]
    with torch.inference_mode():
        for widths in [[1,7,64],[63,257,511],[769,1025,1537]]:
            w=torch.tensor(widths,dtype=torch.int64);x=torch.rand(len(w),1,64,((max(widths)+7)//8)*8)
            for i,n in enumerate(widths):x[i,:,:,n:]=0
            y,n=source(x,w);z,m=model(x,w)
            error=float((y-z).abs().max()); agreement=float((y.argmax(-1)==z.argmax(-1)).float().mean())
            # Near-empty synthetic crops amplify fp32 summation-order noise in SE.
            tolerance = .02 if min(widths)<8 else 2e-4
            assert error < tolerance and agreement == 1.0 and torch.equal(n,m), (widths,error,agreement)
            controls.append(dict(widths=widths,max_abs_logit_error=error,frame_argmax_agreement=agreement))
    identity=dict(checkpoint=str(a.checkpoint.resolve()),sha256=hashlib.sha256(a.checkpoint.read_bytes()).hexdigest(),
        parameters=sum(t.numel()for t in source.parameters()),lowering_controls=controls,
        torch=torch.__version__,output_contract=['tokens:int32','confidence:float32','lengths:int32'])
    (a.out/'source.json').write_text(json.dumps(identity,indent=2)+'\n')
    compact=Compact(model).eval();x=torch.zeros(1,1,64,512);w=torch.tensor([512])
    if a.onnx:
        import onnx
        from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer,DefaultWeightOnlyQuantConfig
        from onnxruntime.quantization.quant_utils import load_model_with_shape_infer
        fp=a.out/'recognizer-fp32.onnx'
        torch.onnx.export(compact,(x,w),str(fp),opset_version=18,dynamo=False,
            input_names=['pixels','widths'],output_names=['tokens','confidence','lengths'],
            dynamic_axes={'pixels':{0:'batch',3:'width'},'widths':{0:'batch'},
                          'tokens':{0:'batch',1:'frames'},'confidence':{0:'batch',1:'frames'},'lengths':{0:'batch'}})
        graph=onnx.load(fp);onnx.checker.check_model(graph)
        for block in map(int,a.blocks.split(',')):
            q=MatMulNBitsQuantizer(load_model_with_shape_infer(fp),
                algo_config=DefaultWeightOnlyQuantConfig(block_size=block,is_symmetric=True,accuracy_level=4))
            q.process();target=a.out/f'recognizer-w4-block{block}.onnx'
            q.model.save_model_to_file(str(target),use_external_data_format=False)
            graph=onnx.load(target);onnx.checker.check_model(graph)
            counts={}
            for node in graph.graph.node:counts[node.op_type]=counts.get(node.op_type,0)+1
            (target.with_suffix('.audit.json')).write_text(json.dumps(dict(bytes=target.stat().st_size,
                ops=counts,weight_bits=4,graph_input_bits=32,accuracy_level=4,retained_float_ops=['stem','depthwise convolutions','biases'],
                format='ORT MatMulNBits W4 with runtime-dependent internal precision; mixed-weight precision, not all-INT4'),indent=2)+'\n')
        if a.hqq:
            from onnxruntime.quantization.matmul_nbits_quantizer import HQQWeightOnlyQuantConfig
            q=MatMulNBitsQuantizer(load_model_with_shape_infer(fp),algo_config=HQQWeightOnlyQuantConfig(block_size=32,bits=4))
            q.process();q.model.save_model_to_file(str(a.out/'recognizer-hqq4-block32.onnx'),use_external_data_format=False)
    if a.coreml:
        import coremltools as ct
        from coremltools.optimize.coreml import OptimizationConfig,OpLinearQuantizerConfig,linear_quantize_weights
        # Dynamic width with exact mask semantics; small bounded batch for mobile.
        traced=torch.jit.trace(compact,(x,w))
        ml=ct.convert(traced,inputs=[ct.TensorType(name='pixels',shape=(ct.RangeDim(1,8,default=1),1,64,ct.RangeDim(8,8192,default=512))),
            ct.TensorType(name='widths',shape=(ct.RangeDim(1,8,default=1),),dtype=np.int32)],
            outputs=[ct.TensorType(name='tokens'),ct.TensorType(name='confidence'),ct.TensorType(name='lengths')],
            minimum_deployment_target=ct.target.iOS18,compute_precision=ct.precision.FLOAT32,
            compute_units=ct.ComputeUnit.CPU_AND_GPU)
        ml.save(str(a.out/'recognizer-fp32.mlpackage'))
        for block in map(int,a.blocks.split(',')):
            config=OptimizationConfig(global_config=OpLinearQuantizerConfig(mode='linear_symmetric',dtype='int4',
                granularity='per_block',block_size=block,weight_threshold=512))
            quant=linear_quantize_weights(ml,config=config)
            quant.save(str(a.out/f'recognizer-w4-block{block}.mlpackage'))
    print(json.dumps(identity),flush=True)

if __name__=='__main__':main()