User:Fleeting Frames/selectSquads

Save as selectSquads.lua into /hack/scripts folder. Call by typing selectSquads or selectSquads R in dfhack console while having base squad menu open.

--Utility script for quickly selecting metal weapon or bolts combatants.
--ver 1.00

if df.global.ui.main.mode ~= 1 or
    df.global.ui.squads.in_move_order or
    df.global.ui.squads.in_kill_order or
    df.global.ui.squads.in_kill_rect or
    df.global.ui.squads.sel_indiv_squad ~= -1
then qerror("Need to be in base squad select menu") end

local args = {...}

local prevSelectIndiv = df.global.ui.squads.in_select_indiv

df.global.ui.squads.in_select_indiv = true

local function repeatfunction(rfunc)
    --repeats a function every frame, until the function returns false or nil
    local myrfunc
    
    myrfunc = function ()
      timeoutid = dfhack.timeout(1, "frames", function(options)
        if rfunc() then myrfunc() end
       end)
    end
    
    myrfunc()
end

--AutomagicTables from lua wiki by ThomasWrensch and RiciLake, copied wholesale
do 
  local auto, assign

  function auto(tab, key)
    return setmetatable({}, {
            __index = auto,
            __newindex = assign,
            parent = tab,
            key = key
    })
  end

  local meta = {__index = auto}

  -- The if statement below prevents the table from being
  -- created if the value assigned is nil. This is, I think,
  -- technically correct but it might be desirable to use
  -- assignment to nil to force a table into existence.

  function assign(tab, key, val)
  if val ~= nil then
    local oldmt = getmetatable(tab)
    oldmt.parent[oldmt.key] = tab
    setmetatable(tab, meta)
    tab[key] = val
  end
  end

  function AutomagicTable()
    return setmetatable({}, meta)
  end
end

local metalt=AutomagicTable()

for i,squad in ipairs(df.global.ui.squads.list) do 
    for j, squaddie in ipairs(squad.positions) do 
        local notset = true 
        for k,item_id in pairs(squaddie.assigned_items) do 
            if notset and df.item.find(item_id)._type == df.item_weaponst then 
                if df.item.find(item_id).mat_type == 0 then 
                    --print(dfhack.TranslateName(df.historical_figure.find(squaddie.occupant).name))
                    metalt[i+1][j+1]=true
                    notset=false
                end
            end
        end
    end
end

local rangedt = AutomagicTable()

for i,squad in ipairs(df.global.ui.squads.list) do 
    for j, squaddie in ipairs(squad.positions) do 
        local notset = true
        local crossbow = false
        local metalbolts = false
        for k,item_id in pairs(squaddie.assigned_items) do 
            if notset and df.item.find(item_id)._type == df.item_weaponst and df.item.find(item_id).subtype.skill_ranged ~= -1 then
                crossbow = true
            end
            if notset and df.item.find(item_id)._type == df.item_quiverst then
                for m, ammunition in pairs(df.item.find(item_id).general_refs) do
                    if ammunition._type == df.general_ref_contains_itemst and df.item.find(ammunition.item_id).mat_type == 0 then 
                        metalbolts = true
                    end
                end
            end
            if notset and crossbow and metalbolts then 
                rangedt[i+1][j+1]=true
                --print(dfhack.TranslateName(df.historical_figure.find(squaddie.occupant).name))
                notset = false
                crossbow = false
                metalbolts = false
            end
        end
    end
end

local indexkey = {
'OPTION1',
'OPTION2',
'OPTION3',
'OPTION4',
'OPTION5',
'OPTION6',
'OPTION7',
'OPTION8',
'OPTION9',
'OPTION10'
}

local unitindexkey = {
'SEC_OPTION1',
'SEC_OPTION2',
'SEC_OPTION3',
'SEC_OPTION4',
'SEC_OPTION5',
'SEC_OPTION6',
'SEC_OPTION7',
'SEC_OPTION8',
'SEC_OPTION9',
'SEC_OPTION10'
}

local inputlist, tin = {}, table.insert

local targett=metalt
if args and #args>0 and args[1]=="R" then targett=rangedt end

for index,squad in pairs(targett) do
    tin(inputlist,{[indexkey[index]]=true})
    for unitindex,_ in pairs(squad) do
        tin(inputlist, {[unitindexkey[unitindex]]=true})
    end
    tin(inputlist,{['LEAVESCREEN']=true})
end

local guiSim = curry((require 'gui').simulateInput, df.global.gview.view.child)

dfhack.timeout((1+#inputlist), "frames", function() df.global.ui.squads.in_select_indiv = prevSelectIndiv end)

repeatfunction(function() local tr=table.remove(inputlist,1) if tr then guiSim(tr) return true end end)