from pathlib import Path

import cv2
import h5py
import numpy as np
import pytest
import torch

from src.demo.infer_single_image import _read_depth_array, load_depth


def test_prompt_depth_loads_plain_npy_and_npz(tmp_path: Path) -> None:
    expected = np.arange(12, dtype=np.float32).reshape(3, 4)
    npy_path = tmp_path / "depth.npy"
    npz_path = tmp_path / "depth.npz"
    np.save(npy_path, expected)
    np.savez(npz_path, depth=expected)

    np.testing.assert_array_equal(_read_depth_array(npy_path), expected)
    np.testing.assert_array_equal(_read_depth_array(npz_path), expected)


def test_prompt_depth_loads_sparse_npz(tmp_path: Path) -> None:
    depth_path = tmp_path / "depth.npz"
    mask = np.array([[False, True], [True, False]])
    np.savez(depth_path, mask=mask, value=np.array([2.0, 4.0], dtype=np.float32))

    actual = _read_depth_array(depth_path)

    np.testing.assert_array_equal(
        actual,
        np.array([[0.0, 2.0], [4.0, 0.0]], dtype=np.float32),
    )


def test_prompt_depth_loads_channel_first_npy(tmp_path: Path) -> None:
    depth_path = tmp_path / "depth.npy"
    expected = np.arange(12, dtype=np.float32).reshape(3, 4)
    np.save(depth_path, expected[None])

    np.testing.assert_array_equal(_read_depth_array(depth_path), expected)


def test_prompt_depth_rejects_png(tmp_path: Path) -> None:
    depth_path = tmp_path / "depth.png"
    depth_path.touch()

    with pytest.raises(ValueError, match="Unsupported prompt depth extension"):
        _read_depth_array(depth_path)


def test_prompt_depth_loads_nested_hdf5(tmp_path: Path) -> None:
    depth_path = tmp_path / "depth.h5"
    expected = np.arange(12, dtype=np.float32).reshape(3, 4)
    with h5py.File(depth_path, "w") as h5_file:
        h5_file.create_dataset("nested/depth", data=expected)

    np.testing.assert_array_equal(_read_depth_array(depth_path), expected)


def test_prompt_depth_loads_exr(tmp_path: Path) -> None:
    depth_path = tmp_path / "depth.exr"
    expected = np.arange(12, dtype=np.float32).reshape(3, 4)
    assert cv2.imwrite(str(depth_path), expected)

    np.testing.assert_array_equal(_read_depth_array(depth_path), expected)


def test_load_depth_preserves_default_sparse_sampling(tmp_path: Path) -> None:
    depth_path = tmp_path / "depth.npy"
    np.save(depth_path, np.full((40, 50), 2.0, dtype=np.float32))
    np.random.seed(0)

    dense_depth, sampled_depth, depth_mask = load_depth(
        depth_path=depth_path,
        tar_size=(40, 50),
    )

    assert dense_depth.shape == (1, 1, 40, 50)
    assert depth_mask.sum().item() == 2000
    assert torch.count_nonzero(sampled_depth).item() == 1500