####################### Wheel Rig ###########################
##
## Create a simple rig for a wheel
##
## Created by Abby O'Malley
## aabrbty.com
## Follow me on instagram! @aabrbty
## Reach out for questions/comments!
##
##############################################################################
import maya.cmds as cmds

### Assign selected wheel to variable and get its shape node ###
curr_whl = cmds.ls(sl = True)[0]
print("Current Wheel: ", curr_whl)#X
whlShape = cmds.listRelatives(curr_whl)[0]
print("Current Wheel Shape: ", whlShape)#X

### Prefix  ###
PFX = curr_whl + "_"
print("Prefix: ", PFX)#X

"""
### RESET (for Testing) ###
if cmds.objExists (PFX+"*"): cmds.delete (PFX+"*")
"""

### Get Size of Wheel ###
sizeMaxX = (cmds.getAttr(curr_whl + '.boundingBoxMaxX'))
sizeMinX = cmds.getAttr(curr_whl + '.boundingBoxMinX')
sizeMaxY = cmds.getAttr(curr_whl + '.boundingBoxMaxY')
sizeMinY = cmds.getAttr(curr_whl + '.boundingBoxMinY')
sizeMaxZ = (cmds.getAttr(curr_whl + '.boundingBoxMaxZ'))
sizeMinZ = cmds.getAttr(curr_whl + '.boundingBoxMinZ')

sizeX = sizeMaxX-sizeMinX
sizeY = sizeMaxY-sizeMinY
sizeZ = sizeMaxZ-sizeMinZ
print("sizeX: ", sizeX)#X
print("sizeY: ", sizeY)#X
print("sizeZ: ", sizeZ)#X


"""
## Test Bounding Box ##
cmds.polyCube(n = "bboxcube")
cmds.scale(sizeX,sizeY,sizeZ)
cmds.matchTransform("bboxcube",f"{curr_whl}",px=True,py=True,pz=True)
"""


#### Create Root CTRL ####
cmds.circle(n=f"{PFX}root_CTRL", nr=(0,1,0), c = (0,0,0))
cmds.group(f"{PFX}root_CTRL", n= f"{PFX}CTRL_GRP")
cmds.matchTransform(f"{PFX}CTRL_GRP", f"{curr_whl}", px=True,pz= True, py=True)
cmds.move((sizeY*.5)*-1,y=True, r=True)
#cmds.move((sizeX),x=True, r=True)


### Size the Root CTRL ###
cmds.select(f"{PFX}root_CTRL")
if sizeZ >= sizeX:
   rootScale = (sizeZ - sizeX)/2.5
   cmds.scale(rootScale,rootScale,rootScale)
   print ("IF one")
elif sizeX >= sizeZ:
   rootScale = (sizeX - sizeZ)*3
   cmds.scale(rootScale,rootScale,rootScale)
   print("IF two")
## Shape Ctrl ##
cmds.select(f"{PFX}root_CTRL.cv[5]")
cmds.move(0,0,rootScale*.63, r= 1)
## Delete History / Freeze Transforms ##
cmds.select(f"{PFX}root_CTRL")
cmds.FreezeTransformations(f"{PFX}root_CTRL")
cmds.DeleteHistory(f"{PFX}root_CTRL")
### Add auto attr to root CTRL ###
cmds.select(f"{PFX}root_CTRL")
cmds.addAttr(ln="auto", at="bool", dv=0)
cmds.setAttr(f"{PFX}root_CTRL.auto",1, k=1)
### Add auto rotation attr to root CTRL ###
cmds.addAttr(ln="autoRot", at="float", dv=0)
cmds.setAttr(f"{PFX}root_CTRL.autoRot",1, k=1)


#### Create Locator ####
cmds.spaceLocator(n=f"{PFX}LOC")
cmds.matchTransform(f"{PFX}LOC", f"{curr_whl}")
cmds.parent(f"{PFX}LOC",f"{PFX}root_CTRL")# put LOC under root_CTRL
cmds.FreezeTransformations(f"{PFX}LOC")
cmds.DeleteHistory(f"{PFX}LOC")


