require "table2"

$define STX = "\x02"
$define ETX = "\x03"

local function bwt(s)
    if s:find(STX, 1, true) or s:find(ETX, 1, true) then return nil end
    s = STX .. s .. ETX
    local len = #s
    local tbl = table.rep(len, "")
    tbl[1] = s
    for i = 2, len do tbl[i] = s:sub(i) .. s:sub(1, i - 1) end
    tbl:sort()
    local last_chars = table.rep(len, "")
    for i = 1, len do last_chars[i] = tbl[i][len] end
    return last_chars:concat()
end

local function ibwt(r)
    local len = #r
    if len == 0 then return "" end
    local tbl = table.rep(len, "")
    for i = 1, len do
        for j = 1, len do tbl[j] = r[j].. tbl[j] end
        tbl:sort()
    end
    for tbl as row do
        if row:endswith(ETX) then return row:sub(2, -2) end
    end
    return ""
end

local function make_printable(s)
    -- Substitute ^ for STX and | for ETX to print results.
    s = s:replace(STX, "^")
    return s:replace(ETX, "|")
end

local tests = {
    "banana",
    "appellee",
    "dogwood",
    "TO BE OR NOT TO BE OR WANT TO BE OR NOT?",
    "SIX.MIXED.PIXIES.SIFT.SIXTY.PIXIE.DUST.BOXES",
    "\x02ABC\x03"
}
for tests as test do
    print(make_printable(test))
    io.write(" --> ")
    local t = bwt(test)
    if !t then
        print("ERROR: String can't contain STX or ETX")
        t = ""
    else
        print(make_printable(t))
    end
    local r = ibwt(t)
    print($" --> {r}\n")
end