"""
Contains classes and functions for storing poses (translate, rotate, and scale)
"""
#--------UTILITIES------#
import maya.cmds as MC
g_uiInstance = None
g_poseInstance = None

def selList():
    """
    Return a list of the selected objects
    """
    return MC.ls(sl=True)

def initUI():
    """
    Creates the Pose Store UI
    """
    #global var
    global g_uiInstance
    
    if not g_uiInstance:
        g_uiInstance = PoseStoreUI()
    else:
        g_uiInstance.show()

#--------CLASSES--------#

class PoseStoreUI:
    """
    Creates a UI for pose storing
     Attributes:
            self.name - window name
            self.poseNameText - name of the "pose name" textbox
            self.restoreCheckboxGrp - name of the pose restore checkboxes
            self.poseList - name of the stored poses list
            self.fileOutText - name of the file output textfield
            self.fileInText - name of the file input textfield
            self.tabLayout - name of the main tab layout
            self.animationLayout - name of the main aniamtion tools layout
            self.riggingLayout - name of the main rigging tools layout
    Functions:
            # functions to glue UI class to backend pose classes 
            self.deletePose
            self.deleteAllPoses
            self.storePose
            self.restorePose
            self.loadFromFile
            self.saveToFile
            self.refreshPoselist - clears and reloads the pose list window to reflect what's currently stored in g_poseInstance.data
            # other class functions
            self.show - build the window
    """

    def __init__(self):
        """
        Builds the MEL UI
        Attributes:
            self.name - window name
            self.poseNameText - name of the "pose name" textbox
            self.restoreCheckboxGrp - name of the pose restore checkboxes
            self.psoeList - name of the stored poses list
            self.fileOutText - name of the file output textfield
            self.fileInText - name of the file input textfield
            self.tabLayout - name of the main tab layout
            self.animationLayout - name of the main aniamtion tools layout
            self.riggingLayout - name of the main rigging tools layout
        
        """
        
        ## Create a PoseManager object if one doesn't exist ##
        
        global g_poseInstance
        if not g_poseInstance:
            g_poseInstance = PoseManager()
        
        ## Build the UI ##
        #Create Window and top layouts
        self.name = MC.window(title ="jpTools", w=350, h=500)
        self.tabLayout = MC.tabLayout()
        self.animationLayout = MC.columnLayout("Pose Guy", w=350, columnAttach=("both",0), adj=True )
        
        #Store Pose Section
        MC.separator( height=10 )
        MC.text(label="Stored Poses:", fn="boldLabelFont", align = "left", h=25)
        #RL houses columns for TSL and Buttons on right
        MC.rowLayout(nc = 3, cw3 = (215, 10, 125), adj=1)
        if True:    #block for whitespace organization
            self.poseList = MC.textScrollList(w=215, h=150, allowMultiSelection=True)
            MC.text(" ")
            #CL for buttons next to TSL
            MC.columnLayout()
            if True:    #block for whitespace organization
                MC.button (label = "Delete Selected Pose", c = self.deletePose, w=125)
                MC.button (label = "Delete All Poses", c = self.deleteAllPoses, w=125)
                MC.setParent("..")# end CL - buttons
            MC.setParent("..")# end RL - TSL and Buttons
        #RL to format "store bar"
        MC.rowLayout(nc=4, cw4 = (40 ,170, 10, 125), adj=2)
        if True:    #block for whitespace organization
            MC.text("Name:")
            self.poseNameText = MC.textField()
            MC.text(" ")
            MC.button(label = "Store Pose", c=self.storePose, w=125)
            MC.setParent("..")#end RL - store bar
        
        #Restore Pose Section
        MC.separator( height=10 )
        MC.text(label="Restore Pose:", fn="boldLabelFont", align = "left", h=25)
        MC.text("  Options:", align="left")
        MC.rowLayout(nc=2, cw2=(25,325), adj=2)
        if True:    #block for whitespace organization
            MC.text(" ")
            MC.columnLayout(adj=True)
            if True:    #block for whitespace organization
                self.restoreCheckboxGrp = MC.checkBoxGrp(numberOfCheckBoxes = 3, cw3= (80,80,80), labelArray3=['Translate', 'Rotate', 'Scale'],
                valueArray3 = (True, True, True), cl3 = ('left', 'left', 'left'))
                MC.setParent("..")
            MC.setParent("..")# BACK to main CL
        MC.rowLayout(nc=3, cw3 = (215, 10, 125), adj=1)
        if True:    #block for whitespace organization
            MC.text("Select pose to restore:")
            MC.text(" ")
            MC.button(label = "Restore Pose", c = self.restorePose, w=125)
            MC.setParent("..")
            
        #File IO Section
        MC.separator( height=10 )
        MC.text(label="Save/Load from File:", fn="boldLabelFont", align = "left", h=25)
        MC.text("  Quicksave File Path", align = "left")
        #Row Layout to organize file load
        MC.rowLayout(nc = 3, cw3 = (215, 10, 125), adj = 1)
        if True:    #block for whitespace organization
            self.fileInText = MC.textField()
            MC.text(" ")
            MC.button(label = "Load from File", c=self.loadFromFile, w=125)
            MC.setParent("..")#end RL - Loading
        #Row Layout to organize file save
        MC.rowLayout(nc = 3, cw3 = (215, 10, 125), adj = 1)
        if True:
            self.fileOutText = MC.textField()
            MC.text(" ")
            MC.button(label = "Save to File", c = self.saveToFile,  w=125)
            MC.setParent("..")#end RL - Saving
        
        #Show it, baby
        MC.showWindow(self.name)
        ##  END UI Build  ##
        
    
    def deletePose (self, *args):
        """
        Delete selected TSL items
        """
        global g_poseInstance
        
        objs = MC.textScrollList(self.poseList, q=True, si=True)
        try:
            for pose in objs:
                g_poseInstance.deletePose(pose)
            self.refreshPoselist()
        except:
            print "oops"

    def deleteAllPoses (self, *args):
        """
        Clear TSL and delete poses from dictionary.
        """
        objs = MC.textScrollList(self.poseList, q=True, ai=True)
        try:
            MC.textScrollList(self.poseList, e=True, ra=True)
            g_poseInstance.__init__()
        except:
            print "Something went wrong!"
            
    def storePose (self, *args):
        """
        Query self.poseNameText for the new pose name and call the PoseManager.storePose method
        Update self.poseList
        """
        global g_poseInstance
        poseName = MC.textField(self.poseNameText, q=True, text=True)
        g_poseInstance.storePose(poseName)
        self.refreshPoselist()
        
    def restorePose (self, *args):
        """
        Query self.poseList for the selected pose(s)
        Query self.restoreCheckboxGrp for restore options
        Call the PoseManager.restorePose method
        """
        global g_poseInstance
        t=True
        r=True
        s=True
        if not MC.checkBoxGrp(self.restoreCheckboxGrp, q=True, v1=True):
            t=False
        if not MC.checkBoxGrp(self.restoreCheckboxGrp, q=True, v2=True):
            r=False
        if not MC.checkBoxGrp(self.restoreCheckboxGrp, q=True, v3=True):
            s=False
        objs = MC.textScrollList(self.poseList, q=True, si=True)
        if not objs:
            print "Must select pose(s) to restore"
            return
       
        for obj in objs:
            try:
                g_poseInstance.restorePose(obj, t, r, s)
            except TypeError, e:
                print "An object doesn't exist: %s" %e
                print "Continuing to restore other objs in pose..."
            
    def loadFromFile (self, *args):
        """
        Query self.fileInText for a filename 
        """
        global g_poseInstance
        path = MC.textField(self.fileInText, q=True, text=True)
        #try:
        g_poseInstance.loadFromFile(path)
        self.refreshPoselist()
        #except:
        #    print "Oops!  Check the file path"
            
    def saveToFile (self, *args):
        """
        Query self.fileOutText
        """
        global g_poseInstance
        path = MC.textField(self.fileOutText, q=True, text=True)
        try:
            g_poseInstance.saveToFile(path)
            self.refreshPoselist()
        except:
            print "Oops!  Check the file path"
    def refreshPoselist (self):
        """
        Refresh the TSL to display all stored poses
        """
        global g_poseInstance
        MC.textScrollList(self.poseList, e=True, ra=True)
        MC.textScrollList(self.poseList, e=True, append=g_poseInstance.data.keys())
        """
        Reload the UI window
        """
        if not MC.window(self.name, ex=True):
            self.__init__()
    def show(self):
        """
        Reload the UI window
        """
        if not MC.window(self.name, ex=True):
            self.__init__()
    
