"""FLOAT32 CoreML convert + trajectory parity from zero state (warmup→steady)."""

import time
import numpy as np
import torch

import common
import dense_wrapper as dw


def zero_state(wrapper):
    wrapper.mem_bank.zero_()
    wrapper.img_bank.zero_()
    wrapper.ptr_bank.zero_()
    wrapper.mem_valid.zero_()
    wrapper.ptr_valid.zero_()


def main():
    cache = torch.load("eager_cache.pt", weights_only=False)
    wrapper, model = dw.build_wrapper()
    # Export at a mid frame so frame_pos path (dynamic cond tpos) is exercised
    inputs = dw.frame_inputs(model, cache, 3)

    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__}: {e}); retrying strict=False")
            ep = torch.export.export(wrapper, inputs, strict=False)
        ep = ep.run_decompositions({})
    print("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,
        # Convert/load on CPU; GPU JIT of masked SDPA previously ballooned RAM.
        # Deploy path still CPU_AND_GPU via MLModel(..., compute_units=...).
        compute_units=ct.ComputeUnit.CPU_ONLY,
        compute_precision=ct.precision.FLOAT32,
    )
    print(f"fp32 convert OK in {time.time()-t0:.1f}s")
    mlmodel.save("dense_sam3_trackstep_fp32.mlpackage")

    zero_state(wrapper)
    in_names = [i.name for i in mlmodel.input_description._fd_spec]
    state = mlmodel.make_state()
    print(f"{'frame':>5} {'osl_rel':>9} {'iou_rel':>9} {'sign_agree':>10} "
          f"{'lowmask_rel':>11} {'t_nan':>5} {'c_nan':>5}")
    worst = 0.0
    for f in range(1, 11):  # warmup→near-steady; enough to prove state+mask
        fi = dw.frame_inputs(model, cache, f)
        with torch.no_grad():
            t_low, t_high, t_osl, t_ious = [x.clone() for x in wrapper(*fi)]
        feed = {n: v.numpy().astype(np.float32) for n, v in zip(in_names, fi)}
        got = mlmodel.predict(feed, state=state)
        by_shape = {tuple(np.asarray(v).shape): torch.from_numpy(
            np.asarray(v)).float() for v in got.values()}
        c_low = by_shape[tuple(t_low.shape)]
        c_osl = by_shape[tuple(t_osl.shape)]
        c_ious = by_shape[tuple(t_ious.shape)]

        def rel(a, b):
            return ((a - b).abs().max() / b.abs().max().clamp_min(1e-9)).item()

        sign = ((c_low > 0) == (t_low > 0)).float().mean().item()
        r_osl, r_iou, r_low = rel(c_osl, t_osl), rel(c_ious, t_ious), rel(c_low, t_low)
        worst = max(worst, r_osl, r_iou, r_low)
        print(f"{f:>5} {r_osl:>9.2e} {r_iou:>9.2e} "
              f"{sign:>10.4f} {r_low:>11.2e} "
              f"{int(t_low.isnan().sum()):>5} {int(c_low.isnan().sum()):>5}")
    print(f"worst rel: {worst:.3e}")
    print("FP32 TRAJ:", "PASS" if worst < 1e-4 else "CHECK")


if __name__ == "__main__":
    main()