"""Verify pinned-boundary smoothing stays within its fixed edge band."""

from collections import Counter

import bpy
import bmesh
import numpy as np
import pytest

from tests.support.depth_surface import surface, evaluated, assert_closed, component_count


def diagonal_surface(triangles=False):
    obj, set_value = surface(resolution=16)
    if triangles:
        mesh = bmesh.new()
        mesh.from_mesh(obj.data)
        bmesh.ops.triangulate(mesh, faces=list(mesh.faces))
        mesh.to_mesh(obj.data)
        mesh.free()
    image = bpy.data.images["Camera"]
    pixels = np.array(image.pixels[:], np.float32).reshape(32, 64, 4)
    yy, xx = np.mgrid[:32, :64]
    pixels[..., 2] = 1 + (xx > 18 + yy * 0.7) * 6
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Split Threshold", 0.5)
    set_value("Depth Scale", 0)
    return obj, set_value


def boundary_neighbors(faces):
    edges = Counter(tuple(sorted((a, b))) for face in faces for a, b in zip(face, (*face[1:], face[0])))
    neighbors = {}
    for (a, b), count in edges.items():
        if count == 1:
            neighbors.setdefault(a, []).append(b)
            neighbors.setdefault(b, []).append(a)
    return neighbors


def edge_band(faces, seeds, rings=2):
    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)
    selected = set(seeds)
    for _ in range(rings):
        selected |= {neighbor for vertex in selected for neighbor in adjacency[vertex]}
    return selected


@pytest.mark.parametrize("triangles", [False, True])
def test_smoothing_reduces_stairs_without_leaving_two_ring_band(triangles):
    obj, set_value = diagonal_surface(triangles)
    original, faces = evaluated(obj)
    boundary = boundary_neighbors(faces)
    outer = np.isclose(original[:, 0], 0) | np.isclose(original[:, 0], 2)
    outer |= np.isclose(original[:, 1], 0) | np.isclose(original[:, 1], 1)
    cut = {i for i in boundary if not outer[i]}
    band = edge_band(faces, boundary)
    movable = np.array([i in band for i in range(len(original))])
    assert movable.any() and (~movable).any()
    set_value("Boundary Smooth", 20)
    smoothed, result_faces = evaluated(obj)
    assert result_faces == faces
    np.testing.assert_array_equal(smoothed[~movable], original[~movable])
    assert np.max(np.linalg.norm(smoothed - original, axis=1)) > 0.005
    nearby = [i for i in band if i not in boundary and not outer[i]]
    assert np.max(np.linalg.norm(smoothed[nearby] - original[nearby], axis=1)) > 1e-4
    before, after = [], []
    for i in cut:
        adjacent = boundary[i]
        if len(adjacent) != 2:
            continue
        before.append(np.linalg.norm(original[adjacent].mean(axis=0) - original[i]))
        after.append(np.linalg.norm(smoothed[adjacent].mean(axis=0) - smoothed[i]))
    assert sum(after) < 0.8 * sum(before)
    np.testing.assert_array_equal(smoothed[:, 2], original[:, 2])
    set_value("Smooth Weight", 0)
    np.testing.assert_array_equal(evaluated(obj)[0], original)
    set_value("Smooth Weight", 1)
    set_value("Boundary Smooth", 0)
    np.testing.assert_array_equal(evaluated(obj)[0], original)


@pytest.mark.parametrize("source", ["definition", "asset"])
def test_smoothing_without_split_moves_original_boundary_only_within_band(source):
    obj, set_value = diagonal_surface()
    if source == "asset":
        from anyimage.common.node import node_asset_path

        with bpy.data.libraries.load(str(node_asset_path())) as (_, target):
            target.node_groups = ["O Image Depth Cutout"]
        obj.modifiers[0].node_group = target.node_groups[0]
    set_value("Split Threshold", 0)
    original, faces = evaluated(obj)
    set_value("Boundary Smooth", 20)
    result, result_faces = evaluated(obj)
    boundary = boundary_neighbors(faces)
    band = edge_band(faces, boundary)
    outside = [i for i in range(len(original)) if i not in band]
    assert np.max(np.linalg.norm(result[list(boundary)] - original[list(boundary)], axis=1)) > 1e-4
    np.testing.assert_array_equal(result[outside], original[outside])
    assert result_faces == faces


