dfhack/plugins/lua/sort.lua

local _ENV = mkmodule('plugins.sort')

local gui = require('gui')
local overlay = require('plugins.overlay')
local utils = require('utils')
local widgets = require('gui.widgets')

local CH_UP = string.char(30)
local CH_DN = string.char(31)

local MELEE_WEAPON_SKILLS = {
    df.job_skill.AXE,
    df.job_skill.SWORD,
    df.job_skill.MACE,
    df.job_skill.HAMMER,
    df.job_skill.SPEAR,
    df.job_skill.MELEE_COMBAT, --Fighter
}

local RANGED_WEAPON_SKILLS = {
    df.job_skill.CROSSBOW,
    df.job_skill.RANGED_COMBAT,
}

local LEADERSHIP_SKILLS = {
    df.job_skill.MILITARY_TACTICS,
    df.job_skill.LEADERSHIP,
    df.job_skill.TEACHING,
}

local function sort_noop(a, b)
    -- this function is used as a marker and never actually gets called
    error('sort_noop should not be called')
end

local function sort_by_name_desc(unit_id_1, unit_id_2)
    if unit_id_1 == unit_id_2 then return 0 end
    local unit1 = df.unit.find(unit_id_1)
    local unit2 = df.unit.find(unit_id_2)
    if not unit1 then return -1 end
    if not unit2 then return 1 end
    local name1 = dfhack.TranslateName(dfhack.units.getVisibleName(unit1))
    local name2 = dfhack.TranslateName(dfhack.units.getVisibleName(unit2))
    return utils.compare_name(name1, name2)
end

local function sort_by_name_asc(unit_id_1, unit_id_2)
    if unit_id_1 == unit_id_2 then return 0 end
    local unit1 = df.unit.find(unit_id_1)
    local unit2 = df.unit.find(unit_id_2)
    if not unit1 then return -1 end
    if not unit2 then return 1 end
    local name1 = dfhack.TranslateName(dfhack.units.getVisibleName(unit1))
    local name2 = dfhack.TranslateName(dfhack.units.getVisibleName(unit2))
    return utils.compare_name(name2, name1)
end

local function sort_by_migrant_wave_desc(unit_id_1, unit_id_2)
    if unit_id_1 == unit_id_2 then return 0 end
    local unit1 = df.unit.find(unit_id_1)
    local unit2 = df.unit.find(unit_id_2)
    if not unit1 then return -1 end
    if not unit2 then return 1 end
    return utils.compare(unit2.id, unit1.id)
end

local function sort_by_migrant_wave_asc(unit_id_1, unit_id_2)
    if unit_id_1 == unit_id_2 then return 0 end
    local unit1 = df.unit.find(unit_id_1)
    local unit2 = df.unit.find(unit_id_2)
    if not unit1 then return -1 end
    if not unit2 then return 1 end
    return utils.compare(unit1.id, unit2.id)
end

local function get_skill(unit_id, skill, unit)
    unit = unit or df.unit.find(unit_id)
    return unit and
        unit.status.current_soul and
        utils.binsearch(unit.status.current_soul.skills, skill, 'id')
end

local function get_max_skill(unit_id, list)
    local unit = df.unit.find(unit_id)
    if not unit or not unit.status.current_soul then return end
    local max
    for _,skill in ipairs(list) do
        local s = get_skill(unit_id, skill, unit)
        if s then
            if not max then
                max = s
            else
                if max.rating == s.rating and max.experience < s.experience then
                    max = s
                elseif max.rating < s.rating then
                    max = s
                end
            end
        end
    end
    return max
end

local function make_sort_by_skill_list_desc(list)
    return function(unit_id_1, unit_id_2)
        if unit_id_1 == unit_id_2 then return 0 end
        if unit_id_1 == -1 then return -1 end
        if unit_id_2 == -1 then return 1 end
        local s1 = get_max_skill(unit_id_1, list)
        local s2 = get_max_skill(unit_id_2, list)
        if s1 == s2 then return sort_by_name_desc(unit_id_1, unit_id_2) end
        if not s2 then return -1 end
        if not s1 then return 1 end
        if s1.rating ~= s2.rating then
            return utils.compare(s2.rating, s1.rating)
        end
        if s1.experience ~= s2.experience then
            return utils.compare(s2.experience, s1.experience)
        end
        return sort_by_name_desc(unit_id_1, unit_id_2)
    end
