"""Convert-only fp16 export of DenseTrackStep (real SAM3.1 weights) and save.
Splits convert from verify so the RAM-heavy predict loop doesn't OOM the
convert on a 16GB host. Verify separately with verify_coreml_lean.py.
"""

import time
import torch

import common
import dense_wrapper as dw


def main():
    cache = torch.load("eager_cache.pt", weights_only=False)
    wrapper, model = dw.build_wrapper()
    inputs = dw.frame_inputs(model, cache, 3)  # mid frame exercises cond-tpos path

    common.hide_triton_stub()
    with torch.no_grad():
        try:
            ep = torch.export.export(wrapper, inputs)
        except Exception as e:
            print(f"strict export failed ({type(e).__name__}); retry strict=False")
            ep = torch.export.export(wrapper, inputs, strict=False)
        ep = ep.run_decompositions({})
    print("torch.export OK")

    import coremltools as ct
    from coremltools.converters.mil.frontend.torch.torch_op_registry import (
        register_torch_op,
    )
    from coremltools.converters.mil.frontend.torch.ops import _get_inputs
    from coremltools.converters.mil.mil import Builder as mb

    @register_torch_op(torch_alias=["where.scalarother"])
    def where_scalarother(context, node):
        cond, a, b = _get_inputs(context=context, node=node, expected=3)
        context.add(mb.select(cond=cond, a=a, b=b), node.name)

    t0 = time.time()
    mlmodel = ct.convert(
        ep,
        minimum_deployment_target=ct.target.iOS18,
        compute_units=ct.ComputeUnit.CPU_ONLY,  # convert/load only; deploy sets CU
    )
    print(f"fp16 convert OK in {time.time()-t0:.1f}s")
    mlmodel.save("dense_sam3_trackstep.mlpackage")
    print("saved dense_sam3_trackstep.mlpackage")


if __name__ == "__main__":
    main()