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 makeLink(name)
	if not isRealValue(name) then
		return ''
	end
	return string.format('[[%s|%s]]', name, name)
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 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 = {}
		}
	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 = {}
			}
		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)

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

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

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

local function loadUnions(people)
	local results = cargo.query(
		'Unions',
		'UnionID,Partner1,Partner2,UnionType,Status',
		{ 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)
		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

-- =========================================
-- Rendering
-- =========================================

local function renderPersonBox(name, extraClass)
	if not isRealValue(name) then
		return nil
	end

	local box = html.create('div')
	box:addClass('familytree-person')
	if isRealValue(extraClass) then
		box:addClass(extraClass)
	end
	box:wikitext(makeLink(name))
	return box
end

local function renderRow(label, names, rowClass)
	names = uniq(names)
	if #names == 0 then
		return nil
	end

	local row = html.create('div')
	row:addClass('familytree-row')
	if isRealValue(rowClass) then
		row:addClass(rowClass)
	end

	row:tag('div')
		:addClass('familytree-row-label')
		:wikitext(label)

	local items = row:tag('div')
		:addClass('familytree-row-items')

	for _, name in ipairs(names) do
		local node = renderPersonBox(name)
		if node then
			items:node(node)
		end
	end

	return row
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 grandparents = getGrandparents(people, root)
	local parents = uniq(person.parents)
	local siblings = getSiblings(people, root)
	local partners = uniq(person.partners)
	local children = uniq(person.children)

	sortNames(people, grandparents)
	sortNames(people, parents)
	sortNames(people, siblings)
	sortNames(people, partners)
	sortNames(people, children)

	local rootNode = html.create('div')
	rootNode:addClass('familytree-wrapper')

	local gpRow = renderRow('Grandparents', grandparents, 'familytree-grandparents')
	if gpRow then rootNode:node(gpRow) end

	local pRow = renderRow('Parents', parents, 'familytree-parents')
	if pRow then rootNode:node(pRow) end

	local selfRow = html.create('div')
	selfRow:addClass('familytree-row')
	selfRow:addClass('familytree-focus-row')

	selfRow:tag('div')
		:addClass('familytree-row-label')
		:wikitext('Focus')

	selfRow:tag('div')
		:addClass('familytree-row-items')
		:node(renderPersonBox(root, 'familytree-focus-person'))

	rootNode:node(selfRow)

	local sRow = renderRow('Siblings', siblings, 'familytree-siblings')
	if sRow then rootNode:node(sRow) end

	local partnerRow = renderRow('Partners', partners, 'familytree-partners')
	if partnerRow then rootNode:node(partnerRow) end

	local childRow = renderRow('Children', children, 'familytree-children')
	if childRow then rootNode:node(childRow) end

	return tostring(rootNode)
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 rootNode = html.create('div')
	rootNode:addClass('familytree-profile')

	rootNode:tag('div')
		:addClass('familytree-profile-name')
		:wikitext(makeLink(person.name))

	local dl = rootNode:tag('dl')

	local function addField(label, value)
		if isRealValue(value) then
			dl:tag('dt'):wikitext(label)
			dl:tag('dd'):wikitext(value)
		end
	end

	addField('Display Name', person.displayName ~= person.name and person.displayName or nil)

	if #person.parents > 0 then
		local links = {}
		for _, name in ipairs(person.parents) do
			table.insert(links, makeLink(name))
		end
		addField('Parents', table.concat(links, ', '))
	end

	if #person.partners > 0 then
		local links = {}
		for _, name in ipairs(person.partners) do
			table.insert(links, makeLink(name))
		end
		addField('Partners', table.concat(links, ', '))
	end

	if #person.children > 0 then
		local links = {}
		for _, name in ipairs(person.children) do
			table.insert(links, makeLink(name))
		end
		addField('Children', table.concat(links, ', '))
	end

	return tostring(rootNode)
end

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 rootNode = html.create('div')
	rootNode:addClass('familytree-connected')

	rootNode:tag('div')
		:addClass('familytree-connected-title')
		:wikitext('Connected to ' .. makeLink(root))

	local items = rootNode:tag('div')
		:addClass('familytree-row-items')

	for _, name in ipairs(connected) do
		items:node(renderPersonBox(name))
	end

	return tostring(rootNode)
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