import argparse
from dataclasses import asdict
import importlib.metadata
import json
import os
from pathlib import Path
import runpy
import sys

os.environ['HF_HUB_OFFLINE']='1'
os.environ['TRANSFORMERS_OFFLINE']='1'
ROOT=Path('/work')
import torch
torch.set_num_threads(4)
torch.set_num_interop_threads(1)
converted=ROOT/'converted/checkpoint.pth'
converted.parent.mkdir(exist_ok=False)
sys.argv=['convert_release_checkpoint.py','--framework','rfdetr','--input',str(ROOT/'source/checkpoint_best_total.safetensors'),
          '--output',str(converted),'--class-names','person','head']
runpy.run_path(str(ROOT/'source/official/scripts/convert_release_checkpoint.py'),run_name='__main__')
from rfdetr import RFDETRSmall
model=RFDETRSmall(pretrain_weights=str(converted),device='cpu')
assert model.model.resolution==512
checkpoint=torch.load(converted,map_location='cpu',weights_only=False)['model']
state=model.model.model.state_dict()
missing=[k for k in checkpoint if k not in state]
different=[k for k in checkpoint if k in state and (checkpoint[k].shape!=state[k].shape or not torch.equal(checkpoint[k],state[k]))]
assert not missing and not different, {'missing':missing,'different':different}
info=dict(resolution=model.model.resolution,class_names=model.class_names,
          model_config=model.model_config.model_dump(mode='json'),loaded_tensors=len(checkpoint),
          strict_uploaded_tensor_match=True,versions={n:importlib.metadata.version(n) for n in ['torch','rfdetr','safetensors','onnx']})
with (ROOT/'export-contract.json').open('x') as f:json.dump(info,f,indent=2)
path=model.export(output_dir=str(ROOT/'onnx'),shape=(512,512),batch_size=1,dynamic_batch=True,
                  opset_version=17,verbose=False,format='onnx')
import onnx
m=onnx.load(str(path));onnx.checker.check_model(m)
def desc(v):
    t=v.type.tensor_type
    return dict(dtype=t.elem_type,shape=[d.dim_value if d.HasField('dim_value') else d.dim_param for d in t.shape.dim])
with (ROOT/'onnx-contract.json').open('x') as f:
    json.dump(dict(path=str(path),inputs={v.name:desc(v) for v in m.graph.input},outputs={v.name:desc(v) for v in m.graph.output}),f,indent=2)
print('ONNX_EXPORT_PASS',str(path),flush=True)