import sys
from pathlib import Path
from unittest.mock import patch
import numpy as np

SOURCE_ROOT = Path(__file__).parents[1] / "src"

sys.path.insert(0, str(SOURCE_ROOT))


from anyimage.server.models import moge2 as moge2_model
from anyimage.server.models.geometry import GeometryFrame


def test_geometry_frame_requires_matching_float32_fields():
    with np.testing.assert_raises_regex(ValueError, "validity"):
        GeometryFrame(
            depth=np.ones((2, 2), dtype=np.float32),
            validity=np.ones((1, 2), dtype=np.float32),
            intrinsics=np.eye(3, dtype=np.float32),
        )
    with np.testing.assert_raises_regex(ValueError, "normal"):
        GeometryFrame(
            depth=np.ones((2, 2), dtype=np.float32),
            validity=np.ones((2, 2), dtype=np.float32),
            intrinsics=np.eye(3, dtype=np.float32),
            normal=np.ones((2, 2, 2), dtype=np.float32),
        )


def test_moge_infer_returns_geometry_frames_with_model_validity():
    mask = np.asarray(((True, False), (True, True)))
    prediction = {
        "depth": np.ones((2, 2), dtype=np.float32),
        "normal": np.zeros((2, 2, 3), dtype=np.float32),
        "mask": mask,
        "intrinsics": np.asarray(
            ((1.0, 0.0, 0.5), (0.0, 1.0, 0.5), (0.0, 0.0, 1.0)),
            dtype=np.float32,
        ),
        "points": np.ones((2, 2, 3), dtype=np.float32),
    }

    with patch.object(moge2_model, "infer_one", return_value=prediction):
        frames = moge2_model.infer(None, {}, ["image"], include_points=True)

    assert len(frames) == 1
    frame = frames[0]
    assert np.array_equal(frame.validity, mask.astype(np.float32))
    assert np.array_equal(frame.normal, prediction["normal"])
    assert np.array_equal(frame.points, prediction["points"])
    assert np.array_equal(
        frame.intrinsics,
        np.asarray(((2.0, 0.0, 1.0), (0.0, 2.0, 1.0), (0.0, 0.0, 1.0))),
    )


def test_moge_z_depth_frame_does_not_require_points():
    prediction = {
        "depth": np.ones((2, 2), dtype=np.float32),
        "normal": np.zeros((2, 2, 3), dtype=np.float32),
        "mask": np.ones((2, 2), dtype=bool),
        "intrinsics": np.eye(3, dtype=np.float32),
        "points": np.ones((2, 2, 3), dtype=np.float32),
    }

    with patch.object(moge2_model, "infer_one", return_value=prediction):
        frame = moge2_model.infer(None, {}, ["image"], include_points=False)[0]

    assert frame.points is None
    assert frame.validity.all()