"""Behavior checks for face separation and connected Depth Surface shells."""

import sys
from pathlib import Path
from collections import Counter

import bpy
import bmesh
import numpy as np
import pytest

ROOT = Path(__file__).parents[1]
sys.path.insert(0, str(ROOT / "src"))
sys.path.insert(0, str(ROOT / "tools" / "node_assets"))
from build_cutout import build_node_group, build_cutout_modifier_group


def surface(*, step=6.0, size=1.0, resolution=2, wrapper=False):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    groups = {
        shape: build_node_group(shape)
        for shape in ("SURFACE", "BALLOON", "DEPTH_BALLOON", "DEPTH_SURFACE")
    }
    group = build_cutout_modifier_group(groups) if wrapper else groups["DEPTH_SURFACE"]
    # Two columns with a complete depth jump along their common boundary.
    points = [
        (x * size / resolution, y * size / resolution, 0)
        for y in range(resolution + 1)
        for x in range(2 * resolution + 1)
    ]
    stride = 2 * resolution + 1
    faces = [
        (
            y * stride + x,
            y * stride + x + 1,
            (y + 1) * stride + x + 1,
            (y + 1) * stride + x,
        )
        for y in range(resolution)
        for x in range(2 * resolution)
    ]
    mesh = bpy.data.meshes.new("Step")
    mesh.from_pydata(points, [], faces)
    uv = mesh.uv_layers.new(name="UVMap")
    for loop in mesh.loops:
        x, y, _ = points[loop.vertex_index]
        uv.data[loop.index].uv = (x / (2 * size), y / size)
    obj = bpy.data.objects.new("Step", mesh)
    bpy.context.collection.objects.link(obj)
    image = bpy.data.images.new("Camera", width=64, height=32, float_buffer=True)
    image.colorspace_settings.name = "Non-Color"
    image.alpha_mode = "CHANNEL_PACKED"
    yy, xx = np.mgrid[:32, :64]
    pixels = np.ones((32, 64, 4), np.float32)
    pixels[..., 0] = (xx + 0.5) / 32
    pixels[..., 1] = -(yy + 0.5) / 32
    pixels[..., 2] = 1 + (xx >= 32) * step
    image.pixels.foreach_set(pixels.ravel())
    modifier = obj.modifiers.new("Depth Surface", "NODES")
    modifier.node_group = group
    inputs = {
        s.name: s.identifier
        for s in group.interface.items_tree
        if s.item_type == "SOCKET" and s.in_out == "INPUT"
    }

    def set_value(name, value):
        modifier[inputs[name]] = value
        obj.update_tag(refresh={"DATA"})
        bpy.context.view_layer.update()

    for name, value in (
        ("Depth Image", image),
        ("Uniform Scale", size),
        ("Reference Depth", size),
        ("Depth Scale", 1.0),
        ("Thickness", 0.0),
        ("Cleanup", 0.0),
        ("Split", 0.0),
    ):
        set_value(name, value)
    if wrapper:
        set_value("Shape", 3)
        set_value("Inward Axis", 0)
        set_value("Smooth Iterations", 0)
    return obj, set_value


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(p.vertices) for p in mesh.polygons
        ]
    finally:
        result.to_mesh_clear()


def assert_closed(faces):
    edges = Counter(
        tuple(sorted((a, b))) for f in faces for a, b in zip(f, (*f[1:], f[0]))
    )
    assert edges and set(edges.values()) == {2}, Counter(edges.values())


@pytest.mark.parametrize("thickness", [0.1, 1.4])
@pytest.mark.parametrize("split", [0.0, 0.5])
def test_surface_caps_and_segmented_walls_are_connected(thickness, split):
    obj, set_value = surface()
    set_value("Split", split)
    set_value("Thickness", thickness)
    coords, faces = evaluated(obj)
    assert np.isfinite(coords).all()
    assert_closed(faces)
    assert coords[:, 2].min() == pytest.approx(-6 - thickness, abs=1e-5)


