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

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

ROOT = Path(__file__).parents[1]
from tools.nodes.groups.image_plane import build_image_plane_group
from tools.nodes.groups.image_depth_plane import build_image_depth_plane_group
from tools.nodes.groups.image_relief_plane import build_image_relief_plane_group
from node_asset_test_support import validate_node_structure
from tools.nodes.arrangement.flow import origin_socket


@pytest.fixture(params=("definition", "asset"))
def surface(request):
    return create_surface(request.param, "DEPTH")


@pytest.fixture(params=("definition", "asset"))
def relief(request):
    return create_surface(request.param, "RELIEF")


def create_surface(source, plane_type):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    group_name = f"O Image {plane_type.title()} Plane"
    if source == "asset":
        with bpy.data.libraries.load(str(ROOT / "src/anyimage/assets/O_AnyImage.blend")) as (library, target):
            target.node_groups = [group_name]
        group = target.node_groups[0]
        plane = group.nodes["Group"].node_tree
    else:
        plane = build_image_plane_group()
        builder = build_image_depth_plane_group if plane_type == "DEPTH" else build_image_relief_plane_group
        group = builder(plane)
    mesh = bpy.data.meshes.new("Rectangle")
    mesh.from_pydata([(-2, -1, 0), (2, -1, 0), (2, 1, 0), (-2, 1, 0)], [], [(0, 1, 2, 3)])
    material = bpy.data.materials.new("Image")
    mesh.materials.append(material)
    obj = bpy.data.objects.new("Surface", mesh)
    bpy.context.collection.objects.link(obj)
    modifier = obj.modifiers.new("Surface", "NODES")
    modifier.node_group = group
    image = bpy.data.images.new("Camera", width=16, height=8, float_buffer=True)
    image.colorspace_settings.name = "Non-Color"
    image.alpha_mode = "CHANNEL_PACKED"
    yy, xx = np.mgrid[:8, :16]
    # Camera principal point is deliberately off center; Y points down.
    pixels = np.ones((8, 16, 4), dtype=np.float32)
    pixels[..., 0] = ((xx + 0.5) - 5) / 4
    pixels[..., 1] = (3 - (yy + 0.5)) / 4
    pixels[..., 2] = 2 + (xx + 0.5) / 16
    pixels[..., 3] = 1
    image.pixels.foreach_set(pixels.ravel())
    inputs = {item.name: item for item in group.interface.items_tree
              if item.item_type == "SOCKET" and item.in_out == "INPUT"}
    def set_value(name, value):
        modifier[inputs[name].identifier] = value
        obj.update_tag(refresh={"DATA"})
        bpy.context.view_layer.update()

    for name, value in (("Subdivide", 2), ("Depth Image", image),
                        ("Uniform Scale", 0.5), ("Reference Depth", 1.2),
                        ("Depth Scale", 0.0), ("Thickness", 0.0)):
        set_value(name, value)
    return obj, set_value, plane, inputs, image


def evaluated(obj):
    result = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    mesh = result.to_mesh()
    try:
        return (
            np.array([v.co[:] for v in mesh.vertices]),
            [tuple(face.vertices) for face in mesh.polygons],
            np.array([value.uv[:] for value in mesh.uv_layers["UVMap"].data]) if mesh.uv_layers else np.empty((0, 2)),
            [material.name if material else None for material in mesh.materials],
            [face.material_index for face in mesh.polygons],
            [attr.name for attr in mesh.attributes],
        )
    finally:
        result.to_mesh_clear()


@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, attributes = 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}
    assert not any(name.startswith("o_depth_") for name in attributes)


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("normal_smooth", [0, 2, 50])
def test_curved_shell_smooths_normals_without_shrinking_thickness(surface, normal_smooth):
    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()
    set_value("Normal Smooth", normal_smooth)
    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))
    if normal_smooth == 0:
        np.testing.assert_allclose(normals, raw_normals, atol=2e-4)
    else:
        assert roughness(normals) < roughness(raw_normals)


