BallanceBlenderHelper/bbp_ng/UTIL_blender_mesh.py

import bpy, bmesh, mathutils
import typing, array, collections
from . import UTIL_functions, UTIL_virtools_types

## Blender Mesh Usage
#  This module create a universal mesh visitor, including MeshReader, MeshWriter and MeshUVModifier
#  for every other possible module using.
#  Obviously, MeshReader is served for 2 exporter, MeshWriter is served for 2 importer.
#  MeshWriter also served for BMERevenge module and Ballance element loading.
#  MeshUVModifier is used by Flatten UV and Rail UV.
#

#region Assist Functions

class FaceVertexData():
    mPosIdx: int
    mNmlIdx: int
    mUvIdx: int
    
    def __init__(self, pos: int = 0, nml: int = 0, uv: int = 0):
        self.mPosIdx = pos
        self.mNmlIdx = nml
        self.mUvIdx = uv

class FaceData():
    ## @remark List or tuple. List is convenient for adding and removing
    mIndices: tuple[FaceVertexData, ...] | list[FaceVertexData]
    ## Face used material slot index
    #  @remark If material slot is empty, or this face do not use material, set this value to 0.
    mMtlIdx: int
    
    def __init__(self, indices: tuple[FaceVertexData, ...] | list[FaceVertexData] = tuple(), mtlidx: int = 0):
        self.mIndices = indices
        self.mMtlIdx = mtlidx
    
    def conv_co(self) -> None:
        """
        Change indice order between Virtools and Blender
        """
        if isinstance(self.mIndices, list):
            self.mIndices.reverse()
        elif isinstance(self.mIndices, tuple):
            self.mIndices = self.mIndices[::-1]
        else:
            raise UTIL_functions.BBPException('invalid indices container.')
    
    def is_indices_legal(self) -> bool:
        return len(self.mIndices) >= 3

class MeshWriterIngredient():
    mVertexPosition: typing.Iterator[UTIL_virtools_types.VxVector3] | None
    mVertexNormal: typing.Iterator[UTIL_virtools_types.VxVector3] | None
    mVertexUV: typing.Iterator[UTIL_virtools_types.VxVector2] | None
    mFace: typing.Iterator[FaceData] | None
    mMaterial: typing.Iterator[bpy.types.Material | None] | None
    
    def __init__(self):
        self.mVertexPosition = None
        self.mVertexNormal = None
        self.mVertexUV = None
        self.mFace = None
        self.mMaterial = None
    
    def is_valid(self) -> bool:
        if self.mVertexPosition is None: return False
        if self.mVertexNormal is None: return False
        if self.mVertexUV is None: return False
        if self.mFace is None: return False
        if self.mMaterial is None: return False
        return True

def _flat_vxvector3(it: typing.Iterator[UTIL_virtools_types.VxVector3]) -> typing.Iterator[float]:
    for entry in it:
        yield entry.x
        yield entry.y
        yield entry.z

def _flat_vxvector2(it: typing.Iterator[UTIL_virtools_types.VxVector2]) -> typing.Iterator[float]:
    for entry in it:
        yield entry.x
        yield entry.y

def _flat_face_nml_index(nml_idx: array.array, nml_array: array.array) -> typing.Iterator[float]:
    for idx in nml_idx:
        pos: int = idx * 3
        yield nml_array[pos]
        yield nml_array[pos + 1]
        yield nml_array[pos + 2]

def _flat_face_uv_index(uv_idx: array.array, uv_array: array.array) -> typing.Iterator[float]:
    for idx in uv_idx:
        pos: int = idx * 2
        yield uv_array[pos]
        yield uv_array[pos + 1]

def _nest_custom_split_normal(nml_array: array.array) -> typing.Iterator[UTIL_virtools_types.ConstVxVector3]:
    # following statement create a iterator for normals array by `iter(nml_array)`
    # then triple it, because iterator is a reference type, so 3 items of this tuple is pointed to the same iterator and share the same iteration progress.
    # then use star macro to pass it to zip, it will cause zip receive 3 params pointing to the same iterator.
    # now zip() will call 3 param‘s __next__() function from left to right.
    # zip will get following iteration result because all iterator are the same one: (0, 1, 2), (3, 4, 5) and etc (number is index to corresponding value).
    # finally, use tuple to expand it to a tuple, not a generator.
    return tuple(zip(*(iter(nml_array), ) * 3))