end

local function make_sort_by_skill_list_asc(list)
    return function(unit_id_1, unit_id_2)
        if unit_id_1 == unit_id_2 then return 0 end
        if unit_id_1 == -1 then return -1 end
        if unit_id_2 == -1 then return 1 end
        local s1 = get_max_skill(unit_id_1, list)
        local s2 = get_max_skill(unit_id_2, list)
        if s1 == s2 then return sort_by_name_desc(unit_id_1, unit_id_2) end
        if not s2 then return 1 end
        if not s1 then return -1 end
        if s1.rating ~= s2.rating then
            return utils.compare(s1.rating, s2.rating)
        end
        if s1.experience ~= s2.experience then
            return utils.compare(s1.experience, s2.experience)
        end
        return sort_by_name_desc(unit_id_1, unit_id_2)
    end
end

local function make_sort_by_skill_desc(sort_skill)
    return function(unit_id_1, unit_id_2)
        if unit_id_1 == unit_id_2 then return 0 end
        if unit_id_1 == -1 then return -1 end
        if unit_id_2 == -1 then return 1 end
        local s1 = get_skill(unit_id_1, sort_skill)
        local s2 = get_skill(unit_id_2, sort_skill)
        if s1 == s2 then return sort_by_name_desc(unit_id_1, unit_id_2) end
        if not s2 then return -1 end
        if not s1 then return 1 end
        if s1.rating ~= s2.rating then
            return utils.compare(s2.rating, s1.rating)
        end
        if s1.experience ~= s2.experience then
            return utils.compare(s2.experience, s1.experience)
        end
        return sort_by_name_desc(unit_id_1, unit_id_2)
    end
end

local function make_sort_by_skill_asc(sort_skill)
    return function(unit_id_1, unit_id_2)
        if unit_id_1 == unit_id_2 then return 0 end
        if unit_id_1 == -1 then return -1 end
        if unit_id_2 == -1 then return 1 end
        local s1 = get_skill(unit_id_1, sort_skill)
        local s2 = get_skill(unit_id_2, sort_skill)
        if s1 == s2 then return sort_by_name_desc(unit_id_1, unit_id_2) end
        if not s2 then return 1 end
        if not s1 then return -1 end
        if s1.rating ~= s2.rating then
            return utils.compare(s1.rating, s2.rating)
        end
        if s1.experience ~= s2.experience then
            return utils.compare(s1.experience, s2.experience)
        end
        return sort_by_name_desc(unit_id_1, unit_id_2)
    end
end

local SORT_FNS = {
    sort_by_any_melee_desc=make_sort_by_skill_list_desc(MELEE_WEAPON_SKILLS),
    sort_by_any_melee_asc=make_sort_by_skill_list_asc(MELEE_WEAPON_SKILLS),
    sort_by_any_ranged_desc=make_sort_by_skill_list_desc(RANGED_WEAPON_SKILLS),
    sort_by_any_ranged_asc=make_sort_by_skill_list_asc(RANGED_WEAPON_SKILLS),
    sort_by_leadership_desc=make_sort_by_skill_list_desc(LEADERSHIP_SKILLS),
    sort_by_leadership_asc=make_sort_by_skill_list_asc(LEADERSHIP_SKILLS),
    sort_by_axe_desc=make_sort_by_skill_desc(df.job_skill.AXE),
    sort_by_axe_asc=make_sort_by_skill_asc(df.job_skill.AXE),
    sort_by_sword_desc=make_sort_by_skill_desc(df.job_skill.SWORD),
    sort_by_sword_asc=make_sort_by_skill_asc(df.job_skill.SWORD),
    sort_by_mace_desc=make_sort_by_skill_desc(df.job_skill.MACE),
    sort_by_mace_asc=make_sort_by_skill_asc(df.job_skill.MACE),
    sort_by_hammer_desc=make_sort_by_skill_desc(df.job_skill.HAMMER),
    sort_by_hammer_asc=make_sort_by_skill_asc(df.job_skill.HAMMER),
    sort_by_spear_desc=make_sort_by_skill_desc(df.job_skill.SPEAR),
    sort_by_spear_asc=make_sort_by_skill_asc(df.job_skill.SPEAR),
    sort_by_crossbow_desc=make_sort_by_skill_desc(df.job_skill.CROSSBOW),
    sort_by_crossbow_asc=make_sort_by_skill_asc(df.job_skill.CROSSBOW),
}