#### Create Rotate GRP ####
cmds.group(f"{curr_whl}", n=f"{PFX}rotate_GRP")
### Apply Point and Scale Constraint to LOC and rotate_GRP ###
cmds.pointConstraint(f"{PFX}LOC", f"{PFX}rotate_GRP", mo=1)
cmds.scaleConstraint(f"{PFX}LOC", f"{PFX}rotate_GRP", mo=1)
cmds.orientConstraint(f"{PFX}LOC", f"{PFX}rotate_GRP", mo=1)
### Add offset attr to rotate GRP ###
cmds.select(f"{PFX}rotate_GRP")
cmds.addAttr(ln="offset", at="float", dv=0)
cmds.setAttr(f"{PFX}rotate_GRP.offset",1, k=1)


### Offset CTRL ###
cmds.circle(n=f"{PFX}offset_CTRL", nr=(1,0,0), c = (0,0,0))
cmds.parent(f"{PFX}offset_CTRL",f"{PFX}root_CTRL")
cmds.matchTransform(f"{PFX}offset_CTRL", curr_whl, px= True, py=True, pz= True)
cmds.move(sizeX*.7,0,0, r=1)
offsetScale = sizeY*.15
cmds.scale(offsetScale,offsetScale,offsetScale)
## Shape Ctrl ##
cmds.select(f"{PFX}offset_CTRL.cv[0]",f"{PFX}offset_CTRL.cv[2]",f"{PFX}offset_CTRL.cv[4]",f"{PFX}offset_CTRL.cv[6]", add=False)
offsetScale2 = sizeY*.015
cmds.scale(offsetScale2, offsetScale2, offsetScale2, r=True)
## Delete History / Freeze Transforms ##
cmds.select(f"{PFX}offset_CTRL")
cmds.FreezeTransformations(f"{PFX}offset_CTRL")
cmds.DeleteHistory(f"{PFX}offset_CTRL")
## Connect offset CTRL to GEO ##
cmds.parent(curr_whl, f"{PFX}offset_CTRL")
cmds.parent(f"{PFX}offset_CTRL", f"{PFX}rotate_GRP")


#************************* Auto ANIMTAION *************************#
## Set Driven Key ## : DR= Auto DN = AutoRot
cmds.setAttr(f"{PFX}root_CTRL.auto", 1)
cmds.setDrivenKeyframe(f"{PFX}root_CTRL.autoRot", v= 2, cd = f"{PFX}root_CTRL.auto")
cmds.setAttr(f"{PFX}root_CTRL.auto", 0)
cmds.setDrivenKeyframe(f"{PFX}root_CTRL.autoRot", v= 0, cd = f"{PFX}root_CTRL.auto")

### EXPRESSION ###
cmds.expression(s= f"{curr_whl}.rotateX = {PFX}root_CTRL.translateZ*{PFX}root_CTRL.autoRot")


#************************** CLEAN UP **************************#
#### Lock and Hide Attributes for CTRLs ####

## Root CTRL ##
cmds.setAttr(f"{PFX}root_CTRL.visibility", lock=1, k=0,cb=0)

## Offrot CTRL ##
cmds.select(f"{PFX}offset_CTRL")
cmds.setAttr(f"{PFX}offset_CTRL.translateX", lock=1, k=0,cb=0)
cmds.setAttr(f"{PFX}offset_CTRL.translateY", lock=1, k=0,cb=0)
cmds.setAttr(f"{PFX}offset_CTRL.translateZ", lock=1, k=0,cb=0)
cmds.setAttr(f"{PFX}offset_CTRL.rotateY", lock=1, k=0,cb=0)
cmds.setAttr(f"{PFX}offset_CTRL.rotateZ", lock=1, k=0,cb=0)
cmds.setAttr(f"{PFX}offset_CTRL.scaleX", lock=1, k=0,cb=0)
cmds.setAttr(f"{PFX}offset_CTRL.scaleY", lock=1, k=0,cb=0)
cmds.setAttr(f"{PFX}offset_CTRL.scaleZ", lock=1, k=0,cb=0)
cmds.setAttr(f"{PFX}offset_CTRL.visibility", lock=1, k=0,cb=0)