@pytest.mark.parametrize("mode", [0, 1])
def test_smoothed_cut_shells_stay_closed_and_separate(mode):
    from anyimage.common.object import modifier_input_identifier, set_modifier_input

    obj, set_value = diagonal_surface()
    obj.data.attributes.new("o_balloon", "FLOAT", "POINT").data.foreach_set(
        "value", np.full(len(obj.data.vertices), 0.3, dtype=np.float32)
    )
    set_value("Depth Scale", 1)
    set_value("Boundary Smooth", 4)
    original, faces = evaluated(obj)
    count = component_count(faces)
    set_value("Mode", mode)
    modifier = obj.modifiers[0]
    identifier = modifier_input_identifier(modifier.node_group, "Thickness", subtype="NONE" if mode == 0 else "DISTANCE")
    set_modifier_input(modifier, identifier, 0.05)
    obj.update_tag(refresh={"DATA"})
    bpy.context.view_layer.update()
    result, shell_faces = evaluated(obj)
    assert np.isfinite(result).all()
    assert_closed(shell_faces)
    assert component_count(shell_faces) == count


def test_panorama_smoothing_preserves_interior_and_uvs():
    from tests.tools.nodes.test_image_depth_panorama import panorama, evaluated as panorama_mesh

    yy, xx = np.mgrid[:64, :128]
    depth = np.where(xx > 35 + yy * 0.6, 8.0, 2.0)
    obj, set_value = panorama(depth, subdivide=4)
    set_value("Split Threshold", 0.5)
    original, faces, uv = panorama_mesh(obj)
    rim = boundary_neighbors(faces)
    band = edge_band(faces, rim)
    interior = [i for i in range(len(original)) if i not in band]
    set_value("Boundary Smooth", 4)
    result, result_faces, result_uv = panorama_mesh(obj)
    assert np.isfinite(result).all()
    assert result_faces == faces
    np.testing.assert_array_equal(result[interior], original[interior])
    np.testing.assert_array_equal(result_uv, uv)
    assert np.max(np.linalg.norm(result-original, axis=1)) > 1e-4


@pytest.mark.parametrize("source", ["definition", "asset"])
def test_plane_edge_smoothing_preserves_uvs_and_closes_thickness(source):
    from tests.support.planes import create_surface, evaluated as plane_mesh

    obj, set_value, _, inputs, image = create_surface(source, "DEPTH")
    assert inputs["Boundary Smooth"].default_value == 5
    assert inputs["Boundary Smooth"].parent.name == "Options"
    pixels = np.array(image.pixels[:], np.float32).reshape(8, 16, 4)
    yy, xx = np.mgrid[:8, :16]
    pixels[..., 2] = 1 + (xx > 5 + yy * 0.7) * 6
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Subdivide", 5)
    set_value("Depth Scale", 1)
    set_value("Split Threshold", 0.5)
    original, faces, uv, *_ = plane_mesh(obj)
    set_value("Boundary Smooth", 4)
    result, result_faces, result_uv, *_ = plane_mesh(obj)
    assert result_faces == faces
    np.testing.assert_array_equal(result_uv, uv)
    assert np.max(np.linalg.norm(result-original, axis=1)) > 1e-4
    set_value("Thickness", 0.03)
    assert_closed(plane_mesh(obj)[1])


def test_edge_iterations_reuse_boundary_topology_and_mapping():
    from tools.nodes.arrangement.flow import origin_socket
    from tools.nodes.arrangement.zones import node_zones

    obj, _ = diagonal_surface()
    group = obj.modifiers[0].node_group
    entry, _, members = next(
        zone for zone in node_zones(group)
        if origin_socket(zone[0].inputs["Iterations"].links[0].from_socket).name == "Boundary Smooth"
    )
    nearest = next(node for node in group.nodes if node.bl_idname == "GeometryNodeSampleNearest"
                   and any(link.to_node in members for link in node.outputs["Index"].links))
    fixed_position = origin_socket(nearest.inputs["Sample Position"].links[0].from_socket).node
    assert fixed_position.bl_idname == "GeometryNodeSampleIndex"
    source = origin_socket(entry.inputs["Geometry"].links[0].from_socket)
    assert origin_socket(fixed_position.inputs["Geometry"].links[0].from_socket) == source
    assert fixed_position.domain == "POINT"
    assert any(node.bl_idname == "GeometryNodeSeparateGeometry" for node in members)


