New Feature: autotraining

autotraining.lua

@@ -0,0 +1,296 @@
+-- Based on the original code by RNGStrategist (who also got some help from Uncle Danny)
+--@ enable = true
+--@ module = true
+
+local repeatUtil = require('repeat-util')
+local utils=require('utils')
+
+local GLOBAL_KEY  = "autotraining"
+local MartialTraining = df.need_type['MartialTraining']
+local ignore_count = 0
+
+local function get_default_state()
+    return {
+        enabled=false,
+        threshold=-5000,
+        ignored={},
+        ignored_nobles={},
+        training_squads = {},
+    }
+end
+
+state = state or get_default_state()
+
+function isEnabled()
+    return state.enabled
+end
+
+-- persisting a table with numeric keys results in a json array with a huge number of null entries
+-- therefore, we convert the keys to strings for persistence
+local function to_persist(persistable)
+    local persistable_ignored = {}
+    for k, v in pairs(persistable) do
+        persistable_ignored[tostring(k)] = v
+    end
+    return persistable_ignored
+end
+
+-- loads both from the older array format and the new string table format
+local function from_persist(persistable)
+    if not persistable then
+        return
+    end
+    local ret = {}
+    for k, v in pairs(persistable) do
+        ret[tonumber(k)] = v
+    end
+    return ret
+end
+
+function persist_state()
+    dfhack.persistent.saveSiteData(GLOBAL_KEY, {
+        enabled=state.enabled,
+        threshold=state.threshold,
+        ignored=to_persist(state.ignored),
+        ignored_nobles=state.ignored_nobles,
+        training_squads=to_persist(state.training_squads)
+    })
+end
+
+--- Load the saved state of the script
+local function load_state()
+    -- load persistent data
+    local persisted_data = dfhack.persistent.getSiteData(GLOBAL_KEY, {})
+    state.enabled = persisted_data.enabled or state.enabled
+    state.threshold = persisted_data.threshold or state.threshold
+    state.ignored = from_persist(persisted_data.ignored) or state.ignored
+    state.ignored_nobles = persisted_data.ignored_nobles or state.ignored_nobles
+    state.training_squads = from_persist(persisted_data.training_squads) or state.training_squads
+    return state
+end
+
+dfhack.onStateChange[GLOBAL_KEY] = function(sc)
+    if sc == SC_MAP_UNLOADED then
+        state.enabled = false
+        return
+    end
+    -- the state changed, is a map loaded and is that map in fort mode?
+    if sc ~= SC_MAP_LOADED or df.global.gamemode ~= df.game_mode.DWARF then
+        -- no its isnt, so bail
+        return
+    end
+    -- yes it was, so:
+
+    -- retrieve state saved in game. merge with default state so config
+    -- saved from previous versions can pick up newer defaults.
+    load_state()
+    if state.enabled then
+        start()
+    end
+    persist_state()
+end
+
+
+--######
+--Functions
+--######
+local function isIgnoredNoble(unit)
+    local noblePos = dfhack.units.getNoblePositions(unit)
+    if noblePos ~= nil then
+        for _, position in ipairs(noblePos) do
+            if state.ignored_nobles[position.position.code] then
+                return true
+            end
+        end
+    end
+    return false
+end
+
+---@return table<integer, { ['unit']: df.unit, ['need']: integer }>
+function getTrainingCandidates()
+    local ret = {}
+    ignore_count = 0
+    for _, unit in ipairs(dfhack.units.getCitizens(true)) do
+        if not dfhack.units.isAdult(unit) then
+            goto next_unit
+        end
+        local need = getTrainingNeed(unit)
+        if not need or need.focus_level >= state.threshold  then
+            goto next_unit
+        end
+        -- ignored units are those that would like to train but are forbidden from doing so
+        if state.ignored[unit.id] then
+            ignore_count = ignore_count + 1
+            goto next_unit
+        end
+        if isIgnoredNoble(unit) then
+            ignore_count = ignore_count + 1
+            goto next_unit
+        end
+        if unit.military.squad_id ~= -1 then
+            goto next_unit
+        end
+        table.insert(ret, { unit = unit, need = need.focus_level })
+        ::next_unit::
+    end
+    table.sort(ret, function (a, b) return a.need < b.need end)
+    return ret
+end
+
+function getTrainingSquads()
+    local squads = {}
+    for squad_id, _ in pairs(state.training_squads) do
+        local squad = df.squad.find(squad_id)
+        if squad then
+            table.insert(squads, squad)
+        else
+            -- setting to nil during iteration is permitted by lua
+            state.training_squads[squad_id] = nil
+        end
+    end
+    return squads
+end
+
+function getTrainingNeed(unit)
+    if unit == nil then return nil end
+    local needs =  unit.status.current_soul.personality.needs
+    for _, need in ipairs(needs) do
+        if need.id == MartialTraining then
+            return need
+        end
+    end
+    return nil
+end
+
+--######
+--Main
+--######
+
+-- Find all training squads
+-- Abort if no squads found
+function checkSquads()
+    local squads = {}
+    for _, squad in ipairs(getTrainingSquads()) do
+        if squad.entity_id == df.global.plotinfo.group_id then
+            local leader = squad.positions[0].occupant
+            if leader ~= -1 then
+                table.insert(squads,squad)
+            end
+        end
+    end
+
+    if #squads == 0 then
+        return nil
+    end
+
+    return squads
+end
+
+function addTraining(unit,good_squads)
+    if unit.military.squad_id ~= -1 then
+        for _, squad in ipairs(good_squads) do
+            if unit.military.squad_id == squad.id then
+                return true
+            end
+        end
+        return false
+    end
+    for _, squad in ipairs(good_squads) do
+        for i=1,9,1   do
+            if squad.positions[i].occupant  == -1 then
+                return dfhack.military.addToSquad(unit.id,squad.id,i)
+            end
+        end
+    end
+
+    return false
+end
+
+function removeAll()
+    if state.training_squads == nil then return end
+    for _, squad in ipairs(getTrainingSquads()) do
+        for i=1,9,1 do
+            local hf = df.historical_figure.find(squad.positions[i].occupant)
+            if hf ~= nil then
+                dfhack.military.removeFromSquad(hf.unit_id)
+            end
+        end
+    end
+end
+
+
+function check()
+    local squads = checkSquads()
+    local intraining_count = 0
+    local inque_count = 0
+    if squads == nil then return end
+    for _,squad in ipairs(squads) do
+        for i=1,9,1   do
+            if squad.positions[i].occupant  ~= -1 then
+                local hf = df.historical_figure.find(squad.positions[i].occupant)
+                if hf ~= nil then
+                    local unit = df.unit.find(hf.unit_id)
+                    local training_need = getTrainingNeed(unit)
+                    if not training_need or training_need.focus_level >= state.threshold then
+                        dfhack.military.removeFromSquad(unit.id)
+                    end
+                end
+            end
+        end
+    end
+    for _, p in ipairs(getTrainingCandidates()) do
+        local added = addTraining(p.unit, squads)
+        if added then
+            intraining_count = intraining_count +1
+        else
+            inque_count = inque_count +1
+        end
+    end
+    print(("%s: %d training, %d waiting, and %d excluded units with training needs"):
+        format(GLOBAL_KEY, intraining_count, inque_count, ignore_count))
+end
+
+function start()
+    repeatUtil.scheduleEvery(GLOBAL_KEY, 1, 'days', check)
+end
+
+function stop()
+    repeatUtil.cancel(GLOBAL_KEY)
+end
+
+function enable()
+    state.enabled = true
+    persist_state()
+    start()
+end
+
+function disable()
+    state.enabled = false
+    persist_state()
+    stop()
+    removeAll()
+end
+
+if dfhack_flags.module then
+    return
+end
+
+validArgs = utils.invert({
+    't'
+})
+
+local args = utils.processArgs({...}, validArgs)
+
+if dfhack_flags.enable then
+    if dfhack_flags.enable_state then
+        enable()
+    else
+        disable()
+    end
+else
+    -- called on the command-line
+    if args.t then
+        state.threshold = 0-tonumber(args.t)
+    end
+    print(("autotraining is %s"):format(state.enabled and "enabled" or "disabled"))
+end

