Module:FamilyTree

Documentation for this module may be created at Module:FamilyTree/doc

local p = {}

local cargo = mw.ext.cargo
local html = mw.html

-- =========================================
-- Helpers
-- =========================================

local function trim(s)
	if s == nil then
		return nil
	end
	s = tostring(s)
	s = mw.text.trim(s)
	if s == '' then
		return nil
	end
	return s
end

local function isRealValue(v)
	v = trim(v)
	if not v then
		return false
	end
	local lowered = mw.ustring.lower(v)
	return lowered ~= 'unknown' and lowered ~= 'none' and lowered ~= 'n/a'
end

local function addUnique(list, value)
	if not isRealValue(value) then
		return
	end
	for _, existing in ipairs(list) do
		if existing == value then
			return
		end
	end
	table.insert(list, value)
end

local function uniq(list)
	local out = {}
	local seen = {}
	for _, v in ipairs(list or {}) do
		if isRealValue(v) and not seen[v] then
			seen[v] = true
			table.insert(out, v)
		end
	end
	return out
end

local function getArg(frame, key)
	local v = frame.args[key]
	if isRealValue(v) then
		return trim(v)
	end

	local parent = frame:getParent()
	if parent then
		v = parent.args[key]
		if isRealValue(v) then
			return trim(v)
		end
	end

	return nil
end

local function getRoot(frame)
	return getArg(frame, 'root') or getArg(frame, 'person') or mw.title.getCurrentTitle().text
end

local function makeLink(name, displayName)
	if not isRealValue(name) then
		return ''
	end
	displayName = trim(displayName) or name
	return string.format('[[%s|%s]]', name, displayName)
end

local function ensurePerson(people, name)
	name = trim(name)
	if not isRealValue(name) then
		return nil
	end

	if not people[name] then
		people[name] = {
			name = name,
			displayName = name,
			parents = {},
			children = {},
			partners = {},
			unions = {},
			childLinks = {}
		}
	end

	return people[name]
end

local function sortNames(people, names)
	table.sort(names, function(a, b)
		local ad = (people[a] and people[a].displayName) or a
		local bd = (people[b] and people[b].displayName) or b
		return mw.ustring.lower(ad) < mw.ustring.lower(bd)
	end)
end

-- =========================================
-- Data loading
-- =========================================

local function loadCharacters()
	local results = cargo.query(
		'Characters',
		'Page,DisplayName',
		{ limit = 5000 }
	)

	local people = {}

	for _, row in ipairs(results) do
		local page = trim(row.Page)
		local displayName = trim(row.DisplayName)

		if isRealValue(page) then
			people[page] = {
				name = page,
				displayName = displayName or page,
				parents = {},
				children = {},
				partners = {},
				unions = {},
				childLinks = {}
			}
		end
	end

	return people
end

local function loadParentChild(people)
	local results = cargo.query(
		'ParentChild',
		'Child,Parent1,Parent2,UnionID,RelationshipType,BirthOrder',
		{ limit = 5000 }
	)

	for _, row in ipairs(results) do
		local child = trim(row.Child)
		local p1 = trim(row.Parent1)
		local p2 = trim(row.Parent2)
		local unionID = trim(row.UnionID)
		local relationshipType = trim(row.RelationshipType)
		local birthOrder = tonumber(trim(row.BirthOrder)) or 999

		if isRealValue(child) then
			ensurePerson(people, child)

			if isRealValue(p1) then
				ensurePerson(people, p1)
				addUnique(people[child].parents, p1)
				addUnique(people[p1].children, child)

				table.insert(people[p1].childLinks, {
					child = child,
					otherParent = p2,
					unionID = unionID,
					relationshipType = relationshipType,
					birthOrder = birthOrder
				})
			end

			if isRealValue(p2) then
				ensurePerson(people, p2)
				addUnique(people[child].parents, p2)
				addUnique(people[p2].children, child)

				table.insert(people[p2].childLinks, {
					child = child,
					otherParent = p1,
					unionID = unionID,
					relationshipType = relationshipType,
					birthOrder = birthOrder
				})
			end
		end
	end
end

