-- Renders all the available render targets of the loaded project.
-- Asks the user for an output directory and saves all the images into
-- that directory (<rendertargetname>_01.png, <rendertargetname>_02.png, ...)
--
-- @author      Mark, Thomas
-- @version     0.2
-- @shortcut    ctrl + b

local version = "v1.0 - for OctaneRender 1.5"

-- helper to pop-up an OkCancel dialog and return the button clicked
function showWarning(text,msg)

 	local ret = octane.gui.showDialog

		{
            type  = octane.gui.dialogType.BUTTON_DIALOG,
			buttonTexts =  {"Continue"}, 
	        icon  = octane.gui.dialogIcon.WARNING,
            title = "WARNING:\n"..tostring(text),
            text  = tostring(msg),
		}
        
end


-- global table with our settings
local settings =
{
    ["projectPath"]     = octane.project.getCurrentProject(),
    ["renderTargets"]   = octane.project.getSceneGraph():findNodes(octane.NT_RENDERTARGET, true),
    ["outputDirectory"] = nil,
    ["enabled"]         = {},
    ["fileType"]        = {}
}

-- filter out render targets with no kernel
local filtered = {}
local strSkippedRenderTargets = ""

for ix, renderTargetNode in ipairs(settings.renderTargets) do

    if renderTargetNode:getInputNode(octane.P_KERNEL) ~= nil then
        table.insert(filtered, renderTargetNode) 
    else
        strSkippedRenderTargets =  strSkippedRenderTargets.. "\n" .. renderTargetNode.name
    end
    
end

if strSkippedRenderTargets ~= "" then

    strSkippedRenderTargets =  strSkippedRenderTargets.."\n\n\nTo process these Render Targets exit the script and wire a kernel node to them."
    showWarning("The following Render Targets have missing kernels and will not be processed.",strSkippedRenderTargets)
    
end

settings.renderTargets = filtered

