---
--- 随机集合对象
--- Created by colin.
---

local tinsert = table.insert
local tremove = table.remove
local tmove = table.move
local tconcat = table.concat
local mrandom = math.random
local ipairs = ipairs
local pairs = pairs
local tostring = tostring

---@class RandSet
---@field list table
---@field idxs table
local RandSet = {}

local RandSetMT = {
    __index = RandSet,
    __pairs = function(t)
        local itr = function(t, k)
            return next(t.idxs, k)
        end
        return itr, t, nil
    end,
    __len = function(t)
        return #t.list
    end,
    __tostring  = function(t)
        return string.format("RandSet: (%s)", RandSet.concat(t, ", "))
    end
}

---增加元素
---@return boolean 是否增加成功
function RandSet:add(e)
    if self.idxs[e] then
        return false
    else
        tinsert(self.list, e)
        self.idxs[e] = #self.list
        return true
    end
end

---删除元素
function RandSet:remove(e)
    local idx = self.idxs[e]
    if idx then
        local ln = #self.list
        if idx > 0 and idx <= ln then
            local re
            if idx == ln then
                re = tremove(self.list)
                self.idxs[re] = nil
            else
                re = self.list[idx]
                self.idxs[re] = nil
                local le = tremove(self.list)
                self.list[idx] = le
                self.idxs[le] = idx
            end
        end
    end
end

---判断集合中是否存在元素e
---@return boolean 是否存在
function RandSet:has(e)
    return self.idxs[e] ~= nil
end

---从集合随机选择n个元素返回，如果n为nil，则返回元素，否则返回元素列表
---@return any|any[]
function RandSet:choice(n)
    local ln = #self.list
    if not n then
        return self.list[mrandom(1, ln)]
    elseif n >= ln then
        return tmove(self.list, 1, ln, 1, {})
    else
        local s = 1
        local i, ei, es
        for _ = 1, n do
            i = mrandom(s, ln)
            ei = self.list[i]
            es = self.list[s]
            self.list[s] = ei
            self.idxs[ei] = s
            self.list[i] = es
            self.idxs[es] = i
            s = s + 1
        end
        return tmove(self.list, 1, s -1, 1, {})
    end
end

---将集合中的元素连接成字符串返回
---@param sep string 分隔符
---@param i number 起始索引，默认为1
---@param j number 结束索引，默认为#set
---@return string
function RandSet:concat(sep, i, j)
    i = i or 1
    j = j or #self.list
    local t = {}
    for _, e in ipairs(self.list) do
        t[#t+1] = tostring(e)
    end
    return tconcat(t, sep)
end

---新建一个集合对象返回
---@return RandSet
function RandSet.new()
    return setmetatable({
        list = {},
        idxs = {},
    }, RandSetMT)
end

-------------------------------------------------------------------
local function test()
    local set = RandSet.new()
    print("add---------------")
    for i = 1, 10 do
        set:add(i)
    end
    print(set)
    print("len: ", #set)
    print("remove---------------")
    set:remove(8)
    set:remove(7)
    set:remove(3)
    print(set)
    print("has---------------")
    print(set:has(8), set:has(1), set:has(10))
    print("choice---------------")
    print(set:choice())
    print(tconcat(set:choice(2), ", "))
    print(tconcat(set:choice(5), ", "))
    print(tconcat(set:choice(10), ", "))
    print(tconcat(set:choice(0), ", "))
    print("for---------------")
    for e in pairs(set) do
        print(e)
    end
end
--test()
-------------------------------------------------------------------

return RandSet