# "simple" exporter for BCK anim files from SMG
# file format information --> https://humming-owl.neocities.org/smg-stuff/pages/tutorials/bck

import bpy, math
from . import blender_funcs
from . import bck_funcs
from . import file_ops
from . import math_funcs
import mathutils

# export BCK animation from the selected armature object
def export_bck_func(options, context):
  
  # this thing is always needed for stuff
  scene = context.scene
  
  # checking stage
  
  # if no armature is selected
  if (scene.objects.active == None):
    blender_funcs.disp_msg("No Armature selected. Select one and try again.")
    return {"FINISHED"}
  elif (scene.objects.active.type != "ARMATURE"):
    blender_funcs.disp_msg("No Armature selected. Currently selecting: \"%s\"" % (scene.objects.active.name))
    return {"FINISHED"}
  
  # select the armature object
  armature = scene.objects.active
  blender_funcs.select_obj(armature, False, "OBJECT")
  
  # if the armature has no bones (lmao)
  if (len(armature.data.bones) == 0):
    blender_funcs.disp_msg("Armature selected \"%s\" does not have any bones." % (armature.name))
    return {"FINISHED"}
  
  # if the armature has no animation data
  if (armature.animation_data == None
      or armature.animation_data.action == None):
    blender_funcs.disp_msg("Armature selected \"%s\" does not have an animation active." % (armature.name))
    return {"FINISHED"}
  
  # start gathering the animation information
  bck_anim = bck_funcs.smg_bck_anim()
  
  # loop mode (dark python string magic - ascii math)
  bck_anim.loop_mode = options.loop_mode.encode()[-1] - "A".encode()[0]
  bck_anim.anim_length = options.anim_length
  bck_anim.bone_count = len(armature.data.bones)
  for i in range(bck_anim.bone_count):
    # append the bone component animation data
    bck_anim.anim_data.append(bck_funcs.smg_bck_anim.anim_data())    
  
  # start getting the actual animation data
  for i in range(len(armature.data.bones)):
    data_bone = armature.data.bones[i] # correct bone index order
    pose_bone = armature.pose.bones[data_bone.name] # pose matrix is got from here
    
    # check if the bone has animation data (1 or more keyframes)
    # if not just add its rest pose value to the structure
    
    # gather the existing fcurves for a bone
    bone_fcurves = [None, None, None, None, None, None, None, None, None]
    # ^ sx, rx, tx, sy (24!), ry, ty, sz, rz and tz in that order (bck order)
    bone_data_path_str = "pose.bones[\"%s\"]." % (data_bone.name)
    foreign_fcurves = [None, None, None, None] # quaternions/axis angles have 4 components
    temp_fcurves = [None, None, None]
    # ^ for rotation mode conversion to euler selected order mode
    # from different euler order, quaternion or axis angle
    
    # detect the rotation mode string (it will be the rotation mode active on the bone)
    bone_rot_mode_str = pose_bone.rotation_mode
    if ("Z" in bone_rot_mode_str): # awful but fast check
      bone_rot_mode_str = "rotation_euler"
    elif (bone_rot_mode_str == "QUATERNION"):
      bone_rot_mode_str = "rotation_quaternion"
    elif (bone_rot_mode_str == "AXIS_ANGLE"):
      bone_rot_mode_str = "rotation_axis_angle"
    
    # now get the animation data
    for fcurve in armature.animation_data.action.fcurves:
      # scaling
      if (fcurve.data_path == bone_data_path_str + "scale"):
        bone_fcurves[int((3 * fcurve.array_index) + 0)] = fcurve
      # rotation (Euler [all its combinations], Quaternions, Axis angle)
      elif (fcurve.data_path == bone_data_path_str + bone_rot_mode_str):
        # check if it is Euler in the euler order selected, if not, select the curves for later conversion
        if (bone_rot_mode_str == "rotation_euler" and bone_rot_euler_order_str == options.euler_mode):
          bone_fcurves[int((3 * fcurve.array_index) + 1)] = fcurve  
        else:
          foreign_fcurves[fcurve.array_index] = fcurve
      # translation
      elif (fcurve.data_path == bone_data_path_str + "location"):
        bone_fcurves[int((3 * fcurve.array_index) + 2)] = fcurve
    
    # get the rest pose matrix
    rest_mat = data_bone.matrix_local.copy()
    if (pose_bone.parent != None):
      rest_mat = data_bone.parent.matrix_local.copy().inverted() * rest_mat.copy()
    else:
      rest_mat = mathutils.Matrix.Identity(4)
    
    # if there are foreign rotation modes used
    # convert them to the euler order mode requested
    # and assign those fcurves to the bone_fcurves list
    # (this is tuff)
    temp_fcurve_group = "temp"
    temp_fcurve_data_path = "temp_fcurves"
    if (options.euler_mode != pose_bone.rotation_mode):
      # generate the temp fcurves but first check if they are animation tracks named that way already
      for fcurve in list(armature.animation_data.action.fcurves):
        if (fcurve.group.name == temp_fcurve_group
            and fcurve.data_path == temp_fcurve_data_path): # delete this animation track
          armature.animation_data.action.fcurves.remove(fcurve)
      for j in range(3):
        temp_fcurves[j] = armature.animation_data.action.fcurves.new(temp_fcurve_data_path, j, temp_fcurve_group)      
      og_rot_values = [None, None, None, None] # to hold original data
      
      # other euler order
      if ("Z" in pose_bone.rotation_mode):
        # go through all animation frames
        for j in range(options.first_frame, options.anim_length):
          # get the og data, if there is no fcurve, get the rest pose values
          for k in range(3):
            if (og_rot_values[k] != None):
              og_rot_values[k] = foreign_fcurves[k].evaluate(j)
            else:
              og_rot_values[k] = rest_mat.to_euler(pose_bone.rotation_mode)
              
          # compile into a matrix, get the desired angles from it in the new euler order
          tmp = mathutils.Euler(og_rot_values[0:3], pose_bone.rotation_mode)
          tmp = tmp.to_matrix().to_euler(options.euler_mode)
          # assign the respective points to the temp fcurves
          for k in range(3):
            temp_fcurves[k].keyframe_points.insert(j, tmp[k])
        
      # quaternions
      elif (pose_bone.rotation_mode == "QUATERNION"):
        # go through all animation frames
        for j in range(options.first_frame, options.anim_length):
          # get the og data, if there is no fcurve, get the rest pose values
          for k in range(4):
            if (foreign_fcurves[k] != None):
              og_rot_values[k] = foreign_fcurves[k].evaluate(j)
            else:
              og_rot_values[k] = rest_mat.to_quaternion()[k]
              
          # compile into a quaternion, get the desired angles from it in the new euler order
          tmp = mathutils.Quaternion(og_rot_values).to_euler(options.euler_mode)
          # assign the respective points to the temp fcurves
          for k in range(3):
            temp_fcurves[k].keyframe_points.insert(j, tmp[k])
        
      # axis angle
      elif (pose_bone.rotation_mode == "AXIS_ANGLE"):
        # go through all animation frames
        for j in range(options.first_frame, options.anim_length):
          # get the og data, if there is no fcurve, get the rest pose values
          for k in range(4):
            if (og_rot_values[k] != None):
              og_rot_values[k] = foreign_fcurves[k].evaluate(j)
            else:
              og_rot_values[k] = rest_mat.to_quaternion().to_axis_angle()[k]
          # compile into a matrix, get the desired angles from it in the new euler order
          tmp = mathutils.Matrix.Rotation(og_rot_values[0], 3, og_rot_values[1:]).to_euler(options.euler_mode)
          # assign the respective points to the temp fcurves
          for k in range(3):
            temp_fcurves[k].keyframe_points.insert(j, tmp[k])
      
      # assign the resulting fcurves to the bone_fcurves list
      # rotation fcurves
      for j in range(3):
        bone_fcurves[int((j * 3) + 1)] = temp_fcurves[j]
    
    # all the fcurves with the correct units were selected
    # generate all the new animation points, interpolation stuff will be done later
    
    # get the points on all frames, only the points
    for j in range(bck_anim.anim_length):
      # find the values respect to rest pose
      scale = [1, 1, 1]
      rot = [0, 0, 0]
      transl = [0, 0, 0]
      
      # all components
      for k in range(9):
        # components with fcurve 
        if (bone_fcurves[k] != None and len(bone_fcurves[k].keyframe_points) >= 1):
          value = bone_fcurves[k].evaluate(options.first_frame + j)
          # check which is the component to get
          if (k == 0 or k == 3 or k == 6):
            scale[int((k - 0) / 3)] = value
          elif (k == 1 or k == 4 or k == 7):
            rot[int((k - 1) / 3)] = value
          elif (k == 2 or k == 5 or k == 8):
            transl[int((k - 2) / 3)] = value
        
      # convert the values to be respect to parent
      new_mat = rest_mat.copy() * math_funcs.calc_transf_mat(scale, rot, transl,
                                                             options.euler_mode,
                                                             options.mult_order).copy()
      for k in range(9):
        value = None
        # check which is the component to get
        if (k == 0 or k == 3 or k == 6):
          value = round(new_mat.to_scale()[int((k - 0) / 3)], options.rounding_vec[0])
        elif (k == 1 or k == 4 or k == 7):
          value = round(new_mat.to_euler(options.euler_mode)[int((k - 1) / 3)], options.rounding_vec[1])
        elif (k == 2 or k == 5 or k == 8):
          value = round(100 * new_mat.to_translation()[int((k - 2) / 3)], options.rounding_vec[2])
          # 100 times because of blenxy's coordinates
        bck_anim.anim_data[i].comp[k].value.append(value)
    
    # delete the temp_fcurves generated
    for fcurve in temp_fcurves:
      if (fcurve != None):
        armature.animation_data.action.fcurves.remove(fcurve)
        
  # got all the animation points
    
  # delete constant value animation tracks
  for i in range(bck_anim.bone_count):
    for j in range(9):
      anim_track_constant = True
      for k in range(bck_anim.anim_length):
        if (k == 0):
          continue
        # check if the whole animation track is the same
        if (bck_anim.anim_data[i].comp[j].value[k - 1] != bck_anim.anim_data[i].comp[j].value[k]):
          anim_track_constant = False
          break
      if (anim_track_constant == True):
        bck_anim.anim_data[i].comp[j].kf_count = 1
        bck_anim.anim_data[i].comp[j].interp_mode = 0
        bck_anim.anim_data[i].comp[j].time = [None]
        bck_anim.anim_data[i].comp[j].value = [bck_anim.anim_data[i].comp[j].value[0]]
        bck_anim.anim_data[i].comp[j].in_slope = [None]
        bck_anim.anim_data[i].comp[j].out_slope = [None]
  
  # ~ print(bck_anim)
  
  # keep all the samples intact and calculate the slopes
  # using linear interpolation between consecutive frames
  if (options.export_type == "OPT_A"):
    
    # assign the rest of the variables
    for i in range(bck_anim.bone_count):
      for j in range(9):
        # skip 1 keyframe animations
        if (bck_anim.anim_data[i].comp[j].kf_count == 1):
          continue
        bck_anim.anim_data[i].comp[j].kf_count = bck_anim.anim_length
        bck_anim.anim_data[i].comp[j].interp_mode = 1 # has to be like this
        for k in range(bck_anim.anim_length):
          bck_anim.anim_data[i].comp[j].time.append(k)
          in_slope = 0
          out_slope = 0
          if (k > 0):
            in_slope = bck_anim.anim_data[i].comp[j].value[k] - bck_anim.anim_data[i].comp[j].value[k - 1]
          if (k < bck_anim.anim_length - 1):
            out_slope = bck_anim.anim_data[i].comp[j].value[k + 1] - bck_anim.anim_data[i].comp[j].value[k]
          # set the rounding digit
          rounding = options.rounding_vec[0] # scale
          if (j == 1 or j == 4 or j == 7): # rotation
            rounding = options.rounding_vec[1]
          elif (j == 2 or j == 5 or j == 8): # translation
            rounding = options.rounding_vec[2]
          bck_anim.anim_data[i].comp[j].in_slope.append(round(in_slope, rounding))
          bck_anim.anim_data[i].comp[j].out_slope.append(round(out_slope, rounding))
  
  # find "best" interpolator fits for the samples
  elif (options.export_type == "OPT_B"):
    
    # assign the rest of the variables
    for i in range(bck_anim.bone_count):
      # assign the best fit for each animation component
      for j in range(9):
        # skip 1 keyframe animations
        if (bck_anim.anim_data[i].comp[j].kf_count == 1):
          continue
        # get the best fit interpolation result
        interp_result = math_funcs.find_best_cubic_hermite_spline_fit(options.first_frame,
                                                                      bck_anim.anim_data[i].comp[j].value,
                                                                      options.angle_limit)      
        # set the rounding digit
        rounding = options.rounding_vec[0] # scale
        if (j == 1 or j == 4 or j == 7): # rotation
          rounding = options.rounding_vec[1]
        elif (j == 2 or j == 5 or j == 8): # translation
          rounding = options.rounding_vec[2]        
        # round all the values returned by the interpolation fit
        for k in range(interp_result.kf_count):
          interp_result.value[k] = round(interp_result.value[k], rounding)
          if (k > 0):
            interp_result.in_slope[k] = round(interp_result.in_slope[k], rounding)
          if (k < interp_result.kf_count - 1):
            interp_result.out_slope[k] = round(interp_result.out_slope[k], rounding)
        
        # check if the fit can be made in interpolation mode == 0 (in_slope = out_slope)
        # assign the best fit for each animation component
        can_use_smooth_interp = True
        for k in range(interp_result.kf_count):
          if (k == 0 or k == interp_result.kf_count - 1):
            continue
          if (interp_result.in_slope[k] != interp_result.out_slope[k]):
            can_use_smooth_interp = False
            break
        
        # nice, adjust in_slope[0] and out_slope[-1]
        if (can_use_smooth_interp == True):
          interp_result.in_slope[0] = interp_result.out_slope[0]
          interp_result.out_slope[-1] = interp_result.in_slope[-1]
        else:
          interp_result.in_slope[0] = 0
          interp_result.out_slope[-1] = 0
          
        # overwrite the old animation track
        bck_anim.anim_data[i].comp[j].kf_count = interp_result.kf_count
        bck_anim.anim_data[i].comp[j].interp_mode = 1
        if (can_use_smooth_interp == True):
          bck_anim.anim_data[i].comp[j].interp_mode = 0
        bck_anim.anim_data[i].comp[j].time = interp_result.time
        bck_anim.anim_data[i].comp[j].value = interp_result.value
        bck_anim.anim_data[i].comp[j].in_slope = interp_result.in_slope
        bck_anim.anim_data[i].comp[j].out_slope = interp_result.out_slope
  
  # hopefully everything went okay
  # ~ print(bck_anim)
  
  # create a raw bck struct and write the BCK file
  raw = bck_funcs.create_smg_bck_raw(bck_anim)
  # ~ print(raw)  
  endian_ch = ">" # big endian character for struct.unpack()   
  if (options.endian == "OPT_B"): # little character
    endian_ch = "<"
  bck_funcs.write_smg_bck_raw(raw, options.filepath, endian_ch)
  
  # done!
  blender_funcs.disp_msg("BCK animation \"%s\" written" % (file_ops.get_file_name(options.filepath)))  
  return {"FINISHED"}

