A family tree generator for Dwarf Fortress.

family.lua

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