#!BPY
"""
Name: 'BSoD Rigs'
Blender: 248
Group: 'Animation'
Tooltip: 'Create Rig for Humanoid Character'
"""
__author__ = "Rasmus Lerchedahl Petersen (rusmus(at)itu.dk)"
__url__ = "http://wiki.blender.org/index.php/Extensions:Py/Scripts/Manual/Animation/BSoD_Rigs", "http://itu.dk/people/rusmus/blender/BSoD_rigs.py"

__version__ = "0.1"

__bpydoc__ = """ BSoD_rigs.py v0.1

Script for generating rigs for humanoid charachters.

This script generates the rigs discussed in "BSoD Introduction to Rigging"
http://mediawiki.blender.org/index.php/BSoD/Introduction_to_Rgging

Usage:<br>
 - Create guides (press "Create Guides");<br>
 - Adjust guides to fit character;<br>
 - Choose a name for the armature;<br>
 - Choose rig types for arms, legs, spine feet and hands;<br>  
 - Create armature (press "Create Armature"), the guides will be deleted.
"""

 #  ***** BEGIN GPL LICENSE BLOCK *****
 #
 #  This program is free software: you can redistribute it and/or modify
 #  it under the terms of the GNU General Public License as published by
 #  the Free Software Foundation, either version 3 of the License, or
 #  (at your option) any later version.
 #
 #  This program is distributed in the hope that it will be useful,
 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 #  GNU General Public License for more details.
 #
 #  You should have received a copy of the GNU General Public License
 #  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
 #  The Original Code is Copyright (C) <date>-<date> by <Rasmus Lerchedahl Petersen>	###
 #  All rights reserved.
 #
 #  Contact:      <name>@<domain.ext>	###
 #  Information:  http://<domain>.<ext>	###
 #
 #  The Original Code is: all of this file.
 #
 #  Contributor(s): none yet.
 #
 #  ***** END GPL LICENSE BLOCK *****
 
import Blender
#from Blender import Pose

#############
# Stage one:
#   Create Guides
#####################

# create/delete empties
def createEmpty(name, x, y, z):
    empty = Blender.Object.New("Empty", name)
    empty.loc = [x, y, z]
    scene = Blender.Scene.GetCurrent()
    scene.link(empty)

def deleteEmpty(name):
    empty = Blender.Object.Get(name)
    scene = Blender.Scene.GetCurrent()
    scene.unlink(empty)

guides = [
    ["COGGuide",        0,  0,  1],
    ["bodyGuide",       0,  0,  2],
    ["pelvisGuide",     0,  0,  0.2],
    ["rHipGuide",      -1,  0,  0],
    ["rKneeGuide",     -1, -1, -4],
    ["rAnkleGuide",    -1,  0, -8],
    ["rBallGuide",     -1, -1, -9],
    ["rToeGuide",      -1, -2, -9],    
    ["lHipGuide",       1,  0,  0],
    ["lKneeGuide",      1, -1, -4],
    ["lAnkleGuide",     1,  0, -8],
    ["lBallGuide",      1, -1, -9],
    ["lToeGuide",       1, -2, -9],    
    ["chestGuide",      0,  0,  3],
    ["rShoulderGuide", -1,  0,  4],
    ["rElbowGuide",    -4,  1  ,4],
    ["rVristGuide",    -7,  0,  4],
    ["rPalmGuide",     -8,  0,  4],
    ["lShoulderGuide",  1,  0,  4],
    ["lElbowGuide",     4,  1,  4],
    ["lVristGuide",     7,  0,  4],
    ["lPalmGuide",      8,  0,  4],
    ["neckGuide",       0,  0,  4],
    ["headGuide",       0,  0,  5]
    ]

# create/delete guides
def createGuides():
    for g in guides:
        createEmpty(g[0], g[1], g[2], g[3])
    Blender.Redraw()

def deleteGuides():
    for g in guides:
        deleteEmpty(g[0])
    Blender.Redraw()
    
#############
# Stage two:
#   Create Armature
#######################

# bone hooks
COGBone = None
hipParent = None
hipChild = {}
footParent = {}
footChild = {}
shoulderParent = {}
shoulderChild = {}
handParent = {}
handChild = {}
headParent = None
headChild = None

IKConstraints = []

# create bones
def addBone_old(armatureData, name, headEmpty, tailEmpty, parent = None):
    bone = Blender.Armature.Bone.New(name)
    bone.setHead(headEmpty.getMatrix("worldspace").translationPart())
    bone.setTail(tailEmpty.getMatrix("worldspace").translationPart())
    if parent != None:
        bone.setParent(parent)
    armatureData.addBone(bone)
    return bone