@pytest.mark.parametrize("source", ["definition", "asset"])
def test_every_split_node_enables_boundary_smoothing_by_default(source):
    from tools.nodes.build import OUTPUT_PATH, build_node_groups
    from tools.nodes.arrangement.flow import origin_socket

    bpy.ops.wm.read_factory_settings(use_empty=True)
    if source == "definition":
        groups = build_node_groups()
    else:
        with bpy.data.libraries.load(str(OUTPUT_PATH)) as (library, target):
            target.node_groups = list(library.node_groups)
        groups = target.node_groups
    split_groups = []
    for group in groups:
        inputs = {item.name: item for item in group.interface.items_tree
                  if item.item_type == "SOCKET" and item.in_out == "INPUT"}
        assert "Edge Smooth" not in inputs
        if "Split Threshold" not in inputs:
            continue
        split_groups.append(group.name)
        smooth = inputs["Boundary Smooth"]
        assert smooth.default_value == 5
        assert (smooth.min_value, smooth.max_value) == (0, 20)
        assert smooth.parent.name == "Options"
        assert not {"Depth Blur Iterations", "Depth Band Rings", "Depth Blur Strength", "Skip Inactive Geometry"} & inputs.keys()
        depth_mix = next(n for n in group.nodes if n.bl_idname == "ShaderNodeMix"
                         and n.data_type == "FLOAT" and n.inputs[3].is_linked
                         and origin_socket(n.inputs[3].links[0].from_socket).node.bl_idname == "GeometryNodeBlurAttribute")
        depth_blur = origin_socket(depth_mix.inputs[3].links[0].from_socket).node
        assert origin_socket(depth_blur.inputs["Iterations"].links[0].from_socket).name == "Boundary Smooth"
        depth_bypass = [n for n in group.nodes if n.bl_idname == "GeometryNodeSwitch"
                        and n.input_type == "FLOAT" and n.inputs["Switch"].is_linked
                        and origin_socket(n.inputs["Switch"].links[0].from_socket).name == "Boundary Smooth"]
        assert len(depth_bypass) == 1
    assert set(split_groups) == {
        "O Image Depth Plane", "O Image Depth Cutout", "O Image Depth Panorama",
    }


@pytest.mark.parametrize("source", ["definition", "asset"])
def test_validity_cut_smoothing_preserves_original_outline_and_interior(source):
    from tests.support.planes import create_surface, evaluated as plane_mesh

    obj, set_value, _, _, image = create_surface(source, "DEPTH")
    set_value("Subdivide", 5)
    set_value("Split Threshold", 0)
    set_value("Valid Only", False)
    whole = plane_mesh(obj)[0]
    low, high = whole.min(axis=0), whole.max(axis=0)
    pixels = np.array(image.pixels[:], np.float32).reshape(8, 16, 4)
    yy, xx = np.mgrid[:8, :16]
    pixels[..., 3] = xx < 5 + yy * 0.7
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Valid Only", True)
    before, faces, uv, *_ = plane_mesh(obj)
    boundary = boundary_neighbors(faces)
    outer = np.any(np.isclose(before[:, :2], low[:2]) | np.isclose(before[:, :2], high[:2]), axis=1)
    cut = {i for i in boundary if not outer[i]}
    band = edge_band(faces, cut)
    fixed = np.array([i not in band or outer[i] for i in range(len(before))])
    set_value("Boundary Smooth", 5)
    after, result_faces, result_uv, *_ = plane_mesh(obj)
    assert result_faces == faces
    np.testing.assert_array_equal(result_uv, uv)
    np.testing.assert_array_equal(after[fixed], before[fixed])
    assert np.max(np.linalg.norm(after-before, axis=1)) > 1e-4
    before_bends, after_bends = [], []
    for i in cut:
        if len(boundary[i]) == 2:
            before_bends.append(np.linalg.norm(before[boundary[i]].mean(axis=0)-before[i]))
            after_bends.append(np.linalg.norm(after[boundary[i]].mean(axis=0)-after[i]))
    assert sum(after_bends) < sum(before_bends) * 0.8
    set_value("Thickness", 0.03)
    assert_closed(plane_mesh(obj)[1])
    set_value("Thickness", 0)
    set_value("Valid Only", False)
    np.testing.assert_array_equal(plane_mesh(obj)[0], whole)


