"""Evaluate complete-image projection and rectangular surface topology."""

from collections import Counter
import bpy
import numpy as np
import pytest

from tests.support.planes import surface, relief, evaluated


@pytest.mark.parametrize("subdivide", [0, 1, 3])
def test_zero_depth_matches_plane_geometry_and_uv(surface, subdivide):
    obj, set_value, plane, inputs, _ = surface
    set_value("Subdivide", subdivide)
    projected = evaluated(obj)
    modifier = obj.modifiers[0]
    modifier.node_group = plane
    for item in plane.interface.items_tree:
        if item.item_type == "SOCKET" and item.in_out == "INPUT":
            if item.name == "Subdivide":
                modifier[item.identifier] = subdivide
            elif item.name == "Thickness":
                modifier[item.identifier] = 0.0
    obj.update_tag(refresh={"DATA"})
    bpy.context.view_layer.update()
    regular = evaluated(obj)
    np.testing.assert_allclose(projected[0], regular[0], atol=1e-6)
    assert projected[1] == regular[1]
    np.testing.assert_allclose(projected[2], regular[2], atol=1e-6)
    assert projected[3:5] == regular[3:5]


@pytest.mark.parametrize("strength", [0, 0.5, 1, 2])
def test_camera_projection_and_reference_depth(surface, strength):
    obj, set_value, _, _, _ = surface
    base = evaluated(obj)
    uv = (base[0][:, :2] + (2, 1)) / (4, 2)
    texel = np.clip(uv * (16, 8), (0.5, 0.5), (15.5, 7.5))
    camera = np.column_stack(((texel[:, 0] - 5) / 8,
                              (texel[:, 1] - 3) / 8,
                              1.2 - (2 + texel[:, 0] / 16) * 0.5))
    expected = base[0] * (1 - strength) + camera * strength
    set_value("Depth Scale", float(strength))
    projected = evaluated(obj)
    np.testing.assert_allclose(projected[0], expected, atol=1e-6)
    assert projected[1] == base[1]
    np.testing.assert_allclose(projected[2], base[2])
    set_value("Reference Depth", 1.7)
    shifted = evaluated(obj)[0]
    np.testing.assert_allclose(shifted - projected[0],
                               np.tile((0, 0, 0.5 * strength), (len(shifted), 1)), atol=1e-6)


@pytest.mark.parametrize("thickness", [0.000001, 0.2, 2.0])
def test_shell_keeps_front_and_closes_rectangular_boundary(surface, thickness):
    obj, set_value, _, _, _ = surface
    set_value("Depth Scale", 1.0)
    front = evaluated(obj)[0]
    set_value("Thickness", thickness)
    points, faces, uv, materials, indices, _ = evaluated(obj)
    assert len(points) == 2 * len(front)
    edges = Counter(tuple(sorted((a, b))) for f in faces for a, b in zip(f, (*f[1:], f[0])))
    assert set(edges.values()) == {2}
    normal = np.array((0.25, 0, 1.0))
    normal /= np.linalg.norm(normal)
    expected = np.vstack((front, front - normal * thickness))
    from scipy.spatial import cKDTree
    assert cKDTree(points).query(expected)[0].max() < 1e-6
    assert np.isfinite(points).all() and np.isfinite(uv).all()
    assert materials[0] == "Image" and set(indices) == {0}


@pytest.mark.parametrize("valid_only", [False, True])
def test_depth_split_preserves_faces_and_closes_each_shell(surface, valid_only):
    from tests.support.depth_surface import component_count

    obj, set_value, _, _, image = surface
    pixels = np.asarray(image.pixels[:], dtype=np.float32).reshape(8, 16, 4)
    pixels[:, 8:, 2] += 10
    pixels[:2, :, 3] = 0
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Depth Scale", 1.0)
    set_value("Depth Mask", valid_only)
    base = evaluated(obj)
    set_value("Depth Split", 1.0)
    points, faces, uv, materials, indices, _ = evaluated(obj)
    assert len(points) > len(base[0])
    assert len(faces) == len(base[1])
    assert component_count(faces) > component_count(base[1])
    np.testing.assert_allclose(uv, base[2])
    assert (materials, indices) == base[3:5]
    set_value("Thickness", 0.2)
    shell_points, shell_faces, shell_uv, *_ = evaluated(obj)
    assert len(shell_points) == 2 * len(points)
    edges = Counter(tuple(sorted((a, b))) for face in shell_faces
                    for a, b in zip(face, (*face[1:], face[0])))
    assert set(edges.values()) == {2}
    assert np.isfinite(shell_points).all() and np.isfinite(shell_uv).all()
    set_value("Thickness", 0.0)
    set_value("Depth Split", 0.0)
    restored = evaluated(obj)
    np.testing.assert_allclose(restored[0], base[0])
    assert restored[1] == base[1]


def test_depth_jump_keeps_rectangular_faces(surface):
    obj, set_value, _, _, image = surface
    base = evaluated(obj)
    pixels = np.ones((8, 16, 4), dtype=np.float32)
    pixels[:, 8:, 2] = 10
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Depth Scale", 1.0)
    result = evaluated(obj)
    assert result[1] == base[1]
    assert len(result[0]) == len(base[0])
    assert np.ptp(result[0][:, 2]) > 4


