import math, numpy
import mathutils
import warnings
warnings.simplefilter("ignore", numpy.RankWarning)

# file containning math related functions

# function to return a 0 filled matrix 3x3
def get_zero_mat3x3():
  return mathutils.Matrix(([0, 0, 0], [0, 0, 0], [0, 0, 0]))
      
# function to return a 0 filled matrix 4x4
def get_zero_mat4x4():
  return mathutils.Matrix(([0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]))

# function to return an identity matrix 3x3
def get_id_mat3x3():
  return mathutils.Matrix(([1, 0, 0], [0, 1, 0], [0, 0, 1]))
      
# function to return an identity matrix 4x4
def get_id_mat4x4():
  return mathutils.Matrix(([1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]))  

# calc_scale_matrix function
# function to build the scale matrix
def calc_scale_matrix(sx, sy, sz):

    mat = mathutils.Matrix(([sx, 0, 0, 0],
                            [0, sy, 0, 0],
                            [0, 0, sz, 0],
                            [0, 0, 0, 1]))
    return mat

# calc_rotation_matrix function
# function to calculate the rotation matrix
# for a Extrinsic Euler XYZ system (radians)
def calc_rotation_matrix(rx, ry, rz, order):
  # calculate the axis rotation matrices
  x_rot = mathutils.Matrix(([1, 0, 0, 0],
                            [0, math.cos(rx), -math.sin(rx), 0],
                            [0, math.sin(rx), math.cos(rx), 0],
                            [0, 0, 0, 1]))
  y_rot = mathutils.Matrix(([math.cos(ry), 0, math.sin(ry), 0],
                            [0, 1, 0, 0],
                            [-math.sin(ry), 0, math.cos(ry), 0],
                            [0, 0, 0, 1]))
  z_rot = mathutils.Matrix(([math.cos(rz), -math.sin(rz), 0, 0],
                            [math.sin(rz), math.cos(rz), 0, 0],
                            [0, 0, 1, 0],
                            [0, 0, 0, 1]))               
  # check the rotation order
  if (order == "XYZ"):
    return z_rot * y_rot * x_rot
  elif (order == "XZY"):
    return y_rot * z_rot * x_rot
  elif (order == "YXZ"):
    return z_rot * x_rot * y_rot
  elif (order == "YZX"):
    return x_rot * z_rot * y_rot
  elif (order == "ZXY"):
    return y_rot * x_rot * z_rot
  elif (order == "ZYX"):
    return x_rot * y_rot * z_rot
  return None

# calc_translation_matrix function
# function to build the translation matrix
def calc_translation_matrix(tx, ty, tz):

    mat = mathutils.Matrix(([1, 0, 0, tx],
                            [0, 1, 0, ty],
                            [0, 0, 1, tz],
                            [0, 0, 0, 1]))
    return mat

# calculate a transformation matrix (euler)
def calc_transf_mat(scale, rotation, translation, rot_order, mult_order):
  # get the 3 matrices
  transl = calc_translation_matrix(translation[0], translation[1], translation[2])
  rot = calc_rotation_matrix(rotation[0], rotation[1], rotation[2], rot_order)
  scale = calc_scale_matrix(scale[0], scale[1], scale[2])
  # check the multiplication order
  if (mult_order == "TRS"):
    return scale * rot * transl
  elif (mult_order == "TSR"):
    return rot * scale * transl
  elif (mult_order == "RTS"):
    return scale * transl * rot
  elif (mult_order == "RST"):
    return transl * scale * rot
  elif (mult_order == "STR"):
    return rot * transl * scale
  elif (mult_order == "SRT"):
    return transl * rot * scale
  return None

