"""Check projected clipping, fixed-plane symmetry and straight closed walls."""
from collections import Counter

import bpy
import bmesh
import numpy as np
import pytest
from mathutils import Euler, Vector
from scipy.spatial import cKDTree

from anyimage.common.object import modifier_input_identifier, set_modifier_input
from anyimage.operators.cutout_tool.object import load_cutout_node_group
from anyimage.operators.cutout_tool.shape import CUTOUT_Z_TO_X
from tests.support.depth_surface import evaluated
from tools.nodes.groups.image_depth_symmetry import build_image_depth_symmetry_group
from tools.nodes.common.normal_map import ROTATION_ATTRIBUTE, FACE_ATTRIBUTE, AXIS_ATTRIBUTE


@pytest.fixture(params=['definition', 'asset'])
def symmetry(request):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    group = (build_image_depth_symmetry_group() if request.param == 'definition'
             else load_cutout_node_group('DEPTH_SYMMETRY'))
    count = 17
    y, x = np.meshgrid(np.linspace(-1, 1, count), np.linspace(-1, 1, count), indexing='ij')
    points = np.column_stack((x.ravel(), y.ravel(), np.zeros(x.size)))
    faces = []
    for row in range(count - 1):
        for col in range(count - 1):
            a = row * count + col
            faces.extend(((a, a + 1, a + count + 1), (a, a + count + 1, a + count)))
    mesh = bpy.data.meshes.new('Curved Depth')
    mesh.from_pydata(points, [], faces)
    uv = mesh.uv_layers.new(name='UVMap')
    for loop in mesh.loops:
        row, col = divmod(loop.vertex_index, count)
        uv.data[loop.index].uv = ((col + .5) / count, (row + .5) / count)
    profile = .3 * np.maximum(0, 1 - np.maximum(abs(x), abs(y)))
    mesh.attributes.new('o_balloon', 'FLOAT', 'POINT').data.foreach_set('value', profile.ravel())
    mesh.attributes.new('user_probe', 'FLOAT', 'POINT').data.foreach_set('value', np.ones(x.size))
    image = bpy.data.images.new('Depth', width=count, height=count, float_buffer=True)
    image.colorspace_settings.name = 'Non-Color'
    image.alpha_mode = 'CHANNEL_PACKED'
    depth = 1 + .6 * x ** 2 + .3 * y ** 2
    image.pixels.foreach_set(np.dstack((x, -y, depth, np.ones_like(x))).astype(np.float32).ravel())
    obj = bpy.data.objects.new('Symmetry', mesh)
    bpy.context.collection.objects.link(obj)
    modifier = obj.modifiers.new('Symmetry', 'NODES')
    modifier.node_group = group

    def set_value(name, value):
        set_modifier_input(modifier, modifier_input_identifier(group, name), value)
        obj.update_tag(refresh={'DATA'})
        bpy.context.view_layer.update()

    for name, value in [('Mode', 1), ('Depth Image', image), ('Reference Depth', 1.45),
                        ('Depth Scale', 1.), ('Uniform Scale', 1.), ('Boundary Smooth', 0),
                        ('Smooth', 0), ('Double Sided', True)]:
        set_value(name, value)
    yield obj, set_value, depth.ravel(), faces
    bpy.ops.wm.read_factory_settings(use_empty=True)


@pytest.mark.parametrize('direction,smooth,scale', [((0, 0, 1), 0, .001), ((-.4, .2, 1), 0, 1),
                                                 ((.2, -.3, 1), 3, 100)])
