import json
from types import SimpleNamespace
import numpy as np


from anyimage.common.selection import SelectionMask
from anyimage.common.depth import (
    depth_uniform_scale,
    load_depth_metadata,
)
from anyimage.operators.cutout_tool import geometry as cutout_geometry


def test_depth_metadata_loader_uses_one_schema_for_plane_and_cutout(tmp_path):
    path = tmp_path / "depth.json"
    path.write_text(
        json.dumps(
            {
                "image_size": [3, 2],
                "intrinsics": np.eye(3).tolist(),
                "reference_depth": 1.0,
            }
        ),
        encoding="utf-8",
    )

    metadata = load_depth_metadata(path)

    assert metadata["image_size"] == (3, 2)
    assert metadata["intrinsics"] == ((1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0))
    assert metadata["reference_depth"] == 1.0


def test_base_plane_fit_is_stable_for_a_slender_diagonal_domain():
    along = np.linspace(-1.0, 1.0, 101)
    across = np.resize(np.asarray((-0.01, 0.0, 0.01)), len(along))
    coordinates = np.column_stack(
        (
            along - across,
            along + across,
        )
    )
    values = 0.4 * along + np.sign(across) * 0.05

    plane = cutout_geometry._stable_plane(coordinates, values)
    selected_plane = coordinates @ plane[:2] + plane[2]
    corners = np.asarray(((-1.0, -1.0), (-1.0, 1.0), (1.0, -1.0), (1.0, 1.0)))
    corner_plane = corners @ plane[:2] + plane[2]

    assert np.ptp(corner_plane) <= np.ptp(selected_plane) * 1.25
    assert np.corrcoef(selected_plane, along)[0, 1] > 0.99


def test_depth_balloon_plane_calibration_matches_cutout_axes():
    vertices = []
    uv = []
    camera_depth = np.empty((3, 3), dtype=np.float32)
    for row in range(3):
        local_y = 1.0 - row
        for column in range(3):
            local_x = column - 1.0
            vertices.append((local_x, local_y, 0.0))
            uv.append((column * 0.5, 1.0 - row * 0.5))
            camera_depth[row, column] = 2.0 + 0.2 * local_x + 0.1 * local_y

    reference, direction = cutout_geometry.depth_balloon_projection_parameters(
        vertices,
        uv,
        camera_depth,
        np.ones((3, 3), dtype=bool),
        0.5,
    )
    expected_direction = np.asarray((0.1, 0.05, 1.0), dtype=np.float64)
    expected_direction /= np.linalg.norm(expected_direction)

    assert np.isclose(
        reference,
        np.percentile(camera_depth, 95.0) * 0.5,
        atol=1e-6,
    )
    assert np.allclose(direction, expected_direction, atol=1e-6)


def test_depth_balloon_reference_uses_the_shared_upper_depth_percentile():
    size = 10
    row, column = np.indices((size, size))
    camera_depth = np.full((size, size), 2.0, dtype=np.float32)
    camera_depth[(row + column) % 2 == 0] = 2.2
    vertices = []
    uv = []
    for y in range(size):
        for x in range(size):
            vertices.append((float(x), float(size - 1 - y), 0.0))
            uv.append((x / (size - 1), 1.0 - y / (size - 1)))

    reference, direction = cutout_geometry.depth_balloon_projection_parameters(
        vertices,
        uv,
        camera_depth,
        np.ones((size, size), dtype=bool),
        1.0,
    )

    assert reference > float(np.median(camera_depth))
    assert np.isclose(reference, 2.2, atol=1e-6)
    assert np.allclose(direction, (0.0, 0.0, 1.0), atol=1e-6)


def test_depth_surface_projection_uses_one_scale_for_all_camera_axes():
    points = np.zeros((51, 101, 3), dtype=np.float32)
    points[..., 2] = np.linspace(1.0, 3.0, 51 * 101).reshape(51, 101)
    metadata = {
        "intrinsics": (
            (100.0, 0.0, 50.0),
            (0.0, 100.0, 25.0),
            (0.0, 0.0, 1.0),
        ),
        "image_size": (101, 51),
        "reference_depth": 2.0,
    }

    uniform_scale = depth_uniform_scale(
        SimpleNamespace(),
        metadata,
        bounds=(-2.0, 2.0, -1.0, 1.0),
    )

    assert uniform_scale == 2.0
    reference_depth = cutout_geometry.metric_depth_reference(
        points[..., 2],
        np.ones((51, 101), dtype=bool),
        uniform_scale,
    )
    assert np.isclose(reference_depth, 5.8, atol=1e-6)


def test_cutout_threshold_controls_contour_and_calibration():
    alpha = np.zeros((12, 16), dtype=np.float32)
    alpha[2:10, 2:14] = 0.2
    alpha[4:8, 5:11] = 1.0
    for threshold, lower, upper in [
        (0.0, (2, 2), (14, 10)),
        (0.1, (2, 2), (14, 10)),
        (0.5, (5, 4), (11, 8)),
        (1.0, (5, 4), (11, 8)),
    ]:
        contours = cutout_geometry.extract_cutout_contours(
            alpha, 64, threshold=threshold
        )
        contour = np.asarray(contours.components[0][0])
        np.testing.assert_allclose(contour.min(axis=0), lower)
        np.testing.assert_allclose(contour.max(axis=0), upper)
        valid, _ = cutout_geometry.selection_calibration(
            alpha, alpha.shape, threshold=threshold
        )
        np.testing.assert_array_equal(valid, (alpha > 0) & (alpha >= threshold))


def test_cutout_depth_calibration_uses_five_percent_transparency_threshold():
    assert not cutout_geometry.use_adaptive_depth_calibration(0.0499)
    assert cutout_geometry.use_adaptive_depth_calibration(0.05)


def test_selection_calibration_uses_the_final_selection_mask():
    selection_mask = SelectionMask(
        np.asarray(((1.0, 1.0), (1.0, 0.0)), dtype=np.float32),
        (0, 0, 2, 2),
    )

    valid, transparent_fraction = cutout_geometry.selection_calibration(
        selection_mask.values,
        (2, 2),
    )

    assert np.array_equal(valid, ((True, True), (True, False)))
    assert transparent_fraction == 0.25