'''

Description: This tool will let you set your object(s) to either templete mode, reference or normal mode

'''
#!user/bin/env python

import maya.cmds as cmds

class DisplaySet:
    def __init__(self):
        #Creat the window, but first check to see if its created, if so, then delete it first
        if cmds.window('OG_SetDisplay_v1.0', q=True, exists=True):
            cmds.deleteUI('OG_SetDisplay_v1.0', window=True)
        #Create the actual gui window
        dspSet = cmds.window( title='OG_SetDisplay_v1.0', iconName='OG_DSPSET', w=300, h=100, s=False)
        #delete the old prefs soo we can see the actual updated window
        if cmds.windowPref( posGUI, exists=True):
            cmds.windowPref( posGUI, remove=True)


        cmds.showWindow( dspSet )


def setDisplayType(normal=0, temp=0, ref=0):
    #Lets collete all the objects first and also how many objects selected
    objsel = cmds.ls(sl=True, type='transform')
    intSet = len(objsel)

    #Lets go through an error check to see if anything was selected
    if not objsel:
        raise RuntimeError('Nothing was selected. Please select your object(s) first!')
    #For this script, lets hit every object selected
    for i in objsel:
        #Lets check to see if there are any incomming connetions for the Display type
        chklst = cmds.listConnections(i + '.overrideDisplayType', scn=True, p=True)
        #if there is a connection to that, lets break that
        if chklst:
            killAttr = cmds.disconnectAttr(chklst[0], i + '.overrideDisplayType')
        #Once error check is complete, lets set the display up by turning it on first
        setEnable = cmds.setAttr(i + '.overrideEnabled', 1)
        #Now, we will let the user choose which mode to set it at
        if normal:
            setType = cmds.setAttr(i + '.overrideDisplayType', 0)
            finishTyp = 'normal mode'
        if temp:
            setType = cmds.setAttr(i + '.overrideDisplayType', 1)
            finishTyp = 'template mode'
        if ref:
            setType = cmds.setAttr(i + '.overrideDisplayType', 2)
            finishTyp = 'referenced mode'
    #Finlly, for fancy mode, lets confirm it all with which mode its on and how many objects were hit
    finMes = '%s objects has been switched to %s' % (intSet, finishTyp)
    #With the info stored, lets show it in a small confirmation GUI
    confrBox = cmds.confirmDialog( title='OG_Display_Complete!', message=finMes, button=['Close'], defaultButton='Yes', cancelButton='No', dismissString='No')

def selectType(normalType=0, tempType=0, refType=0):

    selALL = cmds.ls(g=True)
    sepObj = cmds.listRelatives(selALL, allParents=True)

    normal, ref, temp = [], [], []

    for n in sepObj:
        chkType = cmds.getAttr(n + '.overrideDisplayType')

        if chkType == 0:
            normal.append(n)
        if chkType == 1:
            temp.append(n)
        if chkType == 2:
            ref.append(n)
    if normalType:
        try:
            highlight = cmds.select(normal)
        except TypeError:
            errorWindow = cmds.confirmDialog( title='OG_SelectType_Error!', message='There are no objects in the scene in normal mode!', button=['Close'], defaultButton='Yes', cancelButton='No', dismissString='No')
    if tempType:
        try:
            highlight = cmds.select(temp)
        except TypeError:
            errorWindow = cmds.confirmDialog( title='OG_SelectType_Error!', message='There are no objects in the scene in template mode!', button=['Close'], defaultButton='Yes', cancelButton='No', dismissString='No')
    if refType:
        try:
            highlight = cmds.select(ref)
        except TypeError:
            errorWindow = cmds.confirmDialog( title='OG_SelectType_Error!', message='There are no objects in the scene in reference mode!', button=['Close'], defaultButton='Yes', cancelButton='No', dismissString='No')