View source for Module:TypeEffectivenessGenerator

local p = {}

local types = {
"Normal","Fire","Water","Electric","Grass","Ice","Fighting","Poison",
"Ground","Flying","Psychic","Bug","Rock","Ghost","Dragon","Dark",
"Steel","Fairy","Shadow"
}

local chart = {

Normal = {Rock=0.5, Ghost=0, Steel=0.5},
Fire = {Fire=0.5, Water=0.5, Grass=2, Ice=2, Bug=2, Rock=0.5, Dragon=0.5, Steel=2},
Water = {Fire=2, Water=0.5, Grass=0.5, Ground=2, Rock=2, Dragon=0.5},
Electric = {Water=2, Electric=0.5, Grass=0.5, Ground=0, Flying=2, Dragon=0.5},
Grass = {Fire=0.5, Water=2, Grass=0.5, Poison=0.5, Ground=2, Flying=0.5, Bug=0.5, Rock=2, Dragon=0.5, Steel=0.5},
Ice = {Fire=0.5, Water=0.5, Grass=2, Ground=2, Flying=2, Dragon=2, Steel=0.5},
Fighting = {Normal=2, Ice=2, Poison=0.5, Flying=0.5, Psychic=0.5, Bug=0.5, Rock=2, Ghost=0, Dark=2, Steel=2, Fairy=0.5},
Poison = {Grass=2, Poison=0.5, Ground=0.5, Rock=0.5, Ghost=0.5, Steel=0, Fairy=2},
Ground = {Fire=2, Electric=2, Grass=0.5, Poison=2, Flying=0, Bug=0.5, Rock=2, Steel=2},
Flying = {Electric=0.5, Grass=2, Fighting=2, Bug=2, Rock=0.5, Steel=0.5},
Psychic = {Fighting=2, Poison=2, Psychic=0.5, Dark=0, Steel=0.5},
Bug = {Fire=0.5, Grass=2, Fighting=0.5, Poison=0.5, Flying=0.5, Psychic=2, Ghost=0.5, Dark=2, Steel=0.5, Fairy=0.5},
Rock = {Fire=2, Ice=2, Fighting=0.5, Ground=0.5, Flying=2, Bug=2, Steel=0.5},
Ghost = {Normal=0, Psychic=2, Ghost=2, Dark=0.5},
Dragon = {Dragon=2, Steel=0.5, Fairy=0},
Dark = {Fighting=0.5, Psychic=2, Ghost=2, Dark=0.5, Fairy=0.5},
Steel = {Fire=0.5, Water=0.5, Electric=0.5, Ice=2, Rock=2, Fairy=2, Steel=0.5},
Fairy = {Fire=0.5, Fighting=2, Poison=0.5, Dragon=2, Dark=2, Steel=0.5},
Shadow = {
    Normal=2, Fire=2, Water=2, Electric=2, Grass=2, Ice=2, Fighting=2, Poison=2,
    Ground=2, Flying=2, Psychic=2, Bug=2, Rock=2, Ghost=2, Dragon=2, Dark=2,
    Steel=2, Fairy=2, Shadow=0.5,
}
}

local function multiplier(atk, def)
    if not def or def == "" then return 1 end
    if def == "Shadow" then
        if atk == "Shadow" then return 0.5 end
        return 1
    end
    if chart[atk] and chart[atk][def] then
        return chart[atk][def]
    end
    return 1
end

local function renderChart(frame, type1, type2)

    local results = {
        ["4"]={},["2"]={},["1"]={},["0.5"]={},["0.25"]={},["0"]={}
    }

    for _,atk in ipairs(types) do

        local m = multiplier(atk,type1) * multiplier(atk,type2)

        if m==4 then table.insert(results["4"],atk)
        elseif m==2 then table.insert(results["2"],atk)
        elseif m==1 then table.insert(results["1"],atk)
        elseif m==0.5 then table.insert(results["0.5"],atk)
        elseif m==0.25 then table.insert(results["0.25"],atk)
        elseif m==0 then table.insert(results["0"],atk)
        end

    end

    local function join(t) return table.concat(t,", ") end

    return frame:expandTemplate{
        title="TypeEffectiveness",
        args={
            Color=type1,
            SuperWeaknesses=join(results["4"]),
            Weaknesses=join(results["2"]),
            Neutralities=join(results["1"]),
            Resistances=join(results["0.5"]),
            SuperResistances=join(results["0.25"]),
            Immunities=join(results["0"])
        }
    }

end

function p.generate(frame)
    local type1 = mw.text.trim(frame.args.type1 or "")
    local type2 = mw.text.trim(frame.args.type2 or "")
    return renderChart(frame, type1, type2)
end

local function getArgs(frame)
    local args = {}
    local parent = frame:getParent()

    if parent then
        for k, v in pairs(parent.args) do
            if v ~= '' then
                args[k] = v
            end
        end
    end

    for k, v in pairs(frame.args) do
        if v ~= '' then
            args[k] = v
        end
    end

    return args