-- ----------------------
-- SquadAssignmentOverlay
--

SquadAssignmentOverlay = defclass(SquadAssignmentOverlay, overlay.OverlayWidget)
SquadAssignmentOverlay.ATTRS{
    default_pos={x=-33, y=40},
    default_enabled=true,
    viewscreens='dwarfmode/UnitSelector/SQUAD_FILL_POSITION',
    frame={w=63, h=7},
    frame_style=gui.FRAME_PANEL,
    frame_background=gui.CLEAR_PEN,
}

function SquadAssignmentOverlay:init()
    self.dirty = true

    self:addviews{
        widgets.CycleHotkeyLabel{
            view_id='sort',
            frame={l=0, t=0, w=29},
            label='Sort by:',
            key='CUSTOM_SHIFT_S',
            options={
                {label='any melee skill'..CH_DN, value=SORT_FNS.sort_by_any_melee_desc, pen=COLOR_GREEN},
                {label='any melee skill'..CH_UP, value=SORT_FNS.sort_by_any_melee_asc, pen=COLOR_YELLOW},
                {label='any ranged skill'..CH_DN, value=SORT_FNS.sort_by_any_ranged_desc, pen=COLOR_GREEN},
                {label='any ranged skill'..CH_UP, value=SORT_FNS.sort_by_any_ranged_asc, pen=COLOR_YELLOW},
                {label='any leader skill'..CH_DN, value=SORT_FNS.sort_by_leadership_desc, pen=COLOR_GREEN},
                {label='any leader skill'..CH_UP, value=SORT_FNS.sort_by_leadership_asc, pen=COLOR_YELLOW},
                {label='name'..CH_DN, value=sort_by_name_desc, pen=COLOR_GREEN},
                {label='name'..CH_UP, value=sort_by_name_asc, pen=COLOR_YELLOW},
                {label='migrant wave'..CH_DN, value=sort_by_migrant_wave_desc, pen=COLOR_GREEN},
                {label='migrant wave'..CH_UP, value=sort_by_migrant_wave_asc, pen=COLOR_YELLOW},
                {label='axe skill'..CH_DN, value=SORT_FNS.sort_by_axe_desc, pen=COLOR_GREEN},
                {label='axe skill'..CH_UP, value=SORT_FNS.sort_by_axe_asc, pen=COLOR_YELLOW},
                {label='sword skill'..CH_DN, value=SORT_FNS.sort_by_sword_desc, pen=COLOR_GREEN},
                {label='sword skill'..CH_UP, value=SORT_FNS.sort_by_sword_asc, pen=COLOR_YELLOW},
                {label='mace skill'..CH_DN, value=SORT_FNS.sort_by_mace_desc, pen=COLOR_GREEN},
                {label='mace skill'..CH_UP, value=SORT_FNS.sort_by_mace_asc, pen=COLOR_YELLOW},
                {label='hammer skill'..CH_DN, value=SORT_FNS.sort_by_hammer_desc, pen=COLOR_GREEN},
                {label='hammer skill'..CH_UP, value=SORT_FNS.sort_by_hammer_asc, pen=COLOR_YELLOW},
                {label='spear skill'..CH_DN, value=SORT_FNS.sort_by_spear_desc, pen=COLOR_GREEN},
                {label='spear skill'..CH_UP, value=SORT_FNS.sort_by_spear_asc, pen=COLOR_YELLOW},
                {label='crossbow skill'..CH_DN, value=SORT_FNS.sort_by_crossbow_desc, pen=COLOR_GREEN},
                {label='crossbow skill'..CH_UP, value=SORT_FNS.sort_by_crossbow_asc, pen=COLOR_YELLOW},
            },
            initial_option=SORT_FNS.sort_by_any_melee_desc,
            on_change=self:callback('refresh_list', 'sort'),
        },
        widgets.EditField{
            view_id='search',
            frame={l=32, t=0},
            label_text='Search: ',
            on_char=function(ch) return ch:match('[%l _-]') end,
            on_change=function() self:refresh_list() end,
        },
        widgets.Panel{
            frame={t=2, l=0, r=0, b=0},
            subviews={
                widgets.CycleHotkeyLabel{
                    view_id='sort_any_melee',
                    frame={t=0, l=0, w=10},
                    options={
                        {label='any melee', value=sort_noop},
                        {label='any melee'..CH_DN, value=SORT_FNS.sort_by_any_melee_desc, pen=COLOR_GREEN},
                        {label='any melee'..CH_UP, value=SORT_FNS.sort_by_any_melee_asc, pen=COLOR_YELLOW},
                    },
                    initial_option=SORT_FNS.sort_by_any_melee_desc,
                    option_gap=0,
                    on_change=self:callback('refresh_list', 'sort_any_melee'),
                },
                widgets.CycleHotkeyLabel{
                    view_id='sort_any_ranged',
                    frame={t=0, l=13, w=11},
                    options={
                        {label='any ranged', value=sort_noop},
                        {label='any ranged'..CH_DN, value=SORT_FNS.sort_by_any_ranged_desc, pen=COLOR_GREEN},
                        {label='any ranged'..CH_UP, value=SORT_FNS.sort_by_any_ranged_asc, pen=COLOR_YELLOW},
                    },
                    option_gap=0,
                    on_change=self:callback('refresh_list', 'sort_any_ranged'),
                },
                widgets.CycleHotkeyLabel{
                    view_id='sort_leadership',
                    frame={t=0, l=27, w=11},
                    options={
                        {label='leadership', value=sort_noop},
                        {label='leadership'..CH_DN, value=SORT_FNS.sort_by_leadership_desc, pen=COLOR_GREEN},
                        {label='leadership'..CH_UP, value=SORT_FNS.sort_by_leadership_asc, pen=COLOR_YELLOW},
                    },
                    option_gap=0,
                    on_change=self:callback('refresh_list', 'sort_leadership'),
                },
                widgets.CycleHotkeyLabel{
                    view_id='sort_name',
                    frame={t=0, l=41, w=5},
                    options={
                        {label='name', value=sort_noop},
                        {label='name'..CH_DN, value=sort_by_name_desc, pen=COLOR_GREEN},
                        {label='name'..CH_UP, value=sort_by_name_asc, pen=COLOR_YELLOW},
                    },
                    option_gap=0,
                    on_change=self:callback('refresh_list', 'sort_name'),
                },
                widgets.CycleHotkeyLabel{
                    view_id='sort_migrant_wave',
                    frame={t=0, l=48, w=13},
                    options={
                        {label='migrant wave', value=sort_noop},
                        {label='migrant wave'..CH_DN, value=sort_by_migrant_wave_desc, pen=COLOR_GREEN},
                        {label='migrant wave'..CH_UP, value=sort_by_migrant_wave_asc, pen=COLOR_YELLOW},
                    },
                    option_gap=0,
                    on_change=self:callback('refresh_list', 'sort_migrant_wave'),
                },
                widgets.CycleHotkeyLabel{
                    view_id='sort_axe',
                    frame={t=2, l=2, w=4},
                    options={
                        {label='axe', value=sort_noop},
                        {label='axe'..CH_DN, value=SORT_FNS.sort_by_axe_desc, pen=COLOR_GREEN},
                        {label='axe'..CH_UP, value=SORT_FNS.sort_by_axe_asc, pen=COLOR_YELLOW},
                    },
                    option_gap=0,
                    on_change=self:callback('refresh_list', 'sort_axe'),
                },
                widgets.CycleHotkeyLabel{
                    view_id='sort_sword',
                    frame={t=2, l=9, w=6},
                    options={
                        {label='sword', value=sort_noop},
                        {label='sword'..CH_DN, value=SORT_FNS.sort_by_sword_desc, pen=COLOR_GREEN},
                        {label='sword'..CH_UP, value=SORT_FNS.sort_by_sword_asc, pen=COLOR_YELLOW},
                    },
                    option_gap=0,
                    on_change=self:callback('refresh_list', 'sort_sword'),
                },
                widgets.CycleHotkeyLabel{
                    view_id='sort_mace',
                    frame={t=2, l=18, w=5},
                    options={
                        {label='mace', value=sort_noop},
                        {label='mace'..CH_DN, value=SORT_FNS.sort_by_mace_desc, pen=COLOR_GREEN},
                        {label='mace'..CH_UP, value=SORT_FNS.sort_by_mace_asc, pen=COLOR_YELLOW},
                    },
                    option_gap=0,
                    on_change=self:callback('refresh_list', 'sort_mace'),
                },
                widgets.CycleHotkeyLabel{
                    view_id='sort_hammer',
                    frame={t=2, l=25, w=7},
                    options={
                        {label='hammer', value=sort_noop},
                        {label='hammer'..CH_DN, value=SORT_FNS.sort_by_hammer_desc, pen=COLOR_GREEN},
                        {label='hammer'..CH_UP, value=SORT_FNS.sort_by_hammer_asc, pen=COLOR_YELLOW},
                    },
                    option_gap=0,
                    on_change=self:callback('refresh_list', 'sort_hammer'),
                },
                widgets.CycleHotkeyLabel{
                    view_id='sort_spear',
                    frame={t=2, l=36, w=6},
                    options={
                        {label='spear', value=sort_noop},
                        {label='spear'..CH_DN, value=SORT_FNS.sort_by_spear_desc, pen=COLOR_GREEN},
                        {label='spear'..CH_UP, value=SORT_FNS.sort_by_spear_asc, pen=COLOR_YELLOW},
                    },
                    option_gap=0,
                    on_change=self:callback('refresh_list', 'sort_spear'),
                },
                widgets.CycleHotkeyLabel{
                    view_id='sort_crossbow',
                    frame={t=2, l=45, w=9},
                    options={
                        {label='crossbow', value=sort_noop},
                        {label='crossbow'..CH_DN, value=SORT_FNS.sort_by_crossbow_desc, pen=COLOR_GREEN},
                        {label='crossbow'..CH_UP, value=SORT_FNS.sort_by_crossbow_asc, pen=COLOR_YELLOW},
                    },
                    option_gap=0,
                    on_change=self:callback('refresh_list', 'sort_crossbow'),
                },
            }
        },
    }
