"""Verify saved panorama assets using synthetic validity and depth inputs."""

import json

import bpy
import numpy as np


def socket_id(group, name):
    return next(item.identifier for item in group.interface.items_tree
                if item.item_type == "SOCKET" and item.in_out == "INPUT" and item.name == name)


def evaluated_mesh(obj):
    obj.update_tag()
    bpy.context.view_layer.update()
    evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    return evaluated, evaluated.to_mesh()


def evaluated_points(obj):
    evaluated, mesh = evaluated_mesh(obj)
    points = np.empty(len(mesh.vertices) * 3, dtype=np.float32)
    mesh.vertices.foreach_get("co", points)
    counts = (len(mesh.vertices), len(mesh.edges), len(mesh.polygons))
    evaluated.to_mesh_clear()
    return points.reshape(-1, 3), counts


def test_image(name, alpha=1):
    image = bpy.data.images.new(name, width=16, height=8, alpha=True, float_buffer=True)
    image.colorspace_settings.name = "Non-Color"
    image.alpha_mode = "CHANNEL_PACKED"
    pixels = np.full((8, 16, 4), 3, dtype=np.float32)
    pixels[..., 3] = alpha
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    return image


def verify(output):
    full, empty, partial = test_image("Full"), test_image("Empty", 0), test_image("Combined Alpha", 0.25)
    half = test_image("Half")
    pixels = np.full((8, 16, 4), 3, dtype=np.float32)
    pixels[..., 3] = 1
    pixels[:, :8, 3] = 0
    half.pixels.foreach_set(pixels.ravel())
    half.update()
    results = {}
    try:
        for obj in list(bpy.data.objects):
            if obj.type != "MESH" or not obj.modifiers:
                continue
            modifier = obj.modifiers[0]
            group = modifier.node_group
            image_inputs = [item.name for item in group.interface.items_tree if item.item_type == "SOCKET"
                            and item.in_out == "INPUT" and item.socket_type == "NodeSocketImage"]
            assert image_inputs == ["Distance"]
            names = ("Subdivide", "Depth Scale", "Validity Threshold", "Distance")
            keys = {name: socket_id(group, name) for name in names}
            original = {name: modifier[key] for name, key in keys.items()}

            def set_inputs(**values):
                for name, value in values.items():
                    modifier[keys[name.replace("_", " ")]] = value

            try:
                set_inputs(Distance=full, Depth_Scale=0, Validity_Threshold=0.9)
                face_counts = []
                for level in (0, 2, 4):
                    set_inputs(Subdivide=level)
                    evaluated, mesh = evaluated_mesh(obj)
                    try:
                        assert len(mesh.polygons) == 6 * 4 ** level
                        assert all(len(face.vertices) == 4 for face in mesh.polygons)
                        assert len(mesh.vertices) - len(mesh.edges) + len(mesh.polygons) == 2
                        use = np.zeros(len(mesh.edges), dtype=int)
                        for loop in mesh.loops:
                            use[loop.edge_index] += 1
                        assert np.all(use == 2)
                        points = np.array([vertex.co[:] for vertex in mesh.vertices])
                        np.testing.assert_allclose(np.linalg.norm(points, axis=1), 1, atol=1e-5)
                        if level >= 2:
                            uv = mesh.attributes["UVMap"].data
                            assert all(np.ptp([uv[i].vector.x for i in face.loop_indices]) < 0.51
                                       for face in mesh.polygons), "UV seam crosses the full image"
                        face_counts.append(len(mesh.polygons))
                    finally:
                        evaluated.to_mesh_clear()
                sphere, _ = evaluated_points(obj)
                set_inputs(Depth_Scale=1)
                surface, _ = evaluated_points(obj)
                np.testing.assert_allclose(surface, sphere * 3, atol=1e-5)
                set_inputs(Depth_Scale=0.5)
                middle, _ = evaluated_points(obj)
                np.testing.assert_allclose(middle, sphere * 2, atol=1e-5)
                set_inputs(Distance=empty)
                assert evaluated_points(obj)[1][2] == 0
                set_inputs(Distance=half)
                retained = evaluated_points(obj)[1][2]
                assert 0 < retained < face_counts[-1]
                set_inputs(Distance=partial, Validity_Threshold=0.3)
                assert evaluated_points(obj)[1][2] == 0, "Mask times alpha must be below threshold"
                set_inputs(Validity_Threshold=0.2)
                assert evaluated_points(obj)[1][2] == face_counts[-1]
            finally:
                for name, value in original.items():
                    modifier[keys[name]] = value
                obj.update_tag()
            points, counts = evaluated_points(obj)
            assert np.isfinite(points).all()
            assert counts[2] > 0
            evaluated, mesh = evaluated_mesh(obj)
            try:
                directions = np.empty((counts[0], 3), dtype=np.float32)
                uv = np.empty_like(directions)
                mesh.attributes["panorama_direction"].data.foreach_get("vector", directions.ravel())
                mesh.attributes["panorama_sample_uv"].data.foreach_get("vector", uv.ravel())
                expected_u = np.arctan2(directions[:, 1], directions[:, 0]) / (2 * np.pi) % 1
                delta_u = np.abs(uv[:, 0] - expected_u)
                assert np.max(np.minimum(delta_u, 1 - delta_u)) < 1e-6
                np.testing.assert_allclose(uv[:, 1], 1 - np.arccos(np.clip(directions[:, 2], -1, 1)) / np.pi, atol=1e-6)
                image = original["Distance"]
                width, height = image.size
                pixels = np.array(image.pixels[:]).reshape(height, width, 4)
                data_path = output / ("reference.npz" if obj.name == "Reference Room" else "prediction.npz")
                with np.load(data_path) as data:
                    np.testing.assert_allclose(pixels[::-1, :, 3], data["alpha"], atol=0.0005)
                x = np.clip((uv[:, 0] * width).astype(int), 0, width - 1)
                y = np.clip((uv[:, 1] * height).astype(int), 0, height - 1)
                radius = 1 + original["Depth Scale"] * (pixels[y, x, 0] - 1)
                np.testing.assert_allclose(points, directions * radius[:, None], rtol=2e-6, atol=2e-5)
            finally:
                evaluated.to_mesh_clear()
            results[obj.name] = {"subdivide_levels": [0, 2, 4], "quad_face_counts": face_counts,
                                 "closed_quad_sphere": True, "uv_seam_verified": True,
                                 "depth_scale_verified": True, "mask_alpha_verified": True,
                                 "image_sampling_verified": True,
                                 "single_depth_image_verified": True,
                                 "packed_depth_alpha_verified": True,
                                 "vertices": counts[0], "faces": counts[2]}
        assert results
    finally:
        for image in (full, empty, partial, half):
            bpy.data.images.remove(image)
    (output / "blender_validation.json").write_text(json.dumps(results, indent=2), encoding="utf-8")
    print(json.dumps(results, indent=2), flush=True)