import_fbx.py

# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####

# <pep8 compliant>

# Script copyright (C) Blender Foundation

# FBX 7.1.0 -> 7.4.0 loader for Blender

# Not totally pep8 compliant.
#   pep8 import_fbx.py --ignore=E501,E123,E702,E125

if "bpy" in locals():
    import importlib
    if "parse_fbx" in locals():
        importlib.reload(parse_fbx)
    if "fbx_utils" in locals():
        importlib.reload(fbx_utils)

import bpy

# -----
# Utils
from . import parse_fbx, fbx_utils

from .parse_fbx import data_types, FBXElem
from .fbx_utils import (
    units_convertor_iter,
    array_to_matrix4,
    similar_values,
    similar_values_iter,
)

# global singleton, assign on execution
fbx_elem_nil = None

# Units convertors...
convert_deg_to_rad_iter = units_convertor_iter("degree", "radian")

MAT_CONVERT_BONE = fbx_utils.MAT_CONVERT_BONE.inverted()
MAT_CONVERT_LAMP = fbx_utils.MAT_CONVERT_LAMP.inverted()
MAT_CONVERT_CAMERA = fbx_utils.MAT_CONVERT_CAMERA.inverted()


def elem_find_first(elem, id_search, default=None):
    for fbx_item in elem.elems:
        if fbx_item.id == id_search:
            return fbx_item
    return default


def elem_find_iter(elem, id_search):
    for fbx_item in elem.elems:
        if fbx_item.id == id_search:
            yield fbx_item


def elem_find_first_string(elem, id_search):
    fbx_item = elem_find_first(elem, id_search)
    if fbx_item is not None:
        assert(len(fbx_item.props) == 1)
        assert(fbx_item.props_type[0] == data_types.STRING)
        return fbx_item.props[0].decode('utf-8')
    return None


def elem_find_first_bytes(elem, id_search, decode=True):
    fbx_item = elem_find_first(elem, id_search)
    if fbx_item is not None:
        assert(len(fbx_item.props) == 1)
        assert(fbx_item.props_type[0] == data_types.STRING)
        return fbx_item.props[0]
    return None


def elem_repr(elem):
    return "%s: props[%d=%r], elems=(%r)" % (
        elem.id,
        len(elem.props),
        ", ".join([repr(p) for p in elem.props]),
        # elem.props_type,
        b", ".join([e.id for e in elem.elems]),
        )


def elem_split_name_class(elem):
    assert(elem.props_type[-2] == data_types.STRING)
    elem_name, elem_class = elem.props[-2].split(b'\x00\x01')
    return elem_name, elem_class


def elem_name_ensure_class(elem, clss=...):
    elem_name, elem_class = elem_split_name_class(elem)
    if clss is not ...:
        assert(elem_class == clss)
    return elem_name.decode('utf-8')


def elem_split_name_class_nodeattr(elem):
    assert(elem.props_type[-2] == data_types.STRING)
    elem_name, elem_class = elem.props[-2].split(b'\x00\x01')
    assert(elem_class == b'NodeAttribute')
    assert(elem.props_type[-1] == data_types.STRING)
    elem_class = elem.props[-1]
    return elem_name, elem_class


def elem_uuid(elem):
    assert(elem.props_type[0] == data_types.INT64)
    return elem.props[0]


def elem_prop_first(elem, default=None):
    return elem.props[0] if (elem is not None) and elem.props else default


# ----
# Support for
# Properties70: { ... P:
def elem_props_find_first(elem, elem_prop_id):

    # support for templates (tuple of elems)
    if type(elem) is not FBXElem:
        assert(type(elem) is tuple)
        for e in elem:
            result = elem_props_find_first(e, elem_prop_id)
            if result is not None:
                return result
        assert(len(elem) > 0)
        return None

    for subelem in elem.elems:
        assert(subelem.id == b'P')
        if subelem.props[0] == elem_prop_id:
            return subelem
    return None


def elem_props_get_color_rgb(elem, elem_prop_id, default=None):
    elem_prop = elem_props_find_first(elem, elem_prop_id)
    if elem_prop is not None:
        assert(elem_prop.props[0] == elem_prop_id)
        if elem_prop.props[1] == b'Color':
            # FBX version 7300
            assert(elem_prop.props[1] == b'Color')
            assert(elem_prop.props[2] == b'')
            assert(elem_prop.props[3] in {b'A', b'A+', b'AU'})
        else:
            assert(elem_prop.props[1] == b'ColorRGB')
            assert(elem_prop.props[2] == b'Color')
        assert(elem_prop.props_type[4:7] == bytes((data_types.FLOAT64,)) * 3)
        return elem_prop.props[4:7]
    return default