def connectBone(bone, parent, options = [Blender.Armature.CONNECTED]):
    if parent != None:
        bone.parent = parent
        bone.options = options
    
def addBone(armatureData, name, headEmpty, tailEmpty, parent = None, options = [Blender.Armature.CONNECTED]):
    bone = Blender.Armature.Editbone()
    #bone = Pose.PoseBone()
    bone.head = headEmpty.getMatrix("worldspace").translationPart()
    bone.tail = tailEmpty.getMatrix("worldspace").translationPart()
    connectBone(bone, parent, options)
    armatureData.bones[name] = bone
    return bone

# leg rigs  
def createFKLegRig(armatureData, prefix):
    global hipChild, footParent
    upperLegEmpty = Blender.Object.Get(prefix+"HipGuide")
    lowerLegEmpty = Blender.Object.Get(prefix+"KneeGuide")
    footEmpty = Blender.Object.Get(prefix+"AnkleGuide")

    upperLegBone = addBone(armatureData, prefix+"UpperLeg", upperLegEmpty, lowerLegEmpty)
    lowerLegBone = addBone(armatureData, prefix+"LowerLeg", lowerLegEmpty, footEmpty, parent = upperLegBone)

    hipChild[prefix] = upperLegBone
    footParent[prefix] = lowerLegBone
    
def createChickenLegRig(armatureData, prefix):
    global COGBone, hipChild, footParent
    upperLegEmpty = Blender.Object.Get(prefix+"HipGuide")
    lowerLegEmpty = Blender.Object.Get(prefix+"KneeGuide")
    footEmpty = Blender.Object.Get(prefix+"AnkleGuide")
    toeEmpty = Blender.Object.Get(prefix+"BallGuide")

    legBone = addBone(armatureData, prefix+"Leg", upperLegEmpty, toeEmpty)
    legBone.tail = legBone.head*0.5 + legBone.tail*0.5
    legUVBone = addBone(armatureData, prefix+"UVLeg", upperLegEmpty, toeEmpty, parent = legBone, options = [])
    legUVBone.tail = legUVBone.head*0.75 + legUVBone.tail*0.25
    upperLegBone = addBone(armatureData, prefix+"UpperLeg", upperLegEmpty, lowerLegEmpty, parent = legUVBone, options = [])
    lowerLegBone = addBone(armatureData, prefix+"LowerLeg", lowerLegEmpty, footEmpty, parent = upperLegBone)
    footBone = addBone(armatureData, prefix+"Foot", footEmpty, toeEmpty, parent = lowerLegBone)
    legIKBone = addBone(armatureData, prefix+"IKLeg", toeEmpty, upperLegEmpty, parent = COGBone, options = [])
    legIKBone.tail = legIKBone.head*0.75 + legIKBone.tail*0.25

    IKConstraints.append((footBone.name, legIKBone.name, 2))
    IKConstraints.append((legBone.name, legIKBone.name, 1))

    hipChild[prefix] = legBone
    footParent[prefix] = footBone

# arm rigs  
def createFKArmRig(armatureData, prefix):
    global shoulderChild, handParent, IKConstraints
    upperArmEmpty = Blender.Object.Get(prefix+"ShoulderGuide")
    lowerArmEmpty = Blender.Object.Get(prefix+"ElbowGuide")
    handEmpty = Blender.Object.Get(prefix+"VristGuide")

    upperArmBone = addBone(armatureData, prefix+"UpperArm", upperArmEmpty, lowerArmEmpty)
    lowerArmBone = addBone(armatureData, prefix+"LowerArm", lowerArmEmpty, handEmpty, parent = upperArmBone)

    shoulderChild[prefix] = upperArmBone
    handParent[prefix] = lowerArmBone

def createIKArmRig(armatureData, prefix):
    global COGBone, shoulderChild, handParent
    upperArmEmpty = Blender.Object.Get(prefix+"ShoulderGuide")
    lowerArmEmpty = Blender.Object.Get(prefix+"ElbowGuide")
    handEmpty = Blender.Object.Get(prefix+"VristGuide")

    armBone = addBone(armatureData, prefix+"Arm", upperArmEmpty, handEmpty)
    armBone.tail = armBone.head*0.7 + armBone.tail*0.3
    armUVBone = addBone(armatureData, prefix+"ArmUV", upperArmEmpty, lowerArmEmpty, parent = armBone, options = [])
    armVector = armBone.tail - armBone.head
    uvVector = Blender.Mathutils.CrossVecs(armVector, armUVBone.tail - armUVBone.head)
    uvVector = Blender.Mathutils.CrossVecs(uvVector, armVector)
    uvVector.normalize()
    #uv = 
    armUVBone.tail = armUVBone.head + uvVector
    upperArmBone = addBone(armatureData, prefix+"UpperArm", upperArmEmpty, lowerArmEmpty, parent = armUVBone, options = [])
    lowerArmBone = addBone(armatureData, prefix+"LowerArm", lowerArmEmpty, handEmpty, parent = upperArmBone)
    armIKBone = addBone(armatureData, prefix+"ArmIK", handEmpty, lowerArmEmpty, parent = COGBone, options = [])
    armIKBone.tail = armIKBone.head + Blender.Mathutils.Vector(0, 0, 1)

    IKConstraints.append((lowerArmBone.name, armIKBone.name, 2))
    IKConstraints.append((armBone.name, armIKBone.name, 1))

    shoulderChild[prefix] = armBone
    handParent[prefix] = lowerArmBone

