# SPDX-License-Identifier: Apache-2.0
"""Device tests of the Python API: ``model(image)`` vs the HTTP server's /predict, and call latency.

    python -m pytest tt_moge/tests/test_api_device.py -q -s        # one chip; device_id from MOGE_DEVICE (0)

Run this file on its own (it opens the device through ``MoGeModel.from_pretrained``, the published
configuration: ETH dispatch, 1 command queue, 12x10 grid -- the p150 configuration), not in one session with the tests
that use the ``device`` fixture of conftest.py.

* equivalence: for the demo image (``media/source.png``, 1920x1080) and an 800x600 copy, every array of
  ``model(image, **opts)`` equals the server's ``npz`` response (``tt_moge.server.app.predict`` on the
  same TtMoGe object), bit for bit (``normal`` compared after the server's float16 cast).
* speed: warm repeated calls; ``model.forward(image)`` is the published ``tt(image)`` call
  (README / OPT_REPORT.md "p150 ETH-dispatch compliance 2026-10-05", chip 1: median 20.31-21.06 ms, min
  19.24-19.49 ms). Its median must not be slower than 21.06 ms x (1 + MOGE_PERF_TOL), MOGE_PERF_TOL default 0.05
  (shared host).
"""
from __future__ import annotations

import base64
import io
import json
import os
import statistics
import time

import numpy as np
import pytest
import torch
from PIL import Image

if os.environ.get("TT_VISIBLE_DEVICES", "").lower() == "none":
    pytest.skip("device test (TT_VISIBLE_DEVICES=none)", allow_module_level=True)
ttnn = pytest.importorskip("ttnn")

from tt_moge import MoGeModel  # noqa: E402

REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
DEMO = os.path.join(REPO, "media", "source.png")
PUBLISHED_CALL_MS = {"median": (20.31, 21.06), "min": (19.24, 19.49)}   # tt(image), ETH + 1 CQ, chip 1, 2026-10-05
ITERS = int(os.environ.get("MOGE_PERF_ITERS", "60"))


@pytest.fixture(scope="module")
def model():
    m = MoGeModel.from_pretrained(device_id=int(os.environ.get("MOGE_DEVICE", "0")), verbose=True)
    yield m
    m.close()


def _png_b64(img: np.ndarray) -> str:
    buf = io.BytesIO()
    Image.fromarray(img).save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


def _server(model, img: np.ndarray, **opts) -> dict:
    from tt_moge.server import app, fastpath

    saved = dict(app.STATE)
    try:
        app.STATE.clear()
        app.STATE.update(config=app.Config(), model=model._tt, fast=fastpath.enabled())
        resp = app.predict(app.PredictRequest(image=_png_b64(img), output_format="npz_binary", **opts))
    finally:
        app.STATE.clear()
        app.STATE.update(saved)
    return dict(np.load(io.BytesIO(bytes(resp.body))))


def _assert_same(out, srv):
    np.testing.assert_array_equal(out.depth, srv["depth"])
    np.testing.assert_array_equal(out.points, srv["points"])
    np.testing.assert_array_equal(out.normal.astype(np.float16), srv["normal"])
    np.testing.assert_array_equal(out.mask.astype(np.uint8), srv["mask"])
    np.testing.assert_array_equal(out.intrinsics, srv["intrinsics"])
    assert np.float32(out.metric_scale) == srv["metric_scale"]


@pytest.mark.parametrize("opts", [{}, {"fit": "stretch", "fov_x": 60.0},
                                  {"apply_mask": False, "force_projection": False}])
def test_demo_matches_server(model, opts):
    img = np.asarray(Image.open(DEMO).convert("RGB"))
    out = model(DEMO, **opts)
    srv = _server(model, img, **opts)
    _assert_same(out, srv)
    print(f"\n[api] {opts or 'defaults'}: {out!r}")


def test_resized_matches_server(model):
    img = np.asarray(Image.open(DEMO).convert("RGB").resize((800, 600), Image.BILINEAR))
    out = model(img)
    assert out.depth.shape == (600, 800)
    _assert_same(out, _server(model, img))


def test_input_types_agree(model):
    img = np.array(Image.open(DEMO).convert("RGB"))
    ref = model(DEMO)
    f = torch.from_numpy(img).permute(2, 0, 1).float() / 255           # upstream-style float (3, H, W)
    for x in (img, Image.open(DEMO), f, f[None]):
        o = model(x)
        np.testing.assert_array_equal(o.points, ref.points)
        np.testing.assert_array_equal(o.normal, ref.normal)
    many = model.predict_many([DEMO, img])
    for o in many:
        np.testing.assert_array_equal(o.depth, ref.depth)
    d = model.infer(f)
    assert torch.equal(d["depth"], torch.from_numpy(ref.depth))


def _times(fn, n):
    ts = []
    for _ in range(n):
        t0 = time.perf_counter()
        fn()
        ts.append(1e3 * (time.perf_counter() - t0))
    return statistics.median(ts), min(ts)


def test_call_latency(model):
    from tt_moge.benchmark import load_image

    img_f = load_image(DEMO).float()                          # the bench's input: float32 (1, 3, 1080, 1920)
    img_u8 = np.asarray(Image.open(DEMO).convert("RGB"))
    for _ in range(5):
        model.forward(img_f)
        model(img_u8)
    res = {
        "forward_float": _times(lambda: model.forward(img_f), ITERS),       # == tt(image) of bench_breakdown e2e
        "forward_u8": _times(lambda: model.forward(img_u8), ITERS),         # the served u8 canvas input
        "call_u8_array": _times(lambda: model(img_u8), ITERS),              # + focal/shift solve + finish
        "call_path": _times(lambda: model(DEMO), ITERS // 2),               # + PNG decode
    }
    n = 8
    t0 = time.perf_counter()
    model.predict_many([img_u8] * n)
    res["predict_many_per_image"] = (1e3 * (time.perf_counter() - t0) / n,) * 2
    for k, (med, mn) in res.items():
        print(f"\n[api perf] {k:24s} median {med:7.2f} ms  min {mn:7.2f} ms")
    print(f"[api perf] published tt(image): median {PUBLISHED_CALL_MS['median']} ms, min {PUBLISHED_CALL_MS['min']} ms")
    if os.environ.get("MOGE_API_PERF_JSON"):
        with open(os.environ["MOGE_API_PERF_JSON"], "w") as f:
            json.dump({k: {"median_ms": v[0], "min_ms": v[1]} for k, v in res.items()}
                      | {"published_tt_image_ms": PUBLISHED_CALL_MS, "iters": ITERS}, f, indent=1)
    tol = float(os.environ.get("MOGE_PERF_TOL", "0.05"))
    assert res["forward_float"][0] <= PUBLISHED_CALL_MS["median"][1] * (1 + tol), res