def test_surface_interface_and_node_structure(surface):
    obj, _, _, inputs, _ = surface
    assert list(inputs) == ["Geometry", "Subdivide", "Thickness", "Depth Scale", "Smooth", "Valid Only",
                            "Reference Depth", "Normal Smooth", "Smooth Weight", "Validity Threshold", "Uniform Scale", "Depth Image"]
    assert inputs["Validity Threshold"].default_value == pytest.approx(0.9)
    assert inputs["Validity Threshold"].parent.name == "Options"
    assert inputs["Validity Threshold"].subtype == "FACTOR"
    assert inputs["Valid Only"].default_value is True
    assert inputs["Smooth"].default_value == 0
    assert inputs["Smooth"].min_value == 0 and inputs["Smooth"].max_value == 20
    assert inputs["Smooth Weight"].default_value == 1
    assert inputs["Smooth Weight"].subtype == "FACTOR"
    assert inputs["Smooth Weight"].parent.name == "Options"
    assert inputs["Normal Smooth"].default_value == 50
    assert inputs["Normal Smooth"].min_value == 0
    assert inputs["Normal Smooth"].max_value == 50
    assert inputs["Normal Smooth"].parent.name == "Options"
    assert inputs["Thickness"].default_value == 0
    assert inputs["Subdivide"].default_value == 6
    assert inputs["Depth Scale"].default_value == 1
    assert all(item.structure_type == "SINGLE" for item in inputs.values())
    assert inputs["Thickness"].subtype == inputs["Reference Depth"].subtype == "DISTANCE"
    assert inputs["Uniform Scale"].hide_in_modifier and inputs["Depth Image"].hide_in_modifier
    validate_node_structure(obj.modifiers[0].node_group)


def test_relief_interface_and_node_structure(relief):
    obj, _, _, inputs, _ = relief
    assert list(inputs) == ["Geometry", "Subdivide", "Thickness", "Depth Scale",
                            "Reference Depth", "Uniform Scale", "Depth Image"]
    assert inputs["Thickness"].default_value == 0
    assert inputs["Subdivide"].default_value == 6
    assert inputs["Depth Scale"].default_value == 1
    assert all(item.structure_type == "SINGLE" for item in inputs.values())
    assert inputs["Thickness"].subtype == inputs["Reference Depth"].subtype == "DISTANCE"
    assert inputs["Uniform Scale"].hide_in_modifier and inputs["Depth Image"].hide_in_modifier
    validate_node_structure(obj.modifiers[0].node_group)


def test_conversion_preserves_offset_world_placement_and_initializes_inputs(surface, monkeypatch):
    from anyimage.operators.convert_to_plane import object as conversion
    from anyimage.common.image import image_empty_bounds
    from mathutils import Matrix, Vector

    obj, _, _, inputs, image = surface
    group = obj.modifiers[0].node_group
    monkeypatch.setattr(conversion, "load_node_group", lambda name: group)
    source = bpy.data.objects.new("Offset Image", None)
    source.empty_display_type = "IMAGE"
    source.data = image
    source.empty_display_size = 4
    source.empty_image_offset = (-0.2, -0.8)
    source.matrix_world = Matrix.Translation((3, -2, 5)) @ Matrix.Rotation(0.7, 4, "Y") @ Matrix.Diagonal((2, 3, 1, 1))
    bpy.context.collection.objects.link(source)
    bounds = image_empty_bounds(source)
    matrix = source.matrix_world.copy()
    material = bpy.data.materials.new("Conversion")
    metadata = {"image_size": (16, 8), "reference_depth": 2.0,
                "intrinsics": ((4, 0, 5), (0, 4, 3), (0, 0, 1))}
    result = conversion.create_depth_plane_object(
        bpy.context, source, 2, material, image, metadata, "DEPTH")
    assert result.name == "Offset Image"
    assert result["anyimage_mesh_shape"] == "DEPTH_PLANE"
    assert bpy.context.view_layer.objects.active == result
    expected = sorted(tuple(matrix @ Vector((x, y, 0)))
                      for x in bounds[:2] for y in bounds[2:])
    actual = sorted(tuple(result.matrix_world @ v.co) for v in result.data.vertices)
    np.testing.assert_allclose(actual, expected, atol=1e-6)
    center = matrix @ Vector(((bounds[0] + bounds[1]) / 2, (bounds[2] + bounds[3]) / 2, 0))
    np.testing.assert_allclose(result.matrix_world.translation, center, atol=1e-6)
    modifier = result.modifiers[0]
    assert modifier[inputs["Thickness"].identifier] == 0
    assert modifier[inputs["Depth Image"].identifier] == image
    assert modifier[inputs["Valid Only"].identifier] is True
    assert modifier[inputs["Reference Depth"].identifier] == pytest.approx(
        2 * modifier[inputs["Uniform Scale"].identifier])
    assert result.data.materials[0] == material


