Assignment 03 Shi xinyu & Liu

################################################################
## Assignment03 Crowd System
## Professor: Daniel da Rocha
## Students: Liu, ko-cheng & Shi, xinyu
##
## Logic
##   Create 3 kinds of objs
##     1. agent
##       a) active movement
##     2. Good Angel (helps increase the size of agent)
##       a) passive movement
##     3. Bad Evil (shrink the size of agent)
##   a) non-movable
##   When the agents hit the Good Angel, they increase the size.
##   On ther other hand, they shrink when they hit the Bad Evil
##
## Initial Set up:
##   1. Create a cyclinder as wall that holds all the objs inside 
###################################################################

import maya.cmds as cmds
import maya.mel as mel
import random
import math

def hit(hit, agent, GoodAngel, BadEvil):
 ### define what to do when it some collision is detected
 #all transformations are made only on the agent, as the hit object will perform the same queries on this agent
 #if it hit the cube, ignore
 if hit == "wall_rigid":
  return
 agentRole = cmds.getAttr(agent + ".role")
 hitRole = cmds.getAttr(hit + ".role")
 GoodAngelRole= cmds.getAttr(GoodAngel+".role")
 BadEvilRole= cmds.getAttr(BadEvil+".role")
 if hitRole == GoodAngelRole:
  if hitRole== "good":
   cmds.scale(1.1,1.1,1.1, agent, r=1)
 elif hitRole == GoodAngelRole:
  if hitRole == "bad":
   cmds.scale(.1,.1,.1, agent, r=1)
 elif hitRole == agentRole:
  #check if both are bad
  if hitRole == "bad":
   #then shrink
   cmds.scale(.9,.9,.9, agent, r=1)
  else:
   #if good, increase size
   cmds.scale(1.1,1.1,1.1, agent, r=1)
 elif hitRole != agentRole:
  if hitRole == "bad":
   #means agent role is good
   cmds.scale(.90,.90,.90, agent, r=1)
  elif hitRole == "good":
   #means agent was bad
   cmds.scale(1.1,1.1,1.1, agent, r=1)

   
#function to create shader
def createShader(color, transparency=0, type="lambert"):
 #steps to create a new material
 shadingNode1 = cmds.shadingNode(type, asShader=1)
 shadingNode2 = cmds.sets(renderable=1, noSurfaceShader = 1, empty=1, name=shadingNode1+"SG")
 cmds.connectAttr(shadingNode1+".outColor", shadingNode2 + ".surfaceShader", force=1)
 #steps to modify material attributes
 cmds.setAttr(shadingNode1 + ".color", color[0],color[1], color[2]) #color in r g b values
 cmds.setAttr(shadingNode1 + ".transparency", transparency, transparency, transparency)
 
 return shadingNode2
 
def applyShaderOnObject(obj, shader):
 cmds.sets(obj, e=1, forceElement=shader)