# Stuff down is for the menu appending
# of the importer to work plus some setting stuff
# comes from a Blender importer template
from bpy_extras.io_utils import ExportHelper
from bpy.props import StringProperty, BoolProperty, EnumProperty, FloatProperty, IntProperty, IntVectorProperty
from bpy.types import Operator

# export_bck class 
class export_bck(Operator, ExportHelper):
  """Export the animation data from an Armature into a SMG BCK file"""
  # stuff for blender
  bl_idname = "export_scene.bck"
  bl_label = "Export BCK (SMG)"
  filename_ext = ".bck"
  filter_glob = StringProperty(default = "*.bck", options = {"HIDDEN"}, maxlen = 255)    
  
  # exporter options
  export_type = EnumProperty(
    name = "Export Mode",
    description = "Way in which the animation will be exported",
    default = "OPT_B",
    items = (
      ("OPT_A", "Sample Everything", "Animation will be written completely sampled doing linear interpolation between all the frames of the animation. Some cleanup will be done while reading. Fast and accurate but takes a lot of space"),
      ("OPT_B", "Find Best Interpolator", "Will find the best interpolator fits for all the animation curves involved in the animation. Slow and not that accurate but can take less space")
    )
  )
  angle_limit = FloatProperty(
    name = "Derivative angle limit",
    description = "Value used to specify a keyframe generation at curve points at which sudden slope changes occur. Useful to adjust several straight lines. The angle comes from scaling the vertical axis of the animation track so that the \"visual derivative changes\" become visible. For export option \"Find Best Interpolator\"",
    default = 45,
    min = 0,
    max = 180,
  )
  first_frame = IntProperty(
    name = "First frame",
    description = "Value used to specify the first frame of the animation",
    default = 0,
  )
  anim_length = IntProperty(
    name = "Animation length",
    description = "Value used to specify the number of frames of the BCK animation after the first frame specified",
    default = 30,
  )
  loop_mode = EnumProperty(
    name = "Loop mode",
    description = "Way in which the animation be played in-game",
    default = "OPT_C",
    items = (
      ("OPT_A", "Play once - Stop at last frame", "Animation will start playing forwards and, when the animation data finishes, the last frame will be kept loaded into the model"),
      ("OPT_B", "Play once - Stop at first frame", "Animation will start playing forwards and, when the animation data finishes, the first frame will be kept loaded into the model"),
      ("OPT_C", "Repeat - Play forwards always", "Animation will start playing forwards and, when the animation data finishes, will play again from the beginning forwards"),
      ("OPT_D", "Play once - First forwards then backwards", "Animation will start playing forwards and, when the animation data finishes, the animation will be played backwards. This occurs only once"),
      ("OPT_E", "Repeat - Play forwards then backwards always", "Animation will start playing forwards and, when the animation data finishes, the animation will be played backwards. This repeats infinitely")
    )
  )
  endian = EnumProperty(
    name = "Endian order",
    description = "Way in which the animation data will be written",
    default = "OPT_A",
    items = (
      ("OPT_A", "Big", "Write data in the big endian byte ordering"),
      ("OPT_B", "Little", "Write data in the little endian byte ordering")
    )
  )
  rounding_vec = IntVectorProperty(
    name = "Round SRT to digit",
    description = "Round the scale, rotation and translation values (in that order) to the specified decimal position",
    default = (3, 5, 1),
    min = 0,
    max = 9
  )
  euler_mode = EnumProperty(
    name = "Euler order",
    description = "Export rotation animations in the specified Euler angles order",
    default = "XYZ",
    items = (
      ("XYZ", "XYZ", "X rotation first, Y rotation second, Z rotation last"),
      ("XZY", "XZY", "X rotation first, Z rotation second, Y rotation last"),
      ("YXZ", "YXZ", "Y rotation first, X rotation second, Z rotation last"),
      ("YZX", "YZX", "Y rotation first, Z rotation second, X rotation last"),
      ("ZXY", "ZXY", "Z rotation first, X rotation second, Y rotation last"),
      ("ZYX", "ZYX", "Z rotation first, Y rotation second, X rotation last")
    )
  )
  mult_order = EnumProperty(
    name = "Scale/Rot/Transl mult order",
    description = "Export animations in the specified matrix multiplication order",
    default = "SRT",
    items = (
      ("TRS", "TRS", "Translation first, Rotation second, Scaling last"),
      ("TSR", "TSR", "Translation first, Scaling second, Rotation last"),
      ("RTS", "RTS", "Rotation first, Translation second, Scaling last"),
      ("RST", "RST", "Rotation first, Scaling second, Translation last"),
      ("STR", "STR", "Scaling first, Translation second, Rotation last"),
      ("SRT", "SRT", "Scaling first, Rotation second, Translation last")
    )
  )
  # what the importer actually does
  def execute(self, context):
    return export_bck_func(self, context)

# stuff to append the item to the File -> Import/Export menu
def menu_export_bck(self, context):
  self.layout.operator(export_bck.bl_idname, text = "BCK for SMG (.bck)")

# register func
@bpy.app.handlers.persistent
def register(dummy):
  try:
    bpy.utils.register_class(export_bck)
    bpy.types.INFO_MT_file_export.append(menu_export_bck)
  except:
    return
    
# unregister func
def unregister():
  try:
    bpy.utils.unregister_class(export_bck)
    bpy.types.INFO_MT_file_export.remove(menu_export_bck)
  except:
    return