class Xform:
    """
    Class for storing a pose for an object and operating on that pose

    self.translate - local space translation
    self.rotate - local space rotation
    self.scale - local scale

    """

    def __init__(self, object, query=True):
        self.name = object
        if query:
            self.translate = MC.xform(self.name, q=True, t=True)
            self.rotate = MC.xform(self.name, q=True, ro=True)
            self.scale = MC.xform(self.name, q=True, relative = True, s=True)
    
    def apply(self, t=True, r=True, s=True):
        if t:
            MC.xform(self.name, t=self.translate)
        if r:
            MC.xform(self.name, ro=self.rotate)
        if s:
            MC.xform(self.name, s=self.scale)
    
    def setVals(self, t=None, r=None, s=None):
        if t:
            self.translate = [float(item) for item in t] 
        if r:
            self.rotate = [float(item) for item in r]
        if s:
            self.scale = [float(item) for item in s]
            
class PoseManager:
    """
    Stores and restores poses
    """
    
    def __init__(self):
        """
        Set up members
        """
        self.data = {}
        
    def storePose(self, poseName):
        """
        Store all objects in the selection as a pose
        """
        result = []
        sel = MC.ls(sl = True)

        #check for none
        if not sel:
            return result

        #create list of Pose instances
        for obj in sel:
            result.append(Xform(obj))

        self.data[poseName] = result

    def restorePose(self, poseName, t=True, r=True, s=True):
        """
        Restore all objects in a pose set to the pose
        """

        for poseObj in self.data[poseName]:
            poseObj.apply(t, r, s)
    def deletePose(self, poseName):
        """
        Remove a pose from self.data
        """
        try:
            self.data.pop(poseName)
        except KeyError, e:
            print "%s is not a stored pose" % poseName
    def size(self):
        """
        Return number of poses
        """
        return self.data.__len__()
    
    def objsInPose(self, pose):
        """
        Return the number of poses objects in a named pose
        """
        return len(self.data[pose])
        
    def poseNames(self):
        """
        Return names of poses
        """
        return self.data.keys()

    def clear(self):
        """
        Deletes all poses
        """
        self.data.clear()
        
    def saveToFile(self, filePath=None):
        """
        * Advanced *
        Store to file
        """
        if not filePath:
                filePath = MC.fileDialog(m=1)
        if filePath == "":
            return
        #open file to write
        f = open(filePath, "w")
        for poseName, poseList in self.data.items():
            for poseObj in poseList:
                tList = " ".join([str(val) for val in poseObj.translate])
                rList = " ".join([str(val) for val in poseObj.rotate])
                sList = " ".join([str(val) for val in poseObj.scale])
                f.write("%s.%s = %s %s %s\n" % (poseName, poseObj.name, tList, rList, sList))
        f.close()

    def loadFromFile(self, filePath=None):
        """
        * Advanced *
        Load from file
        """
        if not filePath:
            filePath = MC.fileDialog(m=0)
        if filePath == "":
            return
        f = open(filePath, "r")
        for line in f:
            #split up each line so that nameToks is a 2-item list with [posename, objname]
            #and transToks is a 9-item list of translations
            temp = line.split("=")
            nameToks = temp[0].split(".")
            transToks = temp[1].split()
            poseName = nameToks[0]
            objName = nameToks[1]

            #if the pose name on the current line isn't in the data attr, add it
            if not self.checkIfMember(poseName):
                self.data[poseName] =[]
            
            #create the new Pose object and add it to the current "pose" list
            self.data[poseName].append(Xform(objName, query=False))
            #get a reference to the newly created pose in order to set it's attributes
            currentPose = self.data[poseName][(self.objsInPose(poseName))-1]
            currentPose.setVals(t=[transToks[0], transToks[1], transToks[2]], r=[transToks[3], transToks[4], transToks[5]], s=[transToks[6], transToks[7], transToks[8]])
        f.close()
        print ("Poses Loaded!")
        
    def checkIfMember(self, poseName):
        """
        Check if the pose name is already used
        """
        for pose in self.data.keys():
            if poseName == pose:
                return True
        return False