Untitled

def AG_RemoveAll():
    #disable undo recording
    undo = nuke.Undo()
    undo.disable()

    if nuke.ask('are you sure you want to remove all areas?'):
        node = nuke.thisNode()
        knobs = ['Start', 'End', 'Area', 'Grow', 'Color', 'Remove', '_panelDropped']

        for knob in node.allKnobs():
            kn = knob.name()
            if kn.startswith('A') and kn.endswith(tuple(knobs)):
                node.removeKnob(knob)

        #delete radial nodes
        with node:
            for n in nuke.allNodes():
                if n.Class() == 'Radial':
                    nuke.delete(n)

        node['count'].setValue(0)

    #re-anable undo recording
    undo.enable()

def AG_RemoveArea():
    #disable undo recording
    undo = nuke.Undo()
    undo.disable()

    node = nuke.thisNode()
    name = nuke.thisKnob().name().replace('Remove', '')

    #delete radial node
    with node:
        nuke.delete(nuke.toNode(name))

    knobs = ['Start', 'End', 'Area', 'Grow', 'Color', 'Remove', 'Color_panelDropped']

    for i in knobs:
        try:
            node.removeKnob(node.knobs()[name + i])
        except:
            pass
    #re-anable undo recording
    undo.enable()

def AG_AddArea():
    #disable undo recording
    undo = nuke.Undo()
    undo.disable()

    node = nuke.thisNode()

    colorMode = node['colormode'].getValue()

    count = int(node['count'].value() + 1)
    node['count'].setValue(count)

    #Open Group Knob
    gk = nuke.Tab_Knob('A' + str(count) + 'Start', 'area ' + str(count), 1)
    node.addKnob(gk)

    #Add Area Knobs
    removeKnob = nuke.PyScript_Knob('A' + str(count) + 'Remove', 'remove area', 'AG_RemoveArea()')
    node.addKnob(removeKnob)

    areaKnob = nuke.BBox_Knob('A' + str(count) + 'Area', 'area')
    areaKnob.setValue((-50,-50,50,50))
    node.addKnob(areaKnob)

    growKnob = nuke.Double_Knob('A' + str(count) + 'Grow', 'grow')
    node.addKnob(growKnob)

    colorKnob = nuke.Color_Knob('A' + str(count) + 'Color', 'color')
    if colorMode == 0:
        colorKnob.setFlag(0x00040000)
    node.addKnob(colorKnob)

    #Close Group Knob
    gk = nuke.Tab_Knob('A' + str(count) + 'End', ' ', -1)
    node.addKnob(gk)

    #stuff inside main node
    with node:
        merge = nuke.toNode('RadialMerge')
        #add radial node
        radialName = 'A' + str(count)
        radial = nuke.nodes.Radial(name = radialName, cliptype = 0)

        radial['area'].setExpression('[python {AG_GrowArea("%s")[0]}]' % (radialName), 0)
        radial['area'].setExpression('[python {AG_GrowArea("%s")[1]}]' % (radialName), 1)
        radial['area'].setExpression('[python {AG_GrowArea("%s")[2]}]' % (radialName), 2)
        radial['area'].setExpression('[python {AG_GrowArea("%s")[3]}]' % (radialName), 3)

        radial['color'].setValue([0,0,0,1])
        radial['color'].setExpression('colormode ? %sColor : [python AG_SampleSource("%s", 0)]' % (radialName,radialName), 0)
        radial['color'].setExpression('colormode ? %sColor : [python AG_SampleSource("%s", 1)]' % (radialName,radialName), 1)
        radial['color'].setExpression('colormode ? %sColor : [python AG_SampleSource("%s", 2)]' % (radialName,radialName), 2)

        #connect to merge node
        inputs = merge.inputs()
        if inputs == 2:
            inputs += 1

        merge.setInput(inputs,radial)

    #re-anable undo recording
    undo.enable()

def AG_SampleSource(areaName, indx):
    try:
        inputNode = nuke.thisGroup().input(0)
        area = list(nuke.thisGroup()[areaName + 'Area'].value())

        xpos = (area[0] + area[2])/2
        ypos = (area[1] + area[3])/2
        xsize = max(area[0], area[2]) - min(area[0], area[2])
        ysize = max(area[1], area[3]) - min(area[1], area[3])

        if indx == 0:
            return inputNode.sample('rgb.red', xpos+.5, ypos+.5, xsize, ysize)
        if indx == 1:
            return inputNode.sample('rgb.green', xpos+.5, ypos+.5, xsize, ysize)
        if indx == 2:
            return inputNode.sample('rgb.blue', xpos+.5, ypos+.5, xsize, ysize)
        else:
            return 0
    except:
        return 0

def AG_GrowArea(name):
    node = nuke.thisGroup()

    area = list(node[name + 'Area'].value())
    xsize = max(area[0], area[2]) - min(area[0], area[2])
    ysize = max(area[1], area[3]) - min(area[1], area[3])

    grow = node[name+'Grow'].value()
    grow = grow if grow >= 0 else grow/(1-grow)

    growGlobal = node['growGlobal'].value()
    growGlobal = growGlobal if growGlobal >= 0 else growGlobal/(1-growGlobal)

    total = (grow+1)*(growGlobal+1)

    xg = (xsize*total - xsize)/2
    yg = (ysize*total - ysize)/2

    if area[0] > area[2]:
        area[0] = area[0] + xg
        area[2] = area[2] - xg
    else:
        area[0] = area[0] - xg
        area[2] = area[2] + xg
    if area[1] > area[3]:
        area[1] = area[1] + yg
        area[3] = area[3] - yg
    else:
        area[1] = area[1] - yg
        area[3] = area[3] + yg

    return area

def AG_MakeRenderFarmFriendly(node):
    import inspect
    import AreaGradient
    a = inspect.getsource(AreaGradient.AG_SampleSource)
    b = inspect.getsource(AreaGradient.AG_GrowArea)

    node['onCreate'].setValue(a+b)