def elem_props_get_vector_3d(elem, elem_prop_id, default=None):
    elem_prop = elem_props_find_first(elem, elem_prop_id)
    if elem_prop is not None:
        assert(elem_prop.props_type[4:7] == bytes((data_types.FLOAT64,)) * 3)
        return elem_prop.props[4:7]
    return default


def elem_props_get_number(elem, elem_prop_id, default=None):
    elem_prop = elem_props_find_first(elem, elem_prop_id)
    if elem_prop is not None:
        assert(elem_prop.props[0] == elem_prop_id)
        if elem_prop.props[1] == b'double':
            assert(elem_prop.props[1] == b'double')
            assert(elem_prop.props[2] == b'Number')
        else:
            assert(elem_prop.props[1] == b'Number')
            assert(elem_prop.props[2] == b'')
            assert(elem_prop.props[3] in {b'A', b'A+', b'AU'})

        # we could allow other number types
        assert(elem_prop.props_type[4] == data_types.FLOAT64)

        return elem_prop.props[4]
    return default


def elem_props_get_integer(elem, elem_prop_id, default=None):
    elem_prop = elem_props_find_first(elem, elem_prop_id)
    if elem_prop is not None:
        assert(elem_prop.props[0] == elem_prop_id)
        if elem_prop.props[1] == b'int':
            assert(elem_prop.props[1] == b'int')
            assert(elem_prop.props[2] == b'Integer')
        elif elem_prop.props[1] == b'ULongLong':
            assert(elem_prop.props[1] == b'ULongLong')
            assert(elem_prop.props[2] == b'')

        # we could allow other number types
        assert(elem_prop.props_type[4] in {data_types.INT32, data_types.INT64})

        return elem_prop.props[4]
    return default


def elem_props_get_bool(elem, elem_prop_id, default=None):
    elem_prop = elem_props_find_first(elem, elem_prop_id)
    if elem_prop is not None:
        assert(elem_prop.props[0] == elem_prop_id)
        assert(elem_prop.props[1] == b'bool')
        assert(elem_prop.props[2] == b'')
        assert(elem_prop.props[3] == b'')

        # we could allow other number types
        assert(elem_prop.props_type[4] == data_types.INT32)
        assert(elem_prop.props[4] in {0, 1})

        return bool(elem_prop.props[4])
    return default


def elem_props_get_enum(elem, elem_prop_id, default=None):
    elem_prop = elem_props_find_first(elem, elem_prop_id)
    if elem_prop is not None:
        assert(elem_prop.props[0] == elem_prop_id)
        assert(elem_prop.props[1] == b'enum')
        assert(elem_prop.props[2] == b'')
        assert(elem_prop.props[3] == b'')

        # we could allow other number types
        assert(elem_prop.props_type[4] == data_types.INT32)

        return elem_prop.props[4]
    return default


def elem_props_get_visibility(elem, elem_prop_id, default=None):
    elem_prop = elem_props_find_first(elem, elem_prop_id)
    if elem_prop is not None:
        assert(elem_prop.props[0] == elem_prop_id)
        assert(elem_prop.props[1] == b'Visibility')
        assert(elem_prop.props[2] == b'')
        assert(elem_prop.props[3] in {b'A', b'A+', b'AU'})

        # we could allow other number types
        assert(elem_prop.props_type[4] == data_types.FLOAT64)

        return elem_prop.props[4]
    return default


# ----------------------------------------------------------------------------
# Blender

# ------
# Object
from collections import namedtuple


FBXTransformData = namedtuple("FBXTransformData", (
    "loc",
    "rot", "rot_ofs", "rot_piv", "pre_rot", "pst_rot", "rot_ord", "rot_alt_mat",
    "sca", "sca_ofs", "sca_piv",
))


object_tdata_cache = {}


def blen_read_object_transform_do(transform_data):
    from mathutils import Matrix, Euler

    # translation
    lcl_translation = Matrix.Translation(transform_data.loc)

    # rotation
    to_rot = lambda rot, rot_ord: Euler(convert_deg_to_rad_iter(rot), rot_ord).to_matrix().to_4x4()
    lcl_rot = to_rot(transform_data.rot, transform_data.rot_ord) * transform_data.rot_alt_mat
    pre_rot = to_rot(transform_data.pre_rot, transform_data.rot_ord)
    pst_rot = to_rot(transform_data.pst_rot, transform_data.rot_ord)

    rot_ofs = Matrix.Translation(transform_data.rot_ofs)
    rot_piv = Matrix.Translation(transform_data.rot_piv)
    sca_ofs = Matrix.Translation(transform_data.sca_ofs)
    sca_piv = Matrix.Translation(transform_data.sca_piv)

    # scale
    lcl_scale = Matrix()
    lcl_scale[0][0], lcl_scale[1][1], lcl_scale[2][2] = transform_data.sca

    return (
        lcl_translation *
        rot_ofs *
        rot_piv *
        pre_rot *
        lcl_rot *
        pst_rot *
        rot_piv.inverted() *
        sca_ofs *
        sca_piv *
        lcl_scale *
        sca_piv.inverted()
    )