changelog.txt

@@ -27,6 +27,8 @@ Template for new versions:
 # Future
 
 ## New Tools
+- `autotraining`: new tool to assign citizens to a military squad when they need Martial Training
+- `gui/autotraining`: configuration tool for autotraining
 
 ## New Features
 

docs/autotraining.rst

@@ -0,0 +1,41 @@
+autotraining
+============
+
+.. dfhack-tool::
+    :summary: Assigns citizens to a military squad until they have fulfilled their need for Martial Training
+    :tags: fort auto bugfix units
+
+This script automatically assigns citizens with the need for military training to designated training squads.
+
+You need to have at least one squad that is set up for training. The squad should be set to "Constant Training" in the military screen. The squad doesn't need months off. The members leave the squad once they have satisfied their need for military training.
+
+The configured uniform determines the skills that are acquired by the training dwarves. Providing "No Uniform" is a perfectly valid choice and will make your militarily inclined civilians become wrestlers over time. However, you can also provide weapons and armor to pre-train civilians for future drafts.
+
+Once you have made squads for training use `gui/autotraining` to select the squads and ignored units, as well as the needs threshhold.
+
+Usage
+-----
+
+    ``autotraining [<options>]``
+
+Examples
+--------
+
+``autotraining``
+    Current status of script
+
+``enable autotraining``
+    Checks to see if you have fullfilled the creation of a training squad.
+    If there is no squad marked for training use, a clickable notification will appear letting you know to set one up/
+    Searches your fort for dwarves with a need for military training, and begins assigning them to a training squad.
+    Once they have fulfilled their need they will be removed from their squad to be replaced by the next dwarf in the list.
+
+``disable autotraining``
+    Stops adding new units to the squad.
+
+Options
+-------
+    ``-t``
+        Use integer values. (Default 5000)
+        The negative need threshhold to trigger for each citizen
+        The greater the number the longer before a dwarf is added to the waiting list.

docs/gui/autotraining.rst

@@ -0,0 +1,15 @@
+gui/autotraining
+================
+
+.. dfhack-tool::
+    :summary: GUI interface for ``autotraining``
+    :tags: fort auto interface
+
+This is an in-game configuration interface for `autotraining`. You can pick squads for training, select ignored units, and set the needs threshold.
+
+Usage
+-----
+
+::
+
+    gui/autotraining