@pytest.mark.parametrize("thickness", [0.0, 0.4])
def test_relief_reads_z_and_keeps_the_base(relief, thickness):
    obj, set_value, _, inputs, image = relief
    set_value("Thickness", thickness)
    set_value("Depth Scale", 1.0)
    set_value("Reference Depth", 2.0)
    pixels = np.ones((8, 16, 4), dtype=np.float32)
    pixels[..., 0], pixels[..., 1], pixels[..., 2] = -9, 13, 2
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Uniform Scale", 0.5)
    first = evaluated(obj)
    assert first[0][:, 2].max() == pytest.approx(thickness + 1)
    assert first[0][:, 2].min() == pytest.approx(0 if thickness else 1)
    pixels[..., 0], pixels[..., 1] = 31, -22
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Uniform Scale", 0.5)
    np.testing.assert_allclose(evaluated(obj)[0], first[0])
    assert obj.modifiers[0][inputs["Thickness"].identifier] == pytest.approx(thickness)


@pytest.mark.parametrize("reference", [0.5, 1.25, 2.0])
def test_zero_thickness_relief_keeps_far_depth_at_object_space_zero(relief, reference):
    obj, set_value, _, _, _ = relief
    base = evaluated(obj)
    texel_x = np.clip((base[0][:, 0] + 2) / 4 * 16, 0.5, 15.5)
    depth = (2 + texel_x / 16) * 0.5
    expected = np.maximum((reference - depth) * 1.5, 0)
    set_value("Reference Depth", reference)
    set_value("Depth Scale", 1.5)
    result = evaluated(obj)
    np.testing.assert_allclose(result[0][:, 2], expected, atol=1e-6)
    np.testing.assert_array_equal(result[0][:, :2], base[0][:, :2])
    assert result[1] == base[1]
    np.testing.assert_array_equal(result[2], base[2])


def test_curved_shell_smooths_normals_without_shrinking_thickness(surface):
    from scipy.spatial import cKDTree

    obj, set_value, _, _, image = surface
    yy, xx = np.mgrid[:8, :16]
    pixels = np.ones((8, 16, 4), dtype=np.float32)
    pixels[..., 0] = (xx + 0.5 - 8) / 4
    pixels[..., 1] = (4 - yy - 0.5) / 4
    pixels[..., 2] = 2 + ((xx + 0.5 - 8) / 8) ** 2 + 0.12 * (xx % 2)
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Subdivide", 3)
    set_value("Depth Scale", 1.0)
    front, faces, *_ = evaluated(obj)
    result = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    mesh = result.to_mesh()
    try:
        raw_normals = np.array([v.normal[:] for v in mesh.vertices])
    finally:
        result.to_mesh_clear()
    thickness = 0.001
    set_value("Thickness", thickness)
    points = evaluated(obj)[0]
    distance, source = cKDTree(front).query(points)
    rear = distance > thickness * 0.5
    assert rear.sum() == len(front)
    assert len(set(source[rear])) == len(front)
    np.testing.assert_allclose(distance[rear], thickness, atol=1e-6)
    normals = np.empty_like(front)
    normals[source[rear]] = (front[source[rear]] - points[rear]) / thickness
    assert np.all(np.sum(normals * raw_normals, axis=1) > 0)
    edges = np.array(list({tuple(sorted((a, b))) for f in faces
                          for a, b in zip(f, (*f[1:], f[0]))}))
    roughness = lambda n: np.mean(np.linalg.norm(n[edges[:, 0]] - n[edges[:, 1]], axis=1))
    assert roughness(normals) < roughness(raw_normals)



@pytest.mark.parametrize("thickness", [0.0, 0.4])
@pytest.mark.parametrize("direction", [(0.2, -0.1, 1), (-0.2, 0.1, 1), (0.4, -0.2, 2), (0, 0, 0), (1, 0, 0)])
def test_relief_direction_adjusts_depth_and_keeps_base(relief, thickness, direction):
    obj, set_value, _, _, image = relief
    pixels = np.ones((8, 16, 4), dtype=np.float32)
    pixels[..., 2] = 2
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Thickness", thickness)
    set_value("Reference Depth", 2.0)
    layers = evaluated(obj)[0]
    weights = np.clip(layers[:, 2] / thickness, 0, 1) if thickness else np.ones(len(layers))
    set_value("Depth Scale", 0.5)
    base = evaluated(obj)
    set_value("Depth Direction", direction)
    result = evaluated(obj)
    front = base[0][:, 2] > 0
    expected = base[0].copy()
    if direction[2]:
        expected[front, 2] += 0.5 * weights[front] * (expected[front, :2] @ np.array(direction[:2])) / direction[2]
    np.testing.assert_allclose(result[0], expected, atol=1e-6)
    assert result[1] == base[1]
    np.testing.assert_array_equal(result[2], base[2])


