# SPDX-License-Identifier: Apache-2.0
"""First-call latency of the Python API, one variant per fresh process.

    # 1) host only: write the inputs of every variant (new real images, not the warm-up input)
    python code/bench_first_call.py prepare --out /tmp/fc_inputs [--src 'dir/*.png' ...]
    # 2) one fresh process per variant (on a device shell)
    python code/bench_first_call.py run --inputs /tmp/fc_inputs --variant jpeg_1600x900 \
        [--from-pretrained-kwargs '{"warmup_variants": ...}'] [--dump out.npz]

``run`` times ``SuperPoint.from_pretrained(...)``, then one call per image (pass 0: every call
gets a NEW image), then ``--passes`` more passes over the same images. The latency depends on the
image (keypoint count), so call k of pass 0 is compared with the median of the same image in the
later passes ("steady_same_image_ms"). It prints one JSON line: from_pretrained seconds, calls
1 / 2 / 10 with their steady reference, the steady median over all images, and every latency.
``--dump`` writes the pass-0 outputs (keypoints, scores, descriptors) to compare two runs bit for
bit; the later passes are checked against pass 0 in the process.

The process does not import Pillow / ttnn / the port before ``from_pretrained``: the inputs are
read with ``numpy.load`` or passed as file paths, so lazy imports and other host first-use costs
land where a user sees them.
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import statistics
import sys
import time
from pathlib import Path

CODE = Path(__file__).resolve().parent
SAMPLE = CODE / "sample_data" / "house_in_field_1080p.jpg"

#: variant -> (input kind, source size (W, H) | None, call kwargs). kinds: jpeg / png = file path,
#: rgb = (H, W, 3) uint8 array, plane = (H, W) uint8 array, list4 = list of 4 JPEG paths.
VARIANTS = {
    "jpeg_1600x900": ("jpeg", (1600, 900), {}),
    "rgb_1920x1080": ("rgb", (1920, 1080), {}),
    "rgb_1280x720": ("rgb", (1280, 720), {}),
    "plane_480x640": ("plane", (640, 480), {}),
    "rgb_1024x768": ("rgb", (1024, 768), {}),           # a size outside the default list
    "png_1600x900": ("png", (1600, 900), {}),
    "nms_r3": ("rgb", (1600, 900), {"nms_radius": 3}),  # per-radius device NMS variant
    "nms_r0": ("rgb", (1600, 900), {"nms_radius": 0}),  # host NMS
    "nms_r12": ("rgb", (1600, 900), {"nms_radius": 12}),
    "all_kp": ("rgb", (1600, 900), {"max_keypoints": -1}),
    "no_desc": ("rgb", (1600, 900), {"return_descriptors": False}),
    "list4_jpeg": ("list4", (1600, 900), {}),
    "rgb_8000x4500": ("rgb", (8000, 4500), {}),         # beyond the device resize: host Pillow resize
}


def _sources(globs):
    files = []
    for g in globs:
        files += sorted(glob.glob(g))
    return files or [str(SAMPLE)]


def prepare(args) -> None:
    import numpy as np
    from PIL import Image, ImageOps

    srcs = _sources(args.src)
    ims = []
    for f in srcs:
        with Image.open(f) as im:
            im = im.convert("RGB")
        ims += [im, ImageOps.mirror(im)]
    out = Path(args.out)
    sizes = {VARIANTS[v][1] for v in VARIANTS}
    for (w, h) in sorted(sizes):
        d = out / f"{w}x{h}"
        d.mkdir(parents=True, exist_ok=True)
        for i, im in enumerate(ims[:12] if w * h > 4_000_000 else ims):  # 12 large images (disk)
            r = im.resize((w, h), Image.BILINEAR)
            np.save(d / f"{i:03d}.npy", np.asarray(r))
            if (w, h) == (1600, 900):
                r.save(d / f"{i:03d}.jpg", quality=92)
                r.save(d / f"{i:03d}.png")
    (out / "sources.json").write_text(json.dumps(srcs, indent=1))
    print(f"wrote {len(ims)} images x {len(sizes)} sizes to {out}")


def run(args) -> None:
    import numpy as np

    kind, (w, h), kw = VARIANTS[args.variant]
    d = Path(args.inputs) / f"{w}x{h}"
    n_img = len(sorted(d.glob("*.npy")))
    inputs = []
    for k in range(n_img):
        if kind == "jpeg":
            inputs.append(str(d / f"{k:03d}.jpg"))
        elif kind == "png":
            inputs.append(str(d / f"{k:03d}.png"))
        elif kind == "list4":
            inputs.append([str(d / f"{(4 * k + j) % n_img:03d}.jpg") for j in range(4)])
        else:
            a = np.load(d / f"{k:03d}.npy")
            inputs.append(np.ascontiguousarray(a[..., 0]) if kind == "plane" else a)
    fp_kw = json.loads(args.from_pretrained_kwargs) if args.from_pretrained_kwargs else {}

    t0 = time.perf_counter()
    from tt_superpoint import SuperPoint

    t_imp = time.perf_counter() - t0
    t0 = time.perf_counter()
    model = SuperPoint.from_pretrained(device_id=int(os.environ.get("TT_DEVICE_ID", "0")), **fp_kw)
    t_fp = time.perf_counter() - t0

    def flat(out):
        r = []
        for o in out if isinstance(out, list) else [out]:
            r += [o.keypoints.numpy(), o.scores.numpy(), None if o.descriptors is None else o.descriptors.numpy()]
        return r

    # pass 0: every call sees a new image; passes 1..P: the same images again (steady reference)
    lat = [[0.0] * n_img for _ in range(1 + args.passes)]
    first, same = [], True
    try:
        for p in range(1 + args.passes):
            for i, x in enumerate(inputs):
                t0 = time.perf_counter()
                out = model(x, **kw)
                lat[p][i] = (time.perf_counter() - t0) * 1e3
                f = flat(out)
                if p == 0:
                    first.append(f)
                else:
                    same &= all((a is None and b is None) or np.array_equal(a, b) for a, b in zip(f, first[i]))
        cfg = dict(model.config)
    finally:
        model.close()
    if args.dump:
        dump = {}
        for i, f in enumerate(first):
            for j in range(0, len(f), 3):
                for name, a in zip(("kp", "sc", "de"), f[j:j + 3]):
                    if a is not None:
                        dump[f"{i}_{j // 3}_{name}"] = a
        np.savez(args.dump, **dump)
    ref = [statistics.median(lat[p][i] for p in range(1, 1 + args.passes)) for i in range(n_img)]
    med = statistics.median(ref)

    def c(i):  # call i+1 of pass 0 and the steady latency of the same image
        return {"ms": round(lat[0][i], 2), "steady_same_image_ms": round(ref[i], 3), "ratio": round(lat[0][i] / ref[i], 2)}

    res = {"variant": args.variant, "kwargs": kw, "from_pretrained_kwargs": fp_kw, "import_s": round(t_imp, 3),
           "from_pretrained_s": round(t_fp, 2), "call1": c(0), "call2": c(1), "call10": c(min(9, n_img - 1)),
           "steady_median_ms": round(med, 3), "repeat_outputs_identical": bool(same), "distinct_images": n_img,
           "warmup": cfg.get("warmup_s"), "lat_ms": [[round(v, 3) for v in row] for row in lat]}
    print("RESULT " + json.dumps(res))


def main(argv=None) -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("prepare")
    p.add_argument("--out", required=True)
    p.add_argument("--src", nargs="*", default=[], help="image globs (default: the demo image)")
    r = sub.add_parser("run")
    r.add_argument("--inputs", required=True)
    r.add_argument("--variant", required=True, choices=sorted(VARIANTS))
    r.add_argument("--passes", type=int, default=3, help="repeat passes over the same images (steady reference)")
    r.add_argument("--from-pretrained-kwargs", default="")
    r.add_argument("--dump", default="")
    args = ap.parse_args(argv)
    {"prepare": prepare, "run": run}[args.cmd](args)


if __name__ == "__main__":
    sys.exit(main())