class TemporaryMesh():
    """
    Create a temporary mesh for convenient exporting.
    When exporting mesh, we need evaluate it first, then triangulate it.
    We create a temporary mesh to hold the evaluated and triangulated mesh result.
    So that original object will not be affected and keep its original geometry.
    Please note passed bpy.types.Object must be an object which can be converted into mesh.
    You can use this class provided static method to check it.
    """

    __cAllowedObjectType: typing.ClassVar[set[str]] = {
        'MESH', 'CURVE', 'SURFACE', 'FONT', 'META'
    }

    @staticmethod
    def has_geometry(obj: bpy.types.Object):
        """
        Check whether given Blender object has geometry.
        If it has, it can safely utilize this class for visiting mesh.
        """
        return obj.type in TemporaryMesh.__cAllowedObjectType
    
    __mBindingObject: bpy.types.Object
    __mEvaluatedObject: bpy.types.Object
    __mTempMesh: bpy.types.Mesh
    
    def __init__(self, binding_obj: bpy.types.Object):
        depsgraph = bpy.context.evaluated_depsgraph_get()
        self.__mBindingObject = binding_obj
        self.__mEvaluatedObject = self.__mBindingObject.evaluated_get(depsgraph)
        self.__mTempMesh = self.__mEvaluatedObject.to_mesh()
        
        if self.__mTempMesh is None:
            raise UTIL_functions.BBPException('try getting mesh from an object without mesh.')
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.dispose()
    
    def is_valid(self) -> bool:
        if self.__mBindingObject is None: return False
        if self.__mBindingObject is None: return False
        if self.__mTempMesh is None: return False
        return True
    
    def dispose(self) -> None:
        if self.is_valid():
            self.__mTempMesh = None
            self.__mEvaluatedObject.to_mesh_clear()
            self.__mEvaluatedObject = None
            self.__mBindingObject = None
    
    def get_temp_mesh(self) -> bpy.types.Mesh:
        if not self.is_valid():
            raise UTIL_functions.BBPException('try calling invalid TemporaryMesh.')
        return self.__mTempMesh

#endregion