# XXX This might be weak, now that we can add vgroups from both bones and shapes, name collisions become
#     more likely, will have to make this more robust!!!
def add_vgroup_to_objects(vg_indices, vg_weights, vg_name, objects):
    assert(len(vg_indices) == len(vg_weights))
    if vg_indices:
        for obj in objects:
            # We replace/override here...
            vg = obj.vertex_groups.get(vg_name)
            if vg is None:
                vg = obj.vertex_groups.new(vg_name)
            for i, w in zip(vg_indices, vg_weights):
                vg.add((i,), w, 'REPLACE')


def blen_read_object_transform_preprocess(fbx_props, fbx_obj, rot_alt_mat):
    # This is quite involved, 'fbxRNode.cpp' from openscenegraph used as a reference
    const_vector_zero_3d = 0.0, 0.0, 0.0
    const_vector_one_3d = 1.0, 1.0, 1.0

    loc = list(elem_props_get_vector_3d(fbx_props, b'Lcl Translation', const_vector_zero_3d))
    rot = list(elem_props_get_vector_3d(fbx_props, b'Lcl Rotation', const_vector_zero_3d))
    sca = list(elem_props_get_vector_3d(fbx_props, b'Lcl Scaling', const_vector_one_3d))

    rot_ofs = elem_props_get_vector_3d(fbx_props, b'RotationOffset', const_vector_zero_3d)
    rot_piv = elem_props_get_vector_3d(fbx_props, b'RotationPivot', const_vector_zero_3d)
    sca_ofs = elem_props_get_vector_3d(fbx_props, b'ScalingOffset', const_vector_zero_3d)
    sca_piv = elem_props_get_vector_3d(fbx_props, b'ScalingPivot', const_vector_zero_3d)

    is_rot_act = elem_props_get_bool(fbx_props, b'RotationActive', False)

    if is_rot_act:
        pre_rot = elem_props_get_vector_3d(fbx_props, b'PreRotation', const_vector_zero_3d)
        pst_rot = elem_props_get_vector_3d(fbx_props, b'PostRotation', const_vector_zero_3d)
        rot_ord = {
            0: 'XYZ',
            1: 'XYZ',
            2: 'XZY',
            3: 'YZX',
            4: 'YXZ',
            5: 'ZXY',
            6: 'ZYX',
            }.get(elem_props_get_enum(fbx_props, b'RotationOrder', 0))
    else:
        pre_rot = const_vector_zero_3d
        pst_rot = const_vector_zero_3d
        rot_ord = 'XYZ'

    return FBXTransformData(loc,
                            rot, rot_ofs, rot_piv, pre_rot, pst_rot, rot_ord, rot_alt_mat,
                            sca, sca_ofs, sca_piv)


def blen_read_object(fbx_tmpl, fbx_obj, object_data):
    elem_name_utf8 = elem_name_ensure_class(fbx_obj)

    # Object data must be created already
    obj = bpy.data.objects.new(name=elem_name_utf8, object_data=object_data)

    fbx_props = (elem_find_first(fbx_obj, b'Properties70'),
                 elem_find_first(fbx_tmpl, b'Properties70', fbx_elem_nil))
    assert(fbx_props[0] is not None)

    # ----
    # Misc Attributes

    obj.color[0:3] = elem_props_get_color_rgb(fbx_props, b'Color', (0.8, 0.8, 0.8))
    obj.hide = not bool(elem_props_get_visibility(fbx_props, b'Visibility', 1.0))

    # ----
    # Transformation

    from mathutils import Matrix
    from math import pi

    # rotation corrections
    if obj.type == 'CAMERA':
        rot_alt_mat = MAT_CONVERT_CAMERA
    elif obj.type == 'LAMP':
        rot_alt_mat = MAT_CONVERT_LAMP
    else:
        rot_alt_mat = Matrix()

    transform_data = object_tdata_cache.get(obj)
    if transform_data is None:
        transform_data = blen_read_object_transform_preprocess(fbx_props, fbx_obj, rot_alt_mat)
        object_tdata_cache[obj] = transform_data
    obj.matrix_basis = blen_read_object_transform_do(transform_data)

    return obj


# --------
# Armature

