---------------------------------------------------------------------------------------------------
-- 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 
-- (e.g. <rendertargetname>_01.png, <rendertargetname>_02.png, ...)
--
-- @author      Mark Basset, Thomas Loockx and others
-- @description Batch rendering for Octane.
-- @version     0.3


-- Global table with our settings. All global variables should be here to keep an overview.
-- TODO: a property table here would make everything more compact
local gSettings =
{
    -- absolute path of the current project
    ["projectPath"]     = octane.project.getCurrentProject(),
    -- list of render target nodes in the current project together with their enable state
    -- and export file format
    ["renderTargets"]   = {},
    -- absolute path to the output directory of the rendered images
    ["outputDirectory"] = nil,
    -- max samples (if 0 we use the max samples of the node)
    ["maxSamples"]      = 0,
    -- filename template for the output files
    ["template"]        = "%n.%e",
    -- true if the rendering was cancelled
    ["cancelled"]       = false,
    -- handle for the progress update function (takes value and text)
    ["progress"]        = nil,
    -- handle for the batch render function
    ["batchRender"]     = nil,
    -- handle for the window
    ["window"]          = nil,
}


-- 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


-- 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 


-- Cancels rendering.
function cancelRendering()
    gSettings.cancelled = true
    octane.render.callbackStop()
end



--------------------------------------------------------------------------------------------------
-- Main window creation

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


-- Creates a tabular overview of all the render targets in the current project.
local function initRenderTargetOverview()
    local tableChildren = {}
    -- collect all the render targets from the current project
    for ix, renderTarget in ipairs(gSettings.renderTargets) do
        local description = string.format("%5d: %s", ix, renderTarget.node.name)
        local lbl = octane.gui.create
        {
            type   = octane.gui.componentType.LABEL,
            width  = WIDE_LBL_WIDTH,
            height = LBL_HEIGHT,
            text   = description,
        }
        local enableBox = octane.gui.create
        {
            type     = octane.gui.componentType.CHECK_BOX,
            text     = "render",
            enable   = true,
            width    = 80,
            callback = function(box)
                -- toggle the for the render target enable state in the settings
                renderTarget.enabled = box.checked
            end
        }
        local fileTypeCombo = octane.gui.create
        {
            type  = octane.gui.componentType.COMBO_BOX,
            items =
            { 
                "PNG (8-bit)",
                "PNG (16-bit)",
                "EXR",
            },
            selectedIx = 1,
            width      = 120,
            callback   = function(combo)
                -- update the file type for this render target in the settings      
                local fileTypes =
                {
                    octane.render.imageType.PNG8,
                    octane.render.imageType.PNG16,
                    octane.render.imageType.EXR,
                }
                renderTarget.fileType = fileTypes[combo.selectedIx]
            end
        }

        -- add components in the list
        table.insert(tableChildren, lbl)
        table.insert(tableChildren, enableBox)
        table.insert(tableChildren, fileTypeCombo)
    end

    -- add a disable all button
    local disableAllButton = octane.gui.create
    {
        type     = octane.gui.componentType.BUTTON,
        width    = BUTTON_WIDTH,
        height   = BUTTON_HEIGHT,
        text     = "Disable All",
        callback = function(me)
            local enable = nil 
            if me.text == "Disable All" then
                me.text = "Enable All"
                enable  = false
            else 
                me.text = "Disable All"
                enable  = true
            end
            -- update the render target state
            for _,renderTarget in ipairs(gSettings.renderTargets) do
                renderTarget.enable = enable
            end
            -- update the check boxes
            for _,child in ipairs(tableChildren) do
                if child.type == octane.gui.componentType.CHECK_BOX then
                    child.checked = enable
                end
            end
        end
    }

    -- add the disable button between 2 dummy labels
    local dummyLbl = octane.gui.create
    {
        type   = octane.gui.componentType.LABEL,
        text   = "",
        width  = 1,
        height = 1,
    }
    table.insert(tableChildren, dummyLbl)
    table.insert(tableChildren, disableAllButton)
    table.insert(tableChildren, dummyLbl)

    -- create a group to pack it all up
    local renderTargetsGrp = octane.gui.create
    {
        type     = octane.gui.componentType.GROUP,
        children = tableChildren,
        border   = true,
        cols     = 3,
        rows     = #tableChildren / 3,
        border   = false,
        padding  = { GRP_PAD },
    }

    -- wrap the group in a panel stack to make it scrollable
    local panelStack = octane.gui.create
    {
        type     = octane.gui.componentType.PANEL_STACK,
        height   = 400,
        children = { renderTargetsGrp },
        open     = { true },
        captions = { "Render targets" },
    }

    return panelStack 
end