def test_failed_modifier_initialization_removes_partial_object(surface, monkeypatch):
    from anyimage.operators.convert_to_plane import object as conversion

    obj, _, _, _, image = surface
    monkeypatch.setattr(conversion, "load_node_group", lambda name: obj.modifiers[0].node_group)
    source = bpy.data.objects.new("Source", None)
    source.empty_display_type = "IMAGE"
    source.data = image
    bpy.context.collection.objects.link(source)
    material = bpy.data.materials.new("Unused")
    before_objects, before_meshes = set(bpy.data.objects), set(bpy.data.meshes)

    def fail(*args):
        raise RuntimeError("input failed")

    monkeypatch.setattr(conversion, "set_modifier_input", fail)
    metadata = {"image_size": (16, 8), "reference_depth": 2,
                "intrinsics": ((4, 0, 5), (0, 4, 3), (0, 0, 1))}
    with pytest.raises(RuntimeError, match="input failed"):
        conversion.create_depth_plane_object(bpy.context, source, 2, material, image, metadata, "DEPTH")
    assert set(bpy.data.objects) == before_objects
    assert set(bpy.data.meshes) == before_meshes
    assert material.users == 0


@pytest.mark.parametrize("threshold", [0.0, 0.4, 0.9, 1.0])
@pytest.mark.parametrize("pattern", ["valid", "invalid", "partial", "half", "stripe"])
def test_valid_only_recomputes_per_point_and_preserves_boundary_faces(surface, pattern, threshold):
    obj, set_value, _, inputs, _ = surface
    set_value("Validity 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("Valid Only", 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 any(sampled[index] >= threshold for index in face)]
        set_value("Valid Only", 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()
        if pattern == "half" and subdivisions >= 2 and threshold == 1.0:
            assert 0 <= min(point[0] for point in culled[0]) < 0.125
        set_value("Valid Only", 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()


def test_validity_graph_matches_reference_domains_and_connections(surface):
    obj, _, _, _, _ = surface
    group = obj.modifiers[0].node_group
    delete = next(node for node in group.nodes if node.bl_idname == "GeometryNodeDeleteGeometry")
    assert (delete.domain, delete.mode) == ("FACE", "ALL")
    points = origin_socket(delete.inputs["Selection"].links[0].from_socket).node
    assert (points.bl_idname, points.data_type, points.domain) == ("GeometryNodeFieldOnDomain", "BOOLEAN", "POINT")
    compare = origin_socket(points.inputs["Value"].links[0].from_socket).node
    assert (compare.bl_idname, compare.data_type, compare.operation) == ("FunctionNodeCompare", "FLOAT", "LESS_THAN")
    assert origin_socket(compare.inputs["B"].links[0].from_socket).name == "Validity Threshold"
    alpha = origin_socket(compare.inputs["A"].links[0].from_socket)
    assert alpha.name == "Alpha" and alpha.node.bl_idname == "FunctionNodeSeparateColor"
    color = origin_socket(alpha.node.inputs["Color"].links[0].from_socket).node
    assert origin_socket(color.inputs["Image"].links[0].from_socket).name == "Depth Image"
    assert len([node for node in color.id_data.nodes if node.bl_idname == "GeometryNodeImageTexture"]) == 1
    consumers = [link for link in group.links if link.to_node.bl_idname != "NodeReroute"
                 and origin_socket(link.from_socket) == color.outputs["Color"]]
    assert len(consumers) == 2
    assert origin_socket(delete.inputs["Geometry"].links[0].from_socket).node.node_tree.name == "O Image Plane"


@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)