def blen_read_armatures_add_bone(bl_obj, bl_arm, bones, b_uuid, matrices, fbx_tmpl_model):
    from mathutils import Matrix, Vector

    b_item, bsize, p_uuid, clusters = bones[b_uuid]
    fbx_bdata, bl_bname = b_item
    if bl_bname is not None:
        return bl_arm.edit_bones[bl_bname]  # Have already been created...

    p_ebo = None
    if p_uuid is not None:
        # Recurse over parents!
        p_ebo = blen_read_armatures_add_bone(bl_obj, bl_arm, bones, p_uuid, matrices, fbx_tmpl_model)

    if clusters:
        # Note in some cases, one bone can have several clusters (kind of LoD?), in Blender we'll always
        # use only the first, for now.
        fbx_cdata, meshes, objects = clusters[0]
        objects = {blen_o for fbx_o, blen_o in objects}

        # We assume matrices in cluster are rest pose of bones (they are in Global space!).
        # TransformLink is matrix of bone, in global space.
        # TransformAssociateModel is matrix of armature, in global space (at bind time).
        elm = elem_find_first(fbx_cdata, b'Transform', default=None)
        mmat_bone = array_to_matrix4(elm.props[0]) if elm is not None else None
        elm = elem_find_first(fbx_cdata, b'TransformLink', default=None)
        bmat_glob = array_to_matrix4(elm.props[0]) if elm is not None else Matrix()
        elm = elem_find_first(fbx_cdata, b'TransformAssociateModel', default=None)
        amat_glob = array_to_matrix4(elm.props[0]) if elm is not None else Matrix()

        mmat_glob = bmat_glob * mmat_bone

        # We seek for matrix of bone in armature space...
        bmat_arm = amat_glob.inverted() * bmat_glob

        # Bone correction, works here...
        bmat_loc = (p_ebo.matrix.inverted() * bmat_arm) if p_ebo else bmat_arm
        bmat_loc = bmat_loc * MAT_CONVERT_BONE
        bmat_arm = (p_ebo.matrix * bmat_loc) if p_ebo else bmat_loc
    else:
        # Armature bound to no mesh...
        fbx_cdata, meshes, objects = (None, (), ())
        mmat_bone = None
        amat_glob = bl_obj.matrix_world

        fbx_props = (elem_find_first(fbx_bdata, b'Properties70'),
                     elem_find_first(fbx_tmpl_model, b'Properties70', fbx_elem_nil))
        assert(fbx_props[0] is not None)

        # Bone correction, works here...
        transform_data = blen_read_object_transform_preprocess(fbx_props, fbx_bdata, MAT_CONVERT_BONE)
        bmat_loc = blen_read_object_transform_do(transform_data)
        # Bring back matrix in armature space.
        bmat_arm = (p_ebo.matrix * bmat_loc) if p_ebo else bmat_loc

    # ----
    # Now, create the (edit)bone.
    bone_name = elem_name_ensure_class(fbx_bdata, b'Model')

    ebo = bl_arm.edit_bones.new(name=bone_name)
    bone_name = ebo.name  # Might differ from FBX bone name!
    b_item[1] = bone_name  # since ebo is only valid in Edit mode... :/

    # So that our bone gets its final length, but still Y-aligned in armature space.
    ebo.tail = Vector((0.0, 1.0, 0.0)) * bsize
    # And rotate/move it to its final "rest pose".
    ebo.matrix = bmat_arm.normalized()

    # Connection to parent.
    if p_ebo is not None:
        ebo.parent = p_ebo
        if similar_values_iter(p_ebo.tail, ebo.head):
            ebo.use_connect = True

    if fbx_cdata is not None:
        # ----
        # Add a new vgroup to the meshes (their objects, actually!).
        # Quite obviously, only one mesh is expected...
        indices = elem_prop_first(elem_find_first(fbx_cdata, b'Indexes', default=None), default=())
        weights = elem_prop_first(elem_find_first(fbx_cdata, b'Weights', default=None), default=())
        add_vgroup_to_objects(indices, weights, bone_name, objects)

    # ----
    # If we get a valid mesh matrix (in bone space), store armature and mesh global matrices, we need to set temporarily
    # both objects to those matrices when actually binding them via the modifier.
    # Note we assume all bones were bound with the same mesh/armature (global) matrix, we do not support otherwise
    # in Blender anyway!
    if mmat_bone is not None:
        for obj in objects:
            if obj in matrices:
                continue
            matrices[obj] = (amat_glob, mmat_glob)

    return ebo