local function loadUnions(people)
	local results = cargo.query(
		'Unions',
		'UnionID,Partner1,Partner2,UnionType,Status,StartDate,EndDate,MarriageDate,DivorceDate,EngagementDate',
		{ limit = 5000 }
	)

	for _, row in ipairs(results) do
		local p1 = trim(row.Partner1)
		local p2 = trim(row.Partner2)

		if isRealValue(p1) then
			ensurePerson(people, p1)
		end
		if isRealValue(p2) then
			ensurePerson(people, p2)
		end

		if isRealValue(p1) and isRealValue(p2) then
			addUnique(people[p1].partners, p2)
			addUnique(people[p2].partners, p1)

			table.insert(people[p1].unions, {
				unionID = trim(row.UnionID),
				partner = p2,
				unionType = trim(row.UnionType),
				status = trim(row.Status),
				startDate = trim(row.StartDate),
				endDate = trim(row.EndDate),
				marriageDate = trim(row.MarriageDate),
				divorceDate = trim(row.DivorceDate),
				engagementDate = trim(row.EngagementDate)
			})

			table.insert(people[p2].unions, {
				unionID = trim(row.UnionID),
				partner = p1,
				unionType = trim(row.UnionType),
				status = trim(row.Status),
				startDate = trim(row.StartDate),
				endDate = trim(row.EndDate),
				marriageDate = trim(row.MarriageDate),
				divorceDate = trim(row.DivorceDate),
				engagementDate = trim(row.EngagementDate)
			})
		end
	end
end

local function finalizePeople(people)
	for _, person in pairs(people) do
		person.parents = uniq(person.parents)
		person.children = uniq(person.children)
		person.partners = uniq(person.partners)
	end
end

local function loadData()
	local people = loadCharacters()
	loadParentChild(people)
	loadUnions(people)
	finalizePeople(people)
	return people
end

-- =========================================
-- Relationship helpers
-- =========================================

local function getGrandparents(people, root)
	local out = {}
	local person = people[root]
	if not person then
		return out
	end

	for _, parentName in ipairs(person.parents) do
		local parent = people[parentName]
		if parent then
			for _, gp in ipairs(parent.parents) do
				addUnique(out, gp)
			end
		end
	end

	return uniq(out)
end

local function getSiblings(people, root)
	local out = {}
	local seen = {}
	local person = people[root]

	if not person then
		return out
	end

	for _, parentName in ipairs(person.parents) do
		local parent = people[parentName]
		if parent then
			for _, childName in ipairs(parent.children) do
				if childName ~= root and not seen[childName] then
					seen[childName] = true
					table.insert(out, childName)
				end
			end
		end
	end

	return uniq(out)
end

local function getConnectedPeople(people, root)
	local out = {}
	local person = people[root]
	if not person then
		return out
	end

	for _, v in ipairs(person.parents) do addUnique(out, v) end
	for _, v in ipairs(person.children) do addUnique(out, v) end
	for _, v in ipairs(person.partners) do addUnique(out, v) end

	local siblings = getSiblings(people, root)
	for _, v in ipairs(siblings) do addUnique(out, v) end

	local grandparents = getGrandparents(people, root)
	for _, v in ipairs(grandparents) do addUnique(out, v) end

	return uniq(out)
end

local function relationshipBadge(relType)
	if not isRealValue(relType) then
		return nil
	end

	local t = mw.ustring.lower(relType)

	if t:find('adopt') then
		return 'adopted'
	end
	if t:find('step') then
		return 'step'
	end
	if t:find('bio') then
		return nil
	end

	return relType
end

local function findUnionBetween(people, name1, name2)
	if not isRealValue(name1) or not isRealValue(name2) then
		return nil
	end

	local person = people[name1]
	if not person or not person.unions then
		return nil
	end

	for _, union in ipairs(person.unions) do
		if union.partner == name2 then
			return union
		end
	end

	return nil
end

local function getMarriageYear(union)
	if not union then
		return nil
	end

	local raw = union.marriageDate or union.engagementDate or union.startDate
	if not isRealValue(raw) then
		return nil
	end

	local year = tostring(raw):match('^(%d%d%d%d)')
	return year or tostring(raw)
end