def test_panorama_validity_cut_smoothing_stays_in_two_ring_band():
    from tests.tools.nodes.test_image_depth_panorama import panorama, evaluated as panorama_mesh

    yy, xx = np.mgrid[:64, :128]
    alpha = ((xx-64)**2 + (yy-32)**2 > 15**2).astype(np.float32)
    obj, set_value = panorama(np.full((64, 128), 3), alpha, subdivide=4)
    set_value("Split Threshold", 0)
    before, faces, uv = panorama_mesh(obj)
    band = edge_band(faces, boundary_neighbors(faces))
    outside = [i for i in range(len(before)) if i not in band]
    set_value("Boundary Smooth", 5)
    after, result_faces, result_uv = panorama_mesh(obj)
    assert result_faces == faces
    np.testing.assert_array_equal(result_uv, uv)
    np.testing.assert_array_equal(after[outside], before[outside])
    assert np.max(np.linalg.norm(after-before, axis=1)) > 1e-4


@pytest.mark.parametrize("source", ["definition", "asset"])
@pytest.mark.parametrize("kind", ["cutout", "plane", "panorama"])
@pytest.mark.parametrize("iterations,weight", [(0, 1.0), (8, 1.0), (8, 0.0)])
def test_boundary_loop_bypass_uses_iterations_only(source, kind, iterations, weight):
    from anyimage.common.node import node_asset_path
    from tests.support.planes import create_surface
    from tests.tools.nodes.test_image_depth_panorama import panorama
    from tools.nodes.arrangement.flow import origin_socket

    if kind == "cutout":
        obj, set_value = surface(resolution=4)
    elif kind == "plane":
        obj, set_value, *_ = create_surface("definition", "DEPTH")
    else:
        obj, set_value = panorama(subdivide=2)
    if source == "asset":
        with bpy.data.libraries.load(str(node_asset_path())) as (_, target):
            target.node_groups = [obj.modifiers[0].node_group.name]
        obj.modifiers[0].node_group = target.node_groups[0].copy()
    set_value("Split Threshold", 0)
    set_value("Smooth Weight", weight)
    set_value("Boundary Smooth", iterations)
    group = obj.modifiers[0].node_group
    entry = next(n for n in group.nodes if n.bl_idname == "GeometryNodeRepeatInput"
                 and origin_socket(n.inputs["Iterations"].links[0].from_socket).name == "Boundary Smooth")
    bypass = next(n for n in group.nodes if n.bl_idname == "GeometryNodeSwitch"
                  and n.input_type == "GEOMETRY" and n.inputs["True"].is_linked
                  and origin_socket(n.inputs["True"].links[0].from_socket).node == entry.paired_output)
    control = origin_socket(bypass.inputs["Switch"].links[0].from_socket)
    assert control.name == "Boundary Smooth"
    assert control.identifier == origin_socket(entry.inputs["Iterations"].links[0].from_socket).identifier
    assert origin_socket(bypass.inputs["False"].links[0].from_socket) == origin_socket(entry.inputs["Geometry"].links[0].from_socket)
    original, faces = evaluated(obj)
    group.links.remove(bypass.inputs["Switch"].links[0])
    bypass.inputs["Switch"].default_value = True
    obj.update_tag(refresh={"DATA"})
    bpy.context.view_layer.update()
    unskipped, unskipped_faces = evaluated(obj)
    np.testing.assert_array_equal(original, unskipped)
    assert faces == unskipped_faces