### Group groups ###
cmds.group(f"{PFX}rotate_GRP", f"{PFX}CTRL_GRP", n=f"{PFX}rig")

### Color CTRLS ###
## Root ##
cmds.setAttr(f"{PFX}root_CTRLShape.overrideEnabled", 1)
cmds.setAttr(f"{PFX}root_CTRLShape.overrideColor", 17)
## Offset ##
cmds.setAttr(f"{PFX}offset_CTRLShape.overrideEnabled", 1)
cmds.setAttr(f"{PFX}offset_CTRLShape.overrideColor", 18)


##
## Create a Backdrop, Lights, and Camera for Prop rendering
##     Manually change common render settings
##     Manually adj lights and cam as needed
##
## Created by Abby O'Malley
## Follow me on instagram! @aabrbty
## Reach out for questions/comments!
##
##############################################################################
import maya.cmds as cmds
import mtoa.utils as mutils

### Globals ###
PFX = "PS_"
if cmds.objExists(f"{PFX}*"):
  cmds.delete(f"{PFX}*")
if cmds.objExists("PropSet"):
  cmds.delete("PropSet")   
### Clean up Obj ###    
def cleanup(obj):
  cmds.select(obj)
  cmds.DeleteHistory()
  cmds.FreezeTransformations()
   
### Create Backdrop ###
def Backdrop():
  BackdropName = f"{PFX}Backdrop"
  cmds.polyCube(n= f"{BackdropName}")
  cmds.move(0,-.5,0, f"{BackdropName}.scalePivot", f"{BackdropName}.rotatePivot")
  cmds.move(0,.5,0)
  cmds.scale(117,45,60)
  cmds.delete(f"{BackdropName}.f[0:1]",f"{BackdropName}.f[4:5]")
  cleanup(f"{BackdropName}")
  cmds.polyBevel3(f"{BackdropName}.e[2]", offset=18, segments = 16)
  cmds.scale(1.638, z=True)
  cmds.scale(1.455, y=True)
  cmds.move(4.522, z=True)
  cmds.polySoftEdge(f"{BackdropName}", a=180)

  ## Assign AiSS to Backdrop ##
  BackdropTXT = cmds.shadingNode('aiStandardSurface',n = f"{BackdropName}_TXT", asShader=True)
  AiSSSet =  cmds.sets( renderable=True, noSurfaceShader=True, empty=True, n= f"{BackdropName}set");
  cmds.connectAttr( BackdropTXT+".outColor", AiSSSet+".surfaceShader", force=True)
  cmds.select(f"{BackdropName}")
  cmds.hyperShade(assign = BackdropTXT)