# foot rigs  
def createFKFootRig(armatureData, prefix):
    global footChild
    footEmpty = Blender.Object.Get(prefix+"AnkleGuide")
    ballEmpty = Blender.Object.Get(prefix+"BallGuide")
    toeEmpty = Blender.Object.Get(prefix+"ToeGuide")

    footBone = addBone(armatureData, prefix+"Foot", footEmpty, ballEmpty)
    toeBone = addBone(armatureData, prefix+"Toe", ballEmpty, toeEmpty, parent = footBone)

    footChild[prefix] = footBone

def createChickenFootRig(armatureData, prefix):
    global footChild
    ballEmpty = Blender.Object.Get(prefix+"BallGuide")
    toeEmpty = Blender.Object.Get(prefix+"ToeGuide")

    toeBone = addBone(armatureData, prefix+"Toe", ballEmpty, toeEmpty)

    footChild[prefix] = toeBone

# hand rigs  
def createFKHandRig(armatureData, prefix):
    global handChild
    vristEmpty = Blender.Object.Get(prefix+"VristGuide")
    handEmpty = Blender.Object.Get(prefix+"PalmGuide")

    handBone = addBone(armatureData, prefix+"Hand", vristEmpty, handEmpty)

    handChild[prefix] = handBone

# head rigs  
def createFKHeadRig(armatureData, prefix):
    global headChild
    if prefix == "r":
        neckEmpty = Blender.Object.Get("neckGuide")
        headEmpty = Blender.Object.Get("headGuide")

        headBone = addBone(armatureData, "headd", neckEmpty, headEmpty)

        headChild = headBone

def connectRig():
    global hipChild, hipParent
    global footChild, footParent
    global shoulderChild, shoulderParent
    global handChild, handParent
    global headChild, headParent
    for prefix in ["r", "l"]:
        connectBone(hipChild[prefix], hipParent, options = [])
        connectBone(footChild[prefix], footParent[prefix])
        connectBone(shoulderChild[prefix], shoulderParent[prefix])
        connectBone(handChild[prefix], handParent[prefix])
    connectBone(headChild, headParent)

def createBodyRig(armatureData):
    global COGBone, hipParent, shoulderParent, headParent
    COGEmpty = Blender.Object.Get("COGGuide")
    bodyEmpty = Blender.Object.Get("bodyGuide")
    hipEmpty = Blender.Object.Get("pelvisGuide")
    chestEmpty = Blender.Object.Get("chestGuide")
    shoulderEmpty = {}
    for prefix in ["r", "l"]:
        shoulderEmpty[prefix] = Blender.Object.Get(prefix+"ShoulderGuide")
    neckEmpty = Blender.Object.Get("neckGuide")

    COGBone = addBone(armatureData, "COG", COGEmpty, COGEmpty)
    COGBone.head = COGBone.head - Blender.Mathutils.Vector(0, 1, 0)
    bodyBone = addBone(armatureData, "body", COGEmpty, bodyEmpty, parent = COGBone, options = [])
    hipBone = addBone(armatureData, "hip", COGEmpty, hipEmpty, parent = bodyBone, options = [])
    hipParent = hipBone
    chestBone = addBone(armatureData, "chest", bodyEmpty, chestEmpty, parent = bodyBone)
    shoulderBone = {}
    for prefix in ["r", "l"]:
        shoulderBone[prefix] = addBone(armatureData, prefix+"Shoulder", chestEmpty, shoulderEmpty[prefix], parent = chestBone)
        shoulderParent[prefix] = shoulderBone[prefix]
    neckBone = addBone(armatureData, "neck", chestEmpty, neckEmpty, parent = chestBone)
    headParent = neckBone