@pytest.mark.parametrize("size", [0.01, 1.0, 100.0])
def test_split_mapping_and_uniform_scale(size):
    obj, set_value = surface(step=2.0, size=size)
    # Face centers are size/2 apart: gradient is 4.
    counts = []
    for split in (0.0, 0.25, 0.5, 1.0):
        set_value("Split", split)
        coords, faces = evaluated(obj)
        counts.append(len(coords))
        assert len(faces) == 8
    assert counts[:3] == [15] * 3
    assert counts[3] == 18
    set_value("Depth Scale", 0.0)
    coords, faces = evaluated(obj)
    assert len(coords) == 15 and np.abs(coords[:, 2]).max() < 1e-6


def test_midpoint_selects_gradient_six_and_preserves_faces():
    obj, set_value = surface(step=3.0)
    set_value("Split", 0.5)
    coords, faces = evaluated(obj)
    assert len(coords) == 18 and len(faces) == 8
    assert set(np.round(coords[:, 2], 4)) == {0.0, -3.0}


def test_cleanup_removes_new_islands_before_thickness():
    obj, set_value = surface()
    set_value("Cleanup", 1.1)
    assert len(evaluated(obj)[1]) == 8
    set_value("Split", 0.5)
    assert len(evaluated(obj)[1]) == 0
    set_value("Thickness", 0.2)
    assert len(evaluated(obj)[1]) == 0


def test_full_wrapper_smooth_keeps_closed_shell():
    obj, set_value = surface(wrapper=True)
    set_value("Split", 0.5)
    set_value("Thickness", 0.2)
    set_value("Smooth Iterations", 2)
    coords, faces = evaluated(obj)
    assert np.isfinite(coords).all()
    assert_closed(faces)


def component_count(faces):
    adjacency = {}
    for face in faces:
        for a, b in zip(face, (*face[1:], face[0])):
            adjacency.setdefault(a, set()).add(b)
            adjacency.setdefault(b, set()).add(a)
    remaining = set(adjacency)
    count = 0
    while remaining:
        count += 1
        pending = [remaining.pop()]
        while pending:
            for neighbor in adjacency[pending.pop()] & remaining:
                remaining.remove(neighbor)
                pending.append(neighbor)
    return count


def test_coincident_independent_sheets_do_not_weld_together():
    obj, set_value = surface(step=0.0)
    mesh = obj.data
    vertices = [tuple(v.co) for v in mesh.vertices]
    faces = [tuple(f.vertices) for f in mesh.polygons]
    n = len(vertices)
    duplicate = bpy.data.meshes.new("Close sheets")
    duplicate.from_pydata(
        vertices + [(x, y, z + 1e-8) for x, y, z in vertices],
        [],
        faces + [tuple(v + n for v in face) for face in faces],
    )
    uv = duplicate.uv_layers.new(name="UVMap")
    for loop in duplicate.loops:
        x, y, z = duplicate.vertices[loop.vertex_index].co
        uv.data[loop.index].uv = (x / 2, y)
    obj.data = duplicate
    set_value("Thickness", 0.1)
    coords, faces = evaluated(obj)
    assert_closed(faces)
    assert component_count(faces) == 2


@pytest.mark.parametrize("resolution", [1, 2, 4])
def test_planar_depth_has_no_split_at_any_density(resolution):
    obj, set_value = surface(step=0.0, resolution=resolution)
    original, original_faces = evaluated(obj)
    set_value("Split", 1.0)
    coords, faces = evaluated(obj)
    assert len(coords) == len(obj.data.vertices)
    assert len(faces) == len(obj.data.polygons)
    np.testing.assert_allclose(coords, original, atol=1e-6)


def test_small_positive_split_preserves_unsplit_projection():
    obj, set_value = surface(step=2.0)
    original, original_faces = evaluated(obj)
    set_value("Split", 0.01)
    coords, faces = evaluated(obj)
    np.testing.assert_allclose(coords, original, atol=1e-6)
    assert faces == original_faces


