"""Linux CPU exporter; use an isolated venv, never the training environment."""
import os
os.environ.update(CUDA_VISIBLE_DEVICES='',OMP_NUM_THREADS='2',OPENBLAS_NUM_THREADS='1',TF_NUM_INTRAOP_THREADS='2',TF_NUM_INTEROP_THREADS='1')
import argparse,json,time,sys
from pathlib import Path
import torch
import litert_torch
from ai_edge_quantizer import quantizer,recipe
from build_recognizer import load,R
from portable_model import PortableRecognizer,Compact

def main():
    parser=argparse.ArgumentParser();parser.add_argument('--detector',action='store_true');parser.add_argument('--checkpoint',type=Path);args=parser.parse_args()
    torch.set_num_threads(2);out=R/'litert';out.mkdir(exist_ok=True)
    if args.detector:
        sys.path.append('/home/rezo/bina02-surya-layout-env/lib/python3.10/site-packages')
        from export_detector import EfficientViTConfig,EfficientViTForSemanticSegmentation,BASE,S,Model
        m=EfficientViTForSemanticSegmentation(EfficientViTConfig.from_pretrained(str(BASE))).eval()
        m.load_state_dict(torch.load(args.checkpoint or S/'run/checkpoints/step-001000.pt',map_location='cpu',weights_only=True)['model'])
        converter=litert_torch.signature('detector',Model(m).eval(),(torch.zeros(1,3,1200,1200),))
        prefix='detector'
    else:
        model=Compact(PortableRecognizer(load(args.checkpoint or R/'source/recognizer-095000.pt'))).eval()
        model.model.index_dtype=torch.int32
        converter=None
        for width in [512,1024,2048,4096,8192]:
            x=torch.zeros(1,1,64,width);w=torch.tensor([width])
            if converter is None:converter=litert_torch.signature('width'+str(width),model,(x,w))
            else:converter.signature('width'+str(width),model,(x,w))
        prefix='recognizer'
    with torch.no_grad():edge=converter.convert()
    fp=out/(prefix+'-fp32.tflite');edge.export(str(fp))
    results=[]
    for name,rp in [('weight-only4',recipe.weight_only_wi4_afp32()),('dynamic4-block32',recipe.dynamic_wi4b32_afp32())]:
        t=time.monotonic()
        try:
            q=quantizer.Quantizer(str(fp));q.load_quantization_recipe(rp)
            target=out/f'{prefix}-{name}.tflite';q.quantize().export_model(str(target))
            results.append(dict(name=name,path=str(target),bytes=target.stat().st_size,recipe=rp,seconds=time.monotonic()-t))
        except Exception as e:results.append(dict(name=name,error=repr(e)))
        (out/(prefix+'-build.json')).write_text(json.dumps(dict(converter='litert-torch 0.10.0',quantizer='ai-edge-quantizer 0.10.0',
            input='Five static width signatures512/1024/2048/4096/8192; true widths input; zero right pad, never truncate',results=results),indent=2,default=str)+'\n')
    print(json.dumps(results,default=str),flush=True)

if __name__=='__main__':main()