def blen_read_armatures(fbx_tmpl, armatures, fbx_bones_to_fake_object, scene, global_matrix, arm_parents):
    from mathutils import Matrix

    if global_matrix is None:
        global_matrix = Matrix()

    for a_item, bones in armatures:
        fbx_adata, bl_adata = a_item
        matrices = {}

        # ----
        # Armature data.
        elem_name_utf8 = elem_name_ensure_class(fbx_adata, b'Model')
        bl_arm = bpy.data.armatures.new(name=elem_name_utf8)

        # Need to create the object right now, since we can only add bones in Edit mode... :/
        assert(a_item[1] is None)

        if fbx_adata.props[2] in {b'LimbNode', b'Root'}:
            # rootbone-as-armature case...
            fbx_bones_to_fake_object[fbx_adata.props[0]] = bl_adata = blen_read_object(fbx_tmpl, fbx_adata, bl_arm)
            # reset transform.
            bl_adata.matrix_basis = Matrix()
        else:
            bl_adata = a_item[1] = blen_read_object(fbx_tmpl, fbx_adata, bl_arm)

        # Instantiate in scene.
        obj_base = scene.objects.link(bl_adata)
        obj_base.select = True

        # Switch to Edit mode.
        scene.objects.active = bl_adata
        bpy.ops.object.mode_set(mode='EDIT')

        for b_uuid in bones:
            blen_read_armatures_add_bone(bl_adata, bl_arm, bones, b_uuid, matrices, fbx_tmpl)

        bpy.ops.object.mode_set(mode='OBJECT')

        # Bind armature to objects.
        arm_mat_back = bl_adata.matrix_basis.copy()
        for ob_me, (amat, mmat) in matrices.items():
            # bring global armature & mesh matrices into *Blender* global space.
            amat = global_matrix * amat
            mmat = global_matrix * mmat

            bl_adata.matrix_basis = amat
            me_mat_back = ob_me.matrix_basis.copy()
            ob_me.matrix_basis = mmat

            mod = ob_me.modifiers.new(elem_name_utf8, 'ARMATURE')
            mod.object = bl_adata

            ob_me.parent = bl_adata
            ob_me.matrix_basis = me_mat_back
            # Store the pair for later space corrections (bring back mesh in parent space).
            arm_parents.add((bl_adata, ob_me))
        bl_adata.matrix_basis = arm_mat_back

        # Set Pose transformations...
        for b_item, _b_size, _p_uuid, _clusters in bones.values():
            fbx_bdata, bl_bname = b_item
            fbx_props = (elem_find_first(fbx_bdata, b'Properties70'),
                         elem_find_first(fbx_tmpl, b'Properties70', fbx_elem_nil))
            assert(fbx_props[0] is not None)

            pbo = b_item[1] = bl_adata.pose.bones[bl_bname]
            transform_data = object_tdata_cache.get(pbo)
            if transform_data is None:
                # Bone correction, gives a mess as result. :(
                transform_data = blen_read_object_transform_preprocess(fbx_props, fbx_bdata, MAT_CONVERT_BONE)
                object_tdata_cache[pbo] = transform_data
            mat = blen_read_object_transform_do(transform_data)
            if pbo.parent:
                # Bring back matrix in armature space.
                mat = pbo.parent.matrix * mat
            pbo.matrix = mat


# ----
# Mesh

def blen_read_geom_layerinfo(fbx_layer):
    return (
        elem_find_first_string(fbx_layer, b'Name'),
        elem_find_first_bytes(fbx_layer, b'MappingInformationType'),
        elem_find_first_bytes(fbx_layer, b'ReferenceInformationType'),
        )


def blen_read_geom_array_mapped_vert(
        mesh, blen_data, blend_attr,
        fbx_layer_data, fbx_layer_index,
        fbx_layer_mapping, fbx_layer_ref,
        stride, item_size, descr,
        ):
    # TODO, generic mapping apply function
    if fbx_layer_mapping == b'ByVertice':
        if fbx_layer_ref == b'Direct':
            assert(fbx_layer_index is None)
            # TODO, more generic support for mapping types
            for i, blen_data_item in enumerate(blen_data):
                setattr(blen_data_item, blend_attr,
                        fbx_layer_data[(i * stride): (i * stride) + item_size])
            return True
        else:
            print("warning layer %r ref type unsupported: %r" % (descr, fbx_layer_ref))
    else:
        print("warning layer %r mapping type unsupported: %r" % (descr, fbx_layer_mapping))

    return False


