"""Contracts for smooth depth volume, editable presets and normal masks."""
import bpy
import numpy as np
import pytest

from tools.nodes.groups.image_depth_cutout import build_volume_fields
from anyimage.common.object import modifier_input_identifier, set_modifier_input


def store_test_fields(group, geometry, fields):
    """Expose evaluated fields only in the test output mesh."""
    for name, field in fields.items():
        store = group.nodes.new("GeometryNodeStoreNamedAttribute")
        store.domain = "POINT"
        store.data_type = "FLOAT_VECTOR" if field.type == "VECTOR" else "FLOAT"
        store.inputs["Name"].default_value = name
        group.links.new(geometry, store.inputs["Geometry"])
        group.links.new(field, store.inputs["Value"])
        geometry = store.outputs["Geometry"]
    return geometry


@pytest.mark.parametrize("mode", [0, 1])
def test_normal_smooth_changes_thickness_direction_without_moving_zero_thickness(mode):
    from tests.support.depth_surface import surface, evaluated

    obj, set_value = surface(step=1, resolution=10)
    xy = np.array([v.co[:2] for v in obj.data.vertices])
    profile = .3 * np.sin(np.pi * xy[:, 0] / 2) * np.sin(np.pi * xy[:, 1])
    profile[abs(profile) < 1e-7] = 0
    obj.data.attributes.new("o_balloon", "FLOAT", "POINT").data.foreach_set("value", profile)
    set_value("Mode", mode)
    set_value("Front Inflation", .2)
    modifier = obj.modifiers[0]
    thickness = modifier_input_identifier(modifier.node_group, "Thickness", subtype="DISTANCE" if mode else "NONE")

    def result(amount, smoothing):
        set_modifier_input(modifier, thickness, amount)
        set_value("Normal Smooth", smoothing)
        return evaluated(obj)

    raw, faces = result(0, 16)
    smooth, smooth_faces = result(0, 512)
    np.testing.assert_array_equal(raw, smooth)
    assert faces == smooth_faces
    raw, faces = result(1, 16)
    smooth, smooth_faces = result(1, 512)
    assert faces == smooth_faces
    assert np.isfinite(smooth).all()
    assert np.max(np.abs(raw - smooth)) > 1e-5


def set_test_structure(monkeypatch, value):
    """Control structural amplitude without depending on storage nodes."""
    from tools.nodes.groups import image_depth_cutout
    build = image_depth_cutout.build_front_fields

    def fields(group, volume, normal, normal_weight):
        result = build(group, volume, normal, normal_weight)
        structure = group.nodes.new("ShaderNodeValue")
        structure.outputs[0].default_value = value
        result["structure"] = structure.outputs[0]
        return result

    monkeypatch.setattr(image_depth_cutout, "build_front_fields", fields)


def test_shell_thickness_extends_behind_the_tilted_front():
    from scipy.spatial import cKDTree
    from tests.support.depth_surface import surface, evaluated, assert_closed
    obj, set_value = surface(step=0, resolution=4)
    image = bpy.data.images["Camera"]
    pixels = np.array(image.pixels[:], dtype=np.float32).reshape(32, 64, 4)
    pixels[..., 2] = 1 + 0.4 * pixels[..., 0]
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    obj.update_tag(refresh={"DATA"})
    bpy.context.view_layer.update()
    front, _ = evaluated(obj)
    set_value("Thickness", 0.2)
    set_value("Front Inflation", 0.5)
    points, faces = evaluated(obj)
    expected = front - 0.2 * np.array((0.4, 0, 1)) / np.sqrt(1.16)
    assert cKDTree(points).query(expected)[0].max() < 1e-5
    assert cKDTree(points).query(front)[0].max() < 1e-5
    assert_closed(faces)


def test_balloon_shell_converges_to_front_with_thickness():
    from scipy.spatial import cKDTree
    from tests.support.depth_surface import surface, evaluated, assert_closed
    obj, set_value = surface(step=1, resolution=10)
    xy = np.array([v.co[:2] for v in obj.data.vertices])
    profile = 0.3 * np.sin(np.pi * xy[:, 0] / 2) * np.sin(np.pi * xy[:, 1])
    profile[abs(profile) < 1e-7] = 0
    obj.data.attributes.new("o_balloon", "FLOAT", "POINT").data.foreach_set("value", profile)
    set_value("Mode", 0)
    modifier = obj.modifiers[0]
    thickness = modifier_input_identifier(modifier.node_group, "Thickness", subtype="NONE")
    def amount(value):
        set_modifier_input(modifier, thickness, value)
        obj.update_tag(refresh={"DATA"})
        bpy.context.view_layer.update()
        return evaluated(obj)
    front, _ = amount(0)
    tree = cKDTree(front)
    gaps = []
    for value in (0.01, 0.001, 0.0001):
        points, faces = amount(value)
        assert_closed(faces)
        gaps.append(tree.query(points)[0].max())
    np.testing.assert_allclose(np.array(gaps[1:]) / gaps[:-1], 0.1, rtol=0.02)