@pytest.mark.parametrize("thickness", [0.0, 0.4])
def test_relief_smooth_preserves_topology_uv_and_materials(relief, thickness):
    obj, set_value, _, _, image = relief
    pixels = np.asarray(image.pixels[:], dtype=np.float32).reshape(8, 16, 4)
    pixels[..., 2] += (np.arange(16) % 2) * 0.5
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Thickness", thickness)
    set_value("Reference Depth", 2.0)
    set_value("Depth Scale", 1.0)
    set_value("Depth Direction", (0.1, -0.1, 1))
    base = evaluated(obj)
    set_value("Smooth", 3)
    result = evaluated(obj)
    assert np.isfinite(result[0]).all()
    assert np.max(abs(result[0] - base[0])) > 1e-4
    assert result[1] == base[1]
    np.testing.assert_array_equal(result[2], base[2])
    assert result[3:] == base[3:]
    set_value("Smooth Weight", 0.0)
    np.testing.assert_allclose(evaluated(obj)[0], base[0], atol=1e-6)
    set_value("Smooth Weight", 1.0)
    set_value("Smooth", 0)
    np.testing.assert_allclose(evaluated(obj)[0], base[0], atol=1e-6)


@pytest.mark.parametrize("pattern,threshold", [
    ("valid", 1.0), ("invalid", 0.0), ("invalid", 0.9),
    ("partial", 0.4), ("partial", 0.9), ("half", 0.9), ("stripe", 0.9),
])
def test_valid_only_deletes_invalid_points_and_wire_edges(surface, pattern, threshold):
    obj, set_value, _, inputs, _ = surface
    set_value("Mask Threshold", threshold)
    depth_image = obj.modifiers[0][inputs["Depth Image"].identifier]
    pixels = np.ones((8, 16, 4), dtype=np.float32)
    alpha = np.ones(16, dtype=np.float32)
    if pattern == "invalid":
        alpha[:] = 0
    elif pattern == "partial":
        alpha[:] = 0.5
    elif pattern == "half":
        alpha[:8] = 0
    elif pattern == "stripe":
        alpha[:] = 0
        alpha[8] = 1
    pixels[..., 3] = alpha
    depth_image.pixels.foreach_set(pixels.ravel())
    depth_image.update()
    for subdivisions in (0, 2, 5):
        set_value("Subdivide", subdivisions)
        set_value("Depth Mask", False)
        base = evaluated(obj)
        u = (base[0][:, 0] + 2) / 4
        sampled = np.interp(u * 16 - 0.5, np.arange(16), alpha)
        expected = [face for face in base[1] if np.all(sampled[list(face)] >= threshold)]
        set_value("Depth Mask", True)
        culled = evaluated(obj)
        assert len(culled[1]) == len(expected)
        expected_points = {tuple(base[0][index]) for face in expected for index in face}
        assert {tuple(point) for point in culled[0]} == expected_points
        assert np.isfinite(culled[2]).all()
        set_value("Depth Mask", False)
        np.testing.assert_allclose(evaluated(obj)[0], base[0])


def test_valid_only_precedes_closed_shell_generation(surface):
    obj, set_value, _, inputs, _ = surface
    depth_image = obj.modifiers[0][inputs["Depth Image"].identifier]
    pixels = np.asarray(depth_image.pixels[:], dtype=np.float32).reshape(8, 16, 4)
    pixels[:, :8, 3] = 0
    depth_image.pixels.foreach_set(pixels.ravel())
    depth_image.update()
    set_value("Depth Scale", 1.0)
    front = evaluated(obj)[0]
    set_value("Thickness", 0.2)
    points, faces, uv, *_ = evaluated(obj)
    assert len(points) == 2 * len(front)
    edges = Counter(tuple(sorted((a, b))) for face in faces for a, b in zip(face, (*face[1:], face[0])))
    assert set(edges.values()) == {2}
    assert np.isfinite(points).all() and np.isfinite(uv).all()


@pytest.mark.parametrize("thickness", [0.0, 0.2])
def test_depth_smooth_changes_positions_and_preserves_topology(surface, thickness):
    obj, set_value, _, _, depth_image = surface
    pixels = np.asarray(depth_image.pixels[:], dtype=np.float32).reshape(8, 16, 4)
    pixels[..., 2] += (np.arange(16) % 2) * 0.3
    pixels[:, :4, 3] = 0
    depth_image.pixels.foreach_set(pixels.ravel())
    depth_image.update()
    set_value("Depth Scale", 1.0)
    set_value("Thickness", thickness)
    base = evaluated(obj)
    set_value("Smooth", 3)
    smoothed = evaluated(obj)
    assert smoothed[1] == base[1]
    np.testing.assert_allclose(smoothed[2], base[2])
    assert smoothed[3:5] == base[3:5]
    assert np.isfinite(smoothed[0]).all()
    assert not np.allclose(smoothed[0], base[0])
    set_value("Smooth Weight", 0.0)
    np.testing.assert_allclose(evaluated(obj)[0], base[0], atol=1e-6)
    set_value("Smooth Weight", 1.0)
    set_value("Smooth", 0)
    np.testing.assert_allclose(evaluated(obj)[0], base[0], atol=1e-6)