end

local function normalize_search_key(search_key)
    local out = ''
    for c in dfhack.toSearchNormalized(search_key):gmatch("[%w%s]") do
        out = out .. c:lower()
    end
    return out
end

local function filter_matches(unit_id, search)
    if unit_id == -1 then return true end
    local unit = df.unit.find(unit_id)
    if not unit then return true end
    local search_key = dfhack.TranslateName(dfhack.units.getVisibleName(unit))
    if unit.status.current_soul then
        for _,skill in ipairs(unit.status.current_soul.skills) do
            search_key = (search_key or '') .. ' ' .. (df.job_skill[skill.id] or '')
        end
    end
    return normalize_search_key(search_key):find(dfhack.toSearchNormalized(search))
end

local unit_selector = df.global.game.main_interface.unit_selector

-- this function uses the unused itemid and selected vectors to keep state,
-- taking advantage of the fact that they are reset by DF when the list of units changes
local function filter_vector(search, prev_search)
    local unid_is_filtered = #unit_selector.selected >= 0 and unit_selector.selected[0] ~= 0
    if #search == 0 or #unit_selector.selected == 0 then
        if not unid_is_filtered then
            -- we haven't modified the unid vector; nothing to do here
            return
        end
        -- restore the unid vector
        unit_selector.unid:assign(unit_selector.itemid)
        -- clear our "we meddled" flag
        unit_selector.selected[0] = 0
        return
    end
    if unid_is_filtered and search == prev_search then
        -- prev filter still stands
        return
    end
    if unid_is_filtered then
        -- restore the unid vector
        unit_selector.unid:assign(unit_selector.itemid)
    else
        -- save the unid vector and set our meddle flag
        unit_selector.itemid:assign(unit_selector.unid)
        unit_selector.selected[0] = 1
    end
    -- do the actual filtering
    for idx=#unit_selector.unid-1,0,-1 do
        if not filter_matches(unit_selector.unid[idx], search) then
            unit_selector.unid:erase(idx)
        end
    end