def test_smooth_back_rejects_front_noise_and_preserves_center_depth():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    n = 41
    points = np.array([(x, y) for y in np.linspace(-1, 1, n) for x in np.linspace(-1, 1, n)])
    profile = np.sqrt(np.maximum(0, np.cos(points[:, 0]*np.pi/2)*np.cos(points[:, 1]*np.pi/2))) * 0.6
    boundary = np.max(abs(points), axis=1) > 0.999
    profile[boundary] = 0
    center = 0.3*points[:, 0] + 0.6*points[:, 1]
    noise = 0.08*np.sin(points[:, 0]*23)*np.sin(points[:, 1]*19)
    noise[boundary] = 0
    faces = [(y*n+x, y*n+x+1, (y+1)*n+x+1) for y in range(n-1) for x in range(n-1)]
    faces += [(y*n+x, (y+1)*n+x+1, (y+1)*n+x) for y in range(n-1) for x in range(n-1)]
    mesh = bpy.data.meshes.new("Noisy front")
    mesh.from_pydata(np.column_stack((points, center + profile + noise)), [], faces)
    mesh.attributes.new("o_balloon", "FLOAT", "POINT").data.foreach_set("value", profile.astype(np.float32))
    obj = bpy.data.objects.new("Noisy front", mesh)
    bpy.context.collection.objects.link(obj)
    group = bpy.data.node_groups.new("Volume field test", "GeometryNodeTree")
    for name, kind in (("Geometry", "NodeSocketGeometry"), ("Thickness", "NodeSocketFloat")):
        socket = group.interface.new_socket(name=name, in_out="INPUT", socket_type=kind)
        if name == "Thickness":
            socket.default_value = 1
    group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    source = group.nodes.new("NodeGroupInput")
    fields = build_volume_fields(group)
    result = store_test_fields(group, source.outputs["Geometry"], {"thickness": fields["thickness"]})
    group.links.new(result, group.nodes.new("NodeGroupOutput").inputs["Geometry"])
    obj.modifiers.new("Volume", "NODES").node_group = group
    bpy.context.view_layer.update()
    evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    result = evaluated.to_mesh()
    try:
        thickness = np.array([item.value for item in result.attributes["thickness"].data])
        rear = center + profile + noise - thickness
        error = rear - (center - profile)
        assert max(abs(error)) < 1e-5
        assert np.sqrt(np.mean(error**2)) < np.sqrt(np.mean(noise**2))*0.01
        assert max(abs(thickness[boundary])) == 0
    finally:
        evaluated.to_mesh_clear()


def test_depth_cutout_normal_map_only_affects_front_faces():
    from tests.support.depth_surface import surface
    obj, set_value = surface()
    set_value("Thickness", 0.3)
    bpy.context.view_layer.update()
    evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    mesh = evaluated.to_mesh()
    try:
        attribute = mesh.attributes["o_normal_reduction"]
        assert attribute.domain == "FACE"
        weights = [1.0 - item.value for item in attribute.data]
        assert weights.count(1) == len(obj.data.polygons)
        assert weights.count(0) == len(mesh.polygons) - len(obj.data.polygons)
    finally:
        evaluated.to_mesh_clear()


@pytest.mark.parametrize("thickness,inflation", [
    (0.0, 1.0), (0.02, 0.0), (0.02, 0.0001),
    (0.02, 0.5), (0.1, 0.5), (0.1, 1.0),
])
def test_shell_ignores_front_inflation(thickness, inflation):
    from tests.support.depth_surface import surface, evaluated, assert_closed
    obj, set_value = surface(step=0, resolution=10)
    xy = np.array([v.co[:2] for v in obj.data.vertices])
    profile = 0.3 * np.sin(np.pi * xy[:, 0] / 2) * np.sin(np.pi * xy[:, 1])
    profile[abs(profile) < 1e-7] = 0
    obj.data.attributes.new("o_balloon", "FLOAT", "POINT").data.foreach_set("value", profile)
    set_value("Thickness", thickness)
    set_value("Front Inflation", 0.0)
    baseline, faces = evaluated(obj)
    set_value("Front Inflation", inflation)
    inflated, inflated_faces = evaluated(obj)
    assert faces == inflated_faces
    np.testing.assert_array_equal(inflated, baseline)
    if thickness:
        assert_closed(inflated_faces)


def test_balloon_suppression_is_stronger_without_a_startup_threshold(monkeypatch):
    from tests.support.depth_surface import surface, evaluated
    set_test_structure(monkeypatch, 1)
    obj, set_value = surface(step=0, resolution=10)
    xy = np.array([v.co[:2] for v in obj.data.vertices])
    profile = 0.3 * np.sin(np.pi * xy[:, 0] / 2) * np.sin(np.pi * xy[:, 1])
    profile[abs(profile) < 1e-7] = 0
    obj.data.attributes.new("o_balloon", "FLOAT", "POINT").data.foreach_set("value", profile)
    group = obj.modifiers[0].node_group
    set_value("Mode", 0)
    thickness_id = modifier_input_identifier(group, "Thickness", subtype="NONE")
    previous = 0
    for thickness in (0.001, 0.01, 0.1, 0.5, 1):
        set_modifier_input(obj.modifiers[0], thickness_id, thickness)
        set_value("Front Inflation", 0)
        baseline, _ = evaluated(obj)
        set_value("Front Inflation", 1)
        inflated, _ = evaluated(obj)
        growth = np.linalg.norm(inflated - baseline, axis=1).max() / profile.max()
        assert growth > previous
        previous = growth
    assert growth == pytest.approx(0.323, abs=0.001)