diff --git a/autotraining.lua b/autotraining.lua new file mode 100644 index 0000000000..2f6b03f644 --- /dev/null +++ b/autotraining.lua @@ -0,0 +1,309 @@ +-- 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 need_id = 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() + else + stop() + end + -- start can change the enabled state if the squad cant be found + if state.enabled then + dfhack.print(GLOBAL_KEY .." was persisted with the following data:\nThreshold: ".. state.threshold .. '\n') + end + persist_state() +end + + +--###### +--Functions +--###### +function getTrainingCandidates() + local ret = {} + local citizen = dfhack.units.getCitizens(true) + ignore_count = 0 + for _, unit in ipairs(citizen) do + if dfhack.units.isAdult(unit) then + local noblePos = dfhack.units.getNoblePositions(unit) + local isIgnNoble = false + if ( not state.ignored[unit.id] ) then + 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 not isIgnNoble then + table.insert(ret, unit) + else + removeTraining(unit) + ignore_count = ignore_count +1 + end + else + removeTraining(unit) + ignore_count = ignore_count +1 + end + end + 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 findNeed(unit) + local needs = unit.status.current_soul.personality.needs + for _, need in ipairs(needs) do + if need.id == need_id then + return need + end + end + return nil +end + +--###### +--Main +--###### + +function getByID(id) + for _, unit in ipairs(getTrainingCandidates()) do + if (unit.hist_figure_id == id) then + return unit + end + end + + return nil +end + +-- 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) + if (unit.military.squad_id ~= -1) then + for _, squad in ipairs(getTrainingSquads()) do + if unit.military.squad_id == squad.id then + return true + end + end + return false + end + for _, squad in ipairs(getTrainingSquads()) do + for i=1,9,1 do + if ( squad.positions[i].occupant == -1 ) then + dfhack.military.addToSquad(unit.id,squad.id,i) + -- squad.positions[i].occupant = unit.hist_figure_id + -- unit.military.squad_id = squad.id + -- unit.military.squad_position = i + return true + end + end + end + + return false +end + +function removeTraining(unit) + for _, squad in ipairs(getTrainingSquads()) do + for i=1,9,1 do + if ( unit.hist_figure_id == squad.positions[i].occupant ) then + dfhack.military.removeFromSquad(unit.id) + -- unit.military.squad_id = -1 + -- unit.military.squad_position = -1 + -- squad.positions[i].occupant = -1 + return true + 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 dwarf = getByID(squad.positions[i].occupant) + if (dwarf ~= nil) then + removeTraining(dwarf) + end + end + end +end + + +function check() + local squads = checkSquads() + local intraining_count = 0 + local inque_count = 0 + if ( squads == nil) then return end + for _, unit in ipairs(getTrainingCandidates()) do + local need = findNeed(unit) + if ( need ~= nil ) then + if ( need.focus_level < state.threshold ) then + local bol = addTraining(unit) + if ( bol ) then + intraining_count = intraining_count +1 + else + inque_count = inque_count +1 + end + else + removeTraining(unit) + end + end + end + + dfhack.println(GLOBAL_KEY .. " | IGN: " .. ignore_count .. " TRAIN: " .. intraining_count .. " QUE: " ..inque_count ) +end + +function start() + dfhack.println(GLOBAL_KEY .. " | START") + + if (args.t) then + state.threshold = 0-tonumber(args.t) + end + repeatUtil.scheduleEvery(GLOBAL_KEY, 1, 'days', check) -- 997 is the closest prime to 1000 +end + +function stop() + removeAll() + repeatUtil.cancel(GLOBAL_KEY) + dfhack.println(GLOBAL_KEY .. " | STOP") +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() + dfhack.println(GLOBAL_KEY .." | Enabled") +else + stop() + dfhack.println(GLOBAL_KEY .." | Disabled") +end +persist_state() diff --git a/docs/autotraining.rst b/docs/autotraining.rst new file mode 100644 index 0000000000..c647905516 --- /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 + +Automation script for citizens to hit the gym when they yearn for the gains. Also passively builds military skills and physical stats. + +You need to have at least one squad that is set up for training. This should be a new non-military-use squad. The uniform should be +set to "No Uniform" and the squad should be set to "Constant Training" in the military screen. Edit the squad's schedule to full time training with around 8 units training. +The squad doesn't need months off. The members leave the squad once they have gotten their gains. + +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 gym. + 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 to go to the gym, and begins assigning them to said gym. + Once they have fulfilled their need they will be removed from the gym 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 0000000000..a86b28adf9 --- /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 0000000000..03ac183267 --- /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 37cd56c4e2..6cd4229fc2 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='Automation script for citizens to hit the gym when they yearn for the gains.'}, {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 8af7c2c187..653d3887d5 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.',