Crawl 0.15 autopickup script

-- Autopickup script for DCSS 0.15

-- TL;DR if you pick up armour or weapon that you think shouldn't have been
-- picked up, add an inscription _lock to your current one (with { key)

-- Picks up:
-- all artefacts
-- one blowgun and one buckler per game (shield for troll/ogre)
-- needles if blowgun hasn't been found yet or one is in inventory
-- armour for empty auxillary slots
-- ego armour/shield/weapon that matches current items
-- * can be disabled with _noego inscription
-- * for shields also heavier versions are picked up
-- better weapons of the same school than anything in your inventory
-- * can be disabled with _noup inscription
-- * going from one handed to two handed can be enabled with _2h inscription
-- items that have better plusses and the same brand as the current item (Ash)
-- * brand can be ignored with _plain inscription

-- Options with default values:
-- autopickup_artefacts = true
-- autopickup_blowgun = true
-- autopickup_buckler = true
-- autopickup_upgrade_weapon = true

-- Inscriptions:
-- _noego       don't pick up matching ego items
-- _noup        don't upgrade this item
-- _lock        shortcut for both _noego and _noup
-- _plain       don't consider this item to have an ego or be an artefact
-- _2h          allow replacing this weapon with a two handed version
--              implies _up
-- _noart       don't pick up artefacts of this type

-- TODO:
-- fix small race two handed weapon detection
-- consider all carried armour
-- add inscriptions to picked up items in certain cases (_2h, _noup)

local AUTOPICKUP_ARTEFACTS = true
local AUTOPICKUP_BLOWGUN = true
local AUTOPICKUP_BUCKLER = true
local AUTOPICKUP_UPGRADE_WEAPON = true

local initialized = false
local need_inventory_check = false

local body_armour_data = {
    { "robe",                   2,  0, false },
    { "leather armour",         3,  4, false },
    { "ring mail",              5,  7, false },
    { "scale mail",             6, 11, false },
    { "chain mail",             8, 15, false },
    { "plate armour",          10, 18, false },
    { "crystal plate armour",  14, 24, false },
    { "animal skin",            2,  0, false },
    { "troll leather armour",   4,  4, true },
    { "steam dragon armour",    5,  0, true },
    { "mottled dragon armour",  6,  5, true },
    { "swamp dragon armour",    7,  7, true },
    { "fire dragon armour",     8, 11, true },
    { "ice dragon armour",      9, 11, true },
    { "pearl dragon armour",   10, 11, true },
    { "storm dragon armour",   10, 17, true },
    { "gold dragon armour",    12, 25, true },
    { "troll hide",             4,  4, true },
    { "steam dragon hide",      5,  0, true },
    { "mottled dragon hide",    6,  5, true },
    { "swamp dragon hide",      7,  7, true },
    { "fire dragon hide",       8, 11, true },
    { "ice dragon hide",        9, 11, true },
    { "pearl dragon hide",     10, 11, true },
    { "storm dragon hide",     10, 17, true },
    { "gold dragon hide",      12, 25, true } }
-- scythe handled the same as halberd
local twohanded_weapon_list = { "great sword", "claymore", "battleaxe",
    "executioner's axe", "dire flail", "great mace", "giant club",
    "giant spiked club", "halberd", "glaive", "bardiche", "quaterstaff",
    "lajatang" }
local spriggan_onehanded_weapon_list = { "dagger", "short sword", "cutlass",
    "quick blade", "club", "whip", "mace", "flail", "demon whip", "falchion",
    "hand axe" }
local weapon_upgrade_table = {
    { "dagger", "short sword", "cutlass", "quick blade" },
    { "falchion", "long sword", "blessed falchion", "scimitar",
      "blessed long sword", "blessed scimitar", "demon blade", "bastard sword",
      "eudemon blade", "great sword", "blessed bastard sword",
      "blessed great sword", "claymore", "blessed claymore" },
    { "hand axe", "war axe", "broad axe", "battleaxe", "executioner's axe" },
    { "club", "whip", "hammer", "mace", "flail", "morningstar", "demon whip",
      "dire flail", "sacred scourge", "eveningstar", "great mace",
      "giant club", "giant spiked club" },
    { "spear", "trident", "halberd", "demon trident", "trishula", "glaive",
      "bardiche" },
    { "quarterstaff", "lajatang" },
    { "hunting sling", "greatsling" },
    { "shortbow", "longbow" },
    { "hand crossbow", "arbalest", "triple crossbow" },
    { "blowgun" } }

-- HELPER FUNCTIONS

local function Set(list)
    local set = {}
    for _, l in ipairs(list) do set[l] = true end
    return set
end

local function control(c)
    return string.char(string.byte(c) - string.byte('a') + 1)
end

local function stash_search(search)
    if search:find("dropped") and you.turns() == 0 then
        return false
    end
    crawl.process_keys(control('f') .. search .. "\r" .. string.char(27)
                       .. string.char(27))
    crawl.flush_input()
    return not crawl.messages(2):find("Can't find anything matching that.")
end

local function warning(message)
    --crawl.mpr("Script error: " .. message)
end

local function inventory()
    return iter.invent_iterator:new(items.inventory())
end

local function upgrade_list_for_weapon(name)
    if name == "scythe" then
        name = "halberd"
    end
    local found
    for row = 1, #weapon_upgrade_table do
        for col = 1, #weapon_upgrade_table[row] do
            if name == weapon_upgrade_table[row][col] then
                found = { row=row, col=col }
                break
            end
        end
        if found then
            break
        end
    end
    local from = {}
    local to = {}
    if not found then
        warning("No upgrade data for weapon: " .. name)
    else
        for col = 1, found.col - 1 do
            from[weapon_upgrade_table[found.row][col]] = true
            if weapon_upgrade_table[found.row][col] == "halberd" then
                from["scythe"] = true
            end
        end
        for col = found.col, #weapon_upgrade_table[found.row] do
            to[weapon_upgrade_table[found.row][col]] = true
            if weapon_upgrade_table[found.row][col] == "halberd" then
                to["scythe"] = true
            end
        end
    end
    return from, to
end

local function plus_or_0(item)
    if item.plus == nil then
        return 0
    else
        return item.plus
    end
end

local function basename(item)
    if item.artefact then
        local name = item.name("base")
        if item.fully_identified then
            return name:gsub(" \".-\"$", ""):gsub(" of .-$", "")
        else
            return name:gsub("^.- ", "")
        end
    else
        return item.name("base")
    end
end

local function i_subtype(item)
    if item.class(true) == "jewellery" then
        if basename(item):find("amulet") then
            return "amulet"
        else
            return "ring"
        end
    end
    return item.subtype()
end

local function er(item)
    if item.class(true) ~= "armour" then return nil end
    local name = basename(item)
    local subtype = i_subtype(item)
    if subtype == "shield" then
        if name == "buckler" then
            return 1
        elseif name == "shield" then
            return 3
        elseif name == "large shield" then
            return 5
        else
            warning("Unknown shield: " .. name)
            return -1
        end
    elseif subtype == "body" then
        if body_armour_data[name] ~= nil then
            return body_armour_data[name].er
        else
            warning("Unknown body armour: " .. name)
            return -1
        end
    else
        if name:find("barding") then
            return 6
        else
            return 0
        end
    end
end

-- also does SH
local function base_ac(item)
    if item.class(true) ~= "armour" then return nil end
    local name = basename(item)
    local subtype = i_subtype(item)
    if subtype == "shield" then
        if name == "buckler" then
            return 5
        elseif name == "shield" then
            return 8
        elseif name == "large shield" then
            return 13
        else
            warning("Unknown shield: " .. name)
            return 99
        end
    elseif subtype == "body" then
        if body_armour_data[name] ~= nil then
            return body_armour_data[name].ac
        else
            return 99
        end
    else
        if name:find("barding") then
            return 4
        elseif name == "hat" then
            return 0
        else
            return 1
        end
    end
end

-- also does SH
local function estimate_ac(item)
    if item.class(true) ~= "armour" then return nil end
    return base_ac(item) + plus_or_0(item)
end

local function is_hide(item)
    if item.class(true) ~= "armour" or i_subtype(item) ~= "armour" then
        return nil
    end
    return basename(item):find("hide")
end

local function inventory_match(func, result)
    for item in inventory() do
        if func(item) == result then
            return item
        end
    end
    return nil
end

local function inscription_no_ego(item)
    return item.name("qual"):find("_noego") ~= nil
           or item.name("qual"):find("_lock") ~= nil
end

local function inscription_no_upgrade(item)
    return item.name("qual"):find("_noup") ~= nil
           or item.name("qual"):find("_lock") ~= nil
end

local function inscription_plain(item)
    return item.name("qual"):find("_plain") ~= nil
end

local function inscription_twohanded(item)
    return item.name("qual"):find("_2h") ~= nil
end

local function inscription_no_artefact(item)
    return item.name("qual"):find("_noart") ~= nil
end

local function is_upgradable(item)
    if item.class(true) == "weapon" and AUTOPICKUP_UPGRADE_WEAPON then
        return not inscription_no_upgrade(item)
    end
    return false
end

local function is_special(item)
    return (item.artefact or item.branded) and not inscription_plain(item)
end

local function is_twohanded(item)
    if item.class(true) ~= "weapon" then return nil end
    if you.genus() == "formicid" or inscription_twohanded(item) then
        return true
    end
    if you.genus() == "spriggan" then
        return not (spriggan_onehanded_weapon_list[basename(item)] ~= nil)
    end
    return twohanded_weapon_list[basename(item)] ~= nil
end

local function ego(item)
    if inscription_plain(item) then
        return "plain"
    elseif item.artefact then
        return "artefact"
    elseif item.ego_type_terse ~= "" and item.ego_type_terse ~= "unknown" then
        return item.ego_type_terse
    elseif body_armour_data[basename(item)] ~= nil
           and body_armour_data[basename(item)].ego == true then
        return "innate"
    elseif item.branded then
        return "unknown"
    else
        return "plain"
    end
end

local function same_ego(a, b)
    local ego_a = ego(a)
    if ego_a ~= ego(b) then
        return false
    end
    if ego_a == "unknown" or ego_a == "artefact" then
        return false
    end
    if ego_a == "innate" then
        return basename(a) == basename(b)
    end
    return true
end

local function parse_body_armour_data()
    local parsed = {}
    for _, v in ipairs(body_armour_data) do
        parsed[v[1]] = { ac=v[2], er=v[3], ego=v[4] }
    end
    return parsed
end

local function set_found(key)
    if inventory_match(basename, key) or stash_search(key .. " && dropped") then
        autopickup_found[key] = true
    else
        autopickup_found[key] = false
    end
end

local function smallest_shield(genus)
    if genus == "troll" or genus == "ogre" then
        return "shield"
    else
        return "buckler"
    end
end

local function initialize()
    if you.genus() == "felid" then
        autopickup_found = {}
        return
    end

    local shield_type = smallest_shield(you.genus())
    if autopickup_found == nil then
        autopickup_found = {}
        if AUTOPICKUP_BLOWGUN then
            autopickup_found["blowgun"] = false
        end
        if AUTOPICKUP_BUCKLER then
            autopickup_found[shield_type] = false
        end
        for key, _ in pairs(autopickup_found) do
            set_found(key)
        end
    else
        if not AUTOPICKUP_BLOWGUN then
            autopickup_found["blowgun"] = nil
        elseif autopickup_found["blowgun"]  == nil then
            set_found("blowgun")
        end
        if not AUTOPICKUP_BUCKLER then
            autopickup_found[shield_type] = nil
        elseif autopickup_found[shield_type]  == nil then
            set_found(shield_type)
        end
    end
end

local function check_inventory()
    for key, value in pairs(autopickup_found) do
        if value == false and inventory_match(basename, key) then
            autopickup_found[key] = true
        end
    end
end

-- CALLED BY THE GAME

function set_artefacts(key, value, mode)
    AUTOPICKUP_ARTEFACTS = string.lower(value) ~= "false"
end

function set_blowgun(key, value, mode)
    AUTOPICKUP_BLOWGUN = string.lower(value) ~= "false"
    if you.genus() == "felid" then return end
    if not AUTOPICKUP_BLOWGUN then
        autopickup_found["blowgun"] = nil
    elseif autopickup_found["blowgun"] == nil then
        set_found("blowgun")
    end
end

function set_buckler(key, value, mode)
    AUTOPICKUP_BUCKLER = string.lower(value) ~= "false"
    if you.genus() == "felid" then return end
    local shield_type = smallest_shield(you.genus())
    if not AUTOPICKUP_BUCKLER then
        autopickup_found[shield_type] = nil
    elseif autopickup_found[shield_type] == nil then
        set_found(shield_type)
    end
end

function set_upgrade_weapon(key, value, mode)
    AUTOPICKUP_UPGRADE_WEAPON = string.lower(value) ~= "false"
end

local function autopickup(item)
    local name = basename(item)
    local class = item.class(true)

    if item.artefact and AUTOPICKUP_ARTEFACTS then
        local skip = false
        local classes = {}
        local subtype = nil
        classes[class] = true
        if class == "weapon" then
            classes["rod"] = true
            classes["magical staff"] = true
        elseif class == "jewellery" or class == "armour" then
            subtype = i_subtype(item)
        end
        for inv in inventory() do
            if classes[inv.class(true)] and inscription_no_artefact(inv) then
                if subtype == nil or i_subtype(inv) == subtype then
                    skip = true
                    break
                end
            end
        end
        if not skip then
            return true
        end
    end

    if autopickup_found ~= nil then
        if autopickup_found[name] == false then
            return true
        end
        if autopickup_found["blowgun"] ~= nil and name:find("needle") then
            return autopickup_found["blowgun"] == false
                   or inventory_match(basename, "blowgun")
        end
    end

    if class == "armour" then
        local subtype = i_subtype(item)
        local current = nil
        if subtype == "body" then
            current = items.equipped_at("armour")
        else
            current = items.equipped_at(subtype)
        end
        if subtype == "body" then
            if current ~= nil and er(item) == er(current)
               and base_ac(item) == base_ac(current) then
                if is_hide(current) and not is_hide(item) then
                    return true
                elseif same_ego(item, current) then
                    if estimate_ac(item) > estimate_ac(current) then
                        return true
                    end
                elseif is_special(item)
                       and not inscription_no_ego(current) then
                    return true
                end
            end
        elseif subtype == "shield" then
            if current ~= nil and base_ac(item) >= base_ac(current) then
                if base_ac(item) > base_ac(current) then
                    if inscription_no_upgrade(item) then
                        return false
                    end
                    if same_ego(item, current) then
                        return true
                    elseif is_special(item)
                           and not inscription_no_ego(current) then
                        return true
                    end
                else
                    if same_ego(item, current) then
                        if estimate_ac(item) > estimate_ac(current) then
                            return true
                        end
                    elseif is_special(item)
                           and not inscription_no_ego(current) then
                        return true
                    end
                end
            end
        else
            if estimate_ac(item) >= 0 or is_special(item) then
                if current == nil then
                    return true
                end
                if same_ego(item, current) then
                    if estimate_ac(item) > estimate_ac(current) then
                        return true
                    end
                    if estimate_ac(item) == estimate_ac(current)
                       and base_ac(item) > base_ac(current) then
                        return true
                    end
                elseif is_special(item) and not inscription_no_ego(current) then
                    return true
                end
            end
        end
    end

    if class == "weapon" then
        local match_found = false
        local upgrade_found = false
        local no_upgrade = false
        local upgrade_from, upgrade_to = upgrade_list_for_weapon(name)
        for inv in inventory() do
            if class == inv.class(true) then
                if name == basename(inv) then
                    match_found = true
                    no_upgrade = true
                    if inscription_no_ego(inv) then
                        return false
                    end
                    if is_special(item) then
                        if same_ego(item, inv)
                           and plus_or_0(inv) >= plus_or_0(item) then
                            return false
                        end
                    elseif is_special(inv)
                           or plus_or_0(inv) >= plus_or_0(item) then
                        return false
                    end
                elseif upgrade_to[basename(inv)] then
                    no_upgrade = true
                elseif upgrade_from[basename(inv)] and
                       (not is_twohanded(item) or is_twohanded(inv)) then
                    upgrade_found = true
                    if not is_upgradable(inv) then
                        no_upgrade = true
                    end
                end
            end
        end
        return match_found or (upgrade_found and not no_upgrade)
    end
end
add_autopickup_func(autopickup)

function c_assign_invletter()
    need_inventory_check = true
end

function ready()
    if not initialized then
        crawl.more_autoclear(true)
        initialize()
        crawl.more_autoclear(false)
        initialized = true
    else
        if need_inventory_check then
            check_inventory()
        end
    end
end

chk_lua_option.autopickup_artefacts = set_artefacts
chk_lua_option.autopickup_blowgun = set_blowgun
chk_lua_option.autopickup_buckler = set_buckler
chk_lua_option.autopickup_upgrade_weapon = set_upgrade_weapon
if options.autopickup_artefacts ~= nil then
    set_artefacts(nil, options.autopickup_artefacts, nil);
end
if options.autopickup_blowgun ~= nil then
    set_blowgun(nil, options.autopickup_blowgun, nil);
end
if options.autopickup_buckler ~= nil then
    set_buckler(nil, options.autopickup_buckler, nil);
end
if options.autopickup_upgrade_weapon ~= nil then
    set_upgrade_weapon(nil, options.autopickup_upgrade_weapon, nil);
end

twohanded_weapon_list = Set(twohanded_weapon_list)
spriggan_onehanded_weapon_list = Set(spriggan_onehanded_weapon_list)
body_armour_data = parse_body_armour_data()

function write_bool_table(table, aname)
    local res = aname .. " = { "
    for k, v in pairs(table) do
        res = res .. k .. '=' .. tostring(v) .. ", "
    end
    return res .. "}\n"
end

function autopickup_save()
    return write_bool_table(autopickup_found, "autopickup_found")
end

table.insert(chk_lua_save, autopickup_save)