def createIKConstraints(armatureObj):
    global IKConstraints

    pose = armatureObj.getPose()
    for (sourceName, targetName, chainLength) in IKConstraints:
        poseBone = pose.bones[sourceName]
        constraint = poseBone.constraints.append(Blender.Constraint.Type.IKSOLVER)
        constraint[Blender.Constraint.Settings.TARGET] = armatureObj
        constraint[Blender.Constraint.Settings.BONE] = targetName
        constraint[Blender.Constraint.Settings.CHAINLEN] = chainLength
        constraint.influence = 1.0
    pose.update()

#######
# GUI:
##########
rigOptions = {
    "Leg": [
    ["FK", createFKLegRig],
    ["Chicken", createChickenLegRig]
    ],
    "Arm": [
    ["FK", createFKArmRig],
    ["IK", createIKArmRig]
    ],
    "Foot": [
    ["FK", createFKFootRig],
    ["Chicken", createChickenFootRig]
    ],
    "Hand": [
    ["FK", createFKHandRig]
    ],
    "Head": [
    ["FK", createFKHeadRig]
    ]
}

ArmatureName = Blender.Draw.Create(0)

guideEvent = 1
rigEvent = 2
rigChoiceEvent = 3
rigChoices = {}
legChoice = None
armChoice = None

def createRig(armatureName):
    global legChoice, armChoice
    global rigChoices
#    armatureData = Blender.Armature.New(armatureName+"_data")
    armatureData = Blender.Armature.Armature(armatureName+"_data")
    armatureData.makeEditable()

    createBodyRig(armatureData)

    for rig in rigOptions.keys():
        for prefix in ["r","l"]:
            rigOptions[rig][rigChoices[rig].val][1](armatureData, prefix)

    connectRig()

    armatureObj = Blender.Object.New("Armature", armatureName)
    armatureObj.link(armatureData)    
    armatureData.update()

    scene = Blender.Scene.getCurrent()
    scene.link(armatureObj)

    createIKConstraints(armatureObj)

# events
def key_event(event, val):
    if event == Blender.Draw.QKEY:
        Blender.Draw.Exit()
    return

def button_event(event):
    global ArmatureName

    print event

    if event == guideEvent: # create guides
        createGuides()
    elif event == rigEvent: # create rig
        armatureName = ArmatureName.val
        createRig(armatureName)
        deleteGuides()
    return

# GUI
def menuOptions(title, options):
    optionString = title+" %t"
    n = 0
    for option in options:
        optionString += "|%s %sx%d" % (option[0], "%", n)
        n +=1
    return optionString

def buildMenu((x, y), title, options, event):
    Blender.BGL.glRasterPos2i(x+2,y)
    Blender.Draw.Text("Choose " + title + ":", "normal")
    return Blender.Draw.Menu(menuOptions(title, options), event, x, y-25, 100, 20, 0, "Choice of " + title)

def menuCoordinates(startX, startY, rowBreak, columnBreak):
    width = 120
    height = 50

    x = startX
    y = startY
    n = 0
    colBreak = rowBreak*columnBreak
    while 1:
        yield (x, y)
        n += 1
        x += width
        if n % rowBreak == 0:
            x -= width*rowBreak
            y += height
        if n % colBreak == 0:
            x += (width + 1)*rowBreak
            y -= height*columnBreak

def MakeRigGUI():
    global ArmatureName, guideEvent, rigEvent, rigChoiceEvent, legChoice, armChoice

    Blender.BGL.glClearColor(0.5, 0.5, 0.5, 1)
    Blender.BGL.glClear(Blender.BGL.GL_COLOR_BUFFER_BIT)
    Blender.BGL.glColor3f(1, 1, 1)
    Blender.BGL.glRasterPos2i(20, 200)
    Blender.Draw.Text("Make Rig v0.1 by Rasmus Lerchedahl Petersen", "normal")    
  
    ArmatureName = Blender.Draw.String("Armature Name: ", 0, 20, 175, 220, 20, "armature", 80, "Name of armature to be created")

#    armChoice = buildMenu(20, 100, "Arm Rig", armOptions, rigChoiceEvent)
#    legChoice = buildMenu(140, 100, "Leg Rig", legOptions, rigChoiceEvent)

    menuXY = menuCoordinates(20, 100, 2, 2)
    for rig in rigOptions.keys():
        rigChoices[rig] = buildMenu(menuXY.next(), rig + " Rig", rigOptions[rig], rigChoiceEvent)

    Blender.Draw.Button("Create Guides", guideEvent, 20, 20, 100, 20, "Create Guides")
    Blender.Draw.Button("Create Armature", rigEvent, 140, 20, 100, 20, "Create Armature")

Blender.Draw.Register(MakeRigGUI, key_event, button_event)
