----------------------------------------
-- Partition node.
-- Visits polygon-based meshes in the input and tag the connected components in an attribute.
-- @author      Otoy
-- @version     1.1
-- @script-id   2324ad82-8919-49fc-9e32-bec753d9783c
-- @node-registry-category |Geometry|Partition elements
-----------------------------------------
local inputLinkersInfo = {}
local outputLinkersInfo = {}
local inputLinkers = {}
local outputLinkers = {}
local nodeGeometry = nil

----------------------------------------------------
-- Union-Find.
-- https://en.wikipedia.org/wiki/Disjoint-set_data_structure
----------------------------------------------------
-- union-find nodes indexed by the vertex index.
local ufParents = {}
local ufSizes = {}

-- Union-Find: Find. Finds and returns the root of an element.
local function ufFind(i)

    -- Make.
    if ufParents[i] == nil then
        ufParents[i] = i
        ufSizes[i] = 1
    end

    -- Path compression. Point each entry to its root.
    if ufParents[i] ~= i then
        ufParents[i] = ufFind(ufParents[i])
    end

    return ufParents[i]
end

-- Union-Find: Union. Merges two sets.
local function ufUnion(i, j)

    local pi = ufFind(i)
    local pj = ufFind(j)

    -- Union by size: merge smaller tree into larger tree.
    if ufSizes[pi] < ufSizes[pj] then
        ufParents[pi] = pj
        ufSizes[pj] = ufSizes[pj] + ufSizes[pi]
    else
        ufParents[pj] = pi
        ufSizes[pi] = ufSizes[pi] + ufSizes[pj]
    end

end

-- Find the first empty attribute slot in a series of attributes.
-- prefix: either "A_FLOAT_ATTRIBUTE_NAME" or "A_COLOR_ATTRIBUTE_NAME".
-- max: number of slots for the desired kind of attribute (2 or 4).
local function findEmptyAttribute(node, prefix, max)
    local attrIndex = 1
    local found = false
    while not found and attrIndex < (max+1)
    do
        local attr = prefix..tostring(attrIndex)
        local name = node:getAttribute(octane.attributeId[attr])
        if name == nil or name == "" then
            found = true
        else
            attrIndex = attrIndex + 1
        end
    end

    if found then
        return tostring(attrIndex)
    else
        return nil
    end
