Merge branch 'sci-gym-script' into SquidHack

Author: realSquidCoder

autotraining.lua

@@ -99,38 +99,43 @@ 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 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
+        -- 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 isIgnNoble then
-            ignore_count = ignore_count +1
+        if isIgnoredNoble(unit) then
+            ignore_count = ignore_count + 1
             goto next_unit
         end
-        table.insert(ret, unit)
+        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
 
@@ -229,52 +234,59 @@ function check()
                     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)
+                        dfhack.military.removeFromSquad(unit.id)
                     end
                 end
             end
         end
     end
-    for _, unit in ipairs(getTrainingCandidates()) do
-        local added = addTraining(unit, squads)
+    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
-
-    dfhack.println(GLOBAL_KEY  .. " | IGNORED: " .. ignore_count .. " TRAINING: " .. intraining_count .. " QUEUE: " ..inque_count )
+    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()
-    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
+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
 
-if state.enabled then
-    start()
+if dfhack_flags.enable then
+    if dfhack_flags.enable_state then
+        enable()
+    else
+        disable()
+    end
 else
-    stop()
-    removeAll()
+    -- 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
-persist_state()

gui/autotraining.lua

@@ -32,7 +32,7 @@ function AutoTrain:getSquads()
             goto continue
         end
         table.insert(squads, {
-            text = dfhack.translation.translateName(squad.name, true)..' ('..squad.alias..')',
+            text = dfhack.translation.translateName(squad.name, true)..(squad.alias ~= '' and ' ('..squad.alias..')' or ''),
             icon = self:callback("getSquadIcon", squad.id ),
             id   = squad.id
         })
@@ -192,25 +192,40 @@ function AutoTrain:init()
                 self.subviews.threshold:setText(tostring(entered_number))
             end
         },
-        widgets.Divider{ frame={t=9, h=1}, frame_style_l = false, frame_style_r = false},
+        widgets.ToggleHotkeyLabel {
+            view_id = 'enable_toggle',
+            frame = { t = 9, h = 1 },
+            label = 'Autotraining is',
+            key = 'CUSTOM_E',
+            options = { { value = true, label = 'Enabled', pen = COLOR_GREEN },
+                        { value = false, label = 'Disabled', pen = COLOR_RED } },
+            on_change = function(val)
+                if val then
+                    autotraining.enable()
+                else
+                    autotraining.disable()
+                end
+            end,
+        },
+        widgets.Divider{ frame={t=10, h=1}, frame_style_l = false, frame_style_r = false},
         widgets.Label{
-            frame={ t = 10 , h = 1 },
+            frame={ t = 11 , h = 1 },
             text = "Ignored noble positions:",
         },
         widgets.List{
-            frame = { t = 11 , h = 11},
+            frame = { t = 12 , 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.Divider{ frame={t=23, h=1}, frame_style_l = false, frame_style_r = false},
         widgets.Label{
-            frame={ t = 23 , h = 1 },
+            frame={ t = 24 , h = 1 },
             text = "Select units to exclude from automatic training:"
         },
         widgets.FilteredList{
-            frame = { t = 24 },
+            frame = { t = 25 },
             view_id = "unit_list",
             edit_key = "CUSTOM_CTRL_F",
             icon_width = 2,
@@ -221,6 +236,10 @@ function AutoTrain:init()
     --self.subviews.unit_list:setChoices(unit_choices)
 end
 
+function AutoTrain:onRenderBody(painter)
+    self.subviews.enable_toggle:setOption(autotraining.state.enabled)
+end
+
 function AutoTrain:onDismiss()
     view = nil
 end