class Agent:
 def __init__(self, solver, role, position=(0,0,0), size=(3,3,3), color=(0,0,0), objType="cube", rigidType="active", bounciness=.6, collide=1):
  #set agent atributes
  self.bounciness = bounciness
  self.collide = collide
  self.rigidType = rigidType
  self.solver=solver
  self.role = role
  self.initialScale = size
  #create the agent and scale it
  self.object = cmds.polyCube(ax=(0,0,1))
  cmds.scale(size[0], size[1], size[2])
  #apply a color to the agent
  #set the object color
  sh = createShader(color, transparency=0, type="lambert")
  applyShaderOnObject(self.object[0], sh)
  
  #create the rigid body
  self.rigid = cmds.rigidBody(self.object,
         name=self.object[0] + "_rigid",
         p=position,
         b=self.bounciness,
         impulse=(0,0,0),
         sio=objType, #stand in object
         cc = 1, #contact count
         cp = 1, #contact position
         cn = 1, #contact name
         solver=self.solver)
         
  #add attribute to the rigid body to store the type of agent
  cmds.select(self.rigid, r=1)
  cmds.addAttr(ln="role", dt="string", keyable=1)
  cmds.setAttr(self.rigid + ".role", "good", type="string")

  if self.rigidType == "active":
   cmds.setAttr(self.rigid + ".active", 1)
  else:
   cmds.setAttr(self.rigid + ".active", 0)
   
  #apply the expression
  self.applyExpression()
 
 def applyExpression(self):
  "Function to apply the expression to the agent"
  #first disconnect the rotation attributes because we want to controll it via expressions
  cmds.disconnectAttr (self.rigid + "ry.output", self.object[0] + ".rotateY")
  cmds.disconnectAttr (self.rigid + "rx.output", self.object[0] +  ".rotateX")
  cmds.disconnectAttr (self.rigid + "rz.output", self.object[0] + ".rotateZ")
  # Create expression for VehicleF with a random multiplier added to the expression.
  randX = random.uniform(-3,3)
  randY = random.uniform(3,-3)
  ####COMPOSE THE EXPRESSION STRING
  #this is how it should look like (in MEL)
  '''
  // Set wander motion of vehicle and add random multipliers.
  pCube1_rigid_0.impulseX = sin(time * -2.808801);
  pCube1_rigid_0.impulseY = (noise(time) * -1.041035);

  // Set orientation of Vehicle according to the velocity direction.
  float $fVel[];
  $fVel  = `getAttr pCube1_rigid_0.velocity`;

  pCube1.rotateX = 0;
  pCube1.rotateZ = atan2d( $fVel[0], $fVel[1] );
  pCube1.rotateY = 0;
  
  //checking bumps on other agents
  string $lastHit;
  if (pCube1_rigid_0.contactCount > 0){
   int $contactCount = `getAttr pCube1_rigid_0.contactCount`;
   $lastHit = `getAttr pCube1_rigid_0.contactName[0]`;
   string $rigid = "pCube1_rigid_0";
   string $rigid01 = "pSphere1_rigid_0";
   python("hit('"+$lastHit+"', '"+ $rigid +"','"+rigid01+"')");
  };//endif
  '''
  #now put the text above in a python string
  #the first lines define how the agent will wander
  expString = "// Set wander motion of vehicle and add random multipliers.\n"
  expString += "%s.impulseX" % self.rigid
  expString += " = sin(time * %f);\n" % randX
  expString += "%s.impulseY" % self.rigid
  expString += " = (noise(time) * %f);\n\n" % randY
  #the second part, how it rotates according to the direction it is going (velocity vector)
  expString += "// Set orientation of Vehicle according to the velocity direction.\n"
  expString += "float $fVel[];\n"
  expString += "$fVel  = `getAttr %s.velocity`;\n\n" % (self.rigid)
  expString += "%s.rotateX = 0;\n" % self.object[0]
  expString += "%s.rotateZ = atan2d( $fVel[0], $fVel[1] );\n" % (self.object[0])
  expString += "%s.rotateY = 0;\n\n" % self.object[0]
  #the third part, checking bumps on other agents
  expString += "//checking bumps on other agents\n"
  expString += "string $lastHit;\n"
  expString += "if (%s.contactCount > 0){\n" % self.rigid
  expString += "  int $contactCount = `getAttr %s.contactCount`;\n" % self.rigid
  expString += "  $lastHit = `getAttr %s.contactName[0]`;\n" %self.rigid
  expString += '  string $rigid = "%s";\n' % self.rigid
  expString += '  string $rigid01 = "%s";\n' % self.rigid
  expString += '  string $rigid02 = "%s";\n' % self.rigid
  expString += '  python("hit(\'"+$lastHit+"\', \'"+ $rigid +"\',\'"+ $rigid01 +"\',\'"+ $rigid02 +"\')");\n'
  expString += "};//endif\n"
  
  
  
  self.expString = expString
  cmds.expression(s=expString)  
  
class Goodobj:
 def __init__(self, solver, role, position=(0,0,0), size=(3,3,3), color=(0,0,0), objType="cube", rigidType="active", bounciness=.6, collide=1):
  #set Goodobj atributes
  self.bounciness = bounciness
  self.collide = collide
  self.rigidType = rigidType
  self.solver=solver
  self.role = role
  self.initialScale = size
  #create the agent and scale it
  self.object = cmds.polySphere(ax=(0,0,1))
  cmds.scale(size[0], size[1], size[2])
  #apply a color to the agent
  #set the object color
  sh = createShader(color, transparency=0, type="lambert")
  applyShaderOnObject(self.object[0], sh)
  
  #create the rigid body
  self.rigid01 = cmds.rigidBody(self.object,
         name=self.object[0] + "_rigid",
         p=position,
         b=self.bounciness,
         impulse=(0,0,0),
         sio=objType, #stand in object
         cc = 1, #contact count
         cp = 1, #contact position
         cn = 1, #contact name
         solver=self.solver)
         
  #add attribute to the rigid body to store the type of agent
  cmds.select(self.rigid01, r=1)
  cmds.addAttr(ln="role", dt="string", keyable=1)
  cmds.setAttr(self.rigid01 + ".role", "good", type="string")

  if self.rigidType == "active":
   cmds.setAttr(self.rigid01 + ".active", 1)
  else:
   cmds.setAttr(self.rigid01 + ".active", 0)

   