def blen_read_geom_array_mapped_edge(
        mesh, blen_data, blend_attr,
        fbx_layer_data, fbx_layer_index,
        fbx_layer_mapping, fbx_layer_ref,
        stride, item_size, descr,
        xform=None,
        ):
    if fbx_layer_mapping == b'ByEdge':
        if fbx_layer_ref == b'Direct':
            if stride == 1:
                if xform is None:
                    for i, blen_data_item in enumerate(blen_data):
                        setattr(blen_data_item, blend_attr,
                                fbx_layer_data[i])
                else:
                    for i, blen_data_item in enumerate(blen_data):
                        setattr(blen_data_item, blend_attr,
                                xform(fbx_layer_data[i]))
            else:
                if xform is None:
                    for i, blen_data_item in enumerate(blen_data):
                        setattr(blen_data_item, blend_attr,
                                fbx_layer_data[(i * stride): (i * stride) + item_size])
                else:
                    for i, blen_data_item in enumerate(blen_data):
                        setattr(blen_data_item, blend_attr,
                                xform(fbx_layer_data[(i * stride): (i * stride) + item_size]))
            return True
        else:
            print("warning layer %r ref type unsupported: %r" % (descr, fbx_layer_ref))
    else:
        print("warning layer %r mapping type unsupported: %r" % (descr, fbx_layer_mapping))

    return False


def blen_read_geom_array_mapped_polygon(
        mesh, blen_data, blend_attr,
        fbx_layer_data, fbx_layer_index,
        fbx_layer_mapping, fbx_layer_ref,
        stride, item_size, descr,
        xform=None,
        ):
    if fbx_layer_mapping == b'ByPolygon':
        if fbx_layer_ref == b'IndexToDirect':
            if stride == 1:
                for i, blen_data_item in enumerate(blen_data):
                    setattr(blen_data_item, blend_attr,
                            fbx_layer_data[i])
            else:
                for i, blen_data_item in enumerate(blen_data):
                    setattr(blen_data_item, blend_attr,
                            fbx_layer_data[(i * stride): (i * stride) + item_size])
            return True
        elif fbx_layer_ref == b'Direct':
            # looks like direct may have different meanings!
            assert(stride == 1)
            if xform is None:
                for i in range(len(fbx_layer_data)):
                    setattr(blen_data[i], blend_attr, fbx_layer_data[i])
            else:
                for i in range(len(fbx_layer_data)):
                    setattr(blen_data[i], blend_attr, xform(fbx_layer_data[i]))
            return True
        else:
            print("warning layer %r ref type unsupported: %r" % (descr, fbx_layer_ref))
    else:
        print("warning layer %r mapping type unsupported: %r" % (descr, fbx_layer_mapping))

    return False


def blen_read_geom_array_mapped_polyloop(
        mesh, blen_data, blend_attr,
        fbx_layer_data, fbx_layer_index,
        fbx_layer_mapping, fbx_layer_ref,
        stride, item_size, descr,
        ):
    if fbx_layer_mapping == b'ByPolygonVertex':
        if fbx_layer_ref == b'IndexToDirect':
            assert(fbx_layer_index is not None)
            for i, j in enumerate(fbx_layer_index):
                if j != -1:
                    setattr(blen_data[i], blend_attr,
                            fbx_layer_data[(j * stride): (j * stride) + item_size])
            return True
        else:
            print("warning layer %r ref type unsupported: %r" % (descr, fbx_layer_ref))
    elif fbx_layer_mapping == b'ByVertice':
        if fbx_layer_ref == b'Direct':
            assert(fbx_layer_index is None)
            loops = mesh.loops
            polygons = mesh.polygons
            for p in polygons:
                for i in p.loop_indices:
                    j = loops[i].vertex_index
                    setattr(blen_data[i], blend_attr,
                            fbx_layer_data[(j * stride): (j * stride) + item_size])
        else:
            print("warning layer %r ref type unsupported: %r" % (descr, fbx_layer_ref))
    else:
        print("warning layer %r mapping type unsupported: %r" % (descr, fbx_layer_mapping))

    return False


def blen_read_geom_layer_material(fbx_obj, mesh):
    fbx_layer = elem_find_first(fbx_obj, b'LayerElementMaterial')

    if fbx_layer is None:
        return

    (fbx_layer_name,
     fbx_layer_mapping,
     fbx_layer_ref,
     ) = blen_read_geom_layerinfo(fbx_layer)

    if fbx_layer_mapping == b'AllSame':
        # only to quiet warning
        return

    layer_id = b'Materials'
    fbx_layer_data = elem_prop_first(elem_find_first(fbx_layer, layer_id))

    blen_data = mesh.polygons
    blen_read_geom_array_mapped_polygon(
        mesh, blen_data, "material_index",
        fbx_layer_data, None,
        fbx_layer_mapping, fbx_layer_ref,
        1, 1, layer_id,
        )