end
----------------------------------------------------
-- Process one mesh node: find the geometrically connected components and tag them via a vertex attribute.
local function processMeshNode(script, node)

    local verticesPerPoly = node:getAttribute(octane.attributeId.A_VERTICES_PER_POLY)
    local polyVertexIndices = node:getAttribute(octane.attributeId.A_POLY_VERTEX_INDICES)

    if polyVertexIndices == nil then
        print("polyVertexIndices is nil")
        return
    end

    -- Get rid of the filename attribute,
    -- In the case of an .OBJ based mesh node, this prevents the node evaluation from importing back the OBJ file,
    -- which would overwrite our attributes.
    node:setAttribute(octane.attributeId.A_FILENAME, "")
    node:updateProperties({name=node.name.."_copy"})

    -- Loop over all polygons and group connected vertices together.
    -- At the end of this pass, ufParents contains the root element id for each vertex.
    ufParents = {}
    local index = 1
    for i=1, #verticesPerPoly do
        local count = verticesPerPoly[i]
        local baseIndex = polyVertexIndices[index] + 1
        for j=1, count-1 do
            local vertIndex = polyVertexIndices[index + j] + 1
            ufUnion(baseIndex, vertIndex)
        end

        index = index + count
    end

    -- Fill vertex attribute values.
    -- At this point the root ids are arbitrary as they depend on the order of vertices.
    -- We maintain a sparse map going from roots to more reasonable ids.
    local map = {}
    local attributeValues = {}
    local attributePolyIndices = {}
    index = 1
    for i=1, #verticesPerPoly do
        local count = verticesPerPoly[i]
        for j=0, count-1 do

            local vertexIndex = polyVertexIndices[index + j] + 1
            local root = ufParents[vertexIndex]

            if map[root] == nil then
                map[root] = #attributeValues + 1
                table.insert(attributeValues, #attributeValues + 1)
            end

            attributePolyIndices[index + j] = map[root] - 1
        end

        index = index + count
    end

    -- Optionnally normalize values. From [1..N] to [0..1].
    if script:getInputValue(inputLinkers.normalize) then
        for i=1, #attributeValues do
            attributeValues[i] = (i - 0.5) / #attributeValues
        end
    end

    -- Get the desired name for the new attribute.
    local name = script:getInputValue(inputLinkers.attributeName)
    if name == nil or name == "" then
        print("Error: Attribute name is empty.")
        return
    end

    -- Find an empty attribute slot.
    local attrIndex = findEmptyAttribute(node, "A_FLOAT_ATTRIBUTE_NAME", 4)
    if attrIndex == nil then
        print("Error: Could not find any empty float attribute.")
        return
    end

    -- Fill in values.
    node:setAttribute(octane.attributeId["A_FLOAT_ATTRIBUTE_NAME"..attrIndex], name, false)
    node:setAttribute(octane.attributeId["A_FLOAT_ATTRIBUTE"..attrIndex], attributeValues, false)
    node:setAttribute(octane.attributeId["A_FLOAT_ATTRIBUTE_INDICES"..attrIndex], attributePolyIndices, false)
    node:evaluate()
end

------------------------------------------------------------------------
-- Recursively process items.
local function processItem(script, node)

    if node == nil then
        return
    end

    if node.type == octane.nodeType.NT_GEO_MESH then
        processMeshNode(script, node)
    elseif node.isOutputLinker then
        -- In this case the downstream node was a nodegraph, a geo archive, a scripted graph or a wrapped scripted graph.
        -- We are now inside the graph, walk up and continue.
        processItem(script, node:getConnectedNode(octane.pinId.P_INPUT, true))
    else
        -- In this case the downstream node was some other geometry node.
        -- For example a Geometry group, a scatter node, a volume, a placement, etc.
        -- We'll loop over all the inputs and process them.
        -- There is a potential pitfall here. If we are inside a scripted graph, modifying the input
        -- will trigger `eval()` on the script we are in and the node we are looping over will be
        -- destroyed and recreated. This is the case with the Scatter tools for example.
        for i=1, node:getPinCount() do
            local pinInfo = node:getPinInfoIx(i)
            if pinInfo.type == octane.pinType.PT_GEOMETRY then
                processItem(script, node:getInputNodeIx(i, true))
            end
        end
    end
end

-----------------------------------------
-- Destroy an item and all its inputs recursively.
local function destroyItemTree(item)

    if item == nil then
        return
    end

    if item.isOutputLinker and item.graphOwned then
        item = item.graphOwner
    end

    if item.isGraph then
        local inputNodes = item:getInputNodes()
        for _, inputNode in ipairs(inputNodes) do
            local connectedNode = inputNode:getConnectedNode(octane.pinId.P_INPUT, false)
            destroyItemTree(connectedNode)
        end
    else
        for pinIx = 1, item:getPinCount() do
            local inputNode = item:getInputNodeIx(pinIx, false)
            destroyItemTree(inputNode)
        end
    end

    item:destroy()
end

-----------------------------------------
-- Destroy the internal copy of the input node and disconnect the output.
local function clean()
    outputLinkers.output:disconnect(octane.pinId.P_INPUT)
    destroyItemTree(nodeGeometry)
    nodeGeometry = nil
end

-----------------------------------------
-- Copy the input node into our own graph.
local function copyInput(graph, linker)

    local inputItem = linker:getConnectedNode(octane.pinId.P_INPUT, false)
    if inputItem == nil then
        return nil
    end

    if inputItem.isOutputLinker then
        inputItem = inputItem.graphOwner
    end

    local copy = graph:copyItemTree(inputItem)
    return copy
end

-------------------------------------
-- The input geometry was changed, re-process it.
local function updateInternals(script, graph)

    clean()
    nodeGeometry = copyInput(graph, inputLinkers.inputGeometry)
    if nodeGeometry == nil then
        return
    end

    -- Connect to our output.
    -- We do this first so we can use our output linker as the start of the recursion.
    -- Otherwise there is no way to know if the input node is a normal node or a wrapper node.
    if nodeGeometry.isGraph then
        local outputNode = nodeGeometry:findFirstOutputNode(octane.pinType.PT_GEOMETRY)
        if outputNode ~= nil then
            outputLinkers.output:connectTo(octane.pinId.P_INPUT, outputNode)
        end
    else
        outputLinkers.output:connectTo(octane.pinId.P_INPUT, nodeGeometry)
    end

    local item = outputLinkers.output:getConnectedNode(octane.pinId.P_INPUT, true)
    processItem(script, item)
end

-----------------------------------------------------
local function getInputLinkersInfo()

    local linkerInfoInputGeometry =
    {
        key = "inputGeometry",
        label = "Geometry",
        description = "Geometry that will be duplicated and partitioned into isolated elements.",
        type = octane.pinType.PT_GEOMETRY,
        defaultNodeType = octane.nodeType.NT_GEO_MESH,
    }

    local linkerInfoAttributeName =
    {
        key = "attributeName",
        label = "Attribute name",
        description = "Name of a float attribute that will be created and filled with element IDs.",
        type = octane.pinType.PT_STRING,
        defaultNodeType = octane.nodeType.NT_STRING,
        defaultValue = "elementid",
    }

    local linkerInfoNormalize =
    {
        key = "normalize",
        label = "Normalize",
        description = "Whether to normalize element IDs to be between 0 and 1 or keep integer IDs.",
        type = octane.pinType.PT_BOOL,
        defaultNodeType = octane.nodeType.NT_BOOL,
        defaultValue = false,
    }

    return { linkerInfoInputGeometry, linkerInfoAttributeName, linkerInfoNormalize }
end

local function getOutputLinkersInfo()

    local linkerInfoOutput =
    {
        key = "output",
        label = "Geometry out",
        type = octane.pinType.PT_GEOMETRY
    }

    return { linkerInfoOutput }
end

local function setInputLinkers(graph)

    local inputs = graph:setInputLinkers(inputLinkersInfo)

    for i, info in ipairs(inputLinkersInfo) do
        inputLinkers[info.key] = inputs[i]
    end
end

local function setOutputLinkers(graph)

    local outputs = graph:setOutputLinkers(outputLinkersInfo)

    for i, info in ipairs(outputLinkersInfo) do
        outputLinkers[info.key] = outputs[i]
    end
end

--------------------------------------------
local node = {}
node._name = "Partition elements"
node._icon = octane.image.fromBase64("iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAACF0lEQVQ4jZWSX0hTcRTHP3f3ZqlbLBnoXJhoEc1hy/XHdA+jP64Mi8hQkAoKgtbTiCh6C6KHQIhCe7MSeomIIDAMQqtpD2YJMrOgqczpNJ3byt1NN9eDOSy5G31ffvwO53zOOV+OSHplAQbAAYwAESCZoeYv5YuidL3+3AW/SlrnBLT/UywAu082nhkYC0aSZeY9nYDp3yRVGkAusL2+saHcPTmH5cCxGmAboAbEPw0QFYpFoDQnV33qUfvDvW/6h8kVRT5+cGmWor+0kkBZAsLAvJSm+/725y8vzslLhH0+pnvec/N09dGDdXZ7PL6YbDh79b4nFGlTWqHAtKuirrpynzQ2MY00PoJZF8D54JlQvrNUrDD8lF48vuPcsUnTpDiByWyxxpICfv8Mqlk/thN28HXBvJfZ3kH4PMWhLYWXlADzX4fcPZ8G3ceTkRhZKgFZjpIYcfHj7QC+vnHiagMxObIoKAA2AiXAEZ2+qNagK9gqBDz6d23NuO+2oNJtJhwKcf5V1zUlACxfYTagAcQaU/FopTGbueEE8ag883p04t73hcRTCaCpytjC8rmuqPVJ79BlYAEIAUKeegOHbXpuffHQ+c1bC0wCUyseOGpteVgtWlz9QTq6A46mKuNqIMWF6ynKz0GnyQHoW4mnTLRatKm3ozvA7Ssla3byeOU1sRTA1R9MTQBwo9mj5E3r6o+QwYOM+g29HLcZ+SC1fAAAAABJRU5ErkJggg==", 3)
node._evaluateAllChanges = true

-- Callback called after initializing the scripted graph with the code.
function node.onInit(self, graph)
    graph:setAttribute(octane.attributeId.A_COLOR, { 1.0, 0.5156, 0.8932 }, true)

    inputLinkersInfo = getInputLinkersInfo()
    setInputLinkers(graph)

    outputLinkersInfo = getOutputLinkersInfo()
    setOutputLinkers(graph)

    updateInternals(self, graph)

end

-- Callback called when the connected value of one of the input linker nodes changes.
function node.onEvaluate(self, graph)

    if self:inputWasChanged(inputLinkers.inputGeometry) or
      self:inputWasChanged(inputLinkers.attributeName) or
      self:inputWasChanged(inputLinkers.normalize) then
          updateInternals(self, graph)
    end

end

return node