assign-facets.lua

-- Change the facets of a unit.
--@ module = true

local help = [====[

assign-facets
=============
A script to change the facets (traits) of a unit.

Facets are defined with a token and a number from -3 to 3, which describes
the different levels of facets strength, as explained here:
https://dwarffortresswiki.org/index.php/DF2014:Personality_trait#Facets

========  =========
Strength  Effect
========  =========
3         Highest
2         Very High
1         High
0         Neutral
-1        Low
-2        Very Low
-3        Lowest
========  =========

Resetting a facet means setting it to a level that does not
trigger a report in the "Thoughts and preferences" screen.

Usage:

``-help``:
                    print the help page.

``-unit <UNIT_ID>``:
                    set the target unit ID. If not present, the
                    currently selected unit will be the target.

``-beliefs [ <FACET> <LEVEL> <FACET> <LEVEL> <...> ]``:
                    the facets to modify and their levels. The
                    valid facet tokens can be found in the wiki page
                    linked above; level values range from -3 to 3.
                    There must be a space before and after each square
                    bracket.

``-reset``:
                    reset all facets to a neutral level. If the script is
                    called with both this option and a list of facets/levels,
                    first all the unit facets will be reset and then those
                    facets listed after ``-facets`` will be modified.

Example::

    assign-facets -reset -facets [ HATE_PROPENSITY -2 CHEER_PROPENSITY -1 ]

Resets all the unit facets, then sets the listed facets to the following values:

* Hate propensity: a value between 10 and 24 (level -2);
* Cheer propensity: a value between 25 and 39 (level -1).

The final result (for a dwarf) will be: "She very rarely develops negative
feelings toward things. She is rarely happy or enthusiastic, and she is
conflicted by this as she values parties and merrymaking in the abstract."

Note that the facets are compared to the beliefs, and if conflicts arise they
will be reported.

]====]

local utils = require("utils")

local valid_args = utils.invert({
                                    'help',
                                    'unit',
                                    'facets',
                                    'reset',
                                })

-- ----------------------------------------------- UTILITY FUNCTIONS ------------------------------------------------ --
local function print_yellow(text)
    dfhack.color(COLOR_YELLOW)
    print(text)
    dfhack.color(-1)
end

-- ------------------------------------------------- ASSIGN FACETS -------------------------------------------------- --
-- -- NOTE: in the game data, facets are called traits.

--- Assign the given facets to a unit, resetting all the other facets if requested.
---   :facets: nil, or a table. The fields have the facet name as key and its level as value.
---   :unit: a valid unit id, a df.unit object, or nil. If nil, the currently selected unit will be targeted.
---   :reset: boolean, or nil.
function assign(facets, unit, reset)
    assert(not facets or type(facets) == "table")
    assert(not unit or type(unit) == "number" or type(unit) == "userdata")
    assert(not reset or type(reset) == "boolean")

    facets = facets or {}
    reset = reset or false

    if type(unit) == "number" then
        unit = df.unit.find(tonumber(unit)) --luacheck:retype
    end
    unit = unit or dfhack.gui.getSelectedUnit(true)
    if not unit then
        qerror("No unit found.")
    end

    --- Calculate a random value within the desired facet strength level
    local function calculate_random_facet_strength(facet_level)
        local facets_levels = {
            -- {minimum level value, maximum level value}
            { 0, 9 }, -- -3: lowest
            { 10, 24 }, -- -2: very low
            { 25, 39 }, -- -1: low
            { 40, 60 }, -- 0: neutral
            { 61, 75 }, -- 1: high
            { 76, 90 }, -- 2: very high
            { 91, 100 } -- 3: highest
        }
        local lvl = facets_levels[facet_level + 4]
        return math.random(lvl[1], lvl[2])
    end

    -- reset facets
    if reset then
        for i, _ in pairs(unit.status.current_soul.personality.traits) do
            local value = calculate_random_facet_strength(0)
            unit.status.current_soul.personality.traits[i] = value
        end
    end

    --assign facets
    for facet, level in pairs(facets) do
        assert(type(level) == "number")
        facet = facet:upper()
        -- there's a typo in the game data
        if facet == "PERSEVERANCE" then
            facet = "PERSEVERENCE"
        end
        if df.personality_facet_type[facet] then
            if level >= -3 and level <= 3 then
                local facet_strength = calculate_random_facet_strength(level)
                unit.status.current_soul.personality.traits[facet] = facet_strength
            else
                print_yellow("WARNING: level out of range for facet '" .. level .. "'. Skipping...")
            end
        else
            print_yellow("WARNING: '" .. facet .. "' is not a valid facet token. Skipping...")
        end
    end
end

-- ------------------------------------------------------ MAIN ------------------------------------------------------ --
local function main(...)
    local args = utils.processArgs({...}, valid_args)

    if args.help then
        print(help)
        return
    end

    local unit
    if args.unit then
        unit = tonumber(args.unit)
        if not unit then
            qerror("'" .. args.unit .. "' is not a valid unit ID.")
        end
    end

    local reset = false
    if args.reset then
        reset = true
    end

    -- parse facets list
    local facets = {}
    if args.facets then
        local i = 1
        while i <= #args.facets do
            local v = args.facets[i]
            -- v can be a facet name but it can also be a level value, so we have to check
            if not tonumber(v) then
                -- assume it's a valid facet name, for now
                local facet_name = tostring(v):upper()
                -- then try to get the level value
                local level_str = args.facets[i + 1]
                if not level_str then
                    -- we reached the end of the facets list
                    qerror("Missing level value after '" .. v .. "'.")
                end
                local level_int = tonumber(level_str)
                if not level_int then
                    qerror("'" .. level_str .. "' is not a valid number.")
                end
                -- assume the level value is in range, for now
                facets[facet_name] = level_int
            end
            i = i + 1 -- skip next arg because we already consumed it
        end
    end

    assign(facets, unit, reset)
end

if not dfhack_flags.module then
    main(...)
end