"""Check adaptive outline batches against a numerical graph reference."""

from collections import Counter

import bpy
import numpy as np
import pytest

from tests.support.depth_surface import evaluated
from tools.nodes.common.outline_smoothing import smooth_outline_camera
from tools.nodes.common.depth_surface import CUT_ATTRIBUTE


@pytest.mark.parametrize("steps", [1, 4])
def test_outline_batches_update_weights_preserve_split_and_camera_rays(steps):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    points = np.array([(x, y, 0) for y in range(9) for x in range(9)], dtype=float)
    faces = [(y*9+x, y*9+x+1, (y+1)*9+x+1, (y+1)*9+x) for y in range(8) for x in range(8)]
    deleted = [i % 8 == 4 for i in range(len(faces))]
    result_faces = [face for face, remove in zip(faces, deleted) if not remove]
    cut = np.isin(points[:, 0], [4, 5])
    depth = np.full(81, 2.0)
    depth[2] = 20
    depth[80] = 0
    camera = np.column_stack((points[:, :2]*depth[:, None]/10, depth))
    distance = np.linalg.norm(camera, axis=1)
    mesh = bpy.data.meshes.new("Split outline")
    mesh.from_pydata(points.tolist(), [], faces)
    for name, kind, domain, values, prop in (
        ("test_camera", "FLOAT_VECTOR", "POINT", camera.ravel(), "vector"),
        (CUT_ATTRIBUTE, "FLOAT", "POINT", cut, "value"),
        ("test_deleted", "BOOLEAN", "FACE", deleted, "value"),
    ):
        mesh.attributes.new(name, kind, domain).data.foreach_set(prop, values)
    obj = bpy.data.objects.new("Outline", mesh)
    bpy.context.collection.objects.link(obj)
    group = bpy.data.node_groups.new("Adaptive outline test", "GeometryNodeTree")
    for direction in ("INPUT", "OUTPUT"):
        group.interface.new_socket(name="Geometry", in_out=direction, socket_type="NodeSocketGeometry")
    ids = {}
    for name in ("Outline Depth Fix", "Boundary Smooth"):
        ids[name] = group.interface.new_socket(name=name, in_out="INPUT", socket_type="NodeSocketInt").identifier
    source = group.nodes.new("NodeGroupInput")
    attributes = {}
    for name, kind in (("test_camera", "FLOAT_VECTOR"), ("test_deleted", "BOOLEAN")):
        node = group.nodes.new("GeometryNodeInputNamedAttribute")
        node.data_type = kind
        node.inputs["Name"].default_value = name
        attributes[name] = node.outputs["Attribute"]
    delete = group.nodes.new("GeometryNodeDeleteGeometry")
    delete.domain, delete.mode = "FACE", "ONLY_FACE"
    group.links.new(source.outputs["Geometry"], delete.inputs["Geometry"])
    group.links.new(attributes["test_deleted"], delete.inputs["Selection"])
    scale = group.nodes.new("ShaderNodeValue")
    scale.outputs[0].default_value = 1
    result = smooth_outline_camera(group, source.outputs["Geometry"], delete.outputs[0], attributes["test_camera"], scale.outputs[0])
    position = group.nodes.new("GeometryNodeSetPosition")
    group.links.new(delete.outputs[0], position.inputs["Geometry"])
    group.links.new(result, position.inputs["Position"])
    output = group.nodes.new("NodeGroupOutput")
    group.links.new(position.outputs[0], output.inputs["Geometry"])
    modifier = obj.modifiers.new("Outline", "NODES")
    modifier.node_group = group

    # ONLY_FACE retains the separating column's edges in the current topology.
    edges = np.array([tuple(e.vertices) for e in mesh.edges])
    adjacency = np.zeros((81, 81))
    adjacency[edges[:, 0], edges[:, 1]] = 1
    adjacency += adjacency.T
    degree = adjacency.sum(axis=1)
    def blur(value, weight):
        return (value + weight*(adjacency @ value))/(1 + weight*degree)
    counts = Counter(tuple(sorted((a, b))) for f in result_faces for a, b in zip(f, (*f[1:], f[0])))
    boundary = np.zeros(81)
    for edge, count in counts.items():
        if count == 1:
            boundary[list(edge)] = 1
    protected = boundary.astype(bool) & cut
    seed = boundary*(~protected)
    influence = np.maximum(seed, blur(blur(seed, 1), 1))*(~protected)
    influence *= depth > 1e-8
    for rounds, batch in ((0, steps), (16, 0), (16, steps)):
        modifier[ids["Outline Depth Fix"]] = rounds
        modifier[ids["Boundary Smooth"]] = batch
        obj.update_tag(refresh={"DATA"})
        bpy.context.view_layer.update()
        actual, actual_faces = evaluated(obj)
        expected = distance.copy()
        for _ in range(rounds):
            ratio = abs(expected-blur(expected, .5))
            t = np.clip((ratio-.25)/.25, 0, 1)
            weight = t*t*(3-2*t)*influence
            for _ in range(batch):
                expected = blur(expected, weight)
        np.testing.assert_allclose(np.linalg.norm(actual, axis=1), expected, atol=2e-5)
        assert actual_faces == result_faces
        np.testing.assert_array_equal(actual[influence == 0], camera[influence == 0].astype(np.float32))
        if rounds and batch:
            assert 2 < actual[2, 2] < 15
            valid = depth > 0
            np.testing.assert_allclose(actual[valid, :2]/actual[valid, 2:3], camera[valid, :2]/camera[valid, 2:3], atol=1e-6)
        else:
            np.testing.assert_array_equal(actual, camera.astype(np.float32))
    assert not any(n.bl_idname in {"GeometryNodeCaptureAttribute", "GeometryNodeStoreNamedAttribute"} for n in group.nodes)