Quick Assign module.ttslua

-- ~~~~~~
-- Script by dzikakulka
-- Issues, history at: https://github.com/tjakubo2/TTS_xwing
-- ~~~~~~

-- TODO Use zone.getValue/setValue

#include TTS_lib/Vector/Vector

#include TTS_lib/Util/Util

#include TTS_lib/ObjType/ObjType


ObjType.AddType('ship', function(obj)
    return (obj.tag == 'Figurine')
end)
ObjType.AddType('token', function(obj)
    return (obj.tag == 'Chip' or (obj.getVar('__XW_Token') and obj.getVar('__XW_TokenIdle')))
end)
ObjType.AddType('obj', function(obj)
    return ( (obj.tag ~= 'Fog') and obj.interactable )
end)

Zone = {}

-- Return zone area, zero if size below some lower limit
Zone.Area = function(obj)
    if obj == nil then return 0 end
    local scale = obj.getScale()
    if scale[1] < 10 or scale[3] < 5 then return 0 end
    return scale[1]*scale[3]
end

-- Get all objects clipping a zone
Zone.GetObjects = function(zone, type)
    type = type or 'obj'
    return ObjType.GetBoxOfType(zone.getPosition(), zone.getRotation(), zone.getScale(), type)
end

-- Return true if the zone is empty
Zone.IsEmpty = function(zone)
    return (Zone.GetObjects(zone)[1] == nil)
end

-- Get local position from a zone (scale invariant)
Zone.LocalPos = function(zone, pos)
    local posAdj = Vect.ScaleEach(pos, Vect.Inverse(zone.getScale()))
    return zone.positionToWorld(posAdj)
end

-- Save current zone position and scale
-- Return true if saved, false if there is saved data already
Zone.StoreSetting = function(zone)
    if zone.getTable('extendData') then
        return false
    end
    local extData = {
        oldPos = zone.getPosition(),
        oldScale = zone.getScale()
    }
    zone.setTable('extendData', extData)
    return true
end

-- Restore saved zone position and scale
-- Return true if success, false if no data
Zone.RestoreSetting = function(zone)
    if not zone.getTable('extendData') then
        return false
    end
    local extData = zone.getTable('extendData')
    zone.setPosition(extData.oldPos)
    zone.setScale(extData.oldScale)
    zone.setVar('extendData', nil)
end

-- Target zone height per row of dials
Zone.heightPerRow = 12
-- Save current zone scale and position and adjust it to fit rowNum dial rows
Zone.AdjustHeight = function(zone, rowNum)
    Zone.RestoreSetting(zone)
    Zone.StoreSetting(zone)
    local currHeight = zone.getScale()[3]
    local targetHeight = Zone.heightPerRow*rowNum
    local offset = (targetHeight - currHeight)/2
    zone.setPosition(Zone.LocalPos(zone, {0, 0, -1*offset}))
    local zScale = zone.getScale()
    zone.setScale({zScale[1], zScale[2], targetHeight})
end

-- Pull out an extension tray for zone
Zone.Extend = function(extension)
    extension.call('extend')
end

-- Retract an extension tray for zone
Zone.Retract = function(extension)
    extension.call('retract')
end


AssignModule = {}

AssignModule.maxShipCount = 12
AssignModule.msg = {}
AssignModule.msg.emptyZoneAndSelection = 'Quick Assign: Place ships you want to assign dials for in your hidden zone or select them first'
AssignModule.msg.notShipsOnly = 'Quick Assign: Remove objects that are not ship models from your hidden zone and try again'
AssignModule.msg.tooManyShips = 'Quick Assign: Too many ship models (max ' .. AssignModule.maxShipCount ..')'
AssignModule.msg.unrecognizedShips = 'Quick Assign: Some ship models in your zone have not been recognized. Make sure your models are sourced from squad building tray on this table and contact author if this issue persists'
AssignModule.msg.clickTooFast = 'Quick Assign: Clicking too fast, please wait until assignment finishes'
--AssignModule.msg.zoneTooSmall = 'Quick Assign: Too little horizontal space, please widen your zone (Gizmo tool -> X scale)'
AssignModule.msg.zoneTooSmall = 'Quick Assign: Too many ships!'

-- Sort ships by their positions in the zone, left to right
AssignModule.SortShips = function(shipData, zone)
    local lPosData = {}
    for k,sData in pairs(shipData) do
        local lPos = Vect.Sub(sData.shipRef.getPosition(), zone.getPosition())
        lPos = Vect.RotateDeg(lPos, -1*zone.getRotation()[2])[1]
        table.insert(lPosData, {data = sData, lPos = lPos})
    end
    local sortedData = {}
    repeat
        local lMax = {ind=1, val=lPosData[1].lPos}
        for k,data in pairs(lPosData) do
            if data.lPos < lMax.val then
                lMax.val = data.lPos
                lMax.ind = k
            end
        end
        table.insert(sortedData, lPosData[lMax.ind].data)
        table.remove(lPosData, lMax.ind)
    until lPosData[1] == nil
    return sortedData
