"""Verify split profile shaping and thickness-weighted normals."""

import bpy
import numpy as np
import pytest

from anyimage.common.object import modifier_input_identifier, set_modifier_input
from tests.support.depth_surface import surface, evaluated
from tests.tools.nodes.test_boundary_smoothing import diagonal_surface, boundary_neighbors, edge_band
from tools.nodes.common.cutout_boundary import remove_boundary_triangles
from tools.nodes.groups.image_depth_cutout import build_normal_weight, blend_normal


def profile_values(obj):
    result = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    mesh = result.to_mesh()
    try:
        return np.array([v.value for v in mesh.attributes["o_balloon"].data])
    finally:
        result.to_mesh_clear()


def test_boundary_triangle_cleanup_is_one_pass():
    obj, _ = surface(resolution=2)
    group = obj.modifiers[0].node_group
    mesh = bpy.data.meshes.new("Boundary ears")
    mesh.from_pydata([(0,0,0),(1,0,0),(1,1,0),(0,1,0),(.5,.5,0),(2,0,0)], [],
                     [(0,1,4),(1,2,4),(2,3,4),(3,0,4),(1,5,2)])
    obj.data = mesh
    source = group.nodes.new("NodeGroupInput")
    cleaned = remove_boundary_triangles(group, source.outputs["Geometry"])
    group.links.new(cleaned, next(n for n in group.nodes if n.type == "GROUP_OUTPUT").inputs["Geometry"])
    obj.update_tag(refresh={"DATA"})
    bpy.context.view_layer.update()
    _, faces = evaluated(obj)
    assert len(faces) == 4
    assert all(len(f) == 3 for f in faces)


@pytest.mark.parametrize("iterations", [0, 16])
def test_cleaned_sawtooth_smoothing_preserves_face_orientation(iterations):
    obj, set_value = diagonal_surface(triangles=True)
    original, faces = evaluated(obj)
    set_value("Boundary Smooth", iterations)
    points, actual = evaluated(obj)
    assert actual == faces
    indices = np.array(faces)
    def normals(p):
        return np.cross(p[indices[:,1]]-p[indices[:,0]], p[indices[:,2]]-p[indices[:,0]])
    assert np.all(np.sum(normals(original)*normals(points), axis=1) > 0)


@pytest.mark.parametrize("resolution,aspect", [(12, .5), (24, 2)])
def test_split_profile_tapers_within_boundary_band(resolution, aspect):
    obj, set_value = surface(step=6, resolution=resolution)
    for vertex in obj.data.vertices:
        vertex.co.x *= aspect
    image = bpy.data.images["Camera"]
    pixels = np.array(image.pixels[:], np.float32).reshape(32, 64, 4)
    pixels[..., 0] *= aspect
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    obj.data.attributes.new("o_balloon", "FLOAT", "POINT").data.foreach_set(
        "value", np.full(len(obj.data.vertices), .3, dtype=np.float32))
    set_value("Depth Scale", 0)
    np.testing.assert_allclose(profile_values(obj), .3)
    set_value("Depth Split", .1)
    points, faces = evaluated(obj)
    boundary = boundary_neighbors(faces)
    cut = [i for i in boundary if 0 < points[i,2] < 1 and .5*aspect < points[i,0] < 1.5*aspect]
    assert cut
    band = edge_band(faces, cut)
    outside = [i for i in range(len(points)) if i not in band]
    assert outside
    values = profile_values(obj)
    np.testing.assert_array_equal(values[cut], 0)
    np.testing.assert_allclose(values[outside], .3)
    assert np.isfinite(values).all()
    assert values.max() == pytest.approx(.3)
    set_value("Depth Split", 0)
    np.testing.assert_allclose(profile_values(obj), .3)


@pytest.mark.parametrize("mode", [0, 1])
@pytest.mark.parametrize("reference", [0, 2])
def test_normal_blend_weights_balloon_and_keeps_full_shell_smoothing(mode, reference):
    obj, set_value = surface(resolution=2)
    group = obj.modifiers[0].node_group
    # The production interface renames the Shell thickness socket after building.
    socket = next(s for s in group.interface.items_tree if s.item_type == "SOCKET" and s.in_out == "INPUT"
                  and s.name == "Thickness" and s.subtype == "DISTANCE")
    socket.name = "Shell Thickness"
    is_shell = group.nodes.new("FunctionNodeInputBool")
    is_shell.boolean = bool(mode)
    weight = build_normal_weight(group, is_shell.outputs[0])
    raw = group.nodes.new("ShaderNodeCombineXYZ")
    raw.inputs["Z"].default_value = 1
    blurred = group.nodes.new("ShaderNodeCombineXYZ")
    blurred.inputs["X"].default_value = 1
    result = blend_normal(group, raw.outputs[0], blurred.outputs[0], weight)
    store = group.nodes.new("GeometryNodeStoreNamedAttribute")
    store.data_type, store.domain = "FLOAT_VECTOR", "POINT"
    store.inputs["Name"].default_value = "normal_result"
    source = group.nodes.new("NodeGroupInput")
    group.links.new(source.outputs["Geometry"], store.inputs["Geometry"])
    group.links.new(result, store.inputs["Value"])
    group.links.new(store.outputs[0], next(n for n in group.nodes if n.type == "GROUP_OUTPUT").inputs["Geometry"])
    modifier = obj.modifiers[0]
    identifier = socket.identifier if mode else modifier_input_identifier(group, "Thickness", subtype="NONE")
    set_value("Reference Depth", reference)
    results = []
    for amount in (0, .001, .01, 1):
        set_modifier_input(modifier, identifier, amount * (max(reference, 1e-8) if mode else 1))
        obj.update_tag(refresh={"DATA"}); bpy.context.view_layer.update()
        ev = obj.evaluated_get(bpy.context.evaluated_depsgraph_get()); mesh = ev.to_mesh()
        try:
            results.append(np.array(mesh.attributes["normal_result"].data[0].vector[:]))
        finally:
            ev.to_mesh_clear()
    np.testing.assert_array_equal(results[-1], [1,0,0])
    if mode:
        np.testing.assert_array_equal(results, [[1,0,0]] * len(results))
    else:
        np.testing.assert_array_equal(results[0], [0,0,1])
        assert np.linalg.norm(results[1]-results[0]) < np.linalg.norm(results[2]-results[0]) < .001