-- Initializes a GUI group with all the config settings
local function initSettingGroup(width)
    -- override label
    local overrideLbl = octane.gui.create
    {
        type   = octane.gui.componentType.LABEL,
        text   = "Override s/px",
        width  = NARROW_LBL_WIDTH,
        height = LBL_HEIGHT,
    }
    -- maximum samples label
    local lblMaxSamples = octane.gui.create
    {
        type   = octane.gui.componentType.LABEL,
        text   = "Samples/px",
        width  = NARROW_LBL_WIDTH,
        height = LBL_HEIGHT,
    }
    -- maximum samples numeric box
    local maxSamplexBox = octane.gui.create
    {
        type     = octane.gui.componentType.NUMERIC_BOX,
        width    = BUTTON_WIDTH,
        height   = LBL_HEIGHT,
        enable   = false,
        maxValue = 256000,
        minValue = 1,
        step     = 1,
        value    = 100,
        callback = function(numericBox)
            gSettings.maxSamples = numericBox.value
        end
    }
    -- maximum samples override thickbox
    local overrideCheckBox = octane.gui.create
    {
        type     = octane.gui.componentType.CHECK_BOX,
        width    = 200,
        height   = 24,
        text     = "",
        checked  = false,
        tooltip  = "Override max samples in the kernel node",
        callback = function(box)
            if box.checked then
                maxSamplexBox.enable = box.checked
                gSettings.maxSamples = maxSamplexBox.value
            else
                maxSamplexBox.enable = false
                gSettings.maxSamples = 0
            end
        end
    }
    -- template label
    local templateLbl = octane.gui.create
    {
        type   = octane.gui.componentType.LABEL,
        text   = "Filename template",
        width  = NARROW_LBL_WIDTH,
        height = LBL_HEIGHT,
    }
    -- template text editor
    local templateEditor = octane.gui.create
    {
        type     = octane.gui.componentType.TEXT_EDITOR,
        text     = gSettings.template,
        width    = width - templateLbl.width - 10,
        height   = 20,
        enable   = true,
        tooltip  = "Template parameters: %i render target index %n node name %e extension %t timestamp",
        callback = function(editor)
            gSettings.template = editor.text
        end
    }
    -- create a button to browse for an output folder
    local outputButton = octane.gui.create
    {
        type     = octane.gui.componentType.BUTTON,
        text     = "Output Folder...",
        width    = BUTTON_WIDTH,
        height   = BUTTON_HEIGHT,
        tooltip  = "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   = "",
        width  = templateEditor.width,
        height = 20,
        enable = true,
    }
    -- callback function for the output button
    outputButton.callback = function()
        -- 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(gSettings.projectPath),
            browseDirectory = true,
            save            = false,
        }
        if not ret.result or ret.result == "" then
            gSettings.outputDirectory = nil  
        else 
            gSettings.outputDirectory = ret.result
            txtEditorDirectory.text   = gSettings.outputDirectory
        end
    end

    -- group holding it all together
    local settingsGrp = octane.gui.create
    {
        type     = octane.gui.componentType.GROUP,
        children =
        {
            overrideLbl  , overrideCheckBox  ,
            lblMaxSamples, maxSamplexBox     ,
            templateLbl  , templateEditor    ,
            outputButton , txtEditorDirectory,
        },
        text     = "Settings",
        border   = true,
        cols     = 2,
        rows     = 4,
        inset    = { GRP_PAD },
        padding  = { GRP_PAD },
    }

    return settingsGrp
end


-- Inits the progress bar.
local function initProgressBar(width)
    local progressBar = octane.gui.create
    {
        type   = octane.gui.componentType.PROGRESS_BAR,
        width  = width - 10,
        height = LBL_HEIGHT,
    }
    -- the global progress function needs to update the progress bar
    gSettings.progress = function(progress, text)
        progressBar.progress = progress
        progressBar.text     = text
    end

    return progressBar
end


-- Inits a GUI group for the control buttons.
local function initControlsGroup()
    local cancelButton = octane.gui.create
    {
        type     = octane.gui.componentType.BUTTON,
        text     = "Cancel",
        width    = BUTTON_WIDTH,
        height   = BUTTON_HEIGHT,
        tooltip  = "Cancel batch processing.",
        enable   = false,
        callback = function() cancelRendering() end
    }
    local startButton = octane.gui.create
    {
        type    = octane.gui.componentType.BUTTON,
        text    = "Start",
        width   = BUTTON_WIDTH,
        height  = BUTTON_HEIGHT,
        tooltip = "Start batch processing.",
        callback = function()
            -- we will hang in this callback until rendering is finished, the only way to stay
            -- responsive is via the render callback
            
            -- clear cancel flag
            gSettings.cancelled = false

            -- disable all but the cancel button
            setEnabled(gSettings.window, false)
            cancelButton.enable = true

            -- do the work
            gSettings.batchRender()

            -- enable all but the cancel button
            setEnabled(gSettings.window, true)
            cancelButton.enable = false
            if not gSettings.cancelled then
                gSettings.progress(100, "Finished")
            end
        end
    }
    local controlsGroup = octane.gui.create
    {
        type     = octane.gui.componentType.GROUP,
        children = { startButton, cancelButton },
        rows     = 1,
        cols     = 2,
        border   = false,
        padding  = { 2 },
    }

    return controlsGroup 