def blen_read_geom_layer_uv(fbx_obj, mesh):
    for layer_id in (b'LayerElementUV',):
        for fbx_layer in elem_find_iter(fbx_obj, layer_id):
            # all should be valid
            (fbx_layer_name,
             fbx_layer_mapping,
             fbx_layer_ref,
             ) = blen_read_geom_layerinfo(fbx_layer)

            fbx_layer_data = elem_prop_first(elem_find_first(fbx_layer, b'UV'))
            fbx_layer_index = elem_prop_first(elem_find_first(fbx_layer, b'UVIndex'))

            uv_tex = mesh.uv_textures.new(name=fbx_layer_name)
            uv_lay = mesh.uv_layers[-1]
            blen_data = uv_lay.data[:]

            # some valid files omit this data
            if fbx_layer_data is None:
                print("%r %r missing data" % (layer_id, fbx_layer_name))
                continue

            blen_read_geom_array_mapped_polyloop(
                mesh, blen_data, "uv",
                fbx_layer_data, fbx_layer_index,
                fbx_layer_mapping, fbx_layer_ref,
                2, 2, layer_id,
                )


def blen_read_geom_layer_color(fbx_obj, mesh):
    # almost same as UV's
    for layer_id in (b'LayerElementColor',):
        for fbx_layer in elem_find_iter(fbx_obj, layer_id):
            # all should be valid
            (fbx_layer_name,
             fbx_layer_mapping,
             fbx_layer_ref,
             ) = blen_read_geom_layerinfo(fbx_layer)

            fbx_layer_data = elem_prop_first(elem_find_first(fbx_layer, b'Colors'))
            fbx_layer_index = elem_prop_first(elem_find_first(fbx_layer, b'ColorIndex'))

            color_lay = mesh.vertex_colors.new(name=fbx_layer_name)
            blen_data = color_lay.data[:]

            # some valid files omit this data
            if fbx_layer_data is None:
                print("%r %r missing data" % (layer_id, fbx_layer_name))
                continue

            # ignore alpha layer (read 4 items into 3)
            blen_read_geom_array_mapped_polyloop(
                mesh, blen_data, "color",
                fbx_layer_data, fbx_layer_index,
                fbx_layer_mapping, fbx_layer_ref,
                4, 3, layer_id,
                )


def blen_read_geom_layer_smooth(fbx_obj, mesh):
    fbx_layer = elem_find_first(fbx_obj, b'LayerElementSmoothing')

    if fbx_layer is None:
        return False

    # all should be valid
    (fbx_layer_name,
     fbx_layer_mapping,
     fbx_layer_ref,
     ) = blen_read_geom_layerinfo(fbx_layer)

    layer_id = b'Smoothing'
    fbx_layer_data = elem_prop_first(elem_find_first(fbx_layer, layer_id))

    # udk has 'Direct' mapped, with no Smoothing, not sure why, but ignore these
    if fbx_layer_data is None:
        return False

    if fbx_layer_mapping == b'ByEdge':
        # some models have bad edge data, we cant use this info...
        if not mesh.edges:
            return False

        blen_data = mesh.edges
        ok_smooth = blen_read_geom_array_mapped_edge(
            mesh, blen_data, "use_edge_sharp",
            fbx_layer_data, None,
            fbx_layer_mapping, fbx_layer_ref,
            1, 1, layer_id,
            xform=lambda s: not s,
            )
        return ok_smooth
    elif fbx_layer_mapping == b'ByPolygon':
        blen_data = mesh.polygons
        return blen_read_geom_array_mapped_polygon(
            mesh, blen_data, "use_smooth",
            fbx_layer_data, None,
            fbx_layer_mapping, fbx_layer_ref,
            1, 1, layer_id,
            xform=lambda s: (s != 0),  # smoothgroup bitflags, treat as booleans for now
            )
    else:
        print("warning layer %r mapping type unsupported: %r" % (fbx_layer.id, fbx_layer_mapping))
        return False


def blen_read_geom_layer_normal(fbx_obj, mesh):
    fbx_layer = elem_find_first(fbx_obj, b'LayerElementNormal')

    if fbx_layer is None:
        return False

    (fbx_layer_name,
     fbx_layer_mapping,
     fbx_layer_ref,
     ) = blen_read_geom_layerinfo(fbx_layer)

    layer_id = b'Normals'
    fbx_layer_data = elem_prop_first(elem_find_first(fbx_layer, layer_id))

    blen_data = mesh.vertices

    return blen_read_geom_array_mapped_vert(
        mesh, blen_data, "normal",
        fbx_layer_data, None,
        fbx_layer_mapping, fbx_layer_ref,
        3, 3, layer_id,
        )


