forked from DFHack/dfhack
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdwarfvet.lua
More file actions
257 lines (230 loc) · 7.93 KB
/
dwarfvet.lua
File metadata and controls
257 lines (230 loc) · 7.93 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
local _ENV = mkmodule('plugins.dwarfvet')
local argparse = require('argparse')
local utils = require('utils')
local GLOBAL_KEY = 'dwarfvet'
-- used for reassigning animals to pastures after treatment
tracked_patients = tracked_patients or {}
local function persist_state()
local state = {}
for unit_id, zone_id in pairs(tracked_patients) do
table.insert(state, {unit=unit_id, zone=zone_id})
end
dfhack.persistent.saveSiteData(GLOBAL_KEY, state)
end
dfhack.onStateChange[GLOBAL_KEY] = function(sc)
if sc == SC_MAP_UNLOADED then
tracked_patients = {}
return
end
if sc ~= SC_MAP_LOADED or not dfhack.world.isFortressMode() then
return
end
local state = dfhack.persistent.getSiteData(GLOBAL_KEY, {})
tracked_patients = {}
for _,elem in ipairs(state) do
tracked_patients[elem.unit] = elem.zone
end
end
local function is_valid_animal(unit)
return unit and
dfhack.units.isActive(unit) and
dfhack.units.isAnimal(unit) and
dfhack.units.isFortControlled(unit) and
dfhack.units.isTame(unit) and
not dfhack.units.isDead(unit) and
not dfhack.units.getGeneralRef(unit, df.general_ref_type.BUILDING_CHAIN) and
not dfhack.units.getGeneralRef(unit, df.general_ref_type.CONTAINED_IN_ITEM)
end
local function get_cur_patients()
local cur_patients = {}
for _,job in utils.listpairs(df.global.world.jobs.list) do
if job.job_type ~= df.job_type.Rest then goto continue end
local unit = dfhack.job.getWorker(job)
if is_valid_animal(unit) then
cur_patients[unit.id] = unit
end
::continue::
end
return cur_patients
end
local function get_new_patients(cur_patients)
local new_patients = {}
for _,unit in ipairs(df.global.world.units.active) do
if unit.job.current_job then goto continue end
if not unit.health or not unit.health.flags.needs_healthcare then goto continue end
if cur_patients[unit.id] or not is_valid_animal(unit) then goto continue end
new_patients[unit.id] = unit
::continue::
end
return new_patients
end
local function print_status()
print(('dwarfvet is %srunning'):format(isEnabled() and '' or 'not '))
print()
print('The following animals are receiving treatment:')
local cur_patients = get_cur_patients()
if not next(cur_patients) then
print(' None')
else
for _, unit in pairs(cur_patients) do
print((' %s (%d)'):format(dfhack.units.getReadableName(unit), unit.id))
end
end
print()
print('The following animals are injured and need treatment:')
local new_patients = get_new_patients(cur_patients)
if not next(new_patients) then
print(' None')
else
for _, unit in pairs(new_patients) do
print((' %s (%d)'):format(dfhack.units.getReadableName(unit), unit.id))
end
end
end
HospitalZone = defclass(HospitalZone)
HospitalZone.ATTRS{
zone=DEFAULT_NIL,
}
local ONE_TILE = xy2pos(1, 1)
function HospitalZone:find_spot(unit_pos)
self.x = self.x or self.zone.x1
self.y = self.y or self.zone.y1
local zone = self.zone
for y=self.y,zone.y2 do
for x=self.x,zone.x2 do
local pos = xyz2pos(x, y, zone.z)
if dfhack.maps.canWalkBetween(unit_pos, pos) and
dfhack.buildings.containsTile(zone, x, y) and
dfhack.buildings.checkFreeTiles(pos, ONE_TILE)
then
return pos
end
end
end
end
local function detach_unit_from_pasture(bld, unit)
if not bld or not df.building_civzonest:is_instance(bld) then return end
for idx, id in ipairs(bld.assigned_units) do
if id == unit.id then
bld.assigned_units:erase(idx)
break
end
end
end
local function unassign_and_remember_pasture(unit)
for idx = #unit.general_refs-1, 0, -1 do
local ref = unit.general_refs[idx]
if df.general_ref_building_civzone_assignedst:is_instance(ref) then
detach_unit_from_pasture(df.building.find(ref.building_id), unit)
tracked_patients[unit.id] = ref.building_id
unit.general_refs:erase(idx)
ref:delete()
end
end
persist_state()
end
local function attach_to_pasture(unit, zone)
if not unit or not zone then return end
local ref = df.new(df.general_ref_building_civzone_assignedst)
ref.building_id = zone.id;
unit.general_refs:insert('#', ref)
utils.insert_sorted(zone.assigned_units, unit.id)
end
-- TODO: If health.should_not_move is set, the creature can't move under its own power
-- and a Recover Wounded or Pen/Pasture job must be created by hand
function HospitalZone:assign_spot(unit, unit_pos)
local pos = self:find_spot(unit_pos)
if not pos then return false end
local job = df.job:new()
dfhack.job.linkIntoWorld(job, true)
job.pos.x = pos.x
job.pos.y = pos.y
job.pos.z = pos.z
job.flags.special = true
job.job_type = df.job_type.Rest
local gref = df.new(df.general_ref_unit_workerst)
gref.unit_id = unit.id
job.general_refs:insert('#', gref)
unit.job.current_job = job
unassign_and_remember_pasture(unit)
return true
end
local function get_hospital_zones()
local hospital_zones = {}
local site = dfhack.world.getCurrentSite() or {buildings={}}
for _,location in ipairs(site.buildings) do
if not df.abstract_building_hospitalst:is_instance(location) then goto continue end
for _,bld_id in ipairs(location.contents.building_ids) do
local zone = df.building.find(bld_id)
if zone then
table.insert(hospital_zones, HospitalZone{zone=zone})
end
end
::continue::
end
return hospital_zones
end
local function distance(zone, pos)
return math.abs(zone.x1 - pos.x) + math.abs(zone.y1 - pos.y) + 50*math.abs(zone.z - pos.z)
end
local function reassign_healed_animals_to_pastures(cur_patients)
for unit_id, zone_id in pairs(tracked_patients) do
if cur_patients[unit_id] then goto continue end
tracked_patients[unit_id] = nil
local unit = df.unit.find(unit_id)
if is_valid_animal(unit) and not
dfhack.units.getGeneralRef(unit, df.general_ref_type.BUILDING_CIVZONE_ASSIGNED)
then
attach_to_pasture(unit, df.building.find(zone_id))
end
::continue::
end
persist_state()
end
function checkup()
local cur_patients = get_cur_patients()
reassign_healed_animals_to_pastures(cur_patients)
local new_patients = get_new_patients(cur_patients)
if not next(new_patients) then return end
local hospital_zones = get_hospital_zones()
local count, assigned = 0, 0
for _, unit in pairs(new_patients) do
count = count + 1
local unit_pos = xyz2pos(dfhack.units.getPosition(unit))
table.sort(hospital_zones,
function(a, b) return distance(a.zone, unit_pos) < distance(b.zone, unit_pos) end)
for _,hospital_zone in ipairs(hospital_zones) do
if hospital_zone:assign_spot(unit, unit_pos) then
assigned = assigned + 1
break
end
end
end
print(('dwarfvet scheduled treatment for %d of %d injured animals'):format(assigned, count))
end
local function process_args(opts, args)
if args[1] == 'help' then
opts.help = true
return
end
return argparse.processArgsGetopt(args, {
{'h', 'help', handler=function() opts.help = true end},
})
end
function parse_commandline(args)
local opts = {}
local positionals = process_args(opts, args)
if opts.help or not positionals then
return false
end
local command = positionals[1]
if not command or command == 'status' then
print_status()
elseif command == 'now' then
dwarfvet_cycle()
else
return false
end
return true
end
return _ENV