import bpy
from math import *

def set_kp(kp, co, easing):
    kp.co = co
    kp.handle_left = co
    kp.handle_right = co
    kp.easing = easing
    kp.interpolation = 'SINE'


def wipe_fcurves(action):
    while len(action.fcurves) > 0:
        action.fcurves.remove(action.fcurves[-1])


def keyframe_action1(action):
    wipe_fcurves(action)

    fc = action.fcurves.new("location", 2)
    fc.keyframe_points.add(5)
    set_kp(fc.keyframe_points[0], [0,0], 'EASE_OUT')
    set_kp(fc.keyframe_points[1], [20,1], 'EASE_IN')
    set_kp(fc.keyframe_points[2], [40,0], 'EASE_OUT')
    set_kp(fc.keyframe_points[3], [60,-1], 'EASE_IN')
    set_kp(fc.keyframe_points[4], [80,0], 'EASE_OUT')

def keyframe_action2(action, scale=1):

    wipe_fcurves(action)
    timescale = 400
    fc = action.fcurves.new("location", 2)
    fc.keyframe_points.add(5)
    set_kp(fc.keyframe_points[0], [0,0], 'EASE_IN')
    set_kp(fc.keyframe_points[1], [1*timescale,1*scale], 'EASE_OUT')
    set_kp(fc.keyframe_points[2], [2*timescale,2*scale], 'EASE_IN')
    set_kp(fc.keyframe_points[3], [3*timescale,1*scale], 'EASE_OUT')
    set_kp(fc.keyframe_points[4], [4*timescale,0], 'EASE_IN')

def action_for(u,v):

    name = "high frequency sine"
    action = bpy.data.actions.get(name)
    if action is None:
        action = bpy.data.actions.new(name)
        keyframe_action1(action)

    return action

def action2a():
    name = "low frequency cosine up"
    action = bpy.data.actions.get(name)
    if action is None:
        action = bpy.data.actions.new(name)
        keyframe_action2(action, 2)
    return action

def action2b():
    name = "low frequency cosine down"
    action = bpy.data.actions.get(name)
    if action is None:
        action = bpy.data.actions.new(name)
        keyframe_action2(action, -2)
    return action


def rig_track_one_strip(ad, track_no, strip_start_frame, action, source_start, source_end, repeat_count, strip_scale=1,
                        track_name="high frequency", blend_type='REPLACE'):
    if track_no < len(ad.nla_tracks):
        tr1 = ad.nla_tracks[track_no]
    else:
        tr1 = ad.nla_tracks.new()
    tr1.name = track_name
    if len(tr1.strips) < 1:
        strip1 = tr1.strips.new("bacon", strip_start_frame, action)
    else:
        strip1 = tr1.strips[0]
        strip1.action = action
    strip1.frame_start = strip_start_frame
    strip1.action_frame_start = source_start
    strip1.action_frame_end = source_end
    strip1.repeat = repeat_count
    strip1.scale = strip_scale
    strip1.blend_type = blend_type

    return (tr1, strip1)

def rig_nlas(obj, u,v):
    theta = u/3+v/10 # + frame*pi/40

    if obj.animation_data is None:
        obj.animation_data_create()

    ad = obj.animation_data

    if 0 == u % 2:
        repeat = 40
        scale = 1
    else:
        repeat = 38
        scale = 40 / 38
    rig_track_one_strip(ad, 0, -theta / (pi / 40), action_for(u, v), 0,80, repeat, scale)

    if 0==u%2:
        rig_track_one_strip(ad, 1, 1, action2a(), 0, 1600, 2, 1, "low frequency", 'ADD')
    else:
        rig_track_one_strip(ad, 1, 1, action2b(), 0, 1600, 2, 1, "low frequency", 'ADD')


def fab_object_for(scn, u,v):
    name = "ball2 at %d,%d" % (u, v)
    obj = bpy.data.objects.get(name)
    if obj is None:
        obj = bpy.data.objects.new(name, bpy.data.meshes.get("Icosphere"))
    try:
        scn.objects.link(obj)
    except:
        pass
    obj.location = (u,v,0)
    rig_nlas(obj, u, v)


scn = bpy.context.scene

for v in range(20):
    for u in range(20):
        fab_object_for(scn ,u,v)

keyframe_action1(bpy.data.actions['high frequency sine'])
keyframe_action2(bpy.data.actions['low frequency cosine down'], -2)
keyframe_action2(bpy.data.actions['low frequency cosine up'], 2)