def test_projected_volume_is_closed_symmetric_and_has_straight_walls(symmetry, direction, smooth, scale):
    obj, set_value, _, _ = symmetry
    for name, value in [('Depth Direction', direction), ('Smooth', smooth),
                        ('Uniform Scale', scale), ('Reference Depth', 1.45 * scale)]:
        set_value(name, value)
    result = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    mesh = result.to_mesh()
    try:
        points = np.array([v.co[:] for v in mesh.vertices])
        assert len(points) and np.isfinite(points).all()
        assert cKDTree(points).query(points * (1, 1, -1))[0].max() < scale * 1e-5
        bm = bmesh.new()
        bm.from_mesh(mesh)
        assert all(v.is_manifold for v in bm.verts)
        assert all(e.is_manifold for e in bm.edges)
        bm.free()
        orientation = Counter()
        wall_count = 0
        for polygon in mesh.polygons:
            face = tuple(polygon.vertices)
            for a, b in zip(face, (*face[1:], face[0])):
                orientation[tuple(sorted((a, b)))] += 1 if a < b else -1
            heights = points[list(face), 2]
            if heights.min() < 0 < heights.max():
                wall_count += 1
                assert abs(polygon.normal.z) < 1e-5
                assert not polygon.use_smooth
        assert wall_count and not any(orientation.values())
        assert 'UVMap' in mesh.uv_layers and 'user_probe' in mesh.attributes
    finally:
        result.to_mesh_clear()


def test_projected_deletes_whole_crossing_faces_and_supports_empty_output(symmetry):
    obj, set_value, depth, source_faces = symmetry
    set_value('Double Sided', False)
    points, faces = evaluated(obj)
    assert (points[:, 2] > 0).all()
    expected = [face for face in source_faces if np.all(depth[list(face)] < 1.45 - 1e-7)]
    assert len(faces) == len(expected)
    assert set(np.concatenate(faces)) == set(range(len(points)))
    set_value('Reference Depth', .5)
    points, faces = evaluated(obj)
    assert not len(points) and not faces


def test_mode_switch_and_output_axis_preserve_state(symmetry):
    obj, set_value, _, _ = symmetry
    projected, projected_faces = evaluated(obj)
    set_value('Mode', 0)
    balloon, balloon_faces = evaluated(obj)
    assert len(balloon_faces) != len(projected_faces)
    set_value('Mode', 1)
    restored, restored_faces = evaluated(obj)
    np.testing.assert_array_equal(restored, projected)
    assert restored_faces == projected_faces
    set_value('Depth Axis', 1)
    rotated, rotated_faces = evaluated(obj)
    np.testing.assert_allclose(rotated, projected @ np.array(CUTOUT_Z_TO_X)[:3, :3].T, atol=1e-6)
    assert rotated_faces == projected_faces
    set_value('Depth Axis', 0)
    set_value('Mode', 0)
    restored, restored_faces = evaluated(obj)
    np.testing.assert_array_equal(restored, balloon)
    assert restored_faces == balloon_faces


def test_symmetry_asset_has_direct_mode_branches(symmetry):
    obj, _, _, _ = symmetry
    group = obj.modifiers[0].node_group
    assert not any(node.bl_idname in {'GeometryNodeMeshBoolean', 'GeometryNodeGroup',
                                     'GeometryNodeCaptureAttribute'} for node in group.nodes)
    selector = next(node for node in group.nodes if node.bl_idname == 'GeometryNodeMenuSwitch'
                    and node.inputs['Menu'].links[0].from_socket.name == 'Mode')
    assert [item.name for item in selector.enum_items] == ['Balloon', 'Projected']
    inputs = {item.name: item for item in group.interface.items_tree
              if item.item_type == 'SOCKET' and item.in_out == 'INPUT'}
    for name in ('Depth Direction', 'Depth Offset'):
        assert inputs[name].parent == inputs['Depth Scale'].parent
    for name in ('Double Sided', 'Reference Depth', 'Boundary Smooth'):
        assert inputs[name].parent.name == 'Options'
    assert inputs['Boundary Smooth'].default_value == 1
    assert inputs['Mode'].default_value == 'Projected'
    assert inputs['Fill Sides'].default_value is True
    assert inputs['Depth Offset'].default_value == 0
    assert inputs['Depth Offset'].min_value < 0


