"""Low-RAM CoreML runtime parity: load a saved mlpackage, drive it statefully
frames 1..16 from zero state, compare CoreML outputs to the CACHED eager
targets (warmup f3, steady f16). No torch-eager forward in the loop and no
checkpoint load (synth inputs are weight-independent), so peak RAM stays well
under the 16GB that the in-process double-run in export_fp32.py blows past.

Usage: python verify_coreml_lean.py <mlpackage> [compute_units]
  compute_units: CPU_ONLY (default) | CPU_AND_GPU
"""

import sys
import numpy as np
import torch

import common
import dense_wrapper as dw


def main():
    pkg = sys.argv[1] if len(sys.argv) > 1 else "dense_sam3_trackstep_fp32.mlpackage"
    cu = sys.argv[2] if len(sys.argv) > 2 else "CPU_ONLY"

    cache = torch.load("eager_cache.pt", weights_only=False)
    # Random-weight model ONLY for weight-independent input generation
    # (synth features + embedding size). Avoids a second 457-tensor load.
    model = common.build_model(load_checkpoint=False)

    import coremltools as ct
    mlmodel = ct.models.MLModel(pkg, compute_units=getattr(ct.ComputeUnit, cu))
    in_names = [i.name for i in mlmodel.input_description._fd_spec]
    state = mlmodel.make_state()

    def rel(c, t):
        t = t.float()
        c = c.float()
        return ((c - t).abs().max() / t.abs().max().clamp_min(1e-9)).item()

    checks = {3: cache["warmup"], 16: cache["target"]}
    worst = 0.0
    appear_flips = 0
    print(f"compute_units={cu}  pkg={pkg}")
    for f in range(1, 17):
        fi = dw.frame_inputs(model, cache, f)
        feed = {n: v.numpy().astype(np.float32) for n, v in zip(in_names, fi)}
        got = mlmodel.predict(feed, state=state)
        if f in checks:
            ref = checks[f]
            by_shape = {tuple(np.asarray(v).shape): torch.from_numpy(
                np.asarray(v)).float() for v in got.values()}
            t_low = ref["pred_masks"]
            t_osl = ref["object_score_logits"]
            t_iou = ref["ious"]
            c_low = by_shape[tuple(t_low.shape)]
            c_osl = by_shape[tuple(t_osl.shape)]
            c_iou = by_shape[tuple(t_iou.shape)]
            r_low, r_osl, r_iou = rel(c_low, t_low), rel(c_osl, t_osl), rel(c_iou, t_iou)
            sign = ((c_low > 0) == (t_low > 0)).float().mean().item()
            flip = not torch.equal((c_osl > 0), (t_osl > 0))
            appear_flips += int(flip)
            worst = max(worst, r_low, r_osl, r_iou)
            tag = "warmup" if f == 3 else "steady"
            print(f"frame {f:>2} ({tag}): low_rel={r_low:.2e} osl_rel={r_osl:.2e} "
                  f"iou_rel={r_iou:.2e} sign_agree={sign:.4f} appear_flip={flip}")
    thr = 5e-2 if "fp16" in pkg.lower() or cu == "CPU_AND_GPU" else 1e-3
    print(f"worst rel: {worst:.3e}  appear_flips: {appear_flips}  thr: {thr:.0e}")
    print("COREML RUNTIME:", "PASS" if worst < thr and appear_flips == 0 else "CHECK")


if __name__ == "__main__":
    main()