end

-- Get ships selected by player with some color
AssignModule.ShipsFromSelection = function(playerColor)
    if not Player[playerColor].seated then
        return {}
    end
    local selectedShips = {}
    for k,obj in pairs(Player[playerColor].getSelectedObjects()) do
        if obj.tag == 'Figurine' then
            table.insert(selectedShips, obj)
        end
    end
    return selectedShips
end

function API_AssignDials(args)
    return AssignModule.Assign(table.unpack(args))
end

-- Check zone contents and assign dials if everything is OK
AssignModule.Assign = function(playerColor, zone, buttonPos, extension, playerStuff)
    print('A: ' .. tostring(playerColor) .. ' : ' .. tostring(zone) .. ' : ' .. tostring(buttonPos) .. ' : ' .. tostring(extension) .. ' : ' .. tostring(playerStuff))

    -- If assignment is in progress, notify and return
    if zone.getVar('assignBusy') == true then
        broadcastToColor(AssignModule.msg.clickTooFast, playerColor, {1, 0.45, 0})
        return
    end

    -- Get objects inside the hidden zone or selected if zone empty, notify and return if empty
    local assignObjects = Zone.GetObjects(zone)
    if table.empty(assignObjects) then
        assignObjects = AssignModule.ShipsFromSelection(playerColor)
        if table.empty(assignObjects) then
            broadcastToColor(AssignModule.msg.emptyZoneAndSelection, playerColor, {1, 0.45, 0})
            return
        end
    end

    -- Assert object type validity and count
    local shipsOnly = true
    local tokenObjects = {}
    local shipObjects = {}
    for k,obj in pairs(assignObjects) do
        if ObjType.IsOfType(obj, 'token') then
            table.insert(tokenObjects, obj)
        elseif ObjType.IsOfType(obj, 'ship') then
            table.insert(shipObjects, obj)
        else
            print(obj.getName() .. ' : ' .. obj.tag)
            shipsOnly = false
            break
        end
    end
    if not shipsOnly then
        broadcastToColor(AssignModule.msg.notShipsOnly, playerColor, {1, 0.45, 0})
        return
    end
    --[[if #shipObjects > AssignModule.maxShipCount then
        broadcastToColor(AssignModule.msg.tooManyShips, playerColor, {1, 0.45, 0})
        return
    end]]--

    -- Get and assert ship type data
    local typesOK = true
    local shipData = {}
    -- shipData => Key: numerical, Entry: { shipRef = shipRef, shipType = shipType }
    local bagData = {}
    -- bagData => Key: shipType, Entry: { bagRef = bagRef }
    for k,ship in pairs(shipObjects) do
        local shipType = Global.call('API_ModelDB_GetData', {ship}).type
        if not shipType then
            typesOK = false
            break
        elseif bagData[shipType] == nil then
            bagData[shipType] = {}
        end
        table.insert(shipData, {shipRef=ship, shipType=shipType})
    end
    if not typesOK then
        broadcastToColor(AssignModule.msg.unrecognizedShips, playerColor, {1, 0.45, 0})
        return
    end

    -- Block the zone
    zone.setVar('assignBusy', false)

    -- Sort the ships to match their order
    shipData = AssignModule.SortShips(shipData, zone)

    for k,data in pairs(shipData) do
        data.shipRef = Global.call('API_ModelDB_SetCollider', {data.shipRef, 'full'}) or data.shipRef
    end

    -- Dial bags desired position and its step
    local bagPos = Vect.Sum(Vect.RotateDeg({0, 1.5, 4}, zone.getRotation()[2]), buttonPos)
    local bagStep = Vect.RotateDeg({0, 0, 2.6}, zone.getRotation()[2])

    -- If there are valid dial bags on the table, take them
    playerStuff = playerStuff or {}
    playerStuff = table.join(playerStuff, ObjType.GetBoxOfType({-71, 0, 0}, {0, 0, 0}, {28, 5, 54}, 'obj'))
    for k,obj in pairs(playerStuff) do
        if AssignModule.IsDialBag(obj) then
            local type = AssignModule.DialBagType(obj)
            if bagData[type] ~= nil and bagData[type].bagRef == nil then
                bagData[type].bagRef = obj
            end
        end
    end

    -- Destroy dial bags obstructing target position for current set
    local dialNum = table.size(bagData)
    local angle = Vect.AngleDeg2D(Vect.Between(bagPos, Vect.Sum(bagPos, bagStep)))

    local len = Vect.Length(Vect.Scale(bagStep, dialNum-1))
    local currentDials = ObjType.GetBoxOfType(Vect.Middle(bagPos, Vect.Sum(bagPos, Vect.Scale(bagStep, dialNum-1))), {0, angle, 0}, {len+3, 1, 3}, 'obj')
    for k,obj in pairs(currentDials) do
        if AssignModule.IsDialBag(obj) then
            local type = AssignModule.DialBagType(obj)
            if bagData[type] == nil or ( bagData[type] ~= nil and bagData[type].bagRef ~= obj ) then
                obj.destruct()
            end
        end
    end

    -- If some bags were not on the table, grab them from squad builder module
    local bagsToGrab = {}
    local grabbed = nil
    for type,data in pairs(bagData) do
        if data.bagRef == nil then
            table.insert(bagsToGrab, type)
        end
    end
    if bagsToGrab[1] ~= nil then
        grabbed = AssignModule.GrabDialBags(bagsToGrab)
        for type,data in pairs(bagData) do
            if data.bagRef == nil then
                data.bagRef = grabbed[type]
            end
        end
    end

    -- Move the dial bags to desired position
    for type,data in pairs(bagData) do
        data.bagRef.setPositionSmooth(bagPos)
        data.bagRef.unlock()
        bagPos = Vect.Sum(bagPos, bagStep)
    end

    local scaleTable = {0.625, 0.525, 0.45}    -- Possible dial scales
    local zoneWidth = zone.getScale()[1]
    local dialWidth = 3.1
    local setWidth = dialWidth*5
    local setHeight = dialWidth*5
    local minSlack = 0.3

    -- Get row scale for some dial amount
    local getRowScale = function(shipNum)
        for k,cScale in ipairs(scaleTable) do
            local scaleSlack = (zoneWidth - (setWidth*cScale*shipNum))/(shipNum+1)
            if scaleSlack > (0.25*dialWidth*cScale) then
                return cScale, scaleSlack
            end
        end
        return nil
    end

    -- Split dial set between rows evenly
    local evenRowSplit = function(shipNum, maxRows)
        local split = {}
        local tabSum = function(t)
            local sum = 0
            for _,v in pairs(t) do sum = sum + v end
            return sum
        end
        local evenNum = math.floor(shipNum/maxRows)
        for i=1,maxRows do
            split[i] = evenNum
        end
        local i = 1
        while tabSum(split) < shipNum do
            split[i] = split[i]+1
            i = i+1
        end
        return split
    end

    local shipRows = {}
    local rowScale = {}
    local maxRows = 1
    if extension then
        maxRows = 2
    end

    -- If more than 5 ships split evenly, else cram in one row
    if #shipData > 5 then
        local rowDist = evenRowSplit(#shipData, maxRows)
        local shipNum = 1
        for rowNum, shipCount in pairs(rowDist) do
            shipRows[rowNum] = {}
            for k=1,shipCount do
                table.insert(shipRows[rowNum], shipData[shipNum])
                shipNum = shipNum+1
            end
        end
    else
        shipRows[1] = shipData
    end

    -- Get the scale and leftover space for each row
    -- Assert possible scale (leftover space being above some treshold)
    for row in pairs(shipRows) do
        shipRows[row].scale, shipRows[row].widthSlack = getRowScale(#shipRows[row])
        if shipRows[row].scale == nil then
            broadcastToColor(AssignModule.msg.zoneTooSmall, playerColor, {1, 0.45, 0})
            zone.setVar('assignBusy', false)
            return
        end
    end

    -- Adjust zone height to set amount of dial set rows, use extensions if neccesary
    --Zone.AdjustHeight(zone, #shipRows)
    local zoneHeight = zone.getScale()[3]
    if extension and #shipRows > 1 then
        Zone.Extend(extension)
    elseif extension then
        Zone.Retract(extension)
    end

    -- Target (after assign) ship position and its step
    local shipSpacing = zoneWidth/(#shipData + 1)
    local shipPos = {
        -1*0.5*zoneWidth + shipSpacing,
        0,
        0.5*zoneHeight + 1.5
    }
    local ShipStep = function(pos)
        return Vect.Sum(pos, {shipSpacing, 0, 0})
    end

    for rowNum,rowData in pairs(shipRows) do
        -- Target dial set position in a row and its step
        local setWidthScaled = setWidth*rowData.scale
        local setPos = {
            -1*0.5*zoneWidth + rowData.widthSlack + 0.5*setWidthScaled,
            0,
            0.5*zoneHeight - (0.5*Zone.heightPerRow) - ((rowNum-1)*Zone.heightPerRow)
        }
        local setRot = Vect.Sum(zone.getRotation(), {0, 180, 0})
        local SetStep = function(pos)
            return Vect.Sum(pos, {rowData.widthSlack + setWidthScaled, 0, 0})
        end

        local shipsInZone = Zone.GetObjects(zone, 'ship')
        -- Assignment for each ship in a row
        for shipNum, shipData in ipairs(rowData) do
            -- Take a dial set, layout, assign
            local setDest = Zone.LocalPos(zone, setPos)
            local newStack = bagData[shipData.shipType].bagRef.takeObject({position=setDest, rotation=setRot, callback='expandSet', callback_owner=self, params={ship = shipData.shipRef, scale = {rowData.scale, 1, rowData.scale}, unlock=zone}})
            newStack.setLuaScript(scriptSource.getLuaScript())
            setPos = SetStep(setPos)
            -- Move the ship with its tokens
            if table.find(shipsInZone, shipData.shipRef) then
                local shipDest = Zone.LocalPos(zone, Vect.Sum(shipPos, {0, shipData.shipRef.getPosition()[2], 0}))
                shipData.shipRef.setPositionSmooth(shipDest, false, true)
                shipData.shipRef.setRotationSmooth(Vect.Sum(zone.getRotation(), {0, 180, 0}), false, true)
                shipData.shipRef.highlightOn({0, 1, 0}, 1)
                Global.call('API_QueueShipTokensMove', {ship = shipData.shipRef, noLock = true})
            end
            shipPos = ShipStep(shipPos)
        end
    end

    -- Disable zone mutex
    zone.setVar('assignBusy', false)
    return
end

-- Expand dial set callback
function expandSet(obj, params)
    obj.setScale(params.scale)
    obj.call('init', {ship=params.ship})
    params.unlock.setVar('assignBusy', false)
end

-- Signum
function math.sgn(arg)
    if arg < 0 then
        return -1
    elseif arg > 0 then
        return 1
    end
    return 0
end

-- Is the object a valid dial bag
AssignModule.IsDialBag = function(obj)
    return ( obj.tag == 'Infinite' and obj.getName():find(' Dials') ~= nil )
end

-- Get the ship type that dial bag contains dials for
AssignModule.DialBagType = function(obj)
    return obj.getName():sub(1, obj.getName():find(' Dials')-1)
end

-- Grab dial bags for specified types from spawner and return them
AssignModule.GrabDialBags = function(shipTypesTable)
    -- Get spawner bags template
    local spawnerBagsTemp = nil
    for k,obj in pairs(getAllObjects()) do
        if obj.getName() == 'TEMPLATE{Spawner Bags}' then
            spawnerBagsTemp = obj
        end
    end

    -- CLone it and take the accesories bag (with dials)
    local sPos = self.getPosition()
    local bagClone = spawnerBagsTemp.clone({position = {sPos[1], sPos[2]-1, sPos[3]}})
    local accBagGUID = nil
    for k,data in pairs(bagClone.getObjects()) do
        if data.name == 'Accesories Bag' then
            accBagGUID = data.guid
            break
        end
    end
    local accBag = bagClone.takeObject({guid=accBagGUID, position={sPos[1]+3, sPos[2], sPos[3]}})
    accBag.lock()
    accBag.tooltip = false
    accBag.interactable = false

    -- Grab appropriate dials bags
    local dialBagsInfo = {}
    for k,data in pairs(accBag.getObjects()) do
        if data.name:find('Dials') ~= nil then
            dialBagsInfo[data.name] = data.guid
        end
    end
    local grabbed = {}
    for k,type in pairs(shipTypesTable) do
        grabbed[type] = accBag.takeObject({guid = dialBagsInfo[type .. ' Dials'], position = {sPos[1]-3, sPos[2], sPos[3]}})
        grabbed[type].lock()
    end

    -- Cleaup
    bagClone.destruct()
    accBag.destruct()
    return grabbed
end

function onLoad(save_state)

    -- Get the dial set stack script
    for k,obj in pairs(getAllObjects()) do
        if obj.getName() == 'Dial Set source' then
            scriptSource = obj
        end
    end

    self.setPosition({0, -3, 0})
    self.setRotation({0, 0, 0})
    self.setScale({1, 1, 1})
    self.interactable = false
    self.tooltip = false
    --Zone.UpdateDialZones()
    --AssignModule.UpdateButtons()
end