Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
A family tree generator for Dwarf Fortress.
FamilyNode=defclass(FamilyNode)
FamilyNode.ATTRS {
histfig_id = DEFAULT_NIL,
mother = DEFAULT_NIL,
father = DEFAULT_NIL,
spouse = DEFAULT_NIL,
progenitor = false,
children = {}
}
function FamilyNode:clean()
local remove_table={}
for k,v in pairs(self.children) do
if v.mother~=self and v.father~=self then
remove_table[k]=true
end
end
for k,v in pairs(remove_table) do
self.children[k]=nil
end
self.children_nodes=nil
end
function FamilyNode:getChildrenNum()
local size=0
for k,v in pairs(self.children) do
size=size+1
end
return size
end
function FamilyNode:isVoid()
return (not self.mother and not self.father and self:getChildrenNum()==0)
end
function FamilyNode:isDescendedFrom(histfig_id)
local connected=false
for _,family in ipairs(self.families or {}) do
local otherNode=family:histfigLookup(histfig_id)
connected=connected or family:isDescendedFrom(self,otherNode)
if connected then return connected end
end
return connected
end
FamilyTree=defclass(FamilyTree)
FamilyTree.ATTRS {
histfig_lookup={}
}
function FamilyTree:cleanExplored()
for k,v in pairs(self.histfig_lookup) do
v.explored=false
end
end
function FamilyTree:isDescendedFrom(node1,node2,not_top_level) --i'm pretty sure this is a depth first search
node1.explored=true
if node1==node2 then return true end
for _,parent in ipairs({node1.mother,node1.father}) do
if not parent.explored then
if self:isDescendedFrom(parent,node2,true) then self:cleanExplored() return true end
end
end
if not not_top_level then
self:cleanExplored()
end
return false
end
function FamilyTree:isAncestorOf(node1,node2,not_top_level)
node1.explored=true
if node1==node2 then return true end
for _,child in ipairs(node1.children) do
if not child.explored then
if self:isAncestorOf(child,node2,true) then self:cleanExplored() return true end
end
end
if not not_top_level then
self:cleanExplored()
end
return false
end
function FamilyTree:insertNode(node)
if self:histfigLookup(node.histfig_id) then
return self:histfigLookup(node.histfig_id)
else
self.histfig_lookup[node.histfig_id]=node
node.families=node.families or {}
table.insert(node.families,self)
return node
end
end
function FamilyTree:mergeFamily(other_tree)
for k,v in pairs(other_tree.histfig_lookup) do
self.histfig_lookup[k]=self.histfig_lookup[k] or v
end
other_tree=self
end
function FamilyTree:histfigLookup(histfig_id)
self.histfig_lookup=self.histfig_lookup or {}
return self.histfig_lookup[histfig_id]
end
function FamilyTree:getSize()
local size=0
for k,v in pairs(self.histfig_lookup) do
size=size+1
end
return size
end
function FamilyTree:clean()
for k,v in pairs(self.histfig_lookup) do
v:clean()
if v.mother and not self:histfigLookup(v.mother.histfig_id) then
self:insertNode(v.mother)
end
if v.father and not self:histfigLookup(v.father.histfig_id) then
self:insertNode(v.father)
end
for _,child in pairs(v.children) do
if not self:histfigLookup(child.histfig_id) then
self:insertNode(child)
end
end
if v:isVoid() then
self.histfig_lookup[k]=nil
end
end
end
function FamilyTree:giveFamilyName()
for k,v in pairs(self.histfig_lookup) do
if v.mother and v.father then
local father_fig,mother_fig=df.historical_figure.find(v.father.histfig_id),df.historical_figure.find(v.mother.histfig_id)
local name_words={father_fig.name.words[0],mother_fig.name.words[1]}
local name_parts_of_speech={father_fig.name.parts_of_speech[0],mother_fig.name.parts_of_speech[1]}
local fig=df.historical_figure.find(v.histfig_id)
fig.name.words[0]=name_words[2]
fig.name.words[1]=name_words[1]
fig.name.parts_of_speech[0]=name_parts_of_speech[2]
fig.name.parts_of_speech[1]=name_parts_of_speech[1]
end
end
end
function FamilyTree:findProgenitor()
if self.progenitor then return self.progenitor,self.progenitor.spouse or self.progenitor end
local father,mother
for k,v in pairs(self.histfig_lookup) do --actually just returns after first one, heh
father=v.father or v
mother=v.mother or v
while father.father do
father=father.father or father
end
while mother.mother do
mother=mother.mother or mother
end
return father,mother
end
end
function getLeadingZeroes(num,numDigits)
local zeroes=numDigits-math.floor(math.log(num)/math.log(10))-1
local str=''
for i=1,zeroes do
str=str..'0'
end
return str..num
end
function getMonthDate(num)
return num~=-1 and getLeadingZeroes(math.ceil(num/33600),2)..getLeadingZeroes(math.ceil((num%33600)/1200),2) or '0000'
end
function getMonthDateGramps(num)
return num~=-1 and '-'..math.ceil(num/33600)..'-'..math.floor((num%33600)/1200) or ' (Unknown date)'
end
function cantorBijection(a,b) --thanks cantor; this will give a unique marriage id for any pair of integers
return .5*(a+b)*(a+b+1)+b
end
function FamilyTree:exportToFamilyScript()
local file=""
local counter=1
local size=self:getSize()
for k,v in pairs(self.histfig_lookup) do
print('Exporting member #'..counter..'/'..size..', hist fig number #'..k)
local fig=df.historical_figure.find(v.histfig_id)
file=file..'i'..v.histfig_id..' g'..(fig.sex==0 and 'f' or 'm')
file=file..' p'..dfhack.TranslateName(fig.name)
file=file..' b'..getLeadingZeroes(fig.born_year,4)..getMonthDate(fig.born_seconds)
if fig.died_year~=-1 then
file=file..' z1 d'..getLeadingZeroes(fig.died_year,4)..getMonthDate(fig.died_seconds)
end
if v.mother then
file=file..' m'..v.mother.histfig_id
end
if v.father then
file=file..' f'..v.father.histfig_id
end
if v.spouse then
file=file..' s'..v.spouse.histfig_id
end
file=file..NEWLINE
counter=counter+1
end
local f=assert(io.open('family_'..self:findProgenitor().histfig_id..'.family',"w"))
f:write(file)
f:close()
end
function FamilyTree:exportToCSV()
local file="Person,Surname,Given,Gender,Birth date,Death date,Death place"..NEWLINE
local counter=1
local size=self:getSize()
for k,v in pairs(self.histfig_lookup) do
print('Exporting member #'..counter..'/'..size..', hist fig number #'..k)
local fig=df.historical_figure.find(v.histfig_id)
file=file..'['..v.histfig_id..'],'
do
local name=dfhack.TranslateName(fig.name)
local givenName=fig.name.first_name
local spaceInName=name:find(' ')
local surname=spaceInName and name:sub(spaceInName+1) or name
file=file..surname..','..givenName:sub(1,1):upper()..givenName:sub(2)..','
end
file=file..(fig.sex==0 and 'female,' or fig.sex==1 and 'male,' or 'none,')
file=file..fig.born_year..getMonthDateGramps(fig.born_seconds)..','
if fig.died_year~=-1 then
file=file..fig.died_year..getMonthDateGramps(fig.died_seconds)..','
local unit=df.unit.find(fig.unit_id)
if unit then
local incident=df.incident.find(unit.counters.death_id)
if incident then
local site=df.world_site.find(incident.site)
if site then
file=file..dfhack.TranslateName(site.name,true)
end
end
end
else
file=file..','
end
file=file..NEWLINE
counter=counter+1
end
counter=1
file=file..NEWLINE..NEWLINE..'Marriage,Husband,Wife'..NEWLINE
do
local marriageAlreadyDone={} --abusing the do block for scope a lot here; i'm thinking the memory usage would be simply incredible if this were to stick around for the whole import
for k,v in pairs(self.histfig_lookup) do
if v.spouse then
local marriageID=cantorBijection(math.min(v.histfig_id,v.spouse.histfig_id),math.max(v.histfig_id,v.spouse.histfig_id))
if not marriageAlreadyDone[marriageID] then
--min and max so that a and b are always the same for a given marriage
marriageAlreadyDone[marriageID]=true
do -- http://qntm.org/gay
local husband,wife=0,0
local sex1,sex2=df.historical_figure.find(v.histfig_id).sex,df.historical_figure.find(v.spouse.histfig_id).sex
if sex1==1 or sex1==sex2 then --elegance, or it would be if I couldn't just set it up as partner1 and partner2; DF doesn't have polygamy, so I won't complain FOR NOW
husband,wife=v.histfig_id,v.spouse.histfig_id
else
wife,husband=v.histfig_id,v.spouse.histfig_id
end
file=file..'['..marriageID..']'..','..'['..husband..']'..','..'['..wife..']'..NEWLINE
end
end
end
end
end
file=file..NEWLINE..NEWLINE..'Family,Child'..NEWLINE
do
for k,v in pairs(self.histfig_lookup) do
if v.mother and v.father then
file=file..'['..cantorBijection(math.min(v.mother.histfig_id,v.father.histfig_id),math.max(v.mother.histfig_id,v.father.histfig_id))..']'..','..'['..v.histfig_id..']'..NEWLINE
end
end
end
local f=assert(io.open('family_'..self:findProgenitor().histfig_id..'.csv',"w"))
f:write(file)
f:close()
end
Families=defclass(Families)
Families.ATTRS {
families={}
}
function Families:getFamilyFromHistFig(histfig_id)
self.families=self.families or {}
self.families_histfig_lookup = self.families_histfig_lookup or {}
local family=self.families_histfig_lookup[histfig_id]
if family then return family:histfigLookup(histfig_id),family end
return nil
end
function Families:insertNode(node,family,progenitor)
local existingFamily=self:getFamilyFromHistFig(node.histfig_id)
if existingFamily then
return existingFamily:histfigLookup(node.histfig_id)
else
local newFamily=family or FamilyTree{}
if not family then
newFamily.histfig_lookup={}
if progenitor then
newFamily.progenitor=node
end
newFamily:insertNode(node)
self.families_histfig_lookup[node.histfig_id]=newFamily
table.insert(self.families,newFamily)
else
newFamily:insertNode(node)
self.families_histfig_lookup[node.histfig_id]=newFamily
end
return node,newFamily
end
end
function Families:getFamily(hist_fig_id,family,children_only,parent,is_progenitor)
local hist_fig=df.historical_figure.find(hist_fig_id)
local oldNode,oldfamily=self:getFamilyFromHistFig(hist_fig_id)
if oldNode then
if parent and (not oldNode.mother or not oldNode.father) then
for _,fig_link in pairs(hist_fig.histfig_links) do
if fig_link.target_hf==parent.histfig_id then
if fig_link._type==df.histfig_hf_link_motherst then
oldNode.mother=parent
elseif fig_link._type==df.histfig_hf_link_fatherst then
oldNode.father=parent
end
end
end
family:insertNode(parent)
end
return oldNode,oldfamily
end
local node,family=self:insertNode(FamilyNode{histfig_id=hist_fig_id,progenitor=is_progenitor},family,progenitor)
node.children={}
if children_only and parent then
local parent_fig=df.historical_figure.find(parent.histfig_id)
if parent_fig.sex==0 then
node.mother=parent
else
node.father=parent
end
family:insertNode(parent)
end
for _,fig_link in pairs(hist_fig.histfig_links) do
if fig_link._type==df.histfig_hf_link_motherst and not children_only then
node.mother=self:getFamily(fig_link.target_hf,family,children_only)
elseif fig_link._type==df.histfig_hf_link_fatherst and not children_only then
node.father=self:getFamily(fig_link.target_hf,family,children_only)
elseif fig_link._type==df.histfig_hf_link_childst then
local childNode=self:getFamily(fig_link.target_hf,family,children_only,node)
node.children[childNode.histfig_id]=childNode
elseif fig_link._type==df.histfig_hf_link_spousest then
node.spouse=self:getFamily(fig_link.target_hf,family,children_only)
end
end
return node,family
end
function Families:clean()
local prevAmount=1
for i=#self.families,1,-1 do
local size=self.families[i]:getSize()
if size<2 then table.remove(self.families,i) end
prevAmount=size
end
for k,v in pairs(self.families) do
v:clean()
end
end
function Families:sort()
table.sort(self.families,function(a,b) return a:getSize()>b:getSize() end)
end
function Families:getAllHistFigs(children_only)
print('Gathering all historical figures...')
for k,v in pairs(df.global.world.history.figures) do
if k%100==0 then print('At figure #'..k..'/'..#df.global.world.history.figures) end
self:getFamily(v.id,nil,children_only,nil,true)
end
self:clean()
self:clean() --two passes actually does something, weirdly enough
self:sort()
end
function Families:isDescendedFrom(hist_fig_id1,hist_fig_id2)
local node1,family1=self:getFamily(hist_fig_id1)
local node2,family2=self:getFamily(hist_fig_id2)
return family1==family2 and family1:isDescendedFrom(node1,node2)
end
function Families:isAncestorOf(hist_fig_id1,hist_fig_id2)
local node1,family1=self:getFamily(hist_fig_id1)
local node2,family2=self:getFamily(hist_fig_id2)
return family1==family2 and family1:isAncestorOf(node1,node2)
end
function Families:getDirectDescendentsOnly(hist_fig_id,family,progenitor)
local hist_fig=df.historical_figure.find(hist_fig_id)
local oldNode,oldfamily=self:getFamilyFromHistFig(hist_fig_id)
if oldNode then
return oldNode,oldfamily
end
local node,family=self:insertNode(FamilyNode{histfig_id=hist_fig_id,progenitor=progenitor},family,progenitor)
local progenitor = progenitor or node
node.children={}
for _,fig_link in pairs(hist_fig.histfig_links) do
if fig_link._type==df.histfig_hf_link_childst then
local childNode=self:getDirectDescendentsOnly(fig_link.target_hf,true,progenitor)
node.children[childNode.histfig_id]=childNode
end
end
return node,family
end
function Families:getNumChildren(hist_fig_or_id)
local num=0
local hist_fig=df.historical_figure:is_instance(hist_fig_or_id) and hist_fig_or_id or df.historical_figure.find(hist_fig_or_id)
for _,fig_link in pairs(hist_fig.histfig_links) do
if fig_link._type==df.histfig_hf_link_childst then
num=num+1
end
end
return num
end
local utils=require('utils')
validArgs = validArgs or utils.invert({
'all',
'this',
'help',
'heritage',
'exportToFamilyScript',
'exportToCSV'
})
local args = utils.processArgs({...}, validArgs)
if args.exportToFamilyScript then
if args.all then
local dlg=require('gui.dialogs')
dlg.showYesNoPrompt('Family lineage','Really export all families? (May take very long time and may cause game to hang!)',COLOR_WHITE,function()
local families=Families()
families:getAllHistFigs()
for k,v in ipairs(families.families) do
print('Exporting family of ' .. dfhack.TranslateName(df.historical_figure.find(v:findProgenitor().histfig_id).name) .. '...')
v:exportToFamilyScript()
end
end
)
end
if args.this then
local families=Families()
families:getFamily(dfhack.gui.getSelectedUnit().hist_figure_id)
families:clean()
families:clean()
for k,v in ipairs(families.families) do
print('Exporting family of ' .. dfhack.TranslateName(df.historical_figure.find(v:findProgenitor().histfig_id).name) .. '...')
v:exportToFamilyScript()
end
elseif args.unit then
local families=Families()
families:getFamily(df.unit.find(args.unit).hist_figure_id)
families:clean()
families:clean()
for k,v in ipairs(families.families) do
print('Exporting family of ' .. dfhack.TranslateName(df.historical_figure.find(v:findProgenitor().histfig_id).name) .. '...')
v:exportToFamilyScript()
end
elseif args.histfig then
local families=Families()
families:getFamily(histfig)
families:clean()
families:clean()
for k,v in ipairs(families.families) do
print('Exporting family of ' .. dfhack.TranslateName(df.historical_figure.find(v:findProgenitor().histfig_id).name) .. '...')
v:exportToFamilyScript()
end
end
end
if args.exportToCSV then
if args.all then
local dlg=require('gui.dialogs')
dlg.showYesNoPrompt('Family lineage','Really export all families? (May take very long time and may cause game to hang!)',COLOR_WHITE,function()
local families=Families()
families:getAllHistFigs()
for k,v in ipairs(families.families) do
print('Exporting family of ' .. dfhack.TranslateName(df.historical_figure.find(v:findProgenitor().histfig_id).name) .. '...')
v:exportToCSV()
end
end
)
end
if args.this then
local families=Families()
families:getFamily(dfhack.gui.getSelectedUnit().hist_figure_id)
families:clean()
families:clean()
for k,v in ipairs(families.families) do
print('Exporting family of ' .. dfhack.TranslateName(df.historical_figure.find(v:findProgenitor().histfig_id).name) .. '...')
v:exportToCSV()
end
elseif args.unit then
local families=Families()
families:getFamily(df.unit.find(args.unit).hist_figure_id)
families:clean()
families:clean()
for k,v in ipairs(families.families) do
print('Exporting family of ' .. dfhack.TranslateName(df.historical_figure.find(v:findProgenitor().histfig_id).name) .. '...')
v:exportToCSV()
end
elseif args.histfig then
local families=Families()
families:getFamily(histfig)
families:clean()
families:clean()
for k,v in ipairs(families.families) do
print('Exporting family of ' .. dfhack.TranslateName(df.historical_figure.find(v:findProgenitor().histfig_id).name) .. '...')
v:exportToCSV()
end
end
end
if args.heritage then
if args.all then
local families=Families()
families:getAllHistFigs()
families:clean()
for k,v in ipairs(families.families) do
local progenitor1,progenitor2=v:findProgenitor()
print('Naming family of ' .. dfhack.TranslateName(df.historical_figure.find(progenitor1.histfig_id).name) .. ' and ' .. dfhack.TranslateName(df.historical_figure.find(progenitor2.histfig_id).name)..'...')
v:giveFamilyName()
end
elseif args.this then
local families=Families()
families:getFamily(dfhack.gui.getSelectedUnit().hist_figure_id)
families:clean()
families:clean()
for k,v in ipairs(families.families) do
local progenitor1,progenitor2=v:findProgenitor()
print('Naming family of ' .. dfhack.TranslateName(df.historical_figure.find(progenitor1.histfig_id).name) .. ' and ' .. dfhack.TranslateName(df.historical_figure.find(progenitor2.histfig_id).name)..'...')
v:giveFamilyName()
end
elseif args.unit then
local families=Families()
families:getFamily(df.unit.find(args.unit).hist_figure_id)
families:clean()
families:clean()
for k,v in ipairs(families.families) do
local progenitor1,progenitor2=v:findProgenitor()
print('Naming family of ' .. dfhack.TranslateName(df.historical_figure.find(progenitor1.histfig_id).name) .. ' and ' .. dfhack.TranslateName(df.historical_figure.find(progenitor2.histfig_id).name)..'...')
v:giveFamilyName()
end
elseif args.histfig then
local families=Families()
families:getFamily(histfig)
families:clean()
families:clean()
for k,v in ipairs(families.families) do
local progenitor1,progenitor2=v:findProgenitor()
print('Naming family of ' .. dfhack.TranslateName(df.historical_figure.find(progenitor1.histfig_id).name) .. ' and ' .. dfhack.TranslateName(df.historical_figure.find(progenitor2.histfig_id).name)..'...')
v:giveFamilyName()
end
end
end
if args.help then
print([[scripts/family.lua
arguments
-help
print this help message
-exportToFamilyScript
exports specified families to FamilyScript file
-exportToCSV
exports specified families to a CSV file (compatible with Gramps)
-heritage
gives family name to specified families
-all
specifies all the world's families (may crash game)
-this
specifies selected unit's family
-unit
specifies family of unit with given unit id
-histfig
specifies family of histfig with given unit id
]])
return
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment