assign-beliefs.lua

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

local help = [====[

assign-beliefs
==============
A script to change the beliefs (values) of a unit.

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

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

Resetting a belief 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 [ <BELIEF> <LEVEL> <BELIEF> <LEVEL> <...> ]``:
                    the beliefs to modify and their levels. The
                    valid belief 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 beliefs to a neutral level. If the script is
                    called with both this option and a list of beliefs/levels,
                    first all the unit beliefs will be reset and then those
                    beliefs listed after ``-beliefs`` will be modified.

Example:

``assign-beliefs -reset -beliefs [ TRADITION 2 CRAFTSMANSHIP 3 POWER 0 CUNNING -1 ]``
    Resets all the unit beliefs, then sets the listed
    beliefs to the following values:

    * Tradition: a random value between 26 and 40 (level 2);
    * Craftsmanship: a random value between 41 and 50 (level 3);
    * Power: a random value between -10 and 10 (level 0);
    * Cunning: a random value between -25 and -11 (level -1).

    The final result (for a dwarf) will be:
    ``She personally is a firm believer in
    the value of tradition and sees guile and
    cunning as indirect and somewhat worthless.``

Note that the beliefs aligned with the cultural
values of the unit have not triggered a report.
]====]

local utils = require("utils")

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

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

-- ------------------------------------------------- ASSIGN BELIEFS ------------------------------------------------- --
-- NOTE: in the game data, beliefs are called values.

--- Assign the given beliefs to a unit, clearing all the other beliefs if requested.
---   :beliefs: nil, or a table. The fields have the belief 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(beliefs, unit, reset)
    assert(not beliefs or type(beliefs) == "table")
    assert(not unit or type(unit) == "number" or type(unit) == "userdata")
    assert(not reset or type(reset) == "boolean")

    beliefs = beliefs or {}
    reset = reset or false

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

    --- Calculate a random value within the desired belief level
    local function calculate_random_belief_value(belief_level)
        local beliefs_levels = {
            -- {minimum level value, maximum level value}
            { -50, -41 }, -- -3: lowest
            { -40, -26 }, -- -2: very low
            { -25, -11 }, -- -1: low
            { -10, 10 }, -- 0: neutral
            { 11, 25 }, -- 1: high
            { 26, 40 }, -- 2: very high
            { 41, 51 } -- 3: highest
        }
        local lvl = beliefs_levels[belief_level + 4]
        return math.random(lvl[1], lvl[2])
    end

    -- clear beliefs
    if reset then
        unit.status.current_soul.personality.values = {}
    end

    --assign beliefs
    for belief, level in pairs(beliefs) do
        assert(type(level) == "number")
        belief = belief:upper()
        -- there's a typo in the game data
        if belief == "PERSEVERANCE" then
            belief = "PERSEVERENCE"
        end
        if df.value_type[belief] then
            if level >= -3 and level <= 3 then
                local belief_value = calculate_random_belief_value(level)
                utils.insert_or_update(unit.status.current_soul.personality.values,
                                       { new = true, type = df.value_type[belief], strength = belief_value },
                                       "type")
            else
                print_yellow("WARNING: level out of range for belief '" .. belief .. "'. Skipping...")
            end
        else
            print_yellow("WARNING: '" .. belief .. "' is not a valid belief 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 beliefs list
    local beliefs = {}
    if args.beliefs then
        local i = 1
        while i <= #args.beliefs do
            local v = args.beliefs[i]
            -- v can be a belief name but it can also be a level value, so we have to check
            if not tonumber(v) then
                -- assume it's a valid belief name, for now
                local belief_name = tostring(v):upper()
                -- then try to get the level value
                local level_str = args.beliefs[i + 1]
                if not level_str then
                    -- we reached the end of the beliefs 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
                beliefs[belief_name] = level_int
            end
            i = i + 1 -- skip next arg because we already consumed it
        end
    end

    assign(beliefs, unit, reset)
end

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