import blenderproc as bproc  # noqa
import mathutils  # type: ignore

from typing import List

from blenderproc.python.types.MeshObjectUtility import MeshObject, bmesh, Matrix


def create_bvh_tree_multi_objects(
    mesh_objects: List[MeshObject],
) -> mathutils.bvhtree.BVHTree:
    """Creates a bvh tree which contains multiple mesh objects.

    Such a tree is later used for fast raycasting.

    :param mesh_objects: The list of mesh objects that should be put into the BVH tree.
    :return: The built BVH tree.
    """
    # Create bmesh which will contain the meshes of all objects
    bm = bmesh.new()
    # Go through all mesh objects
    for obj in mesh_objects:
        # Get a copy of the mesh
        mesh = obj.get_mesh().copy()
        # Apply world matrix
        mesh.transform(Matrix(obj.get_local2world_mat()))
        # Add object mesh to bmesh
        bm.from_mesh(mesh)

        # Uncomment below line to avoid leaving orphan meshes
        # bpy.data.meshes.remove(mesh)

    # Create tree from bmesh
    bvh_tree = mathutils.bvhtree.BVHTree.FromBMesh(bm)
    bm.free()
    return bvh_tree


bproc.init()
mesh = bproc.object.create_primitive("SPHERE")
for i in range(500):
    create_bvh_tree_multi_objects([mesh] * 100)