end

local SORT_WIDGET_NAMES = {
    'sort',
    'sort_any_melee',
    'sort_any_ranged',
    'sort_leadership',
    'sort_name',
    'sort_migrant_wave',
    'sort_axe',
    'sort_sword',
    'sort_mace',
    'sort_hammer',
    'sort_spear',
    'sort_crossbow',
}

function SquadAssignmentOverlay:refresh_list(sort_widget, sort_fn)
    sort_widget = sort_widget or 'sort'
    sort_fn = sort_fn or self.subviews.sort:getOptionValue()
    if sort_fn == sort_noop then
        self.subviews[sort_widget]:cycle()
        return
    end
    for _,widget_name in ipairs(SORT_WIDGET_NAMES) do
        self.subviews[widget_name]:setOption(sort_fn)
    end
    local search = self.subviews.search.text
    filter_vector(search, self.prev_search)
    self.prev_search = search
    utils.sort_vector(unit_selector.unid, nil, sort_fn)
end

function SquadAssignmentOverlay:onInput(keys)
    if keys._MOUSE_R_DOWN or
        keys._MOUSE_L_DOWN and not self:getMouseFramePos()
    then
        -- if any click is made outside of our window, we may need to refresh our list
        self.dirty = true
    end
    return SquadAssignmentOverlay.super.onInput(self, keys)
