import bpy
import math

def makeMesh(name, nSegs, z1, z2, dTheta, thetaLag):
    mesh = bpy.data.meshes.new(name)
    verts = []
    faces = []
    for u in range(0,nSegs+1):
        v1=len(verts)
        verts.append( [u,0, 0] )
        verts.append( [u,0,z1] )
        verts.append( [u,0, z2] )
        if (u>0):
            v2 = v1+1
            v3 = v1+2
            v4 = v1-3
            v5 = v4+1
            v6 = v4+2
            faces.append( [ v1, v2, v5, v4] )
            faces.append( [ v2, v3, v6, v5] )
    print(len(verts))
    print(len(faces))
    mesh.from_pydata(verts, [], faces)
    mesh.validate(True)
    mesh.show_normal_face = True

    return mesh

def addVertexGroups(obj, nSegs):
    for u in range(0, nSegs+1):
        for r in range(1,3):
            vgn = "u=%d r=%d"%(u,r)
            vg = obj.vertex_groups.new(vgn)
            vg.add( [u*3+r] , 1, 'REPLACE')

def fabArmature(nSegs, z1):
    arm = bpy.data.armatures.new(name)
    oa = bpy.data.objects.new("armature.%s"%name, arm)

    bpy.context.scene.objects.link(oa)

    oa.select = True
    #oa.update()
    bpy.context.scene.update()

    bpy.context.scene.objects.active = oa

    bpy.ops.object.mode_set(mode='EDIT')

    for u in range(0, nSegs+1):
        for r in range(1,3):
            vgn = "u=%d r=%d"%(u,r)
            bone = arm.edit_bones.new(vgn)
            vgnp = "u=%d r=%d"%(u,r-1)
            if (r==1):
                bone.head = (u,0,0)
                bone.tail = (u,0,1)
            else:
                bone.head = (u,0,z1)
                bone.tail = (u,0,1+z1)
                bone.parent = arm.edit_bones[vgnp]
                bone.use_inherit_rotation=False

    if False:
      for u in range(0, nSegs+1):
        for r in range(2,3):
            vgn = "u=%d r=%d"%(u,r)
            bone = arm.edit_bones[vgn]
            vgnp = "u=%d r=%d"%(u,r-1)
            bone.parent = arm.edit_bones[vgnp]

    bpy.ops.object.mode_set(mode='OBJECT')


    for u in range(0, nSegs+1):
        for r in range(1,3):
            vgn = "u=%d r=%d"%(u,r)
            bone = oa.pose.bones[vgn]
            bone.rotation_mode = 'XYZ'

    return oa

def fabRotor(name, nRotors, z):
    arm = bpy.data.armatures.new(name)
    oa = bpy.data.objects.new(name, arm)

    bpy.context.scene.objects.link(oa)

    bpy.context.scene.objects.active = oa

    bpy.ops.object.mode_set(mode='EDIT')

    cs = arm.edit_bones.new("crankshaft")
    cs.head = (0,0,0)
    cs.tail = (-1,0,0)

    for k in range(nRotors):
        kn = "rotor %d"%k
        bone = arm.edit_bones.new(kn)
        bone.parent = cs
        theta = math.pi*2* k/nRotors
        y = math.sin(theta)
        z = math.cos(theta)
        bone.head = (-1, y,z)
        bone.tail = (-2, y,z)
    
    bpy.ops.object.mode_set(mode='OBJECT')

    oa.pose.bones[0].rotation_mode = 'XYZ'

    return oa

def addDrivers(ao, nSegs, rotor):

    nRotors = len(rotor.pose.bones)-1
    for u in range(0, nSegs+1):
        for r in range(1,3):
            bn = "u=%d r=%d"%(u,r)
            dr = ao.driver_add("pose.bones[\"%s\"].rotation_euler"%bn, 0)
            dr.driver.type = 'AVERAGE'
            var = dr.driver.variables.new()
            var.type = 'TRANSFORMS'
            var.targets[0].id = rotor
            if (r==1):
                ridx = u%nRotors
            else:
                ridx = (u-2 + nRotors)%nRotors
            var.targets[0].bone_target = "rotor %d" % ridx
            var.targets[0].transform_type = 'LOC_Y'

def spinRotor(rotor, period):
    crankshaft = rotor.pose.bones[0]
    crankshaft.rotation_euler.y = 0
    crankshaft.keyframe_insert(data_path = 'rotation_euler', index=1, frame=1)
    crankshaft.rotation_euler.y = math.pi*2
    crankshaft.keyframe_insert(data_path = 'rotation_euler', index=1, frame=1+period*bpy.context.scene.render.fps)

    curve = rotor.animation_data.action.fcurves[0]
    curve.extrapolation = 'LINEAR'
    curve.keyframe_points[0].interpolation = 'LINEAR'

#
#

z1 = 2
nSegs = 40
name ="fin"
obj = bpy.data.objects.get(name)
if (obj is None) :
    mesh = makeMesh(name, nSegs, z1, 3, 0,0)
    obj = bpy.data.objects.new(name, mesh)
    bpy.context.scene.objects.link(obj)

    addVertexGroups(obj, nSegs)


ao = bpy.data.objects.get("armature.%s"%name)
if(ao is None):
    ao = fabArmature(nSegs, z1)
    ao.location = (3,0,0)

    mod = obj.modifiers.new("armature", 'ARMATURE')
    mod.object = ao

    cns = obj.constraints.new('CHILD_OF')
    cns.target = ao


rotor = bpy.data.objects.get("rotor")
if rotor is None:
    rotor = fabRotor("rotor", 20, 0.8)

addDrivers(ao, nSegs, rotor)

spinRotor(rotor, 2)