import fnmatch
import nuke


# martin@keller.io
__version__ = '1.0'

MODES = ['MergeBackIntoStream','SelectOnly', 'ContactSheet']
PATTERNDEFAULT = "*direct* !*indirect* !*default"
INIT_DONE = False

class ChannelSelect(KellerNukePlugin):

    def configurePlugin(self):
        self.menu = nuke.menu('Nodes')
        m = self.menu.addMenu('Workgroup','Toolbar_Workgroup.png')
        m.addCommand('ChannelSelect', 'import ChannelSelect; ChannelSelect.createChannelSelect()',  icon=':qrc/images/ToolbarChannel.png')
        nuke.addKnobChanged(onKnobChanged)

    def unconfigurePlugin(self):
        self.menu.removeItem('ChannelSelect')
        nuke.removeKnobChanged(onKnobChanged)
        pass


# events
def onKnobChanged():
    global INIT_DONE
    if not INIT_DONE:
        # not ready
        return

    k = nuke.thisKnob()
    # do a refresh on these knobs
    refresh_resolved = ['pattern','update']

    if k.name() in refresh_resolved:
        print('ChannelSelect: Info. Event refresh triggered.')
        # Config changed, loading new tasks
        n = nuke.thisNode()

        pattern = n['pattern'].value()
        resolved = resolvePattern(n, pattern)
        if resolved:
            resString = '\n'.join(resolved)

        if not resolved:
            n['resolved'].setValue('NO MATCH')
        else:
            n['resolved'].setValue(resString)
        pass


def createChannelSelect():
    # get selected Node (for connecting our node below)
    try:
        s = nuke.selectedNode()
    except:
        s = None
    # deselect all Nodes
    for a in nuke.allNodes():
        a.knob('selected').setValue(False)

    # create Group
    n = nuke.collapseToGroup(show = False)
    n.setName('ChannelSelect')
    # % (r * 255, g * 255, b * 255, 1), 16): change r g and b to your liking
    hexColour = int('%02x%02x%02x%02x' % (0 * 255, 0 * 255, 0 * 255, 1), 16)
    n['tile_color'].setValue(hexColour)

    if s is not None:
        n.setInput(0, s)
        n.setXYpos(s.xpos(), s.ypos() + 100)

    # tab ChannelSelect
    k = nuke.Tab_Knob('select', 'Select')
    n.addKnob(k)

    # Pattern
    k = nuke.EvalString_Knob('pattern', 'Pattern')
    n.addKnob(k)
    n.knob('pattern').setTooltip('*diffuse* *SSS*\nwith exclude (! or ^):\n*direct* !*indirect*')
    n.knob('pattern').setValue(PATTERNDEFAULT)

    # Update
    k = nuke.PyScript_Knob('update', 'Set', 'ChannelSelect.updatePattern(nuke.thisNode())')
    n.addKnob(k)
    n.knob('update').setTooltip('Creates Nodes in Group based on Search Pattern')

    # Resolved Channels
    k = nuke.Multiline_Eval_String_Knob('resolved', 'Resolved')
    n.addKnob(k)
    n.knob('resolved').setTooltip('These AOVs are found based on your search Pattern above')

    # Divide and Conquer
    k = nuke.Text_Knob('divider', '')
    n.addKnob(k)

    # Exposure
    k = nuke.Double_Knob('exposure', 'Exposure')
    n.addKnob(k)
    n.knob('exposure').setRange(-2,2)

    # Multiply
    k = nuke.Color_Knob('multiply', 'Multiply')
    n.addKnob(k)
    n.knob('multiply').setValue([1, 1, 1])

    # Gamma
    k = nuke.Double_Knob('gamma', 'Gamma')
    n.addKnob(k)
    n.knob('gamma').setValue(1)
    n.knob('gamma').setRange(0.2, 2)

    # Compensate Mult
    k = nuke.Boolean_Knob('compensate', 'Compensate Mult')
    n.addKnob(k)
    n.knob('compensate').setValue(True)
    n.knob('compensate').setFlag(nuke.STARTLINE)

    # Reset Values
    k = nuke.PyScript_Knob('reset', 'Reset', 'ChannelSelect.resetValues(nuke.thisNode())')
    n.addKnob(k)
    n.knob('reset').setTooltip('Reset Color Correct values')
    n.knob('reset').setFlag(nuke.STARTLINE)

    # Divide and Conquer
    k = nuke.Text_Knob('divider', '')
    n.addKnob(k)

    # Mode
    k = nuke.Enumeration_Knob('mode', 'Mode', MODES)
    n.addKnob(k)
    k.setTooltip('MergeBackIntoStream: Shuffles the resolved aovs, applies the grading options and merges the result back to main comp stream\n\nSelectOnly: Just shuffles the resolved aovs and adds them together.\n\nLayerContactSheet: Displays the LayerContactSheet with layer names. Note. You need to Press Set in order for this to work. An empty Pattern is possible')
    n.knob('mode').clearFlag(nuke.STARTLINE)

    # refresh again
    n.knob('pattern').setValue('*key*')
    n.knob('pattern').setValue(PATTERNDEFAULT)

    global INIT_DONE
    INIT_DONE = True
    return None