end


-- Initializes the main window.
local function initMainWindow()
    -- initialize the main pieces of the user interface
    local renderTargetOverview = initRenderTargetOverview()
    local userSettings         = initSettingGroup(renderTargetOverview.width)
    local progressBar          = initProgressBar(userSettings.width)
    local controlsGroup        = initControlsGroup()

    -- these guys go on the window top-to-bottom
    local children = { renderTargetOverview, progressBar, userSettings, controlsGroup }

    -- this group holds everything together
    local layoutGrp = octane.gui.create
    {
        type     = octane.gui.componentType.GROUP,
        children = children,
        cols     = 1,
        rows     = #children,
        border   = false,
        padding  = { 2 },
        centre   = true,
    }
    -- window that holds all components
    local window = octane.gui.create
    { 
        type     = octane.gui.componentType.WINDOW,
        text     = "Batch Rendering",
        children = { layoutGrp },
        width    = layoutGrp.width,
        height   = layoutGrp.height,
        callback = function()
            -- cancel rendering on close
            cancelRendering()
        end
    }

    return window
end



---------------------------------------------------------------------------------------------------
-- Batch rendering


-- Creates a file name based on a template string by substitution
local function createFilename(template, ix, name, extension)
    local s = template
    -- %i -> index of the render target
    s = string.gsub(s, "%%i", string.format("%d", ix))
    -- %n -> name of the node
    s = string.gsub(s, "%%n", name)
    -- %e -> extension
    s = string.gsub(s, "%%e", extension)
    -- %t -> timestamp (h_m_s)
    s = string.gsub(s, "%%t", os.date("%H_%M_%S"))
    return s
end


-- The batch rendering function. Loops over all the render targets in the project and renders those
-- that the user selected.
gSettings.batchRender = function()
    -- calculate the progression step
    local numEnabled = 0  
    for ix, renderTarget in ipairs(gSettings.renderTargets) do
        if renderTarget.enabled then numEnabled = numEnabled + 1 end
    end
    local progressStep = 1 / numEnabled
    -- keeps track of the render progress (in range [0,1])
    local progress = 0

    -- render all the render targets that are enabled
    for ix, renderTarget in ipairs(gSettings.renderTargets) do
        -- only render if the rt is enabled
        if renderTarget.enabled 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",
            }
            -- create a filename based on the template
            local filename = createFilename(gSettings.template, ix, renderTarget.node.name, fileExtensions[renderTarget.fileType])
            -- create an output path for the image
            local path = nil
            if gSettings.outputDirectory then
                path = string.format("%s/%s", gSettings.outputDirectory, filename)
            else
                path = string.format("%s (dry-run)", filename)
            end

            -- update the progress
            gSettings.progress(progress, string.format("rendering %s -> %s", renderTarget.node.name, path))
            progress = progress + progressStep

            -- override max samples if configured
            if gSettings.maxSamples > 0 then
                renderTarget.node:getInputNode(octane.P_KERNEL):setPinValue(octane.P_MAX_SAMPLES, gSettings.maxSamples)
            end

            -- do the rendering of the image
            octane.render.start
            { 
                renderTargetNode = renderTarget.node,
                callback         = function()
                    if gSettings.cancelled then
                        octane.render.callbackStop()
                    end
                end
            }

            -- break out if we're cancelled and set it in the progress bar
            if gSettings.cancelled then
                gSettings.progress(0, "Cancelled")
                break 
            end
 
            -- save out the image
            if gSettings.outputDirectory and path then
                octane.render.saveImage(path, renderTarget.fileType)
            end
        end
    end
end



---------------------------------------------------------------------------------------------------
-- Main script

-- create a copy of the original project
local projectGraph = octane.nodegraph.createRootGraph("Project Copy")
projectGraph:copyFromGraph(octane.project.getSceneGraph())

-- fetch all the render target nodes
local renderTargetNodes = projectGraph:findNodes(octane.NT_RENDERTARGET, true)
-- if no render targets are found -> error out
if #renderTargetNodes == 0 then
    error("No render targets were found in this project.", projectPath)
end

-- sort all the render target nodes alphanumerically
alphanumsort(renderTargetNodes)

-- by default we don't override max samples
gSettings.maxSamples = 0

-- initialize the state for the render targets
for ix, node in ipairs(renderTargetNodes) do
    local state =
    {
        ["node"]     = node,
        ["enabled"]  = true,
        ["fileType"] = octane.render.imageType.PNG8,
    }
    table.insert(gSettings.renderTargets, state)
end

-- the script blocks here until the window is closes
gSettings.window = initMainWindow()
gSettings.window:showWindow()