"""Step 4b: attempt on-device ANE compile of the converted model.
Loads dense_sam3_trackstep.mlpackage with CPU_AND_NE and runs one predict.
The ANECompiler runs at load; failures surface as load errors or CPU fallback.
"""

import sys
import time
import numpy as np
import torch

import common
import dense_wrapper as dw


def main():
    import coremltools as ct

    cache = torch.load("eager_cache.pt", weights_only=False)
    # need model only for input synthesis
    model = common.build_model()
    inputs = dw.frame_inputs(model, cache, cache["target_frame"])

    unit = getattr(ct.ComputeUnit, sys.argv[1] if len(sys.argv) > 1 else "CPU_AND_NE")
    print(f"loading with compute_units={unit} ...")
    t0 = time.time()
    mlmodel = ct.models.MLModel("dense_sam3_trackstep.mlpackage",
                                compute_units=unit)
    print(f"load OK in {time.time()-t0:.1f}s")

    in_names = [i.name for i in mlmodel.input_description._fd_spec]
    feed = {n: v.numpy().astype(np.float32) for n, v in zip(in_names, inputs)}
    state = mlmodel.make_state()
    t0 = time.time()
    out = mlmodel.predict(feed, state=state)
    print(f"predict 1 OK in {time.time()-t0:.1f}s")
    t0 = time.time()
    out = mlmodel.predict(feed, state=state)
    print(f"predict 2 OK in {time.time()-t0:.1f}s")
    print("outputs:", {k: np.asarray(v).shape for k, v in out.items()})

    # which units actually got used?
    try:
        plan = ct.models.compute_plan.MLComputePlan.load_from_path(
            mlmodel.get_compiled_model_path(), compute_units=unit
        )
        counts = {}
        prog = plan.model_structure.program
        fn = prog.functions["main"]
        for op in fn.block.operations:
            usage = plan.get_compute_device_usage_for_mlprogram_operation(op)
            if usage is None:
                continue
            dev = type(usage.preferred_compute_device).__name__
            counts[dev] = counts.get(dev, 0) + 1
        print("op placement:", counts)
    except Exception as e:
        print(f"compute plan inspection failed: {type(e).__name__}: {e}")


if __name__ == "__main__":
    main()