import maya.cmds as cmds
import maya.mel as mel
import glTools.tools.controlBuilder
import glTools.utils.base
import glTools.utils.channelState


class PivotObject(object):
    def __init__(self):
        """
        Initializer for PivotObject class.
        """
        # Pivot Control Names
        self.pivotControlName = 'PivotCONTROL'
        self.pivotControlGrp = 'PivotCONTROL_grp'
        # Pivot Locator Names
        self.pivotLocatorName = 'PivotLOCATOR'
        self.pivotLocatorBuf = 'PivotLOCATOR_buf'
        self.pivotLocatorGrp = 'PivotLOCATOR_grp'
        # Channel Arrays
        self.translateChannels = ['tx', 'ty', 'tz']
        self.rotateChannels = ['rx', 'ry', 'rz']
        self.transforcmdshannels = ['tx', 'ty', 'tz', 'rx', 'ry', 'rz']

    def create(self, ctrlList=[], pivotObj='', orientPivotControl=False, initializeCleanUpJob=True,
               constrainTranslate=True, constrainRotate=True, maintainPivotInfo=False):
        """
        This function will setup a pivotObject control, which will allow you to manipulate
        a group of controls from a common pivot position.
        @param ctrlList: List of controls that will be controled by pivotObject
        @type ctrlList: list
        @param pivotObj: Object whose pivot position will be used for the pivotObject position
        @type pivotObj: str
        @param orientPivotControl: Orient pivotObject to the specified control
        @type orientPivotControl: bool
        @param initializeCleanUpJob: Initialize scriptJob that will cleanup pivotObject control and constraints
        @type initializeCleanUpJob: bool
        @param constrainTranslate: Affect translate channels of controls driven by pivotObject
        @type constrainTranslate: bool
        @param constrainRotate: Affect rotate channels of controls driven by pivotObject
        @type constrainRotate: bool
        @param maintainPivotInfo: Maintain pivotObject information by generating a locator with relevant attribute values
        @type maintainPivotInfo: bool
        """
        # Check for existing Pivot Object
        if cmds.objExists(self.pivotControlName):
            self.kill()

        # Check control list
        for ctrl in ctrlList:
            if not cmds.objExists(ctrl):
                raise UserInputError('Control object ' + ctrl + ' does not exists!')
        # Check pivot object
        if not cmds.objExists(pivotObj):
            raise UserInputError('Pivot object ' + pivotObj + ' does not exists!')

        # Create pivot control
        pivCtrl = glTools.tools.controlBuilder.create('sphere', self.pivotControlName)['control']
        pivGrp = glTools.utils.base.group(pivCtrl, 0)
        # Position pivot control
        cmds.delete(cmds.pointConstraint(pivotObj, pivGrp, mo=False))
        # Orient pivot control
        if orientPivotControl:
            cmds.delete(cmds.orientConstraint(pivotObj, pivGrp, mo=False))

        # Set channelStates for pivot control and group
        glTools.utils.channelState.ChannelState().setFlags([0, 0, 0, 0, 0, 0, 0, 0, 0, 1], [pivCtrl])
        glTools.utils.channelState.ChannelState().setFlags([2, 2, 2, 2, 2, 2, 2, 2, 2, 1], [pivGrp])
        glTools.utils.channelState.ChannelState().set(1, [pivCtrl, pivGrp])

        # Create constraints
        for ctrl in ctrlList:
            # Check translate channels
            skipTranslate = []
            if constrainTranslate:
                for ch in self.translateChannels:
                    if not cmds.getAttr(ctrl + '.' + ch, se=True):
                        skipTranslate.append(ch[1])
            else:
                skipTranslate = ['x', 'y', 'z']
            # Check rotate channels
            skipRotate = []
            if constrainRotate:
                for ch in self.rotateChannels:
                    if not cmds.getAttr(ctrl + '.' + ch, se=True):
                        skipRotate.append(ch[1])
            else:
                skipRotate = ['x', 'y', 'z']
            # Generate parentConstraint
            cmds.parentConstraint(pivCtrl, ctrl, skipTranslate=skipTranslate, skipRotate=skipRotate, mo=True)

        # Select pivot control
        cmds.select(self.pivotControlName)

        # Setup kill scriptJob
        if initializeCleanUpJob:
            self.cleanUpJob(maintainPivotInfo)

    def cleanUp(self, maintainPivotInfo=False):
        """
        This function will remove and cleanup any nodes, controls and constraints generated by PivotObject.create()
        @param maintainPivotInfo: Maintain pivotObject information by generating a locator with relevant attribute values
        @type maintainPivotInfo: bool
        """
        # Check pivot control exists
        if cmds.objExists(self.pivotControlName):

            # Check maintain info
            if maintainPivotInfo:
                self.pivotLocator()

            # Get constraint/control lists
            controlList = []
            constraintList = []
            [constraintList.append(conn) for conn in
             cmds.listConnections('PivotCONTROL', s=False, d=True, type='constraint') if not constraintList.count(conn)]
            for constraint in constraintList:
                constraintSlave = cmds.listConnections(constraint + '.constraintRotatePivot', s=1, d=0)
                if constraintSlave:
                    controlList.append(constraintSlave[0])
            # Record control transform values
            controlTransforms = {}
            for ctrl in controlList:
                controlTransforms[ctrl] = []
                controlTransforms[ctrl].extend(cmds.getAttr(ctrl + '.t')[0])
                controlTransforms[ctrl].extend(cmds.getAttr(ctrl + '.r')[0])

            # Delete constraints and control
            cmds.delete(cmds.listConnections('PivotCONTROL', s=False, d=True, type='constraint'))
            cmds.delete(self.pivotControlName)

            # Reapply control transform values
            transforcmdshannels = self.translateChannels + self.rotateChannels
            for ctrl in controlTransforms.iterkeys():
                for i in range(len(transforcmdshannels)):
                    try:
                        cmds.setAttr(ctrl + '.' + transforcmdshannels[i], controlTransforms[ctrl][i])
                    except:
                        pass

        # Delete control group
        if cmds.objExists(self.pivotControlGrp):
            cmds.delete(self.pivotControlGrp)

    def cleanUpJob(self, maintainPivotInfo=False):
        """
        Initializes a scriptJob that will run a function that will remove and cleanup any nodes, controls and constraints
        generated by PivotObject.create()
        @param maintainPivotInfo: Maintain pivotObject information by generating a locator with relevant attribute values
        @type maintainPivotInfo: bool
        """
        mel.eval('scriptJob -e "SelectionChanged" "python(\\\"PivotObject().cleanUp(' + str(
            maintainPivotInfo) + ')\\\")" -runOnce true')

    def pivotLocator(self):
        """
        Generates a pivotObject locator that will contain all necessary information (via attribute values)
        to recreate a pivotObject.
        """
        # Check pivot locator group
        if not cmds.objExists(self.pivotLocatorGrp):
            cmds.group(n=self.pivotLocatorGrp, em=True)
            glTools.utils.channelState.ChannelState().setFlags([2, 2, 2, 2, 2, 2, 2, 2, 2, 1], [self.pivotLocatorGrp])
            glTools.utils.channelState.ChannelState().set(1, [self.pivotLocatorGrp])

        # Create pivot locator
        # pivLoc = glTools.tools.controlBuilder.create('locator',self.pivotLocatorName)
        pivLoc = cmds.group(n=self.pivotLocatorName, em=True)
        # Create pivot locator annotation
        pivLocAnnot = cmds.createNode('annotationShape', p=pivLoc)
        cmds.setAttr(pivLocAnnot + '.text', self.pivotLocatorName, type='string')
        cmds.setAttr(pivLocAnnot + '.displayArrow', 0)
        cmds.delete(cmds.parentConstraint(self.pivotControlName, pivLoc))
        # Create pivot locator group
        pivLocBuf = cmds.group(n=self.pivotLocatorBuf, em=True)
        cmds.delete(cmds.parentConstraint(self.pivotControlGrp, pivLocBuf))
        # Parent locator to group
        cmds.parent(pivLoc, pivLocBuf)
        cmds.parent(pivLocBuf, self.pivotLocatorGrp)
        # Set pivot locator channel states
        glTools.utils.channelState.ChannelState().setFlags([2, 2, 2, 2, 2, 2, 2, 2, 2, 1], [pivLoc, pivLocBuf])
        glTools.utils.channelState.ChannelState().set(1, [pivLoc, pivLocBuf])

        # Record Control List
        # ---------------------
        # Get affected control list
        controlList = []
        constraintList = []
        constraintConnList = cmds.listConnections(self.pivotControlName, s=False, d=True, type='constraint')
        [constraintList.append(conn) for conn in constraintConnList if not constraintList.count(conn)]
        constraintList.sort()
        for constraint in constraintList:
            constraintSlave = cmds.listConnections(constraint + '.constraintRotatePivot', s=1, d=0)
            if constraintSlave:
                controlList.append(constraintSlave[0])
        # Add control list values to locator string array attribute
        if len(controlList):
            cmds.addAttr(pivLoc, ln='controlList', m=True, dt='string')
            for ctrl in range(len(controlList)):
                cmds.setAttr(pivLoc + '.controlList[' + str(ctrl) + ']', controlList[ctrl], type='string')
        cmds.setAttr(pivLoc + '.controlList', l=True)
        # Record Current Frame Number
        # ----------------------------
        cmds.addAttr(pivLoc, ln='frame', at='float')
        cmds.setAttr(pivLoc + '.frame', cmds.currentTime(q=True))
        cmds.setAttr(pivLoc + '.frame', l=True)