def resetValues(n):
    n['exposure'].setValue(0)
    n['multiply'].setValue([1, 1, 1])
    n['gamma'].setValue(1)
    n['compensate'].setValue(True)


def resolvePattern(n, pattern):
    if pattern == "":
        return None

    channels = n.channels()
    layers = list(set([c.split('.')[0] for c in channels]))
    layers.sort()

    matchChannels = []
    ignoreChannels = []
    for p in pattern.split(' '):
        matchChannels += (fnmatch.filter(layers, p))
        # excludes
        if p[0] == '!' or p[0] == '^':
            ignoreChannels += (fnmatch.filter(layers, p[1:]))

    matchChannels = list(set(matchChannels) - set(ignoreChannels))
    matchChannels.sort()
    return matchChannels


def updatePattern(n):

    V_OFFSET = 250

    # n = node
    pattern = n['pattern'].value()
    resolved = resolvePattern(n, pattern)

    # create nodes
    # input and output of Group
    input = nuke.toNode('Input1')
    output = nuke.toNode('Output1')

    # delete old nodes
    g = nuke.toNode(n.name())
    oldNodes = [i for i in g.nodes() if '__' in i['name'].value()]

    for node in oldNodes:
        nuke.delete(node)

    # dot input
    dot_input = nuke.nodes.Dot(xpos=input.xpos() , ypos=input.ypos()+100)
    dot_input['name'].setValue('dot_input__')
    dot_input['label'].setValue('[value name]')
    #dot_input['note_font_size'].setValue(25)
    # connect to the input1 node

    # dot input 2
    dot_input2 = nuke.nodes.Dot(xpos=dot_input.xpos()+1500 , ypos=dot_input.ypos())
    dot_input2['name'].setValue('dot_input2__')
    dot_input2['label'].setValue('[value name]')
    #dot_input2['note_font_size'].setValue(25)
    # connect to the input1 node

    # dot resolved
    dot_resolved = nuke.nodes.Dot(xpos=dot_input.xpos() , ypos=dot_input.ypos() + V_OFFSET)
    dot_resolved['name'].setValue('dot_resolved__')
    dot_resolved['label'].setValue('[value name]')
    #dot_resolved['note_font_size'].setValue(25)

    # dot resolved2
    dot_resolved2 = nuke.nodes.Dot(xpos=dot_resolved.xpos() - 100, ypos=dot_resolved.ypos())
    dot_resolved2['name'].setValue('dot_resolved2__')
    dot_resolved2['label'].setValue('[value name]')
    #dot_resolved2['note_font_size'].setValue(25)

    # dot resolved3
    dot_resolved3 = nuke.nodes.Dot(xpos=dot_resolved2.xpos(), ypos=dot_resolved.ypos() + V_OFFSET)
    dot_resolved3['name'].setValue('dot_resolved3__')
    dot_resolved3['label'].setValue('[value name]')
    #dot_resolved3['note_font_size'].setValue(25)

    # merge from
    merge_from = nuke.nodes.Merge2(xpos=dot_input.xpos()+500 , ypos=dot_input.ypos() )
    merge_from['name'].setValue('merge_from__')
    merge_from['operation'].setValue("from")

    # exposure
    exposure = nuke.nodes.EXPTool(xpos=dot_resolved.xpos() , ypos=dot_resolved.ypos() + V_OFFSET)
    exposure['name'].setValue('exposure__')

    # grade
    grade = nuke.nodes.Grade(xpos=exposure.xpos() , ypos=exposure.ypos() + V_OFFSET)
    grade['name'].setValue('grade__')

    # merge plus
    merge_plus = nuke.nodes.Merge2(xpos=grade.xpos() , ypos=grade.ypos() + V_OFFSET)
    merge_plus['operation'].setValue("plus")
    merge_plus['name'].setValue('merge_plus__')

    # dot input From
    dot_input_from = nuke.nodes.Dot(xpos=merge_from.xpos() , ypos=merge_plus.ypos())
    dot_input_from['name'].setValue('dot_input_from__')
    dot_input_from['label'].setValue('[value name]')
    #dot_input_from['note_font_size'].setValue(25)

    # contact_sheet_all
    contact_sheet_all = nuke.nodes.LayerContactSheet(xpos=merge_plus.xpos() +250 , ypos=merge_plus.ypos() + 200)
    contact_sheet_all['name'].setValue('contact_sheet__')

    # switch
    switch = nuke.nodes.Switch(xpos=merge_plus.xpos() , ypos=merge_plus.ypos() + V_OFFSET)
    switch['name'].setValue('switch__')

    # dot output
    dot_output = nuke.nodes.Dot(xpos=dot_input.xpos() , ypos=dot_input.ypos() + V_OFFSET)
    dot_output['name'].setValue('dot_output__')
    dot_output['label'].setValue('[value name]')
    #dot_output['note_font_size'].setValue(25)

    # create new nodes
    i = 1
    j = 0
    offset = 100

    merge_resolved = None
    if resolved:
        if len(resolved) > 0:
            for r in resolved:
                # shuffle
                shuffle = nuke.nodes.Shuffle(xpos=dot_input.xpos() - (len(resolved)/2)*offset + (i * offset), ypos=dot_input.ypos() + offset)
                shuffle['in'].setValue(resolved[i - 1])
                shuffle['in2'].setValue('rgba')
                shuffle['out'].setValue('rgb')
                shuffle['name'].setValue('shuffle__' + str(i))
                # connect to the dot dot_input
                shuffle.setInput(0, dot_input)

                # create plus
                if i == 1:
                    merge_resolved = nuke.nodes.Merge2(xpos=dot_input.xpos(), ypos=dot_output.ypos() + 50)
                    merge_resolved['operation'].setValue("plus")
                    merge_resolved['name'].setValue("merge_resolved__")

                # connect plus
                # skip Mask Input
                if j==2:
                    j+=1
                # connect input
                merge_resolved.setInput(j, shuffle)

                # next please
                i += 1
                j += 1

            # connect to the input1 node
            dot_resolved.setInput(0, merge_resolved)

    else:
        # shuffle
        shuffle = nuke.nodes.Shuffle(xpos=dot_input.xpos()  + (i * offset), ypos=dot_input.ypos() + offset)
        # create black in rgb and a
        shuffle['in'].setValue('none')
        shuffle['in2'].setValue('none')
        shuffle['out'].setValue('rgb')
        shuffle['name'].setValue('shuffle_empty__')

        shuffle.setInput(0, dot_input)

        # remove old input connections
        #for r in range(output.inputs()):
        #    output.setInput(r, None)

        # connect to the input1 node
        dot_resolved.setInput(0, shuffle)


    # connect the nodes
    dot_input.setInput(0, input)
    dot_input2.setInput(0, dot_input)
    dot_input_from.setInput(0, merge_from)
    dot_resolved2.setInput(0, dot_resolved)
    dot_resolved3.setInput(0, dot_resolved2)
    merge_plus.setInput(1, grade)
    merge_plus.setInput(0, dot_input_from)
    grade.setInput(0, exposure)
    exposure.setInput(0, dot_resolved)
    merge_from.setInput(0, dot_input2)
    merge_from.setInput(1, dot_resolved)
    contact_sheet_all.setInput(0, dot_input2)
    switch.setInput(0, merge_plus)
    switch.setInput(1, contact_sheet_all)
    switch.setInput(2, dot_resolved3)
    dot_output.setInput(0,switch)
    output.setInput(0, dot_output)

    # link the nodes
    exposure['mode'].setValue('Stops')
    exposure['gang'].setValue(0)
    exposure['red'].setExpression('parent.exposure')
    exposure['green'].setExpression('parent.exposure')
    exposure['blue'].setExpression('parent.exposure')
    # init
    grade['multiply'].setValue([1,1,1,1])
    grade['multiply'].setExpression('parent.multiply.r',0)
    grade['multiply'].setExpression('parent.multiply.g',1)
    grade['multiply'].setExpression('parent.multiply.b',2)
    grade['multiply'].setValue(1, 3)
    # enable Compensate Mult by default
    grade['white'].setExpression('parent.compensate == 1 ? 1/((parent.multiply.r+parent.multiply.g+parent.multiply.b)/3) : 1')
    grade['gamma'].setExpression('parent.gamma')
    contact_sheet_all['width'].setExpression('parent.width')
    contact_sheet_all['height'].setExpression('parent.height')
    contact_sheet_all['showLayerNames'].setValue(1)
    # again. Weird Update Issues
    contact_sheet_all['showLayerNames'].setExpression('parent.mode == 0 ? 0 :1')
    switch['which'].setExpression('parent.mode == 0 ? 0 :(parent.mode == 1 ? 2 : 1)')

    # arrange the nodes
    weird_dot_offset = 0
    dot_input.setXYpos(input.xpos()+weird_dot_offset, input.ypos() + 100)
    if merge_resolved:
        merge_resolved.setXYpos(input.xpos(), input.ypos() + 350)
    dot_resolved.setXYpos(input.xpos()+weird_dot_offset, input.ypos() + 500)
    dot_resolved2.setXYpos(dot_resolved.xpos()-250+weird_dot_offset, dot_resolved.ypos())
    dot_resolved3.setXYpos(dot_resolved2.xpos()+weird_dot_offset, dot_resolved2.ypos() + 500)
    exposure.setXYpos(dot_resolved.xpos(), dot_resolved.ypos()+150)
    grade.setXYpos(exposure.xpos(), exposure.ypos() + 150)
    merge_plus.setXYpos(grade.xpos(), grade.ypos() + 150)
    switch.setXYpos(merge_plus.xpos(), merge_plus.ypos() + 150)
    dot_output.setXYpos(switch.xpos()+weird_dot_offset, switch.ypos() + 150)
    output.setXYpos(dot_output.xpos(), dot_output.ypos() + 150)
    merge_from.setXYpos(dot_input2.xpos(), dot_resolved.ypos())
    dot_input_from.setXYpos(merge_from.xpos()+weird_dot_offset, merge_plus.ypos())

    n['label'].setValue('[value pattern]')
    return None