# calculate a value of a cubic hermite interpolation
# it is assumed t0 and tf are 0 and 1
# t is the parametrization variable
# https://humming-owl.neocities.org/smg-stuff/pages/tutorials/model3#cubic-hermite-spline
def cubic_hermite_spline_time_unitary(p0, m0, m1, p1, t):
  # interpolation not defined here bruh
  if (t < 0 or t > 1):
    return None
  t_pow3 = math.pow(t, 3)
  t_pow2 = math.pow(t, 2)
  t_pow1 = t
  interp_result  = p0 * ((+2 * t_pow3) - (3 * t_pow2) + 1)
  interp_result += m0 * ((+1 * t_pow3) - (2 * t_pow2) + t_pow1)
  interp_result += m1 * ((+1 * t_pow3) - (1 * t_pow2) + 0)
  interp_result += p1 * ((-2 * t_pow3) + (3 * t_pow2) + 0)
  return interp_result
  
# same as the above function but t0 and tf can be any time value (tf > t0)
def cubic_hermite_spline_time_general(p0, m0, m1, p1, t0, tf, t):
  # interpolation not defined here bruh
  if (t < t0 or t > tf):
    return None
  interp_result = cubic_hermite_spline_time_unitary(p0, m0 * (tf - t0), m1 * (tf - t0), p1, (t - t0) / (tf - t0))
  return interp_result
  
# structure to be returned by the function below
class best_chs_fits:
  
  def __init__(self):
    self.kf_count = 0
    self.time = []
    self.value = []
    self.in_slope = []
    self.out_slope = []
  
  def __str__(self):
    rtn =  "keyframe count: %s\n" % (self.kf_count)
    rtn += "times: %s\n" % (self.time)
    rtn += "values: %s\n" % (self.value)
    rtn += "in slopes: %s\n" % (self.in_slope)
    rtn += "out slopes: %s" % (self.out_slope)
    return rtn

# function to calculate the best cubic hermite spline interpolator fits
# for a given a set of points. The set of points is expected to be at 
# each frame of the animation (start_frame indicates the start frame, integer)
# the function will return the above structure
# it will be in the general cubic hermite spline form (t0 < tf)