local function getRootSiblingSequence(people, root)
	local siblings = getSiblings(people, root)
	sortNames(people, siblings)

	local seq = {}
	local inserted = false
	local midpoint = math.floor(#siblings / 2) + 1

	for i, sib in ipairs(siblings) do
		if i == midpoint then
			table.insert(seq, root)
			inserted = true
		end
		table.insert(seq, sib)
	end

	if not inserted then
		table.insert(seq, root)
	end

	return seq
end

local function getFamilyGroupsForRoot(people, root)
	local person = people[root]
	if not person or not person.childLinks then
		return {}
	end

	local groups = {}

	for _, link in ipairs(person.childLinks) do
		local key

		if isRealValue(link.unionID) then
			key = 'union::' .. link.unionID
		elseif isRealValue(link.otherParent) then
			key = 'partner::' .. link.otherParent
		else
			key = 'single::' .. root
		end

		if not groups[key] then
			groups[key] = {
				unionID = link.unionID,
				partner = link.otherParent,
				children = {}
			}
		end

		table.insert(groups[key].children, {
			name = link.child,
			relationshipType = link.relationshipType,
			birthOrder = tonumber(link.birthOrder) or 999
		})
	end

	local out = {}
	for _, group in pairs(groups) do
		table.sort(group.children, function(a, b)
			return a.birthOrder < b.birthOrder
		end)
		table.insert(out, group)
	end

	table.sort(out, function(a, b)
		local aSingle = not isRealValue(a.partner)
		local bSingle = not isRealValue(b.partner)

		if aSingle ~= bSingle then
			return aSingle
		end

		local ap = a.partner or ''
		local bp = b.partner or ''
		local ad = (people[ap] and people[ap].displayName) or ap
		local bd = (people[bp] and people[bp].displayName) or bp
		return mw.ustring.lower(ad) < mw.ustring.lower(bd)
	end)

	return out
end

-- =========================================
-- Rendering helpers
-- =========================================

local function renderCard(people, name, badgeText)
	if not isRealValue(name) then
		return nil
	end

	local person = people[name] or { name = name, displayName = name }

	local card = html.create('div')
	card:addClass('kbft-card')
	card:wikitext(makeLink(person.name, person.displayName))

	if isRealValue(badgeText) then
		card:tag('div')
			:addClass('kbft-years')
			:wikitext(badgeText)
	end

	return card
end

local function renderSingle(people, name)
	local wrap = html.create('div')
	wrap:addClass('kbft-single')
	wrap:node(renderCard(people, name))
	return wrap
end

local function renderCouple(people, leftName, rightName, marriageYear)
	if isRealValue(leftName) and isRealValue(rightName) then
		local wrap = html.create('div')
		wrap:addClass('kbft-couple')

		wrap:node(renderCard(people, leftName))

		local marriage = wrap:tag('div')
		marriage:addClass('kbft-marriage')

		if isRealValue(marriageYear) then
			marriage:tag('div')
				:addClass('kbft-marriage-year')
				:wikitext(marriageYear)
		end

		marriage:tag('div')
			:addClass('kbft-marriage-line')

		wrap:node(renderCard(people, rightName))
		return wrap
	end

	if isRealValue(leftName) then
		return renderSingle(people, leftName)
	end

	if isRealValue(rightName) then
		return renderSingle(people, rightName)
	end

	return nil
end

local function renderRowSingles(people, names)
	local row = html.create('div')
	row:addClass('kbft-row')

	for _, name in ipairs(names) do
		row:node(renderSingle(people, name))
	end

	return row
end

local function renderGrandparentGeneration(people, parents)
	if #parents == 0 then
		return nil
	end

	local gen = html.create('div')
	gen:addClass('kbft-generation')

	local row = gen:tag('div')
	row:addClass('kbft-row')

	for _, parentName in ipairs(parents) do
		local parent = people[parentName]
		local gp1 = nil
		local gp2 = nil

		if parent then
			gp1 = parent.parents[1]
			gp2 = parent.parents[2]
		end

		local couple = renderCouple(people, gp1, gp2)
		if couple then
			row:node(couple)
		end
	end

	return gen
end

local function renderParentGeneration(people, root)
	local person = people[root]
	if not person then
		return nil
	end

	local parents = uniq(person.parents)
	if #parents == 0 then
		return nil
	end

	sortNames(people, parents)

	local gen = html.create('div')
	gen:addClass('kbft-generation')
	gen:node(renderCouple(people, parents[1], parents[2]))
	return gen
end

local function renderSiblingGeneration(people, root)
	local sequence = getRootSiblingSequence(people, root)

	local container = html.create('div')
	container:addClass('kbft-siblings')

	container:tag('div')
		:addClass('kbft-sibling-spine')

	local row = container:tag('div')
	row:addClass('kbft-sibling-row')

	for _, name in ipairs(sequence) do
		local unit = row:tag('div')
		unit:addClass('kbft-sibling-unit')

		unit:tag('div')
			:addClass('kbft-sibling-up')

		unit:node(renderSingle(people, name))
	end

	return container
end

local function renderFamilyGroupsGeneration(people, root)
	local groups = getFamilyGroupsForRoot(people, root)
	if #groups == 0 then
		return nil
	end

	local gen = html.create('div')
	gen:addClass('kbft-generation')

	local row = gen:tag('div')
	row:addClass('kbft-row')

	for _, group in ipairs(groups) do
		local unit = row:tag('div')
		unit:addClass('kbft-sibling-unit')

		local mainWrap = unit:tag('div')
		mainWrap:addClass('kbft-family-main-wrap')

		local marriageYear = nil
		if isRealValue(group.partner) then
			local union = findUnionBetween(people, root, group.partner)
			marriageYear = getMarriageYear(union)
		end

		mainWrap:node(renderCouple(people, root, group.partner, marriageYear))

		if #group.children > 0 then
			unit:tag('div')
				:addClass('kbft-child-down')

			local childRow = unit:tag('div')
			childRow:addClass('kbft-children')

			for _, child in ipairs(group.children) do
				childRow:node(
					renderCard(
						people,
						child.name,
						relationshipBadge(child.relationshipType)
					)
				)
			end
		end
	end

	return gen
end

-- =========================================
-- Public renderers
-- =========================================

local function renderConnectedForRoot(people, root)
	local person = people[root]
	if not person then
		return '<strong>FamilyTree error:</strong> No character found for "' .. tostring(root) .. '".'
	end

	local connected = getConnectedPeople(people, root)
	sortNames(people, connected)

	local node = html.create('div')
	node:addClass('kbft-tree')

	node:tag('div')
		:addClass('kbft-title')
		:wikitext('Connected to ' .. makeLink(person.name, person.displayName))

	local row = node:tag('div')
	row:addClass('kbft-row')

	for _, name in ipairs(connected) do
		row:node(renderSingle(people, name))
	end

	return tostring(node)
end

local function renderProfileForRoot(people, root)
	local person = people[root]
	if not person then
		return '<strong>FamilyTree error:</strong> No character found for "' .. tostring(root) .. '".'
	end

	local node = html.create('div')
	node:addClass('kbft-tree')

	node:tag('div')
		:addClass('kbft-title')
		:wikitext(makeLink(person.name, person.displayName))

	local function addSection(label, values)
		values = uniq(values)
		if #values == 0 then
			return
		end

		node:tag('div')
			:addClass('kbft-title')
			:css('margin-top', '22px')
			:wikitext(label)

		node:tag('div')
			:addClass('kbft-generation')
			:node(renderRowSingles(people, values))
	end

	addSection('Parents', person.parents)
	addSection('Partners', person.partners)
	addSection('Children', person.children)

	return tostring(node)
end

local function renderTreeForRoot(people, root)
	local person = people[root]
	if not person then
		return '<strong>FamilyTree error:</strong> No character found for "' .. tostring(root) .. '".'
	end

	local node = html.create('div')
	node:addClass('kbft-tree')

	node:tag('div')
		:addClass('kbft-title')
		:wikitext('Family Tree: ' .. makeLink(person.name, person.displayName))

	local parents = uniq(person.parents)
	sortNames(people, parents)

	local grandparentGen = renderGrandparentGeneration(people, parents)
	if grandparentGen then
		node:node(grandparentGen)
		node:tag('div'):addClass('kbft-connector')
	end

	local parentGen = renderParentGeneration(people, root)
	if parentGen then
		node:node(parentGen)
		node:tag('div'):addClass('kbft-connector')
	end

	node:node(renderSiblingGeneration(people, root))

	local familyGroupsGen = renderFamilyGroupsGeneration(people, root)
	if familyGroupsGen then
		node:tag('div'):addClass('kbft-connector')
		node:node(familyGroupsGen)
	end

	return tostring(node)
end

-- =========================================
-- Public functions
-- =========================================

function p.tree(frame)
	local root = getRoot(frame)
	local people = loadData()
	return renderTreeForRoot(people, root)
end

function p.profile(frame)
	local root = getRoot(frame)
	local people = loadData()
	return renderProfileForRoot(people, root)
end

function p.connected(frame)
	local root = getRoot(frame)
	local people = loadData()
	return renderConnectedForRoot(people, root)
end

return p