#######################################################################################################################

__author__ = "Boris Martinez Castillo Based on Fernando Arbelaez's Idea"
__version__ = "1.0.1"
__maintainer__ = "Boris Martinez Castillo"
__email__ = "boris.vfx@outlook.com"

########################################################################################################################

#IMPORTS

import nuke
import nukescripts
import itertools

#DEFINITIONS

class CamPullPanelB(nukescripts.PythonPanel):
    def __init__(self):
        nukescripts.PythonPanel.__init__(self,"focus_puller")

    #CREATE KNOBS
        self.fstop = nuke.Double_Knob("fstop:")
        self.fstop.clearFlag(nuke.STARTLINE)        
        self.frame_list = nuke.String_Knob("frame list:")
        self.focus_list = nuke.String_Knob("focus plane values:")                
        self.cameralist = nuke.Enumeration_Knob("cameras", "choose camera",self.get_camera_nodes())
        self.pull_focus = nuke.PyScript_Knob("pull focus","pull focus")
        self.pull_focus.clearFlag(nuke.STARTLINE)
        self.author = nuke.Text_Knob("making focus pulling \n great again")
                
   
    #SET DEFAULTS
        i = nuke.root()['first_frame'].value()
        o = nuke.root()['last_frame'].value()
        self.fstop.setValue(5.6)
        self.frame_list.setValue(str(int(i)) + "," + str(int(o)))
        self.focus_list.setValue("10,20")
       

    #ADD KNOBS
        for i in (self.fstop,self.frame_list,self.focus_list,self.cameralist,self.author):
            self.addKnob(i)

      #METHODS
    def get_camera_nodes(self):
        self.camera_list = []
            
        for node in nuke.allNodes():
            if node.Class() == "HubCamera2" or node.Class() == "Camera2":
                self.camera_list.append(node.name())
        return self.camera_list
       

def delete_fstop_focalpoint_anim(node_name):
    
        node = nuke.toNode(node_name)

        for knob in node.knobs():
            if nuke.Knob.isAnimated(node['focal_point']):
                nuke.Knob.clearAnimated(node['focal_point'])

            elif nuke.Knob.isAnimated(node['fstop']):
                        nuke.Knob.clearAnimated(node['fstop'])                    
            else: 
                print "Nothing to clean up!!"
                

def string_to_list(str):

    lst = str.split(",")
    return lst


def set_keys_focal(camera,frame_list,focal_list):
    
    focal = camera['focal_point']
    focal.setAnimated()
    print frame_list, focal_list
    
    for f,b in itertools.izip(frame_list,focal_list):
        focal.setValueAt(float(b),float(f))


def pull_focus_b():
        
    context= CamPullPanelB()

    if not context.showModalDialog():
        print "script aborted"
        return
    else:

        cam_node = context.cameralist.value()
        
        delete_fstop_focalpoint_anim(cam_node)
        
        frame_list = string_to_list(context.frame_list.value())
        focus_list = string_to_list(context.focus_list.value())    
        fstop = context.fstop.value()
           
        target_cam_node = nuke.toNode(cam_node)

        set_keys_focal(target_cam_node,frame_list,focus_list)

        target_cam_node.knob('fstop').setValue(fstop)

if __name__ == "__main__":
    pull_focus_b()