class MeshReader():
    """
    The passed mesh must be created by bpy.types.Object.to_mesh() and destroyed by bpy.types.Object.to_mesh_clear().
    Because this class must trianglate mesh. To prevent change original mesh, this operations is essential.
    A helper class TemporaryMesh can help you do this.
    """
    
    __mAssocMesh: bpy.types.Mesh ##< The binding mesh for this reader. None if this reader is invalid.
    
    def __init__(self, assoc_mesh: bpy.types.Mesh):
        self.__mAssocMesh = assoc_mesh
        
        # triangulate temp mesh
        if self.is_valid():
            self.__triangulate_mesh()
    
    def is_valid(self) -> bool:
        return self.__mAssocMesh is not None
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.dispose()
    
    def dispose(self) -> None:
        if self.is_valid():
            # reset mesh
            self.__mAssocMesh = None
    
    def get_vertex_position_count(self) -> int:
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshReader.')
        
        return len(self.__mAssocMesh.vertices)
    
    def get_vertex_position(self) -> typing.Iterator[UTIL_virtools_types.VxVector3]:
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshReader.')
        
        cache: UTIL_virtools_types.VxVector3 = UTIL_virtools_types.VxVector3()
        for vec in self.__mAssocMesh.vertices:
            cache.x = vec.co.x
            cache.y = vec.co.y
            cache.z = vec.co.z
            yield cache
    
    def get_vertex_normal_count(self) -> int:
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshReader.')
        
        # return loops count, equaling with face count * 3 in theory.
        return len(self.__mAssocMesh.loops)
    
    def get_vertex_normal(self) -> typing.Iterator[UTIL_virtools_types.VxVector3]:
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshReader.')
        
        cache: UTIL_virtools_types.VxVector3 = UTIL_virtools_types.VxVector3()
        for nml in self.__mAssocMesh.corner_normals:
            cache.x = nml.vector.x
            cache.y = nml.vector.y
            cache.z = nml.vector.z
            yield cache
    
    def get_vertex_uv_count(self) -> int:
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshReader.')
        
        if self.__mAssocMesh.uv_layers.active is None:
            # if no uv layer, we need make a fake one
            # return the same value with normals.
            # it also mean create uv for each face vertex
            return len(self.__mAssocMesh.loops)
        else:
            # otherwise return its size, also equaling with face count * 3 in theory
            return len(self.__mAssocMesh.uv_layers.active.uv)
    
    def get_vertex_uv(self) -> typing.Iterator[UTIL_virtools_types.VxVector2]:
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshReader.')
        
        cache: UTIL_virtools_types.VxVector2 = UTIL_virtools_types.VxVector2()
        if self.__mAssocMesh.uv_layers.active is None:
            # create a fake one
            cache.x = 0.0
            cache.y = 0.0
            for _ in range(self.get_vertex_uv_count()):
                yield cache
        else:
            for uv in self.__mAssocMesh.uv_layers.active.uv:
                cache.x = uv.vector.x
                cache.y = uv.vector.y
                yield cache
    
    def get_material_slot_count(self) -> int:
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshReader.')
        
        return len(self.__mAssocMesh.materials)
    
    def get_material_slot(self) -> typing.Iterator[bpy.types.Material | None]:
        """
        @remark This generator may return None if this slot do not link to may material.
        """
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshReader.')
        
        for mtl in self.__mAssocMesh.materials:
            yield mtl
    
    def get_face_count(self) -> int:
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshReader.')
        
        return len(self.__mAssocMesh.polygons)
    
    def get_face(self) -> typing.Iterator[FaceData]:
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshReader.')
        
        # detect whether we have material
        no_mtl: bool = self.get_material_slot_count() == 0
        
        # use list as indices container for convenient adding and deleting.
        cache: FaceData = FaceData([], 0)
        for face in self.__mAssocMesh.polygons:
            # confirm material use
            # a face without mtl have 2 situations. first is the whole object do not have mtl
            # another is this face use an empty mtl slot.
            if no_mtl:
                cache.mMtlIdx = 0
            else:
                cache.mMtlIdx = face.material_index
            
            # resize indices
            self.__resize_face_data_indices(cache.mIndices, face.loop_total)
            # set indices
            for i in range(face.loop_total):
                cache.mIndices[i].mPosIdx = self.__mAssocMesh.loops[face.loop_start + i].vertex_index
                cache.mIndices[i].mNmlIdx = face.loop_start + i
                cache.mIndices[i].mUvIdx = face.loop_start + i
            
            # return value
            yield cache
    
    def __resize_face_data_indices(self, ls: list[FaceVertexData], expected_size: int) -> None:
        diff: int = expected_size - len(ls)
        if diff > 0:
            # add entry
            for _ in range(diff):
                ls.append(FaceVertexData())
        elif diff < 0:
            # remove entry
            for _ in range(diff):
                ls.pop()
        else:
            # no count diff, pass
            pass
    
    def __triangulate_mesh(self) -> None:
        bm: bmesh.types.BMesh = bmesh.new()
        bm.from_mesh(self.__mAssocMesh)
        bmesh.ops.triangulate(bm, faces = bm.faces)
        bm.to_mesh(self.__mAssocMesh)
        bm.free()