end

function SquadAssignmentOverlay:onRenderFrame(dc, frame_rect)
    SquadAssignmentOverlay.super.onRenderFrame(self, dc, frame_rect)
    if self.dirty then
        self:refresh_list()
        self.dirty = false
    end
end

OVERLAY_WIDGETS = {
    squad_assignment=SquadAssignmentOverlay,
}

--[[
local utils = require('utils')
local units = require('plugins.sort.units')
local items = require('plugins.sort.items')

orders = orders or {}
orders.units = units.orders
orders.items = items.orders

function parse_ordering_spec(type,...)
    local group = orders[type]
    if group == nil then
        dfhack.printerr('Invalid ordering class: '..tostring(type))
        return nil
    end

    local specs = table.pack(...)
    local rv = { }

    for _,spec in ipairs(specs) do
        local nil_first = false
        if string.sub(spec,1,1) == '<' then
            nil_first = true
            spec = string.sub(spec,2)
        end

        local reverse = false
        if string.sub(spec,1,1) == '>' then
            reverse = true
            spec = string.sub(spec,2)
        end

        local cm = group[spec]

        if cm == nil then
            dfhack.printerr('Unknown order for '..type..': '..tostring(spec))
            return nil
        end

        if nil_first or reverse then
            cm = copyall(cm)
            cm.nil_first = nil_first
            cm.reverse = reverse
        end

        rv[#rv+1] = cm
    end

    return rv
end

make_sort_order = utils.make_sort_order
]]

return _ENV