import bpy
import bmesh


def verticesFor(xArray6, yArray4, zArray4):
    verts = []
    x0,x1,x2,x3,x4,x5 = xArray6 
    y0,y1,y2,y3 = yArray4
    z0,z1,z2,z3 = zArray4

    verts.append( [ x0, y0, z0 ] )
    verts.append( [ x5, y0, z0 ] )
    verts.append( [ x5, y3, z0 ] )
    verts.append( [ x0, y3, z0 ] )

    verts.append( [ x0, y0, z1 ] )
    verts.append( [ x5, y0, z1 ] )
    verts.append( [ x5, y3, z1 ] )
    verts.append( [ x0, y3, z1 ] )

    verts.append( [ x1, y1, z1 ] )
    verts.append( [ x2, y1, z1 ] )
    verts.append( [ x2, y2, z1 ] )
    verts.append( [ x1, y2, z1 ] )

    verts.append( [ x1, y1, z2 ] )
    verts.append( [ x2, y1, z2 ] )
    verts.append( [ x2, y2, z2 ] )
    verts.append( [ x1, y2, z2 ] )

    verts.append( [ x3, y1, z1 ] )
    verts.append( [ x4, y1, z1 ] )
    verts.append( [ x4, y2, z1 ] )
    verts.append( [ x3, y2, z1 ] )

    verts.append( [ x3, y1, z3 ] )
    verts.append( [ x4, y1, z3 ] )
    verts.append( [ x4, y2, z3 ] )
    verts.append( [ x3, y2, z3 ] )

    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 createMesh(name, xArray6, yArray4, zArray5):
    z0,z1,z2,z3,z4 = zArray5
    z9 = z3-z2+z1
    verts = verticesFor(xArray6, yArray4, 
                        [z0, z1, z9, z9])
    faces = []
    faces.append([3,2,1,0])
    faces.extend(facesForRings([0,1,2,3], [4,5,6,7]))

    faces.append( [4,5, 17,16, 9,8 ])
    faces.append( [4,8,11,7] )
    faces.append( [5,6,18,17] )
    faces.append( [6,7, 11,10, 19,18] )
    faces.append( [9, 16, 19, 10] )

    faces.extend( facesForRings([8,9,10,11], [12,13,14,15] ) )
    faces.append( [12,13,14,15] )

    faces.extend( facesForRings([16,17,18,19], [20,21,22,23] ) )
    faces.append( [20,21,22,23] )

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

    return mesh


def addShapeKeys(obj, xArray6, yArray4, zArray5):
    obj.shape_key_add("Basis")
    sk1 = obj.shape_key_add("Key1")
    sk2a = obj.shape_key_add("Key2A")
    sk2b = obj.shape_key_add("Key2B")
    print (sk2a)
    sk2a.relative_key = sk1
    sk2b.relative_key = sk1

    z0,z1,z2,z3,z4 = zArray5

    bm = bmesh.new()
    bm.from_mesh(obj.data)

    bm.verts.ensure_lookup_table()

    sl = bm.verts.layers.shape.get("Key1")
    verts = verticesFor(xArray6, yArray4, [ z0, z2, z3, z3] )
    for i in range(len(bm.verts)):
        bm.verts[i][sl] = verts[i]

    sl = bm.verts.layers.shape.get("Key2A")
    verts = verticesFor(xArray6, yArray4, [ z0, z2, z4, z3] )
    for i in range(len(bm.verts)):
        bm.verts[i][sl] = verts[i]

    sl = bm.verts.layers.shape.get("Key2B")
    verts = verticesFor(xArray6, yArray4, [ z0, z2, z3, z4] )
    for i in range(len(bm.verts)):
        bm.verts[i][sl] = verts[i]

    bm.to_mesh(obj.data)

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.L")
    bone.head = (0.5,0.5,0)
    bone.tail = (0.5,0.5,1)

    bone = arm.edit_bones.new("bone.R")
    bone.head = (1.5,0.5,0)
    bone.tail = (1.5,0.5,1)

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

    return oa

# compute the a and b such that a+b*x0 is 0 and a+b*x1 = 1
def ABFor(x0, x1):
    b = 1/(x1-x0)
    a = - b*x0
    return [a,b]


def addDrivers(obj, oa, zArray5):

    mesh = obj.data

    z0,z1,z2,z3,z4 = zArray5

    za = z1 + z3-z2
    zb = z3
    zc = z4

    data_path = "key_blocks[\"Key1\"].value"
    dr = mesh.shape_keys.driver_add(data_path)
    dr.driver.type='MAX'

    var = dr.driver.variables.new()
    var.type = 'TRANSFORMS'
    var.targets[0].id = oa
    var.targets[0].bone_target = "bone.L"
    var.targets[0].transform_type = 'LOC_Z'
    dr.modifiers[0].coefficients = ABFor(za, zb)


    var = dr.driver.variables.new()
    var.type = 'TRANSFORMS'
    var.targets[0].id = oa
    var.targets[0].bone_target = "bone.R"
    var.targets[0].transform_type = 'LOC_Z'
    dr.modifiers[0].coefficients = ABFor(za, zb)

    #

    data_path = "key_blocks[\"Key2A\"].value"
    dr = mesh.shape_keys.driver_add(data_path)
    dr.driver.type='MAX'

    var = dr.driver.variables.new()
    var.type = 'TRANSFORMS'
    var.targets[0].id = oa
    var.targets[0].bone_target = "bone.L"
    var.targets[0].transform_type = 'LOC_Z'
    dr.modifiers[0].coefficients = ABFor(zb, zc)

    #

    data_path = "key_blocks[\"Key2B\"].value"
    dr = mesh.shape_keys.driver_add(data_path)
    dr.driver.type='MAX'

    var = dr.driver.variables.new()
    var.type = 'TRANSFORMS'
    var.targets[0].id = oa
    var.targets[0].bone_target = "bone.R"
    var.targets[0].transform_type = 'LOC_Z'
    dr.modifiers[0].coefficients = ABFor(zb, zc)


xArray6 = [0, 0.2, 0.8, 1.2, 1.8, 2]
yArray4 = [0, 0.2, 0.8, 1 ]
zArray5 = [ 0, 0.2, 1, 1.2, 2]
name = "stacks"

mesh = createMesh(name, xArray6, yArray4 ,zArray5)

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

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

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

addShapeKeys(obj, xArray6, yArray4 ,zArray5)

oa = fabArmature(name)

addDrivers(obj, oa, zArray5)