# ---------------------------------------------------------------------
# Copyright (c) 2026 Qualcomm Technologies, Inc. and/or its subsidiaries.
# SPDX-License-Identifier: BSD-3-Clause
# ---------------------------------------------------------------------

from __future__ import annotations

import numpy as np

from qai_hub_models.utils.asset_loaders import load_image

from .app import DepthProApp
from .demo import INPUT_IMAGE_ADDRESS
from .demo import main as demo_main
from .model import DepthPro


def test_task() -> None:
    """Run torch DepthPro end-to-end on the sample fixture.

    Sanity-checks that the pipeline resolves the HF weights, produces a
    depth map at the network's native 1536x1536 unpadded to the original
    input resolution, and yields a plausible field of view (roughly the
    range Apple demos on natural imagery, 30-100 degrees).
    """
    model = DepthPro.from_pretrained()
    (_, _, height, width) = model.get_input_spec()["image"][0]
    app = DepthProApp(model, height, width)
    image = load_image(INPUT_IMAGE_ADDRESS)
    prediction = app.estimate_depth(image)

    assert prediction.depth.ndim == 2
    assert prediction.depth.shape == (image.size[1], image.size[0])
    assert np.all(np.isfinite(prediction.depth))
    assert prediction.depth.min() > 0

    assert 10.0 < prediction.field_of_view < 170.0
    assert prediction.focal_length_px > 0

    assert prediction.heatmap.size == image.size


def test_demo() -> None:
    demo_main(is_test=True)