from mayaUsd import lib as mayaUsdLib
from mayaUsd.lib import GetPrim
from mayaUsd.lib import proxyAccessor as pa
from pxr import Usd, UsdGeom, Sdf, Tf, Vt
import maya.cmds as cmds
cmds.loadPlugin("mayaUsdPlugin")

def createProxyAndStage():
    """
    Create in-memory stage
    """
    cmds.createNode('mayaUsdProxyShape', name='stageShape')

    shapeNode = cmds.ls(sl=True,l=True)[0]
    shapeStage = mayaUsdLib.GetPrim(shapeNode).GetStage()
    
    cmds.select( clear=True )
    cmds.connectAttr('time1.outTime','{}.time'.format(shapeNode))

    return shapeNode,shapeStage


def makeUfePath(dagPath, sdfPath=None):
    """
    Make ufe item out of dag path and sdfpath
    """
    ufePath = ufe.PathString.path('|world{},{}'.format(dagPath,sdfPath) if sdfPath != None else '|world{}'.format(dagPath))
    ufeItem = ufe.Hierarchy.createItem(ufePath)
    return ufeItem


def createAnimatedHierarchy(stage):
    """
    Create simple hierarchy in the stage:
    /ParentA
        /Sphere
        /Cube
    /ParenB
    
    Entire ParentA hierarchy will receive time samples on translate for time 1 and 100
    """
    parentA      = "/ParentA"
    childSphere  = "/ParentA/Sphere"
    #childSphere1 = "/ParentA/pSphere1"
    #childSphere2 = "/ParentA/pSphere2"
    #childSphere3 = "/ParentA/pSphere3"
    
    parentPrimA      = stage.DefinePrim(parentA, 'Xform')
    childPrimSphere  = stage.DefinePrim(childSphere, 'Sphere')
    #childPrimSphere1 = stage.DefinePrim(childSphere1, 'Sphere')
    #childPrimSphere2 = stage.DefinePrim(childSphere2, 'Sphere')   
    #childPrimSphere3 = stage.DefinePrim(childSphere3, 'Sphere')
    
    #UsdGeom.XformCommonAPI(childPrimSphere2).SetTranslate((-3,0,0))
    #UsdGeom.XformCommonAPI(childPrimSphere3).SetTranslate((3,0,0))

nodeDagPath,stage = createProxyAndStage()
createAnimatedHierarchy(stage)

# Create maya locator
cmds.spaceLocator()
locatorNodeDagPath = cmds.ls(sl=True,l=True)[0]
cmds.setAttr('{}.ty'.format(locatorNodeDagPath), 3)

# Get UFE items
ufeItemParent  = makeUfePath(nodeDagPath,'/ParentA')
#ufeItemSphere  = makeUfePath(nodeDagPath,'/ParentA/pSphere1')
ufeItemSphere  = makeUfePath(nodeDagPath,'/ParentA/Sphere')
ufeItemLocator = makeUfePath(locatorNodeDagPath)

# this breaks
translatePlugParent   = pa.GetOrCreateAccessPlug(ufeItemParent, usdAttrName='xformOp:translate')
rotatePlugParent      = pa.GetOrCreateAccessPlug(ufeItemParent, usdAttrName='xformOp:rotateXYZ')
worldMatrixPlugSphere = pa.GetOrCreateAccessPlug(ufeItemSphere, '', Sdf.ValueTypeNames.Matrix4d)

# Animate proxy shape transform
transform = cmds.listRelatives(nodeDagPath,parent=True)[0]
cmds.setKeyframe( '{}.ry'.format(transform), time=1.0, value=0 )
cmds.setKeyframe( '{}.ry'.format(transform), time=100.0, value=90 )

# Animate ParentA Translate and Rotate input plugs ( screenshot 1 )
cmds.setKeyframe( '{}.{}'.format(nodeDagPath,translatePlugParent), time=1.0, value=0.0 )
cmds.setKeyframe( '{}.{}'.format(nodeDagPath,translatePlugParent), time=100.0, value=10.0)


# Rotate ( screenshot 2 )
cmds.setKeyframe( '{}.{}'.format(nodeDagPath,rotatePlugParent), time=1.0, value=0.0 )
cmds.setKeyframe( '{}.{}'.format(nodeDagPath,rotatePlugParent), time=100.0, value=90.0)

# Parent ( screenshot 3 )
pa.ParentItems([ufeItemLocator],ufeItemSphere)