class Badobj:
 def __init__(self, solver, role, position=(0,0,0), size=(3,3,3), color=(0,0,0), objType="cube", rigidType="active", bounciness=.6, collide=1):
  #set Goodobj atributes
  self.bounciness = bounciness
  self.collide = collide
  self.rigidType = rigidType
  self.solver=solver
  self.role = role
  self.initialScale = size
  #create the agent and scale it
  self.object = cmds.polyCone(ax=(0,0,1))
  cmds.scale(size[0], size[1], size[2])
  #apply a color to the agent
  #set the object color
  sh = createShader(color, transparency=0, type="lambert")
  applyShaderOnObject(self.object[0], sh)
  
  #create the rigid body
  self.rigid02 = cmds.rigidBody(self.object,
         name=self.object[0] + "_rigid",
         p=position,
         b=self.bounciness,
         impulse=(0,0,0),
         sio=objType, #stand in object
         cc = 1, #contact count
         cp = 1, #contact position
         cn = 1, #contact name
         solver=self.solver)
         
  #add attribute to the rigid body to store the type of agent
  cmds.select(self.rigid02, r=1)
  cmds.addAttr(ln="role", dt="string", keyable=1)
  cmds.setAttr(self.rigid02 + ".role", "good", type="string")

  if self.rigidType == "active":
   cmds.setAttr(self.rigid02 + ".active", 0)
  else:
   cmds.setAttr(self.rigid02 + ".active", 0)

##############################################
class Crowd:
 def __init__(self, numAgents=20, numGoodobjs=20, numBadobjs=20,minPosition=-70, maxPosition=70, walls=0):
  #set the playback options to increase the time range
  cmds.playbackOptions(min= 1, max=5000)
  self.numAgents = numAgents
  self.numGoodobjs= numGoodobjs
  self.numBadobjs= numBadobjs
  self.minPosition = minPosition
  self.maxPosition = maxPosition
  #get any selected objects
  self.wallObjs = cmds.ls(sl=1)
  #create the rigid solver
  self.solver = self.createSolver()
  #create the agents
  self.createAgents()
  #create the goodobjs
  self.createGoodobj()
  #create the badobjs
  self.createBadobj()
  #create the walls
  if walls: self.makeWalls()
  
 def createSolver(self):
  solver = cmds.rigidSolver( create=1,
            current=1,
            name="crowdSolver",
            velocityVectorScale=0.5,
            displayVelocity=1,
            sc=1, #showcollision
            ctd=1 #contact data
            )
  cmds.setAttr(solver + ".allowDisconnection", 1)
  return solver
  
 def createAgents(self):
  type = ["good", (3,3,3), (1,1,0)] #name, scale, color
  self.goodAgents = []
  for i in range(self.numAgents):
   #get the agent type randomly
   #get random x and y
   x = random.uniform(self.minPosition,self.maxPosition)
   y = random.uniform(self.minPosition,self.maxPosition)
   #create the agents
   a = Agent(self.solver, type[0], color=type[2], size=type[1], position=(x,y,0))

   self.goodAgents.append(a)
 def createGoodobj(self):
   type = ["good", (3,3,3), (1,0,0)] #name, scale, color
   self.goodobjs = []
   for i in range(self.numGoodobjs):
    #get the agent type randomly
    #get random x and y
    x = random.uniform(self.minPosition,self.maxPosition)
    y = random.uniform(self.minPosition,self.maxPosition)
    #create the agents
    a = Goodobj(self.solver, type[0], color=type[2], size=type[1], position=(x,y,0))
 
    self.goodobjs.append(a)

 def createBadobj(self):
   type = ["bad", (2,2,2), (0,0,0)] #name, scale, color
   self.Badobjs = []
   for i in range(self.numBadobjs):
    #get the agent type randomly
    #get random x and y
    x = random.uniform(self.minPosition,self.maxPosition)
    y = random.uniform(self.minPosition,self.maxPosition)
    #create the agents
    a = Badobj(self.solver, type[0], color=type[2], size=type[1], position=(x,y,0))
    self.Badobjs.append(a)


 def makeWalls(self):
  #get selected object,
  #which should be already with its normals facing to the right directions,
  #and convert it to passive rigid bodies
  if len(self.wallObjs) == 0:
   return "No wall objects were selected."
  else:
   for w in self.wallObjs:
    self.wallRigid = cmds.rigidBody(w, passive=1, name= "wall_rigid", bounciness=.8)
########################################
c=Crowd(numAgents=2,numGoodobjs=10,numBadobjs=20,walls=1)