Smooth out keys on infinite cycled controls in Maya's graph editor

SmoothInfiniteKeys.py

import maya.cmds as cmds
import math

def getVector(val1, val2):

    v1_theta = math.atan2(val1[1], val1[0])
    v2_theta = math.atan2(val2[1], val2[0])

    myAngle = math.degrees((v2_theta - v1_theta))

    return myAngle

def setKeyTangentValue(obj, time, value, attribute):
    # get the outweight so it's preserved
    outWeight = cmds.keyTangent(obj, time=(time,time), attribute=attribute, q = True, outWeight=True)
    cmds.keyTangent(obj, e=True, a=True, time=(time,time), outAngle=value, outWeight=outWeight[0], attribute=attribute)
    # grab average of the in and out weight and set the outweight again
    outAngle = (cmds.keyTangent(obj, time=(time, time), attribute=attribute, q=True, outAngle=True))[0]
    inAngle = (cmds.keyTangent(obj, time=(time, time), attribute=attribute, q=True, inAngle=True))[0]
    outAngleNew = ((outAngle + inAngle) / 2)

    cmds.keyTangent(obj, e=True, a=True, time=(time, time), outAngle=outAngleNew, attribute=attribute)

def smoothHax(obj, time, attribute, type):
    cmds.keyTangent(obj, e=True, a=True, time=(time, time), itt='linear', ott='linear', attribute=attribute)
    outWeight = cmds.keyTangent(obj, time=(time, time), attribute=attribute, q=True, outWeight=True)
    inWeight = cmds.keyTangent(obj, time=(time, time), attribute=attribute, q=True, inWeight=True)
    outAngle = (cmds.keyTangent(obj, time=(time, time), attribute=attribute, q=True, outAngle=True))[0]
    inAngle = (cmds.keyTangent(obj, time=(time, time), attribute=attribute, q=True, inAngle=True))[0]
    print attribute
    cmds.keyTangent(obj, e=True, a=True, time=(time, time), itt='auto', ott='auto', attribute=attribute)
    if type == 'in':
        a = cmds.keyTangent(obj, e=True, a=True, time=(time, time), inAngle=outAngle, inWeight=inWeight[0], attribute=attribute)
    if type == 'out':
        a = cmds.keyTangent(obj, e=True, a=True, time=(time, time), outAngle=inAngle, outWeight=outWeight[0], attribute=attribute)

def getAllKeys(obj, transform, overshootMultiplier):
    # print obj
    # get the first and second keys on the control- their time and their value
    min_time = cmds.playbackOptions(q=True, min=True)
    max_time = cmds.playbackOptions(q=True, max=True)
    allKeys = cmds.keyframe((obj + '.' + transform), time=(min_time, max_time), query=True, valueChange=True, timeChange=True)

    # set all keys to non-weighted, because fucking weighted keys, am I right?
    cmds.keyTangent(obj, e=True, a=True, time=(min_time, max_time), weightedTangents=False)

    # Check to make sure allKeys is valid after trying to get keys on the object.
    if allKeys is not None:
        # allKeys is returned as a poopy array of key time, key value, key time, key value, etc.
        # hence, going through it to by 2 is how we get the first two keys  and last two keys in the timeline
        # also make sure there's at least two keys, bruh
        if len(allKeys) >= 4:

            # there is probably a smarter way to do this but I don't care.
            firstKeyTime = allKeys[0]
            firstKeyValue= allKeys[1]
            secondKeyTime = allKeys[2]
            secondKeyValue = allKeys[3]

            secondToLastKeyTime = allKeys[(len(allKeys) - 4)]
            secondToLastKeyValue = allKeys[(len(allKeys) - 3)]
            lastKeyTime = allKeys[(len(allKeys) - 2)]
            lastKeyValue = allKeys[(len(allKeys) - 1)]

            # normalize the first set of values to 0,0 and adjust the second one relative to that
            # get the angle
            # dividing the time by 2 because we want to "aim" at the middle of the curve, it turns out, not the next key, for best smoothing
            # need better math for this though- dividing the time by 2 doesn't work great in all cases

            # get the relative (to the midpoint) time of the second and second to last frames
            length = (lastKeyTime - firstKeyTime)
            secondKeyRelativeTime = (secondKeyTime / (length / 2))
            secondToLastKeyRelativeTime = ((lastKeyTime - secondToLastKeyTime) / (length / 2))

            firstValue = getVector((0, 0), ((secondKeyTime-firstKeyTime)/2, (secondKeyValue-firstKeyValue)))

            #do the same thing again for the last frame and second to last frame. we will use the lower of these two values
            secondValue = getVector((0, 0), ((lastKeyTime - secondToLastKeyTime)/2, (lastKeyValue - secondToLastKeyValue)))

            value = ((firstValue + secondValue) / 2)

            # overshoot mafs!
            # get the difference between 89 and value (so it can never aim straight up)

            # multiply that by overshootMultiplier, and then add to value
            overshootDiff = 0
            if value > 0:
                overshootDiff = (89 - value)
            if value < 0:
                overshootDiff = (-89 - value)

            modifiedValue = (value + (overshootDiff * overshootMultiplier))

            # set the key tangent values on the first and last keys
            # IDEA - change the scale of the values depending on how far from center the next key is?
            setKeyTangentValue(obj, firstKeyTime, modifiedValue, transform)
            setKeyTangentValue(obj, lastKeyTime, modifiedValue, transform)

            # if overshootMultiplier is over 0.0, we need to aim the second and second to last key tangents at the first and last keys
            # which is easy, just use the negative of the pre-multiplied value. I think.
            # holy shit that worked.
            # do this math better, because the closer to the start or end it is, the worse it looks
            # I guess find the relative time of the range- and modify it by how far from the center it is? There is probably a math term for this.
            # or maybe based on the value difference between the two keys?
            # TODO: Lean math terms
            # lol it colors the "TODO" lines differently. oh.
            # these are shitty comments, back to the issue at hand.
            if overshootMultiplier > 0.0:
                # set second and second to last keys to linear, get values, set back to auto, set values.
                # HAX because fuck math and it's the weekend.
                smoothHax(obj, secondKeyTime, transform, 'in')
                smoothHax(obj, secondToLastKeyTime, transform, 'out')


def smoothCycle(selection):
    reporter = mel.eval('string $tmp = $gCommandReporter;')
    cmds.cmdScrollFieldReporter(reporter, e=True, clear=True)
    rots = ['rotateX', 'rotateY', 'rotateZ']
    #rots = ['rotateZ']
    trans = ['translateX', 'translateY', 'translateZ']
    scales = ['scaleX', 'scaleY', 'scaleZ']
    for obj in selection:
        # SET OVERSHOOT MULTIPLIER HERE - range from 0.0 to 1.0. different per transform type
        for r in rots:
            getAllKeys(obj, r, 0.0)
        for t in trans:
            getAllKeys(obj, t, 0.0)
        #for t in scales:
        #    getAllKeys(obj, t, 0.3)

smoothCycle((cmds.ls(selection=True)))