# make it so that keyframe trigger variables can be modified
def find_best_cubic_hermite_spline_fit(start_frame, values, angle_limit):
  
  # check
  if ((type(start_frame) != int)
      or (type(values) != list)
      or (len(values) == 0)):
    return None
  
  # create the object to return
  rtn = best_chs_fits()
  
  # the cubic hermite spline is just a sub-family of 3rd degree bezier curves
  # that makes this curve to be "able" to represent 3rd degree polynomials
  # these polinomials can have a maximum of 2 concavity changes
  # to avoid a greater data loss I will just make the best fit
  # when single concavity change is reached
  
  # special cases
  if (len(values) == 1): # single keyframe values like in BCK
    rtn.kf_count = 1
    rtn.time = [None]
    rtn.value = values
    rtn.in_slope = [None]
    rtn.out_slope = [None]
    return rtn
  elif (len(values) == 2): # 2 keyframe values (force linear interpolation)
    rtn.kf_count = 2
    rtn.time = [start_frame, start_frame + 1]
    rtn.value = values
    rtn.in_slope = [None, (rtn.value[1] - rtn.value[0]) / (rtn.time[1] - rtn.time[0])]
    rtn.out_slope = [(rtn.value[1] - rtn.value[0]) / (rtn.time[1] - rtn.time[0]), None]
    return rtn
  
  # get the lowest and highest value of the set of values given
  # will be used to scale the points to be able to do derivative difference checking
  # also get the average
  # all of them positive
  lowest_value = 0
  highest_value = 0
  avg_value = 0
  for i in range(len(values)):
    if (abs(values[i]) < lowest_value):
      lowest_value = abs(values[i])
    if (abs(values[i]) > highest_value):
      highest_value = abs(values[i])
    avg_value += abs(values[i])
  print(highest_value)
  scale_factor = (1 / highest_value) * (len(values) - start_frame)
  avg_value /= len(values)
  
  # lets do this bruv
  start_index = 0
  old_value = None
  cur_value = None
  new_value = None
  left_first_derivative = None
  right_first_derivative = None
  old_concavity = None
  new_concavity = None
  generate_keyframe = None
  small_time_dif = 0.000001
  # generate the keyframes
  i = 0
  while (i < len(values)):    
    # first keyframe
    if (i == 0):
      rtn.kf_count += 1
      rtn.time.append(start_frame)
      rtn.value.append(values[0])
      rtn.in_slope.append(None)
      rtn.out_slope.append(None)
    
    # last keyframe / concavity never changes
    elif (i == len(values) - 1):
      # add the last keyframe of the animation
      poly_deg = 3
      if (i - start_index == 1):
        poly_deg = 1
      elif(i - start_index == 2):
        poly_deg = 2
      rtn.kf_count += 1
      rtn.time.append(start_frame + i)
      rtn.value.append(values[-1])
      poly = numpy.polyfit(list(range(start_index, i + 1)),
                           values[start_index : i + 1], poly_deg)
      p0 = numpy.polyval(poly, start_index)
      pa = numpy.polyval(poly, start_index + small_time_dif)
      m0 = (pa - p0) / small_time_dif
      pb = numpy.polyval(poly, i - small_time_dif)
      p1 = numpy.polyval(poly, i)
      m1 = (p1 - pb) / small_time_dif
      rtn.in_slope.append(m1)
      rtn.out_slope[-1] = m0
      rtn.out_slope.append(None)
    
    # the rest
    else:
      
      # get the values
      old_value = values[i - 1]
      cur_value = values[i]
      new_value = values[i + 1]
          
      # determine the first derivative on both sides
      left_first_derivative = cur_value - old_value
      right_first_derivative = new_value - cur_value
      new_concavity = right_first_derivative - left_first_derivative
      
      # check if this is the first time getting the concavity
      if (old_concavity == None):
        old_concavity = new_concavity
      
      # check concavity changes
      if (numpy.sign(old_concavity) != numpy.sign(new_concavity)):
        generate_keyframe = True
      
      # scale the values up and check if the slope change is too violent with vector angles
      # the scaling is done to be able to visually see the angle changes, it is the same thing
      # we do with curves when zoomin in or out (either horizontally or vertically)
      left_first_derivative_scaled = left_first_derivative * scale_factor
      right_first_derivative_scaled = right_first_derivative * scale_factor
      left_vec = mathutils.Vector((1, left_first_derivative_scaled))
      right_vec = mathutils.Vector((1, right_first_derivative_scaled))
      angle = abs(math.degrees(left_vec.angle(right_vec)))
      if (angle > angle_limit):
        generate_keyframe = True        
      
      # store the old concavity for the next loop
      old_concavity = new_concavity
    
      # check if a new keyframe should be added
      if (generate_keyframe == True):
        # find the best 3rd degree polynomial "fit" with the set of points studied
        poly_deg = 3
        if (i - start_index == 1):
          poly_deg = 1
        elif(i - start_index == 2):
          poly_deg = 2
        poly = numpy.polyfit(list(range(start_index, i + 1)),
                             values[start_index : i + 1], poly_deg)
        print(poly)
        p0 = numpy.polyval(poly, start_index)
        pa = numpy.polyval(poly, start_index + small_time_dif)
        m0 = (pa - p0) / small_time_dif
        pb = numpy.polyval(poly, i - small_time_dif)
        p1 = numpy.polyval(poly, i)
        m1 = (p1 - pb) / small_time_dif        
        # write the data into the structure
        rtn.kf_count += 1
        rtn.time.append(start_frame + i)
        rtn.value.append(p1)
        rtn.in_slope.append(m1)
        rtn.out_slope[-1] = m0
        rtn.out_slope.append(None)
      
        # reset these variables for the next iteration
        start_index = i
        generate_keyframe = False
    
    # increment i 
    i += 1
    
  # done!
  return rtn