Source code for ddg.visualization.blender.mesh

"""
Collection of functions for Blender Mesh objects.

The related functions on the level of bmesh are found in bmeshutils.py
module.

"""

import bmesh
import bpy
import numpy as np

import ddg.math.projective as putils


def from_bmesh(bm, name, free=False):
    """Write a BMesh to a new mesh.

    Parameters
    ----------
    bm : BMesh
    name : str
        The `.name` of the new mesh.
    free : bool (default=False)
        If True, `bm` is freed.

    Returns
    -------
    bpy.types.Mesh

    """
    me = bpy.data.meshes.new(name)
    bm.to_mesh(me)

    if free:
        bm.free()

    return me


###############        Create mesh function           #######################
#############################################################################


def add_root(collection=None, location=np.array((0, 0, 0)), empty_type='PLAIN_AXES',
             name="root", parent=None):
    if collection is None:
        collection = bpy.context.scene.collection
    root = bpy.data.objects.new(name, None)
    root.empty_display_type = empty_type
    root.location = location
    root.parent = parent

    collection.objects.link(root)

    return root


###############          Modify mesh functions        #######################
#############################################################################
def transform(mesh, M):
    """
    Transforms vertices of mesh with given 3x3 or 4x4 matrix

    Parameters
    ----------
    mesh : bpy.types.Mesh
        mesh to transform
    M : numpy.ndarray
        3x3 or 4x4 transformation matrix
    """

    if M.shape == (3, 3):
        for vert in mesh.vertices:
            vert.co = np.dot(M, np.array(vert.co))
    elif M.shape == (4, 4):
        for vert in mesh.vertices:
            vert.co = putils.dehomogenize(np.dot(M, putils.homogenize(np.array(vert.co))))
    else:
        raise ValueError("Shape of matrix must be (3,3) or (4,4).")

    return mesh


def shade_smooth(mesh_or_bobj, smooth=True):
    """Apply smooth shading on each polygon of mesh_or_bobj.

    Parameters
    ----------
    mesh_or_bobj: bpy.types.Mesh or bpy.types.Object
        The mesh or the object containign a mesh.
        This can also be a curve and will have no effect.

    Returns
    -------
    bpy.types.Mesh or bpy.types.Object
        The given mesh or blender object.
    """
    # Check if mesh object containing a mesh is given.
    if isinstance(mesh_or_bobj, bpy.types.Mesh):
        mesh = mesh_or_bobj
    elif isinstance(mesh_or_bobj.data, bpy.types.Mesh):
        mesh = mesh_or_bobj.data
    else:
        # Silently fail for objects not containing a mesh as data.
        # This allows to use this function curves.
        return

    for p in mesh.polygons:
        p.use_smooth = smooth
    mesh.update()
    return mesh_or_bobj


###############       Select functions        #######################
#####################################################################

def select_boundary_edges(mesh):
    bm = bmesh.from_edit_mesh(mesh)
    for edge in bm.edges:
        edge.select = edge.is_boundary
    mesh.update()


def select_interior_edges(mesh):
    bm = bmesh.from_edit_mesh(mesh)
    for edge in bm.edges:
        edge.select = not edge.is_boundary
    mesh.update()


def select_polygons(mesh, indices):
    for face in mesh.polygons:
        face.select = face.index in indices
    mesh.update()


def select_edges(mesh, indices):
    bm = bmesh.from_edit_mesh(mesh)
    for edge in bm.edges:
        edge.select = edge.index in indices
    mesh.update()


def select_vertices(mesh, indices):
    for face in mesh.vertices:
        face.select = face.index in indices
    mesh.update()


def _select_mesh_object(context, mesh_object, selector_function):
    bpy.ops.object.select_all(action='DESELECT')
    context.view_layer.objects.active = mesh_object
    bpy.ops.object.mode_set(mode='EDIT')
    selector_function(mesh_object.data)
    bpy.ops.object.mode_set(mode='OBJECT')


def select_interior_edges_mesh_object(context, mesh_object):
    _select_mesh_object(context, mesh_object, mutils.select_interior_edges)


def select_boundary_edges_mesh_object(context, mesh_object):
    _select_mesh_object(context, mesh_object, mutils.select_boundary_edges)


def select_faces_mesh_object(context, mesh_object, indices):
    _select_mesh_object(context, mesh_object,
                        lambda mesh_data: mutils.select_polygons(mesh_data, indices))


def select_edges_mesh_object(context, mesh_object, indices):
    _select_mesh_object(context, mesh_object,
                        lambda mesh_data: mutils.select_edges(mesh_data, indices))


def select_vertices_mesh_object(context, mesh_object, indices):
    _select_mesh_object(context, mesh_object,
                        lambda mesh_data: mutils.select_vertices(mesh_data, indices))