import bpy
import bmesh

def vertsForTube(dxArray, dyArray, dzArray):
    z=0
    verts = []
    for i in range(len(dzArray)):
        verts.append( [ dxArray[i], dyArray[i], z] )
        verts.append( [ -dxArray[i], dyArray[i], z] )
        verts.append( [ -dxArray[i], -dyArray[i], z] )
        verts.append( [ dxArray[i], -dyArray[i], z] )
        z = z+dzArray[i]
        verts.append( [ dxArray[i], dyArray[i], z] )
        verts.append( [ -dxArray[i], dyArray[i], z] )
        verts.append( [ -dxArray[i], -dyArray[i], z] )
        verts.append( [ dxArray[i], -dyArray[i], z] )

    return verts

def facesForRings(ring1, ring2):
    faces = []
    n = len(ring1)
    for u in range(n):
        v0 = ring1[ u ]
        v1 = ring1[ (u+1)%n ]
        v2 = ring2[ u ]
        v3 = ring2[ (u+1)%n ]
        faces.append( [v0, v1, v3, v2])
    return faces

def makeBaseMesh(name, dxArray, dyArray, dzArray):
    verts = vertsForTube(dxArray, dyArray, dzArray)
    faces = []
    faces.append( [0,3,2,1] )
    for u in range(len(dzArray)):
        v0 = u*8
        v1 = v0+1
        v2 = v0+2
        v3 = v0+3
        v4 = v0+4
        v5 = v0+5
        v6 = v0+6
        v7 = v0+7
        faces.extend( facesForRings( [ v0,v1,v2,v3] , [v4,v5,v6,v7] ) )
        if (u+1 < len(dzArray)) :
            v8 = v0+8
            v9 = v0+9
            v10 = v0+10
            v11 = v0+11
            faces.extend( facesForRings( [v4,v5,v6,v7], [v8,v9,v10,v11] ) )
        else:
            faces.append( [v4,v5,v6,v7] )

    mesh = bpy.data.meshes.new(name)
    mesh.from_pydata(verts, [], faces)

    mesh.show_normal_face = True

    return mesh

def zsForStage(stage, total):
    rval = []
    for i in range(total):
        if i<stage:
            rval.append(1)
        else:
            rval.append(0.1)
    return rval

def addShapeKeys(obj, dxArray, dyArray):

    n = len(dxArray)
    for s in range(n+1):
        if (s==0):
            obj.shape_key_add("Basis")
        else:
            kn = "stage %d"%s
            obj.shape_key_add(kn)
            if (s>1):
                obj.data.shape_keys.key_blocks[kn].relative_key = obj.data.shape_keys.key_blocks["stage %d"%(s-1)]
            bm = bmesh.new()
            bm.from_mesh(obj.data)
            bm.verts.ensure_lookup_table()
            sl = bm.verts.layers.shape.get(kn)
            newVerts = vertsForTube(dxArray, dyArray, zsForStage(s, n))
            for u in range(len(newVerts)):
                bm.verts[u][sl] = newVerts[u]
            bm.to_mesh(obj.data)
    obj.data.shape_keys.use_relative = True


def fabArmature(name):
    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')

    bone = arm.edit_bones.new("bone")
    bone.head = (0,0,0)
    # if we don't set the tail!=head, the bone disappears
    bone.tail = (0,0,1)

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

    return oa

def addDriver(obj, armature, nSegs):

    extended = 1
    retracted = 0.1

    mesh = obj.data
    for i in range(nSegs):
        data_path = "key_blocks[\"stage %d\"].value" % (i+1)
        dr = mesh.shape_keys.driver_add(data_path)
        dr.driver.type='MIN'
        var = dr.driver.variables.new()
        var.type = 'TRANSFORMS'
        var.targets[0].id = armature
        var.targets[0].bone_target = "bone"
        var.targets[0].transform_type = 'LOC_Z'
        mod = dr.modifiers[0]
        slope = extended/(extended-retracted)
        base = (nSegs-i)*retracted + (i*extended)
        mod.coefficients = [-base*slope, 
                            slope]

dxArray = [ 1, 0.9, 0.8, 0.7, 0.6, 0.5]
dyArray = [ 2, 1.9, 1.8, 1.7, 1.6, 1.5]

dzArray = zsForStage(0, len(dxArray))

name="tube"
mesh = makeBaseMesh( name, dxArray, dyArray, dzArray)

obj = bpy.data.objects.new(name, mesh)
bpy.context.scene.objects.link(obj)

obj.select = True
bpy.context.scene.objects.active = obj

addShapeKeys(obj, dxArray, dyArray)


oa = fabArmature("armature")
oa.select = True

addDriver(obj, oa, len(dxArray))