import math

from mathutils import Matrix


def construct_bone(name, armature, matrix, scale):
    bpy_bone = armature.data.edit_bones.new(name)
    
    tail, roll = mat3_to_vec_roll(matrix.to_3x3())
    tail *= scale
    
    pos_vector = matrix.to_translation()
    bpy_bone.head = pos_vector
    bpy_bone.tail = pos_vector + tail
    bpy_bone.roll = roll
    
    # # Head/tail/roll sets correctly in Blender 2.83, but not in
    # # Blender 3.4?!
    # # So here we'll just manually set the matrix because... I have no idea
    # # why Blender sets the matrix_local incorrectly later
    # # Can't just set the matrix because that prevents the head/tail being set,
    # # so set the head/tail first and then align the roll by setting the matrix
    # # I feel *extremely* uncomfortable about the fact that two different
    # # roll values are required in two versions of Blender to get the same
    # # matrix - need to find out why.
    
    # bpy_bone.matrix = matrix
    
    return bpy_bone

def vec_roll_to_mat3(vec, roll):
    """
    Code from
    https://blender.stackexchange.com/a/90240
    with minor edits:
        - Removed 'mathutils' prefix from Matrix classes
        - Fixed invalid operation in penultimate line: * -> @
    """
    #port of the updated C function from armature.c
    #https://developer.blender.org/T39470
    #note that C accesses columns first, so all matrix indices are swapped compared to the C version

    nor = vec.normalized()
    THETA_THRESHOLD_NEGY = 1.0e-9
    THETA_THRESHOLD_NEGY_CLOSE = 1.0e-5

    #create a 3x3 matrix
    bMatrix = Matrix().to_3x3()

    theta = 1.0 + nor[1];

    if (theta > THETA_THRESHOLD_NEGY_CLOSE) or ((nor[0] or nor[2]) and theta > THETA_THRESHOLD_NEGY):

        bMatrix[1][0] = -nor[0];
        bMatrix[0][1] = nor[0];
        bMatrix[1][1] = nor[1];
        bMatrix[2][1] = nor[2];
        bMatrix[1][2] = -nor[2];
        if theta > THETA_THRESHOLD_NEGY_CLOSE:
            #If nor is far enough from -Y, apply the general case.
            bMatrix[0][0] = 1 - nor[0] * nor[0] / theta;
            bMatrix[2][2] = 1 - nor[2] * nor[2] / theta;
            bMatrix[0][2] = bMatrix[2][0] = -nor[0] * nor[2] / theta;

        else:
            #If nor is too close to -Y, apply the special case.
            theta = nor[0] * nor[0] + nor[2] * nor[2];
            bMatrix[0][0] = (nor[0] + nor[2]) * (nor[0] - nor[2]) / -theta;
            bMatrix[2][2] = -bMatrix[0][0];
            bMatrix[0][2] = bMatrix[2][0] = 2.0 * nor[0] * nor[2] / theta;

    else:
        #If nor is -Y, simple symmetry by Z axis.
        bMatrix = Matrix().to_3x3()
        bMatrix[0][0] = bMatrix[1][1] = -1.0;

    #Make Roll matrix
    rMatrix = Matrix.Rotation(roll, 3, nor)

    #Combine and output result
    mat = rMatrix @ bMatrix
    return mat


def mat3_to_vec_roll(mat):
    """
    Code from
    https://blender.stackexchange.com/a/38337
    https://blender.stackexchange.com/a/90240
    """
    vec = mat.col[1]
    vecmat = vec_roll_to_mat3(mat.col[1], 0)
    vecmatinv = vecmat.inverted()
    rollmat = vecmatinv @ mat
    roll = math.atan2(rollmat[0][2], rollmat[2][2])
    return vec, roll