### Create Lights ###
def Lights():    
  ## Key Light ##
  mutils.createLocator('aiAreaLight', asLight=True)
  keylight = cmds.rename('aiAreaLight1', f"{PFX}Key")
  cmds.select(f"{PFX}Key")
  cmds.move(21.517,6.564,25.303)
  cmds.rotate(-2.726,45,0)
  cmds.scale(8.902,8.902,8.902)
  cmds.setAttr(f"{PFX}KeyShape.color", 1, .888, .829)
  cmds.setAttr(f"{PFX}KeyShape.intensity", 1135.516)
  cmds.setAttr(f"{PFX}KeyShape.exposure", 1.250)
  cmds.setAttr(f"{PFX}KeyShape.aiSamples", 3)
  cmds.setAttr(f"{PFX}KeyShape.aiTranslator", "disk", type="string")
   
  """
  ## Fill Light - Directional ##
  cmds.directionalLight(rot=(-30,-45,0), intensity=2.111, rgb=(.692,.728,1))
  fill = cmds.rename("directionalLight1", f"{PFX}Fill")
  cmds.move(-17.402,22.710,20.709)
  cmds.scale(10.733,10.733,10.733)
  cmds.setAttr(f"{PFX}FillShape.aiExposure",-0.750)
  cmds.setAttr(f"{PFX}FillShape.aiSamples", 3)
  cmds.setAttr(f"{PFX}FillShape.aiCastShadows",0)
  cmds.setAttr(f"{PFX}FillShape.aiCastVolumetricShadows",0)
  """
   
  ## Fill Light - Area ##
  mutils.createLocator('aiAreaLight', asLight=True)
  area01 = cmds.rename('aiAreaLight1', f"{PFX}Fill")
  cmds.select(f"{PFX}Fill")
  cmds.move(-24.333, 15.946, 25.303)
  cmds.rotate(-12.935, -35.536, -2.759)
  cmds.scale(8.902, 8.902, 8.902)
  cmds.setAttr(f"{PFX}FillShape.color", .860, .876, 1)
  cmds.setAttr(f"{PFX}FillShape.intensity", 1135.516)    
  cmds.setAttr(f"{PFX}FillShape.exposure", 1)
  cmds.setAttr(f"{PFX}FillShape.aiSamples", 3)
  cmds.setAttr(f"{PFX}FillShape.aiTranslator", "disk", type="string")

  ## Rim Light ##
  mutils.createLocator('aiAreaLight', asLight=True)
  rimlight = cmds.rename('aiAreaLight1', f"{PFX}Rim")
  cmds.select(f"{PFX}Rim")
  cmds.move(0, 29.103,-22.901)
  cmds.rotate(-24.143,180, 0)
  cmds.scale(8.902, 8.902, 8.902)
  cmds.setAttr(f"{PFX}RimShape.color", 1, 1, 1)
  cmds.setAttr(f"{PFX}RimShape.intensity", 1431.516)
  cmds.setAttr(f"{PFX}RimShape.exposure", .894)
  cmds.setAttr(f"{PFX}RimShape.aiSamples", 2)
  cmds.setAttr(f"{PFX}RimShape.aiDiffuse", .992)
  cmds.setAttr(f"{PFX}RimShape.aiIndirect", 2.129)
  cmds.lightlink(b=True, light= f"{PFX}Rim", object= f"{PFX}Backdrop")
   
  ## Skydome Light ##
  mutils.createLocator('aiSkyDomeLight', asLight=True)
  area01 = cmds.rename('aiSkyDomeLight1', f"{PFX}Skydome")
  cmds.select(f"{PFX}Skydome")
  cmds.move(0, 0, 0)
  cmds.rotate(0, 0, 0)
  cmds.scale(1, 1, 1)
  cmds.setAttr(f"{PFX}SkydomeShape.color", .664, .687, .750)
  cmds.setAttr(f"{PFX}SkydomeShape.intensity", .2)
  cmds.setAttr(f"{PFX}SkydomeShape.exposure", 0)
  cmds.setAttr(f"{PFX}SkydomeShape.aiSamples", 2)
   
  cmds.select(f"{PFX}Skydome")
  cmds.createDisplayLayer(n=f"{PFX}Skdome_LAY")
   
### Create Camera ###
def Camera():
  cmds.camera(n=f"{PFX}rendCAM", fl=50, hfa=1.417, vfa=.945, fs=5.6, coi=21.401)
  cmds.move(0.607, 15.036, 51.905)
  cmds.rotate(-10.796, 0, 0)
  cmds.scale(2, 2, 2)
  cmds.setAttr ("defaultArnoldRenderOptions.AASamples",4)
  cmds.setAttr ("defaultArnoldRenderOptions.GIDiffuseSamples", 2)

def GroupSet():
  cmds.select(f"{PFX}*")
  cmds.group(n= "PropSet")
  cmds.select(cl=True)
   
### MAIN ###    
def Main():
  Backdrop()
  Lights()
  Camera()
  GroupSet()

Main()