@pytest.mark.parametrize('mode', [0, 1])
def test_depth_offset_changes_depth_after_direction_without_moving_xy(symmetry, mode):
    obj, set_value, _, _ = symmetry
    for name, value in [('Mode', mode), ('Depth Direction', (.3, -.2, 1)),
                        ('Reference Depth', 4.), ('Depth Scale', .6)]:
        set_value(name, value)
    original, faces = evaluated(obj)
    for offset in (.2, -.1):
        set_value('Depth Offset', offset)
        points, changed_faces = evaluated(obj)
        assert changed_faces == faces
        np.testing.assert_allclose(points[:, :2], original[:, :2], atol=1e-6)
        assert cKDTree(points).query(points * (1, 1, -1))[0].max() < 1e-6
        delta = abs(points[:, 2]) - abs(original[:, 2])
        if mode == 1:
            np.testing.assert_allclose(delta, offset, atol=1e-6)
        else:
            if offset > 0:
                assert delta.max() > .1
            else:
                assert delta.min() < -.05
            boundary = abs(original[:, 2]) < 1e-8
            np.testing.assert_allclose(points[boundary, 2], 0, atol=1e-8)
    set_value('Depth Offset', 0.)
    restored, _ = evaluated(obj)
    np.testing.assert_allclose(restored, original, atol=1e-6)
    if mode == 1:
        set_value('Depth Offset', -10.)
        points, faces = evaluated(obj)
        assert not len(points) and not faces


def test_material_attributes_match_evaluated_faces_and_direction(symmetry):
    obj, set_value, _, _ = symmetry
    direction = Vector((.3, -.2, 1)).normalized()
    set_value('Depth Direction', direction)
    for axis in (0, 1):
        set_value('Depth Axis', axis)
        result = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
        mesh = result.to_mesh()
        try:
            points = np.array([v.co[:] for v in mesh.vertices])
            if axis:
                points = points @ np.array(CUTOUT_Z_TO_X)[:3, :3]
            for name in (ROTATION_ATTRIBUTE, FACE_ATTRIBUTE, AXIS_ATTRIBUTE):
                assert mesh.attributes[name].domain == 'FACE'
            for polygon in mesh.polygons:
                index = polygon.index
                rotation = Euler(mesh.attributes[ROTATION_ATTRIBUTE].data[index].vector)
                np.testing.assert_allclose(rotation.to_matrix() @ Vector((0, 0, 1)), direction, atol=1e-6)
                heights = points[list(polygon.vertices), 2]
                kind = 1 if heights.min() > 0 else 2 if heights.max() < 0 else 3
                assert mesh.attributes[FACE_ATTRIBUTE].data[index].value == kind
                assert mesh.attributes[AXIS_ATTRIBUTE].data[index].value == axis
        finally:
            result.to_mesh_clear()
    set_value('Mode', 0)
    result = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    assert ROTATION_ATTRIBUTE not in result.data.attributes
    assert FACE_ATTRIBUTE not in result.data.attributes


def test_fill_sides_only_controls_projected_walls(symmetry):
    obj, set_value, _, _ = symmetry
    closed, closed_faces = evaluated(obj)
    set_value('Fill Sides', False)
    points, faces = evaluated(obj)
    assert len(faces) < len(closed_faces)
    assert cKDTree(points).query(closed)[0].max() < 1e-6
    assert cKDTree(points).query(points * (1, 1, -1))[0].max() < 1e-6
    assert all(not (points[list(face), 2].min() < 0 < points[list(face), 2].max()) for face in faces)
    edges = Counter(tuple(sorted((a, b))) for face in faces
                    for a, b in zip(face, (*face[1:], face[0])))
    assert any(count == 1 for count in edges.values())
    set_value('Fill Sides', True)
    restored, restored_faces = evaluated(obj)
    np.testing.assert_array_equal(restored, closed)
    assert restored_faces == closed_faces
    for name, value in [('Double Sided', False), ('Mode', 0)]:
        set_value(name, value)
        before, before_faces = evaluated(obj)
        set_value('Fill Sides', False)
        after, after_faces = evaluated(obj)
        np.testing.assert_array_equal(after, before)
        assert after_faces == before_faces
        set_value('Fill Sides', True)