def blen_read_geom(fbx_tmpl, fbx_obj):
    # TODO, use 'fbx_tmpl'
    elem_name_utf8 = elem_name_ensure_class(fbx_obj, b'Geometry')

    fbx_verts = elem_prop_first(elem_find_first(fbx_obj, b'Vertices'))
    fbx_polys = elem_prop_first(elem_find_first(fbx_obj, b'PolygonVertexIndex'))
    fbx_edges = elem_prop_first(elem_find_first(fbx_obj, b'Edges'))

    if fbx_verts is None:
        fbx_verts = ()
    if fbx_polys is None:
        fbx_polys = ()

    mesh = bpy.data.meshes.new(name=elem_name_utf8)
    mesh.vertices.add(len(fbx_verts) // 3)
    mesh.vertices.foreach_set("co", fbx_verts)

    if fbx_polys:
        mesh.loops.add(len(fbx_polys))
        poly_loop_starts = []
        poly_loop_totals = []
        poly_loop_prev = 0
        for i, l in enumerate(mesh.loops):
            index = fbx_polys[i]
            if index < 0:
                poly_loop_starts.append(poly_loop_prev)
                poly_loop_totals.append((i - poly_loop_prev) + 1)
                poly_loop_prev = i + 1
                index ^= -1
            l.vertex_index = index

        mesh.polygons.add(len(poly_loop_starts))
        mesh.polygons.foreach_set("loop_start", poly_loop_starts)
        mesh.polygons.foreach_set("loop_total", poly_loop_totals)

        blen_read_geom_layer_material(fbx_obj, mesh)
        blen_read_geom_layer_uv(fbx_obj, mesh)
        blen_read_geom_layer_color(fbx_obj, mesh)

    if fbx_edges:
        # edges in fact index the polygons (NOT the vertices)
        import array
        tot_edges = len(fbx_edges)
        edges_conv = array.array('i', [0]) * (tot_edges * 2)

        edge_index = 0
        for i in fbx_edges:
            e_a = fbx_polys[i]
            if e_a >= 0:
                e_b = fbx_polys[i + 1]
                if e_b < 0:
                    e_b ^= -1
            else:
                # Last index of polygon, wrap back to the start.

                # ideally we wouldn't have to search back,
                # but it should only be 2-3 iterations.
                j = i - 1
                while j >= 0 and fbx_polys[j] >= 0:
                    j -= 1
                e_a ^= -1
                e_b = fbx_polys[j + 1]

            edges_conv[edge_index] = e_a
            edges_conv[edge_index + 1] = e_b
            edge_index += 2

        mesh.edges.add(tot_edges)
        mesh.edges.foreach_set("vertices", edges_conv)

    # must be after edge, face loading.
    ok_smooth = blen_read_geom_layer_smooth(fbx_obj, mesh)

    ok_normals = blen_read_geom_layer_normal(fbx_obj, mesh)

    mesh.validate()

    if not ok_normals:
        mesh.calc_normals()

    if not ok_smooth:
        for p in mesh.polygons:
            p.use_smooth = True

    return mesh


def blen_read_shape(fbx_tmpl, fbx_sdata, fbx_bcdata, meshes, scene, global_matrix):
    from mathutils import Vector

    elem_name_utf8 = elem_name_ensure_class(fbx_sdata, b'Geometry')
    indices = elem_prop_first(elem_find_first(fbx_sdata, b'Indexes'), default=())
    dvcos = tuple(co for co in zip(*[iter(elem_prop_first(elem_find_first(fbx_sdata, b'Vertices'), default=()))] * 3))
    # We completely ignore normals here!
    weight = elem_prop_first(elem_find_first(fbx_bcdata, b'DeformPercent'), default=100.0) / 100.0
    vgweights = tuple(vgw / 100.0 for vgw in elem_prop_first(elem_find_first(fbx_bcdata, b'FullWeights'), default=()))

    assert(len(vgweights) == len(indices) == len(dvcos))
    create_vg = bool(set(vgweights) - {1.0})

    keyblocks = []

    for me, objects in meshes:
        vcos = tuple((idx, me.vertices[idx].co + Vector(dvco)) for idx, dvco in zip(indices, dvcos))
        objects = list({blen_o for fbx_o, blen_o in objects})
        assert(objects)

        if me.shape_keys is None:
            objects[0].shape_key_add(name="Basis", from_mix=False)
        objects[0].shape_key_add(name=elem_name_utf8, from_mix=False)
        me.shape_keys.use_relative = True  # Should already be set as such.

        kb = me.shape_keys.key_blocks[elem_name_utf8]
        for idx, co in vcos:
            kb.data[idx].co[:] = co
        kb.value = weight

        # Add vgroup if necessary.
        if create_vg:
            add_vgroup_to_objects(indices, vgweights, elem_name_utf8, objects)