require "table2"

class bifid
    static function encrypt(polybius, message)
        message = message:upper():replace("J", "I")
        local rows = {}
        local cols = {}
        for i = 1, #message do
            local c = message[i]
            local ix = polybius:find(c, 1, true)
            if !ix then continue end
            rows:insert((ix - 1) // 5 + 1)
            cols:insert((ix - 1) % 5 + 1)
        end
        local s = ""
        for table.joined(rows, cols):chunk(2) as pair do
            local ix = (pair[1] - 1) * 5 + pair[2]
            s ..= polybius[ix]
        end
        return s
    end

    static function decrypt(polybius, message)
        local rows = {}
        local cols = {}
        for i = 1, #message do
            local c = message[i]
            local ix = polybius:find(c, 1, true)
            rows:insert((ix - 1) // 5 + 1)
            cols:insert((ix - 1) % 5 + 1)
        end
        local lines = table.zip(rows, cols):flatten()
        local count = #lines // 2
        rows = lines:slice(1, count)
        cols = lines:slice(count + 1)
        local s = ""
        for i = 1, count do
            local ix = (rows[i] - 1) * 5 + cols[i]
            s ..= polybius[ix]
        end
        return s
    end
end

local poly1 = "ABCDEFGHIKLMNOPQRSTUVWXYZ"
local poly2 = "BGWKZQPNDSIOAXEFCLUMTHYVR"
local poly3 = "PLAYFIREXMBCDGHKNOQSTUVWZ"
local polys = {poly1, poly2, poly2, poly3}
local msg1  = "ATTACKATDAWN"
local msg2  = "FLEEATONCE"
local msg3  = "The invasion will start on the first of January"
local msgs  = {msg1, msg2, msg1, msg3}
for i = 1, #msgs do
    local encrypted = bifid.encrypt(polys[i], msgs[i])
    local decrypted = bifid.decrypt(polys[i], encrypted)
    print($"Message   : {msgs[i]}")
    print($"Encrypted : {encrypted}")
    print($"Decrypted : {decrypted}")
    if i < #msgs then print() end
end