end

local PokemonForms = require('Module:PokemonForms')
local FORM_SUFFIXES = PokemonForms.FORM_SUFFIXES
local STRIP_SUFFIXES = PokemonForms.STRIP_SUFFIXES

local function sentenceCase(s)
    local first = mw.ustring.sub(s, 1, 1)
    local rest = mw.ustring.sub(s, 2)
    return mw.ustring.upper(first) .. mw.ustring.lower(rest)
end

local function hyphenSpaceVariant(s)
    if mw.ustring.find(s, '-', 1, true) then
        return (mw.ustring.gsub(s, '%-', ' '))
    end
    return (mw.ustring.gsub(s, ' ', '-'))
end

local function stripKnownFormSuffix(name)
    for _, suffix in ipairs(STRIP_SUFFIXES) do
        local tail = '-' .. suffix
        if mw.ustring.sub(name, -mw.ustring.len(tail)) == tail then
            return mw.ustring.sub(name, 1, -mw.ustring.len(tail) - 1)
        end
    end
    return nil
end

local pokemonTypeCache = nil

local function loadPokemonTypes()
    if pokemonTypeCache then return pokemonTypeCache end

    local csvData = mw.ext.externalData.getExternalData({
        url    = 'https://wiki.pokemonrevolution.net/index.php?title=Special:GetData/PokemonRawList',
        format = 'CSV with header',
        data = {
            Name = 'Name',
            Type1 = 'Type1',
            Type2 = 'Type2',
        },
    })

    pokemonTypeCache = {}
    if csvData then
        for _, row in ipairs(csvData) do
            pokemonTypeCache[row.Name] = {
                type1 = row.Type1 or '',
                type2 = row.Type2 or '',
            }
        end
    end
    return pokemonTypeCache
end

local function lookupTypes(name)
    local typeData = loadPokemonTypes()
    return typeData[name] or typeData[sentenceCase(name)] or typeData[hyphenSpaceVariant(name)]
end

function p.main(frame)
    local args = getArgs(frame)
    local pokemonName = args[1] or args.pokemon or mw.title.getCurrentTitle().text

    if not pokemonName or pokemonName == '' then
        return "<span style='color:red'>Usage: {{PokemonTypeEffectiveness|PokemonName}}</span>"
    end

    local basePt = lookupTypes(pokemonName)
    local baseName = pokemonName
    if not basePt then
        local stripped = stripKnownFormSuffix(pokemonName)
        if stripped then
            local strippedPt = lookupTypes(stripped)
            if strippedPt then
                baseName, basePt = stripped, strippedPt
            end
        end
    end

    if not basePt then
        return "<span style='color:red'>Data load error: No type data found for " .. pokemonName .. "</span>"
    end

    local forms = {}
    for _, form in ipairs(FORM_SUFFIXES) do
        local candidateName
        if form.learnsetStyle == 'prefix-space' then
            candidateName = form.suffix .. ' ' .. baseName
        else
            candidateName = baseName .. '-' .. form.suffix
        end

        local rawListName
        if form.rawList == 'suffix' then
            rawListName = candidateName
        elseif form.rawList == 'prefix-hyphen' then
            rawListName = form.label .. '-' .. baseName
        elseif form.rawList == 'suffix-label' then
            rawListName = baseName .. '-' .. form.label
        elseif form.rawList == 'suffix-space' then
            rawListName = baseName .. ' ' .. form.label
        else
            rawListName = form.label .. ' ' .. baseName
        end

        local pt = lookupTypes(rawListName)
        if pt then
            table.insert(forms, { label = form.label, type1 = pt.type1, type2 = pt.type2, baseLabel = form.baseLabel })
        end
    end

    if #forms == 0 then
        return renderChart(frame, basePt.type1, basePt.type2)
    end

    local baseLabel = 'Normal'
    for _, f in ipairs(forms) do
        if f.baseLabel then
            baseLabel = f.baseLabel
            break
        end
    end

    local participants = { { label = baseLabel, type1 = basePt.type1, type2 = basePt.type2, isBase = true } }
    for _, f in ipairs(forms) do
        table.insert(participants, f)
    end

    local tabHeaders, infoBoxes = '', ''
    for _, part in ipairs(participants) do
        tabHeaders = tabHeaders .. "{{Tabs:MultiformTab|Form=" .. part.label .. "}}"

        local box = renderChart(frame, part.type1, part.type2)
        if part.isBase then
            infoBoxes = infoBoxes .. "{{Tabs:MultiformTableAdditional|Table=" .. box .. "}}"
        else
            infoBoxes = infoBoxes .. "{{Tabs:MultiformTableAdditional|Form=" .. part.label .. "|Table=" .. box .. "}}"
        end
    end

    local output = "{{TabbedInfoBoxes|Width=85|Class=|TabHeaders=" .. tabHeaders .. "|InfoBoxes=\n"
        .. infoBoxes .. "\n}}"

    return frame:preprocess(output)
end

return p