def test_cut_vertices_preserve_camera_rays_and_other_vertices():
    obj, set_value = surface(step=3.0)
    original, original_faces = evaluated(obj)
    set_value("Split", 0.5)
    coords, faces = evaluated(obj)
    for source_face, result_face in zip(original_faces, faces):
        for source_index, result_index in zip(source_face, result_face):
            before, after = original[source_index], coords[result_index]
            if obj.data.vertices[source_index].co.x != 1.0:
                np.testing.assert_allclose(after, before, atol=1e-6)
            before_depth, after_depth = 1.0 - before[2], 1.0 - after[2]
            np.testing.assert_allclose(
                after[:2] / after_depth, before[:2] / before_depth, atol=1e-6
            )


def test_stored_point_positions_prevent_depth_resampling_on_side_walls():
    obj, set_value = surface(step=3.0)
    set_value("Split", 0.5)
    set_value("Thickness", 0.1)
    coords, faces = evaluated(obj)
    assert_closed(faces)
    assert set(np.round(coords[:, 2], 4)) == {0.0, -0.1, -3.0, -3.1}


@pytest.mark.parametrize("thickness", [0.0, 0.2])
@pytest.mark.parametrize("split", [0.0, 0.5])
def test_temporary_attributes_are_removed_and_user_attributes_survive(thickness, split):
    obj, set_value = surface(wrapper=True)
    obj.data.attributes.new("o_depth_probe", "FLOAT", "POINT")
    attribute = obj.data.attributes.new("o_user_weight", "FLOAT", "POINT")
    for value in attribute.data:
        value.value = 0.75
    set_value("Split", split)
    set_value("Thickness", thickness)
    set_value("Smooth Iterations", 2)
    result = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    mesh = result.to_mesh()
    try:
        assert not any(a.name.startswith("o_depth_") for a in mesh.attributes)
        assert "UVMap" in mesh.attributes
        assert all(
            v.value == pytest.approx(0.75)
            for v in mesh.attributes["o_user_weight"].data
        )
    finally:
        result.to_mesh_clear()


def test_zero_center_distance_does_not_split_or_produce_nan():
    obj, set_value = surface(step=6.0)
    mesh = bpy.data.meshes.new("Coincident face centers")
    mesh.from_pydata(
        [(0, 0, 0), (1, 0, 0), (0.5, 1, 0), (0.5, 1, 0)], [], [(0, 1, 2), (1, 0, 3)]
    )
    uv = mesh.uv_layers.new(name="UVMap")
    for face in mesh.polygons:
        for li in face.loop_indices:
            uv.data[li].uv = (0.25 if face.index == 0 else 0.75, 0.5)
    obj.data = mesh
    set_value("Split", 1.0)
    coords, faces = evaluated(obj)
    assert len(coords) == 4 and len(faces) == 2
    assert np.isfinite(coords).all()


def test_triangle_corner_weights_do_not_open_sloped_side_walls():
    obj, set_value = surface(step=3.0)
    bm = bmesh.new()
    bm.from_mesh(obj.data)
    bmesh.ops.triangulate(bm, faces=list(bm.faces))
    bm.to_mesh(obj.data)
    bm.free()
    image = bpy.data.images["Camera"]
    pixels = np.array(image.pixels[:], np.float32).reshape(32, 64, 4)
    pixels[..., 2] += np.arange(32)[:, None] * 0.02
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Split", 0.5)
    set_value("Thickness", 0.8)
    coords, faces = evaluated(obj)
    assert_closed(faces)
    assert component_count(faces) == 2


@pytest.mark.parametrize("resolution", [1, 2, 4])
def test_linear_depth_slope_has_density_independent_selection(resolution):
    obj, set_value = surface(step=0.0, resolution=resolution)
    image = bpy.data.images["Camera"]
    pixels = np.array(image.pixels[:], np.float32).reshape(32, 64, 4)
    pixels[..., 2] = 1 + 4 * pixels[..., 0]
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Split", 0.5)
    assert len(evaluated(obj)[0]) == len(obj.data.vertices)
    set_value("Split", 1.0)
    assert len(evaluated(obj)[0]) > len(obj.data.vertices)