diff --git a/autotraining.lua b/autotraining.lua new file mode 100644 index 000000000..9e936492c --- /dev/null +++ b/autotraining.lua @@ -0,0 +1,280 @@ +-- 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') + +validArgs = utils.invert({ + 't' +}) + +local args = utils.processArgs({...}, validArgs) +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 +--###### +function getTrainingCandidates() + local ret = {} + ignore_count = 0 + for _, unit in ipairs(dfhack.units.getCitizens(true)) do + if state.ignored[unit.id] then + ignore_count = ignore_count +1 + goto next_unit + end + 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 + local noblePos = dfhack.units.getNoblePositions(unit) + local isIgnNoble = false + if noblePos ~=nil then + for _, position in ipairs(noblePos) do + if state.ignored_nobles[position.position.code] then + isIgnNoble = true + break + end + end + end + if isIgnNoble then + ignore_count = ignore_count +1 + goto next_unit + end + table.insert(ret, unit) + ::next_unit:: + 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) + end + end + end + end + end + for _, unit in ipairs(getTrainingCandidates()) do + local added = addTraining(unit, squads) + if added then + intraining_count = intraining_count +1 + else + inque_count = inque_count +1 + end + end + + dfhack.println(GLOBAL_KEY .. " | IGNORED: " .. ignore_count .. " TRAINING: " .. intraining_count .. " QUEUE: " ..inque_count ) +end + +function start() + if args.t then + state.threshold = 0-tonumber(args.t) + end + repeatUtil.scheduleEvery(GLOBAL_KEY, 1, 'days', check) +end + +function stop() + repeatUtil.cancel(GLOBAL_KEY) +end + +if dfhack_flags.enable then + if dfhack_flags.enable_state then + state.enabled = true + else + state.enabled = false + end + persist_state() +end + +if dfhack_flags.module then + return +end + +if state.enabled then + start() +else + stop() + removeAll() +end +persist_state() diff --git a/docs/autotraining.rst b/docs/autotraining.rst new file mode 100644 index 000000000..360f4d08e --- /dev/null +++ b/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 []`` + +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. diff --git a/docs/gui/autotraining.rst b/docs/gui/autotraining.rst new file mode 100644 index 000000000..a86b28adf --- /dev/null +++ b/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 diff --git a/gui/autotraining.lua b/gui/autotraining.lua new file mode 100644 index 000000000..03ac18326 --- /dev/null +++ b/gui/autotraining.lua @@ -0,0 +1,245 @@ +---@diagnostic disable: missing-fields + +local gui = require('gui') +local widgets = require('gui.widgets') + +local autotraining = reqscript('autotraining') + +local training_squads = autotraining.state.training_squads +local ignored_units = autotraining.state.ignored +local ignored_nobles = autotraining.state.ignored_nobles + +AutoTrain = defclass(AutoTrain, widgets.Window) +AutoTrain.ATTRS { + frame_title='Training Setup', + frame={w=55, h=45}, + resizable=true, -- if resizing makes sense for your dialog + resize_min={w=55, h=20}, -- try to allow users to shrink your windows +} + +local SELECTED_ICON = dfhack.pen.parse{ch=string.char(251), fg=COLOR_LIGHTGREEN} +function AutoTrain:getSquadIcon(squad_id) + if training_squads[squad_id] then + return SELECTED_ICON + end + return nil +end + +function AutoTrain:getSquads() + local squads = {} + for _, squad in ipairs(df.global.world.squads.all) do + if not (squad.entity_id == df.global.plotinfo.group_id) then + goto continue + end + table.insert(squads, { + text = dfhack.translation.translateName(squad.name, true)..' ('..squad.alias..')', + icon = self:callback("getSquadIcon", squad.id ), + id = squad.id + }) + + ::continue:: + end + return squads +end + +function AutoTrain:toggleSquad(_, choice) + training_squads[choice.id] = not training_squads[choice.id] + autotraining.persist_state() + self:updateLayout() +end + +local IGNORED_ICON = dfhack.pen.parse{ch='x', fg=COLOR_RED} +function AutoTrain:getUnitIcon(unit_id) + if ignored_units[unit_id] then + return IGNORED_ICON + end + return nil +end + +function AutoTrain:getNobleIcon(noble_code) + if ignored_nobles[noble_code] then + return IGNORED_ICON + end + return nil +end + +function AutoTrain:getUnits() + local unit_choices = {} + for _, unit in ipairs(dfhack.units.getCitizens(true,false)) do + if not dfhack.units.isAdult(unit) then + goto continue + end + + table.insert(unit_choices, { + text = dfhack.units.getReadableName(unit), + icon = self:callback("getUnitIcon", unit.id ), + id = unit.id + }) + ::continue:: + end + return unit_choices +end + +function AutoTrain:toggleUnit(_, choice) + ignored_units[choice.id] = not ignored_units[choice.id] + autotraining.persist_state() + self:updateLayout() +end + +local function to_title_case(str) + return dfhack.capitalizeStringWords(dfhack.lowerCp437(str:gsub('_', ' '))) +end + +function toSet(list) + local set = {} + for _, v in ipairs(list) do + set[v] = true + end + return set +end + +local function add_positions(positions, entity) + if not entity then return end + for _,position in pairs(entity.positions.own) do + positions[position.id] = { + id=position.id+1, + code=position.code, + } + end +end + +function AutoTrain:getPositions() + local positions = {} + local excludedPositions = toSet({ + 'MILITIA_CAPTAIN', + 'MILITIA_COMMANDER', + 'OUTPOST_LIAISON', + 'CAPTAIN_OF_THE_GUARD', + }) + + add_positions(positions, df.historical_entity.find(df.global.plotinfo.civ_id)) + add_positions(positions, df.historical_entity.find(df.global.plotinfo.group_id)) + + -- Step 1: Extract values into a sortable array + local sortedPositions = {} + for _, val in pairs(positions) do + if val and not excludedPositions[val.code] then + table.insert(sortedPositions, val) + end + end + + -- Step 2: Sort the positions (optional, adjust sorting criteria) + table.sort(sortedPositions, function(a, b) + return a.id < b.id -- Sort alphabetically by code + end) + + -- Step 3: Rebuild the table without gaps + positions = {} -- Reset positions table + for i, val in ipairs(sortedPositions) do + positions[i] = { + text = to_title_case(val.code), + value = val.code, + pen = COLOR_LIGHTCYAN, + icon = self:callback("getNobleIcon", val.code), + id = val.id + } + end + + return positions +end + + + +function AutoTrain:toggleNoble(_, choice) + ignored_nobles[choice.value] = not ignored_nobles[choice.value] + autotraining.persist_state() + self:updateLayout() +end + +function AutoTrain:init() + self:addviews{ + widgets.Label{ + frame={ t = 0 , h = 1 }, + text = "Select squads for automatic training:", + }, + widgets.List{ + view_id = "squad_list", + icon_width = 2, + frame = { t = 1, h = 5 }, + choices = self:getSquads(), + on_submit=self:callback("toggleSquad") + }, + widgets.Divider{ frame={t=6, h=1}, frame_style_l = false, frame_style_r = false}, + widgets.Label{ + frame={ t = 7 , h = 1 }, + text = "General options:", + }, + widgets.EditField { + view_id = "threshold", + frame={ t = 8 , h = 1 }, + key = "CUSTOM_T", + label_text = "Need threshold for training: ", + text = tostring(-autotraining.state.threshold), + on_char = function (char, _) + return tonumber(char,10) + end, + on_submit = function (text) + -- still necessary, because on_char does not check pasted text + local entered_number = tonumber(text,10) or 5000 + autotraining.state.threshold = -entered_number + autotraining.persist_state() + -- make sure that the auto correction is reflected in the EditField + self.subviews.threshold:setText(tostring(entered_number)) + end + }, + widgets.Divider{ frame={t=9, h=1}, frame_style_l = false, frame_style_r = false}, + widgets.Label{ + frame={ t = 10 , h = 1 }, + text = "Ignored noble positions:", + }, + widgets.List{ + frame = { t = 11 , h = 11}, + view_id = "nobles_list", + icon_width = 2, + choices = self:getPositions(), + on_submit=self:callback("toggleNoble") + }, + widgets.Divider{ frame={t=22, h=1}, frame_style_l = false, frame_style_r = false}, + widgets.Label{ + frame={ t = 23 , h = 1 }, + text = "Select units to exclude from automatic training:" + }, + widgets.FilteredList{ + frame = { t = 24 }, + view_id = "unit_list", + edit_key = "CUSTOM_CTRL_F", + icon_width = 2, + choices = self:getUnits(), + on_submit=self:callback("toggleUnit") + } + } + --self.subviews.unit_list:setChoices(unit_choices) +end + +function AutoTrain:onDismiss() + view = nil +end + +AutoTrainScreen = defclass(AutoTrainScreen, gui.ZScreen) +AutoTrainScreen.ATTRS { + focus_path='autotrain', +} + +function AutoTrainScreen:init() + self:addviews{AutoTrain{}} +end + +function AutoTrainScreen:onDismiss() + view = nil +end + +if not dfhack.world.isFortressMode() or not dfhack.isMapLoaded() then + qerror('gui/autotraining requires a fortress map to be loaded') +end + +view = view and view:raise() or AutoTrainScreen{}:show() diff --git a/internal/control-panel/registry.lua b/internal/control-panel/registry.lua index 76fbee5c1..08b721f8b 100644 --- a/internal/control-panel/registry.lua +++ b/internal/control-panel/registry.lua @@ -34,6 +34,8 @@ COMMANDS_BY_IDX = { desc='Automatically shear creatures that are ready for shearing.', params={'--time', '14', '--timeUnits', 'days', '--command', '[', 'workorder', 'ShearCreature', ']'}}, {command='autoslab', group='automation', mode='enable'}, + {command='autotraining', group='automation', mode='enable', + desc='Automatically assign units with training needs to training squads. '}, {command='ban-cooking all', group='automation', mode='run'}, {command='buildingplan set boulders false', group='automation', mode='run', desc='Enable if you usually don\'t want to use boulders for construction.'}, diff --git a/internal/notify/notifications.lua b/internal/notify/notifications.lua index 8af7c2c18..653d3887d 100644 --- a/internal/notify/notifications.lua +++ b/internal/notify/notifications.lua @@ -366,6 +366,24 @@ NOTIFICATIONS_BY_IDX = { dlg.showMessage('Rescue stuck squads', message, COLOR_WHITE) end, }, + { + name='auto_train', + desc='Notifies when there are no squads set up for training', + default=true, + dwarf_fn=function() + local at = reqscript('autotraining') + if (at.isEnabled() and at.checkSquads() == nil) then + return {{text="autotraining: no squads selected",pen=COLOR_LIGHTRED}} + end + end, + on_click=function() + local message = + "You have no squads selected for training.\n".. + "You should have a squad set up to be constantly training with about 8 units needed for training.\n".. + "Then you can select that squad for training in the config.\n\nWould you like to open the config? Alternatively, simply close this popup to go create a squad." + dlg.showYesNoPrompt('Training Squads not configured', message, COLOR_WHITE, function () dfhack.run_command('gui/autotraining') end) + end, + }, { name='traders_ready', desc='Notifies when traders are ready to trade at the depot.',