-- sorts a table alpha numerically
-- (snippet from http://notebook.kulchenko.com/algorithms/alphanumeric-natural-sorting-for-humans-in-lua)
function alphanumsort(nodes)
    local function padnum(d)
        local dec, n = string.match(d, "(%.?)0*(.+)")
        return #dec > 0 and ("%.12f"):format(d) or ("%s%03d%s"):format(dec, #n, n)
    end
    local function compare(n0, n1) 
        local a, b = n0.name, n1.name
        return a:gsub("%.?%d+",padnum)..("%3d"):format(#b)
             < b:gsub("%.?%d+",padnum)..("%3d"):format(#a)
    end
    table.sort(nodes, compare)
end 

-- sort the render targets alphabetically
alphanumsort(settings.renderTargets)

-- If no render targets are found, error out
if #settings.renderTargets == 0 then
    error("No render targets were found in this project.", projectPath)
end


--------------------------------------------------------------------------------------------------
-- Gui helpers

-- Creates a text label and returns it
function createLabel(text, width, height)
    return octane.gui.create
    { 
        type   = octane.gui.componentType.LABEL,            -- type of the component
        text   = text,                                      -- text that appears on the label
        width  = width,                                     -- width of the label in pixels
        height = height,                                    -- height of the label in pixels
    }
end


-- Creates a button and returns it
function createButton(name, text, width, height, tooltip)
    return octane.gui.create
    {
        type        = octane.gui.componentType.BUTTON,      -- type of the component
        name        = name,                                 -- name of the component
        text        = text,                                 -- button text
        width       = width,                                -- width in pixels
        height      = height,                               -- height in pixels
        tooltip     = tooltip,                              -- tooltip text
        x           = x,
        y           = y,
    }
end


-- Creates a slider and returns it
function createSlider(name, value, min, max, step, width, height, logarithmic)
    return octane.gui.create
    { 
        type        = octane.gui.componentType.SLIDER,      -- type of the component
        name        = name,                                 -- name of the component
        value       = value,                                -- value of the slider
        minValue    = min,                                  -- minimum value of the slider
        maxValue    = max,                                  -- maximum value of the slider
        step        = step,                                 -- interval between 2 discrete slider values
        width       = width,                                -- width of the slider in pixels
        height      = height,                               -- height of the slider in pixels
        logarithmic = logarithmic,                          -- Make slider logarithmic
    }
end

    
-- Creates a group and retuns it
function createGroup(name, children, border, rows, cols, width, height, text, padding, inset, center)
    return octane.gui.create
    {
        type        = octane.gui.componentType.GROUP,           -- type of the component
        name        = name,                                     -- name of component
        children    = children,                                 -- list of children components
        border      = border,                                   -- boolean flag to show border
        rows        = rows,                                     -- number of rows in group grid
        cols        = cols,                                     -- number of cols in group grid
        width       = width,                                    -- width in pixels
        height      = height,                                   -- height in pixels
        text        = text,                                     -- text to display at top of group
        padding     = padding,                                  -- internal padding in each cell
        inset       = inset,                                    -- inset of the group component
        centre      = center,                                   -- center the group
    }
end


-- Recursively enables or disables all components on a window
function setEnabled(component, enable)

    if component.type == octane.gui.componentType.WINDOW or
       component.type == octane.gui.componentType.GROUP then
        for _, childComponent in ipairs(component.children) do
            setEnabled(childComponent, enable)
        end
    else
        component.enable = enable
    end
    
end
    

-- layout constants
local BUTTON_WIDTH     = 100
local BUTTON_HEIGHT    = 20
local WIDE_LBL_WIDTH   = 226
local NARROW_LBL_WIDTH = BUTTON_WIDTH
local LBL_HEIGHT       = 20
local GRP_PAD          = 2

-- create a table with:
-- * render target sequence number
-- * render target name
-- * render target enable box
-- * render target output format box
local tableChildren = {}

for ix, renderTarget in ipairs(settings.renderTargets) do

    -- Check if a kernel is connected to the render target
    local kernelNode = renderTarget:getInputNode(octane.P_KERNEL)      
    local description = string.format("%d: %s", ix, renderTarget.name .. " - " ..kernelNode:getPinValue(octane.P_MAX_SAMPLES) .. " spp")
    local lbl         = createLabel(description, WIDE_LBL_WIDTH, LBL_HEIGHT)
    local enableBox   = octane.gui.create
    {
        type   = octane.gui.componentType.CHECK_BOX,
        text   = "render",
        enable = true,
        width  = 60,
    }
    
    local fileTypeCombo = octane.gui.create
    {
        type  = octane.gui.componentType.COMBO_BOX,
        items =
        { 
            "PNG 8-bpc"       ,
            "PNG 16-bpc"      ,
            "EXR untonemapped",
            "EXR tonemapped"  ,
        },
        selectedIx = 1,
        width      = 120,
    }
    
    table.insert(tableChildren    , lbl)
    table.insert(tableChildren    , enableBox)
    table.insert(tableChildren    , fileTypeCombo)
    table.insert(settings.enabled , true)
    table.insert(settings.fileType, octane.render.imageType.PNG8)
    
    -- on click, toggle the enabled flag in the list
    enableBox.callback = function()
        settings.enabled[ix] = not settings.enabled[ix]
    end
    
    -- on combo change -> modify the image file type
    fileTypeCombo.callback = function()
        local fileTypes =
        {
            octane.render.imageType.PNG8,
            octane.render.imageType.PNG16,
            octane.render.imageType.EXR,
            octane.render.imageType.EXRTONEMAPPED,
        }
        settings.fileType[ix] = fileTypes[fileTypeCombo.selectedIx]
    end
    
end


-- add a disable all button
local disableAllButton = octane.gui.create
{
    type     = octane.gui.componentType.BUTTON,
    width    = BUTTON_WIDTH,
    height   = BUTTON_HEIGHT,
    text     = "Deselect All",
    x        = 5,
    y        = 5,
    
    callback = function(me)
    
        local enable = nil 
        
        if me.text == "Deselect All" then
            me.text = "Select All"
            enable  = false
        else 
            me.text = "Deselect All"
            enable  = true
        end
        
        for ix=1,#settings.renderTargets do
            settings.enabled[ix] = enable
        end
        
        for _,child in ipairs(tableChildren) do
            if child.type == octane.gui.componentType.CHECK_BOX then
                child.checked = enable
            end
        end
        
    end
}

table.insert(tableChildren, 1,disableAllButton)
table.insert(tableChildren, 2,createLabel("", 1, 1))
table.insert(tableChildren, 3,createLabel("", 1, 1))

createButton("output", "Output Folder...", BUTTON_WIDTH, BUTTON_HEIGHT, "Specify the folder to receive the rendered output files.")

local renderTargetsGrp = octane.gui.create
{
    type     = octane.gui.componentType.GROUP,
    children = tableChildren,
    border   = true,
    cols     = 3,
    rows     = #tableChildren / 3,
    text     = "Available Render Targets",
    padding  =  { GRP_PAD + 0 },
}

-- maximum samples numericbox
--local lblMaxSamples = createLabel("Samples/Pixel for render targets without a kernel node", NARROW_LBL_WIDTH, LBL_HEIGHT)
local lblMaxSamples = createLabel("Samples/Pixel (for render targets with a default kernel)", 300, LBL_HEIGHT)
local boxMaxSamples = octane.gui.create
{
    type     = octane.gui.componentType.NUMERIC_BOX,
    width    = NARROW_LBL_WIDTH,
    height   = LBL_HEIGHT,
    maxValue = 256000,
    minValue = 1,
    step     = 100,
    value    = 500,
}
grpSettings1 = createGroup("grpEsettings1", { lblMaxSamples,boxMaxSamples}, false, 1, 2, nil, nil, "Execute", { 5 }, nil, false)

-- output location
local btnOutput = createButton("output", "Output Folder...", BUTTON_WIDTH, BUTTON_HEIGHT, "Specify the folder to receive the rendered output files.")

-- create an editor that will show the chosen file path
local txtEditorDirectory = octane.gui.create
{
    type    = octane.gui.componentType.TEXT_EDITOR,
    text    = ""                                  ,
    x       = 0                                   ,
    -- calculate the width to make the 2 stacked groups the same size
    -- no magic formula just trial and error ;)
    width   = renderTargetsGrp.width - btnOutput.width - 30,
    height  = 20                                           ,
    enable  = false                                        ,
}
grpSettings2 = createGroup("grpEsettings1", { btnOutput, txtEditorDirectory}, false, 1, 2, nil, nil, "Execute", { 5 }, nil, false)
    
local settingsGrp = octane.gui.create
{
    type     = octane.gui.componentType.GROUP,
    children =
    {
        --grpSettings1,
        grpSettings2,
    },
    text    = "Settings",
    border  = false,
    cols    = 1,
    rows    = 1,
    width = renderTargetsGrp.width,
    padding = { 0 },
}

local progressBar = octane.gui.create
{
    type   = octane.gui.componentType.PROGRESS_BAR,
    width  = settingsGrp.width - 8,
    height = LBL_HEIGHT,
}

-- buttons (ok not used in this case)
local btnApply  = createButton("apply", "Apply", BUTTON_WIDTH, BUTTON_HEIGHT, "Start processing all selected render targets. if no output directory is specified, a dry run is performed")
local btnCancel = createButton("cancel", "Cancel", BUTTON_WIDTH, BUTTON_HEIGHT, "Cancel.")
--local btnPause     = createButton("pause", "Pause", BUTTON_WIDTH, BUTTON_HEIGHT, "Suspend processing the batch.")
--local btnContinue = createButton("continue", "Continue", BUTTON_WIDTH, BUTTON_HEIGHT, "Continue processing a suspended batch.")

grpExecute = createGroup("grpExecute", { btnApply,  btnCancel}, false, 1, 2, nil, nil, "Execute", { 5 }, nil, false)


-- this group holds everything together
local layoutGrp = octane.gui.create
{
    type     = octane.gui.componentType.GROUP,
    rows     = 4,
    cols     = 1,
    children = {renderTargetsGrp, settingsGrp, progressBar, grpExecute},
    border   = false,
    padding  = { 2 },
    centre   = true,
}

-- window that holds all components
local wndMain = octane.gui.create
{ 
    type     = octane.gui.componentType.WINDOW,
    text     = "Batch Rendering",
    children = { layoutGrp },
    width    = layoutGrp.width,
    height   = layoutGrp.height,
}


-- Callback function for the GUI Controls
local function guiCallback(component, event)

    if component == btnOutput then
        -- ask the user for an output directory for the results
        local ret = octane.gui.showDialog
        {
            type            = octane.gui.dialogType.FILE_DIALOG,
            title           = "Select output directory for the render results",
            path            = octane.file.getParentDirectory(settings.projectPath),
            browseDirectory = true,
            save            = false,
        }
        if not ret.result or ret.result == "" then
            settings.outputDirectory = nil  
        else 
            settings.outputDirectory = ret.result
            txtEditorDirectory.text  = settings.outputDirectory
        end
        
    elseif component == btnApply then
        -- Clear the cancel flag
        IS_CANCELLED = false

        btnCancel.enable = true
        --btnPause.enable  = true
        --btnContinue.enable = false
        
        batchRender()

        btnApply.enable  = false
        --btnPause.enable  = false
        btnCancel.enable = true
        progressBar.text = "finished"
        
    elseif component == btnPause then
        --btnApply.enable  = false
        --btnPause.enable  = false
        --btnContinue.enable = true
        --btnCancel.enable = true
        --IS_PAUSED = true
        --cancelRender()
        --progressBar.text  = "paused"

    elseif component == btnContinue then
        --btnPause.enable  = false
        --btnContinue.enable = true
        --IS_PAUSED = false

    elseif component == btnCancel then
        cancelRender()
        wndMain:closeWindow() 
            
    elseif component == wndMain then
        cancelRender()
    end
    
end

-- Hookup GUI controls to their Callback function
boxMaxSamples.callback = guiCallback
btnOutput.callback     = guiCallback
btnApply.callback      = guiCallback
btnCancel.callback     = guiCallback
--btnPause.callback      = guiCallback
--btnContinue.callback   = guiCallback


function cancelRender()
    IS_CANCELLED = true
    octane.render.callbackStop()
end


------------------------------------------------------------------------------------------------------------------------------------
-- MAIN  ROUTINE -------------------------------------------------------------------------------------------------------------------
------------------------------------------------------------------------------------------------------------------------------------

function batchRender()

    progressBar.progress = 0
    
    -- Calculate the progress bar step
    local numEnabled = 0
    local path = nil
    local ms =  boxMaxSamples.value
    local KernelType = ""
    
    -- String to hold log info
    local strLog = ""
    
    for ix, enabled in ipairs(settings.enabled) do
        if enabled then numEnabled = numEnabled + 1 end
    end
    
    local progressBarStep = 1 / numEnabled

    -- Render all the render targets that are enabled
    for ix, renderTarget in ipairs(settings.renderTargets) do
    
        -- Break out if we're cancelled and set it in the progress bar
        if IS_CANCELLED then
            progressBar.text  = "cancelled"
            break 
        end
        
        -- Only render if the rt is enabled
        if settings.enabled[ix] then
        
            -- Check if a kernel is connected to the render target
            local kernelNode = renderTarget:getInputNode(octane.P_KERNEL)
        
            print(kernelNode)

            -- Figure out the type of kernel (you can use .type here)
            if kernelNode ~= nil then
                if kernelNode.type == octane.NT_KERN_DIRECTLIGHTING
                or kernelNode.type == octane.NT_KERN_PATHTRACING 
                or kernelNode.type == octane.NT_KERN_PMC
                or kernelNode.type == octane.NT_KERN_INFO
                or kernelNode.type == octane.NT_KERN_MATPREVIEW then
                
                print(string.format("kernel with %d samples/px", kernelNode:getPinValue(octane.P_MAX_SAMPLES)))
                
                ms = kernelNode:getPinValue(octane.P_MAX_SAMPLES)
                
                if kernelNode:getProperties().type == octane.NT_KERN_DIRECTLIGHTING then KernelType = "Direct Lighting Kernel" end
                if kernelNode:getProperties().type == octane.NT_KERN_PATHTRACING then KernelType = "Path Tracing Kernel" end
                if kernelNode:getProperties().type == octane.NT_KERN_PMC then KernelType = "PMC Kernel" end
                if kernelNode:getProperties().type == octane.NT_KERN_INFO then KernelType = "Info Kernel" end
                if kernelNode:getProperties().type == octane.NT_KERN_MATPREVIEW then KernelType = "Material Preview Kernel" end
            
            else
               
                print("unknown kernel", kernelNode.type)
            
            end
        
        else
    
            print("render target has no connected kernel node")
        
        end
        
            -- create an output path for the image
            if settings.outputDirectory then
                -- common extension for our image output types
                local fileExtensions = 
                {
                    [octane.render.imageType.PNG8]          = "png",
                    [octane.render.imageType.PNG16]         = "png",
                    [octane.render.imageType.EXR]           = "exr",
                    [octane.render.imageType.EXRTONEMAPPED] = "exr",
                }
                path = string.format("%s/%s_%dspp_%03d.%s",
                                     settings.outputDirectory,
                                     renderTarget.name,
                                     ms,
                                     ix,
                                     fileExtensions[settings.fileType[ix]])
            end

            -- update the progress
            progressBar.text = string.format("rendering %s (%s)", renderTarget.name, path or "dry-run")
            
            -- Store the start time
            local ImageStartTime = os.clock()
            
            -- Render the image
            octane.render.start
            { 
                renderTargetNode = renderTarget,
                maxSamples       = ms,
            }
            
            ms =  boxMaxSamples.value
            
            -- Calculate the elapsed time
            local ImageElapsedTime = string.format("%.02f", os.clock() - ImageStartTime)
            -- Add image info to log
            local strImageLog = string.format(renderTarget.name .. " exported with the " .. KernelType .." @" .. boxMaxSamples.value .." samples/pixel in " .. ImageElapsedTime .." secs")
            strLog = strLog .. "\n" .. strImageLog
        
           -- save out the image
            if path then
                octane.render.saveImage(path, settings.fileType[ix])
            end

            -- update the progress bar
            progressBar.progress = progressBar.progress + progressBarStep
        end
    end
    
    if path ~= nil then
    
        local logFileName = "Octane Batch Render.log"
        path = string.format("%s/%s",
                             settings.outputDirectory,
                             logFileName
                             )
                                     
        -- Write log info to a text file.
        local file = io.open(path, "w")
        file:write(strLog)
        file:close()
        
    end
    
end


-- Initialise GUI
setEnabled(wndMain, true)
btnCancel.enable = false
--btnPause.enable  = false
--btn-Continue.enable = false

-- Ths must be placed such that all required functions are read prior to execution.
-- The script will hold here until the window closes
wndMain:showWindow()