import maya.cmds as cmds
from random import uniform
  
# In the "RenderManProceduralShape->Scripts->Pre Shape Python Script" copy the following text:
  
# import rfm2.api.strings as apistr;import pp_place_archives;pp_place_archives.setDataStr(apistr.expand_string("<shape>"))
def setDataStr(shapeName):
    tnode = cmds.getAttr(shapeName + '.target')
    # If the name of the transform node of the target shape has
    # not been specified we assume the RenderManProgram has been
    # parented to the polymesh that will "receive" the spheres.
    if tnode == None or len(tnode) == 0:
        try:
            tnode = cmds.listRelatives(shapeName, parent=True)[0]
            tnode = cmds.listRelatives(tnode, parent=True)[0]
        except:
            cmds.error("Cannot find a transform node to use. Enter the name of a shape in the 'target' field.")
            return
            
    scale = str(cmds.getAttr(shapeName + '.scale'))
    use_local_space = cmds.getAttr(shapeName + '.use_local_space')
  
    probability = cmds.getAttr(shapeName + '.probability')
    jitter =  cmds.getAttr(shapeName + '.jitter') 
    archives_path =  cmds.getAttr(shapeName + '.archives_path') 
    #num_particles =  cmds.getAttr(shapeName + '.num_particles')
    
    coords = get_coordinates(tnode, use_local_space)
    jittered_coords = []
    for n in range(0, len(coords), 3):
        if uniform(0,1) < probability:
            continue
        x = coords[n+0] + uniform(-jitter, jitter)
        y = coords[n+1] + uniform(-jitter, jitter)
        z = coords[n+2] + uniform(-jitter, jitter)
        jittered_coords.extend( [x,y,z] )
    rounded = []
    # Reduce precision to 3 decimal places - less text to pass.
    for coord in jittered_coords:
        rounded.append(round(coord,3))
    coords_str = " ".join(map(str, rounded))
    
    text = archives_path + ' ' + scale + ' ' + str(len(jittered_coords)) + ' ' + ' ' + coords_str
    cmds.setAttr(shapeName + '.data', text, type='string')
    
# Rreturns a flat list of coords ie. [x0,y0,z0,x1,y1,z1....]
def get_coordinates(tnode, useLocalSpace=True):
    print("get_coordinates has been give tnode = %s\n" % tnode)
    verts = []
    listOfShapes = cmds.listRelatives(tnode, shapes=True)
    if listOfShapes == None:
        print("Sorry but %s does not have a shape!" % tnode)
        return []
    shape = listOfShapes[0]
        
    num_verts = cmds.polyEvaluate(tnode, vertex=True)
    for n in range(num_verts):
        vert_str = shape + '.vtx[%d]' % n;
        vert_pos = cmds.pointPosition(vert_str, local=useLocalSpace) 
        verts.extend(vert_pos)    
    return verts