Module:Stats

From https://kenshi.fandom.com/wiki/Module:Stats

---

local p = {}

local groupKeys = { 'Humans', 'Shek', 'Hive', 'Fogmen', 'Skeletons', 'Cannibals', 'Fishmen' }
local raceKeys = { 'greenlander', 'scorchlander', 'shek', 'hive prince', 'hive soldier drone', 'hive worker drone', 'hive prince south hive', 'hive soldier drone south hive', 'hive worker drone south hive', 'hive queen', 'deadhive prince', 'deadhive soldier', 'deadhive worker', 'p4 unit', 'skeleton p4mkii', 'screamer mki', 'skeleton mkii screamer', 'skeleton', 'skeleton log-head mkii', 'skeleton no-head mkii', 'soldierbot', 'cannibal', 'cannibal skav', 'fishman', 'alpha fishman' }

local races = {
    ['greenlander'] = {
        group = 'Humans',
        name = 'Greenlander'
    },
    ['scorchlander'] = {
        group = 'Humans',
        name = 'Scorchlander'
    },
    ['shek'] = {
        group = 'Shek',
        name = 'Shek'
    },
    ['hive soldier drone'] = {
        group = 'Hive',
        name = 'Hive Soldier Drone'
    },
    ['hive prince'] = {
        group = 'Hive',
        name = 'Hive Prince'
    },
    ['hive worker drone'] = {
        group = 'Hive',
        name = 'Hive Worker Drone'
    },
    ['hive prince south hive'] = {
        group = 'Hive',
        name = 'Southern Hive Prince'
    },
    ['hive soldier drone south hive'] = {
        group = 'Hive',
        name = 'Southern Hive Soldier Drone'
    },
    ['hive worker drone south hive'] = {
        group = 'Hive',
        name = 'Southern Hive Worker Drone'
    },
    ['hive queen'] = {
        group = 'Hive',
        name = 'Hive Queen'
    },
    ['deadhive prince'] = {
        group = 'Fogmen',
        name = 'Deadhive Prince'
    },
    ['deadhive soldier'] = {
        group = 'Fogmen',
        name = 'Deadhive Soldier'
    },
    ['deadhive worker'] = {
        group = 'Fogmen',
        name = 'Deadhive Worker'
    },
    ['p4 unit'] = {
        group = 'Skeletons',
        name = 'P4 Unit'
    },
    ['skeleton p4mkii'] = {
        group = 'Skeletons',
        name = 'Skeleton P4MkII'
    },
    ['screamer mki'] = {
        group = 'Skeletons',
        name = 'Screamer MkI'
    },
    ['skeleton mkii screamer'] = {
        group = 'Skeletons',
        name = 'Skeleton MKII Screamer'
    },
    ['skeleton'] = {
        group = 'Skeletons',
        name = 'Skeleton'
    },
    ['skeleton log-head mkii'] = {
        group = 'Skeletons',
        name = 'Skeleton Log-Head MkII'
    },
    ['skeleton no-head mkii'] = {
        group = 'Skeletons',
        name = 'Skeleton No-Head MkII '
    },
    ['soldierbot'] = {
        group = 'Skeletons',
        name = 'Soldierbot'
    },
    ['cannibal'] = {
        group = 'Cannibals',
        name = 'Cannibal'
    },
    ['cannibal skav'] = {
        group = 'Cannibals',
        name = 'Cannibal Skav'
    },
    ['fishman'] = {
        group = 'Fishmen',
        name = 'Fishman'
    },
    ['alpha fishman'] = {
        group = 'Fishmen',
        name = 'Alpha Fishman'
    }
}

function p.main(frame)
    local args = frame:getParent().args
    
    return p._main(frame, args)
end

function p._main(frame, args)
    local template = ''
    if string.trim(args[1]) == 'calculated' then
        template = 'Stats/Calculated'
    else
        template = 'Stats/Base'
    end
    
    if args['Races'] == nil then
        return frame:expandTemplate { title = template, args = args }
    end
    
    return p._generate_tabber(frame, args, template)
end

function p._generate_tabber(frame, args, template)
    local tabberArgs = {}
    
    local racesArg = args['Races'] or ''
    racesArg = string.lower(racesArg)

    local _, racesCount = string.gsub(racesArg, ",", "")
    if racesCount > 4 then
        tabberArgs = p._generate_categorized_tabber(frame, args, template, racesArg)
    else
        tabberArgs = p._generate_uncategorized_tabber(frame, args, template, racesArg)
    end

    return frame:expandTemplate { title = 'Tabber', args = tabberArgs }
end

function p._generate_categorized_tabber(frame, args, template, racesArg)
    local groups = {}

    for i, raceKey in ipairs(raceKeys) do
    	for raceArg in string.gmatch(racesArg, '([^,]+)') do
    		if raceKey == raceArg then
    			local details = races[raceKey]
		    	local templateArgs = table.shallow_copy(args)
		    	
		    	templateArgs['Race'] = raceKey
		    	
		    	local statsTemplate = frame:expandTemplate { title = template, args = templateArgs }

                if groups[details.group] == nil then
                    groups[details.group] = {}
                end

                table.insert(groups[details.group], details.name)
                table.insert(groups[details.group], statsTemplate)
    		end
		end
    end
    
    local results = {}
    
    for i, groupKey in ipairs(groupKeys) do
        if groups[groupKey] ~= nil then
            local groupArgs = groups[groupKey]
            table.insert(results, groupKey)
            table.insert(results, frame:expandTemplate { title = 'Tabber', args = groupArgs })
        end
    end

    return results
end

function p._generate_uncategorized_tabber(frame, args, template, racesArg)
    local results = {}

    for i, raceKey in ipairs(raceKeys) do
    	for raceArg in string.gmatch(racesArg, '([^,]+)') do
    		if raceKey == raceArg then
    			local details = races[raceKey]
		    	local templateArgs = table.shallow_copy(args)
		    	
		    	templateArgs['Race'] = raceKey
		    	
		    	local statsTemplate = frame:expandTemplate { title = template, args = templateArgs }
		    	table.insert(results, details.name)
		    	table.insert(results, statsTemplate)
    		end
		end
    end

    return results
end

function table.shallow_copy(t)
    local t2 = {}
    for k,v in pairs(t) do
        t2[k] = v
    end
    return t2
end

function string.trim(arg)
	if arg == nil then
		return nil
	end
	
	return string.gsub(arg, "%s+", "")
end

return p