class MeshWriter():
    """
    If face do not use material, pass 0 as its material index.
    If face do not have UV becuase it doesn't have material, you at least create 1 UV vector, eg. (0, 0),
    then refer it to all face uv.
    """
    
    __mAssocMesh: bpy.types.Mesh ##< The binding mesh for this writer. None if this writer is invalid.
    
    __mVertexPos: array.array ##< Array item is float(f). Length must be an integer multiple of 3.
    __mVertexNormal: array.array ##< Array item is float(f). Length must be an integer multiple of 3.
    __mVertexUV: array.array ##< Array item is float(f). Length must be an integer multiple of 2.
    
    ## Array item is int32(L).
    #  Length must be the sum of each items in __mFaceVertexCount.
    #  Item is face vertex position index, based on 0, pointing to __mVertexPos (visiting need multiple it with 3 because __mVertexPos is flat struct).
    __mFacePosIndices: array.array
    ## Same as __mFacePosIndices, but store face vertex normal index.
    #  Array item is int32(L). Length is equal to __mFacePosIndices
    __mFaceNmlIndices: array.array
    ## Same as __mFacePosIndices, but store face vertex uv index.
    #  Array item is int32(L). Length is equal to __mFacePosIndices
    __mFaceUvIndices: array.array
    ## Array item is int32(L).
    #  Length is the face count.
    #  It indicate how much vertex need to be consumed in __mFacePosIndices, __mFaceNmlIndices and __mFaceUvIndices for one face.
    __mFaceVertexCount: array.array
    __mFaceMtlIdx: array.array  ##< Array item is int32(L). Length is equal to __mFaceVertexCount.
    
    ## Material Slot.
    #  Each item is unique make sure by __mMtlSlotMap
    __mMtlSlot: list[bpy.types.Material | None]
    ## The map to make sure every item in __mMtlSlot is unique.
    #  Key is bpy.types.Material
    #  Value is key's index in __mMtlSlot.
    __mMtlSlotMap: dict[bpy.types.Material | None, int]
    
    ## The attribute name storing temporary normals data inside mesh.
    __cTempNormalAttrName: typing.ClassVar[str] = 'temp_custom_normals'

    def __init__(self, assoc_mesh: bpy.types.Mesh):
        self.__mAssocMesh = assoc_mesh
        
        self.__mVertexPos = array.array('f')
        self.__mVertexNormal = array.array('f')
        self.__mVertexUV = array.array('f')
        
        self.__mFacePosIndices = array.array('L')
        self.__mFaceNmlIndices = array.array('L')
        self.__mFaceUvIndices = array.array('L')
        self.__mFaceVertexCount = array.array('L')
        self.__mFaceMtlIdx = array.array('L')
        
        self.__mMtlSlot = []
        self.__mMtlSlotMap = {}
    
    def is_valid(self) -> bool:
        return self.__mAssocMesh is not None
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.dispose()
    
    def dispose(self):
        if self.is_valid():
            # write mesh
            self.__write_mesh()
            # reset mesh
            self.__mAssocMesh = None
    
    def add_ingredient(self, data: MeshWriterIngredient):
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshWriter.')
        if not data.is_valid():
            raise UTIL_functions.BBPException('invalid mesh part data.')
        
        # add vertex data
        prev_vertex_pos_count: int = len(self.__mVertexPos) // 3
        self.__mVertexPos.extend(_flat_vxvector3(data.mVertexPosition))
        prev_vertex_nml_count: int = len(self.__mVertexNormal) // 3
        self.__mVertexNormal.extend(_flat_vxvector3(data.mVertexNormal))
        prev_vertex_uv_count: int = len(self.__mVertexUV) // 2
        self.__mVertexUV.extend(_flat_vxvector2(data.mVertexUV))
        
        # add material slot data and create mtl remap
        mtl_remap: list[int] = []
        for mtl in data.mMaterial:
            idx: int | None = self.__mMtlSlotMap.get(mtl, None)
            if idx is not None:
                mtl_remap.append(idx)
            else:
                self.__mMtlSlotMap[mtl] = len(self.__mMtlSlot)
                mtl_remap.append(len(self.__mMtlSlot))
                self.__mMtlSlot.append(mtl)
        
        # add face data
        for face in data.mFace:
            # check indices count
            if not face.is_indices_legal():
                raise UTIL_functions.BBPException('face must have at least 3 vertex.')
            
            # add indices
            for vec_index in face.mIndices:
                self.__mFacePosIndices.append(vec_index.mPosIdx + prev_vertex_pos_count)
                self.__mFaceNmlIndices.append(vec_index.mNmlIdx + prev_vertex_nml_count)
                self.__mFaceUvIndices.append(vec_index.mUvIdx + prev_vertex_uv_count)
            self.__mFaceVertexCount.append(len(face.mIndices))
            
            # add face mtl with remap
            mtl_idx: int = face.mMtlIdx
            if mtl_idx < 0 or mtl_idx >= len(mtl_remap):
                # fall back. add 0
                self.__mFaceMtlIdx.append(0)
            else:
                self.__mFaceMtlIdx.append(mtl_remap[mtl_idx])
    
    def __write_mesh(self):
        # detect status
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshWriter.')
        # and clear mesh
        self.__clear_mesh()
        
        # push material data
        for mtl in self.__mMtlSlot:
            self.__mAssocMesh.materials.append(mtl)
        
        # add corresponding count for vertex position
        self.__mAssocMesh.vertices.add(len(self.__mVertexPos) // 3)
        # add loops data, it is the sum count of indices
        # we use face pos indices size to get it
        self.__mAssocMesh.loops.add(len(self.__mFacePosIndices))
        # set face count
        self.__mAssocMesh.polygons.add(len(self.__mFaceVertexCount))
        # create uv layer
        self.__mAssocMesh.uv_layers.new(do_init = False)
        
        # add vertex position data
        self.__mAssocMesh.vertices.foreach_set('co', self.__mVertexPos)
        # add face vertex pos index data
        self.__mAssocMesh.loops.foreach_set('vertex_index', self.__mFacePosIndices)
        # add face vertex nml by function via mesh custom attribute
        # NOTE: Blender 4.0 / 4.1 changed. I copy these code from FBX Importer.
        temp_normal_attribute: bpy.types.FloatVectorAttribute
        temp_normal_attribute = typing.cast(
            bpy.types.FloatVectorAttribute,
            self.__mAssocMesh.attributes.new(MeshWriter.__cTempNormalAttrName, 'FLOAT_VECTOR', 'CORNER')
        )
        temp_normal_attribute.data.foreach_set('vector',
            tuple(_flat_face_nml_index(self.__mFaceNmlIndices, self.__mVertexNormal))
        )
        # add face vertex uv by function
        self.__mAssocMesh.uv_layers.active.uv.foreach_set('vector',
            tuple(_flat_face_uv_index(self.__mFaceUvIndices, self.__mVertexUV))
        )   # NOTE: blender 3.5 changed. UV must be visited by .uv, not the .data
        
        # iterate face to set face data
        f_vertex_idx: int = 0
        for fi in range(len(self.__mFaceVertexCount)):
            # set start loop
            # NOTE: blender 3.6 changed. Loop setting in polygon do not need set loop_total any more.
            # the loop_total will be auto calculated by the next loop_start.
            # loop_total become read-only
            self.__mAssocMesh.polygons[fi].loop_start = f_vertex_idx
            
            # set material index
            self.__mAssocMesh.polygons[fi].material_index = self.__mFaceMtlIdx[fi]
            
            # set auto smooth. it is IMPORTANT
            # because it related to whether custom split normal can work
            self.__mAssocMesh.polygons[fi].use_smooth = True
            
            # inc vertex idx
            f_vertex_idx += self.__mFaceVertexCount[fi]
        
        # validate mesh.
        # it is IMPORTANT that do NOT delete custom data
        # because we need use these data to set custom split normal later
        self.__mAssocMesh.validate(clean_customdata = False)
        # update mesh without mesh calc
        self.__mAssocMesh.update(calc_edges = False, calc_edges_loose = False)
        
        # set custom split normal data
        # this operation must copy preserved normal data from loops, not the array data in this class,
        # because the validate() may change the mesh and if change happended, an error will occur when applying normals (not matched loops count).
        # this should not happend in normal case, for testing, please load "Level_1.NMO" (Ballance Level 1).
        
        # copy data from loops preserved in validate().
        # NOTE: Blender 4.0 / 4.1 changed. I copy these code from FBX Importer.
        loops_normals = array.array('f', [0.0] * (len(self.__mAssocMesh.loops) * 3))
        temp_normal_attribute = typing.cast(
            bpy.types.FloatVectorAttribute, 
            self.__mAssocMesh.attributes[MeshWriter.__cTempNormalAttrName]
        )
        temp_normal_attribute.data.foreach_get("vector", loops_normals)
        # apply data
        self.__mAssocMesh.normals_split_custom_set(
            tuple(_nest_custom_split_normal(loops_normals))
        )
        self.__mAssocMesh.attributes.remove(
            # MARK: idk why I need fucking get this attribute again.
            # But if I were not, this function must raise bullshit exception!
            self.__mAssocMesh.attributes[MeshWriter.__cTempNormalAttrName]
        )

    def __clear_mesh(self):
        if not self.is_valid():
            raise UTIL_functions.BBPException('try to call an invalid MeshWriter.')
        
        # clear geometry
        self.__mAssocMesh.clear_geometry()
        # clear mtl slot because clear_geometry will not do this.
        self.__mAssocMesh.materials.clear()