--!native
--!optimize 2
--!strict

local bit32 = require("bit32")
local cpx = require("cpx")
local mtx = require("mtx")

local crypmethods = {}

type Vector = {number}
type Matrix = {Vector}
type EllipticCurve = {a: number, b: number}
type EllipticCurvePoint = {x: number, y: number}

local function keccak(message: string, rate: number, capacity: number, output_length: number): string
    assert(rate + capacity == 1600, "Invalid rate and capacity combination")
    local state = table.create(200, 0)
    local message_length = #message
    local block_size = rate // 8
    local rounds = 24

    -- Constants
    local RC = {
        0x0000000000000001, 0x0000000000008082, 0x800000000000808A, 0x8000000080008000,
        0x000000000000808B, 0x0000000080000001, 0x8000000080008081, 0x8000000000008009,
        0x000000000000008A, 0x0000000000000088, 0x0000000080008009, 0x000000008000000A,
        0x000000008000808B, 0x800000000000008B, 0x8000000000008089, 0x8000000000008003,
        0x8000000000008002, 0x8000000000000080, 0x000000000000800A, 0x800000008000000A,
        0x8000000080008081, 0x8000000000008080, 0x0000000080000001, 0x8000000080008008
    }

    local function rot(x: number, n: number): number
        return bit32.bor(bit32.lshift(x, n), bit32.rshift(x, 64 - n))
    end

    local function keccak_f(A: {number})
        for round = 0, rounds - 1 do
            local C = table.create(5, 0)
            for x = 0, 4 do
                C[x + 1] = bit32.bxor(A[x + 1], A[x + 6], A[x + 11], A[x + 16], A[x + 21])
            end

            local D = table.create(5, 0)
            for x = 0, 4 do
                D[x + 1] = bit32.bxor(C[((x + 4) % 5) + 1], rot(C[((x + 1) % 5) + 1], 1))
            end

            for x = 0, 4 do
                for y = 0, 4 do
                    A[5 * x + y + 1] = bit32.bxor(A[5 * x + y + 1], D[x + 1])
                end
            end

            local B = table.create(25, 0)
            for x = 0, 4 do
                for y = 0, 4 do
                    local Y = (2 * x + 3 * y) % 5
                    B[Y * 5 + x + 1] = rot(A[5 * x + y + 1], ((x + 3 * y) * (x + 3 * y) + x + 3 * y) % 64)
                end
            end

            for x = 0, 4 do
                for y = 0, 4 do
                    A[5 * x + y + 1] = bit32.bxor(B[5 * x + y + 1], bit32.band(bit32.bnot(B[5 * ((x + 1) % 5) + y + 1]), B[5 * ((x + 2) % 5) + y + 1]))
                end
            end

            A[1] = bit32.bxor(A[1], RC[round + 1])
        end
        return A
    end

    local padded_message = message .. string.char(0x06) .. string.rep("\0", block_size - (message_length + 1) % block_size - 1) .. string.char(0x80)

    for i = 1, #padded_message, block_size do
        local block = padded_message:sub(i, i + block_size - 1)
        for j = 1, #block do
            state[j] = bit32.bxor(state[j], string.byte(block, j))
        end
        state = keccak_f(state)
    end

    local output = ""
    while #output < output_length do
        for i = 1, math.min(block_size, output_length - #output) do
            output = output .. string.char(state[i])
        end
        if #output < output_length then
            state = keccak_f(state)
        end
    end

    return output
end

local function f(m: string): string
    return keccak(m, 1088, 512, 32)  -- SHA3-256
end

local function h(m1: string, m2: string): string
    return f(m1 .. m2)
end

local function prf(seed: string, addr: number): string
    return f(seed .. string.pack(">I8", addr))
end

local function wots_chain(x: string, i: number, w: number, public_seed: string): string
    for j = 0, i - 1 do
        x = f(public_seed .. x .. string.pack(">I2I2", w, j))
    end
    return x
end

local function wots_sign(message: string, secret_seed: string, public_seed: string, w: number): string
    local n = 32  -- Hash output length
    local l1 = math.ceil(8 * n / math.log(w, 2))
    local l2 = math.floor(math.log(l1 * (w - 1), 2) / math.log(w, 2)) + 1
    local l = l1 + l2

    local b = 0
    for i = 1, l1 do
        b = b + bit32.extract(string.unpack(">I4", message, (i-1)*4 + 1), 0, 8) * (w - 1)
    end

    local signature = ""
    for i = 1, l do
        local sk_i = f(secret_seed .. string.pack(">I2", i))
        local a = i <= l1 and bit32.extract(string.unpack(">I4", message, (i-1)*4 + 1), 0, 8) or bit32.extract(b, (i-l1-1)*8, 8)
        signature = signature .. wots_chain(sk_i, a, w, public_seed)
    end
    return signature
end

local function tree_hash(seed: string, leaf_idx: number, height: number): string
    local stack = {}
    for i = 0, 2^height - 1 do
        local node = f(seed .. string.pack(">I8", leaf_idx + i))
        local node_height = 0
        while #stack > 0 and stack[#stack].height == node_height do
            local right = node
            local left = table.remove(stack).node
            node = f(left .. right)
            node_height = node_height + 1
        end
        table.insert(stack, {node = node, height = node_height})
    end
    return stack[1].node
end

local function gen_leaf_wots(secret_seed: string, public_seed: string, addr: number): string
    local leaf_secret = f(secret_seed .. string.pack(">I8", addr))
    return wots_sign(leaf_secret, leaf_secret, public_seed, 16)
end

local function secure_random(n: number): string
    local result = ""
    for _ = 1, n do
        result = result .. string.char(math.random(0, 255))
    end
    return result
end

local function constant_time_compare(a: string, b: string): boolean
    if #a ~= #b then
        return false
    end
    local result = 0
    for i = 1, #a do
        result = bit32.bor(result, bit32.bxor(string.byte(a, i), string.byte(b, i)))
    end
    return result == 0
end

local crypmethods = {}

function crypmethods.generate_key_pair(): (string, string)
    local secret_seed = secure_random(32)
    local public_seed = f(secret_seed .. "public")
    return secret_seed, public_seed
end

function crypmethods.sphincs_sign(message: string, secret_seed: string, public_seed: string): string
    local n = 32  -- Hash output length
    local w = 16  -- Winternitz parameter
    local h = 64  -- Height of the hypertree
    local d = 8   -- Number of layers in the hypertree

    local r = secure_random(n)
    local R = f(r .. public_seed .. message)

    local signature = r
    local root = tree_hash(public_seed, 0, h)

    for i = 1, d do
        local leaf_idx = bit32.band(bit32.rshift(string.unpack(">I4", R, 1), h - (i-1)*h//d), 2^(h//d) - 1)
        local auth_path = {}
        for j = 0, h//d - 1 do
            local sibling_idx = bit32.bxor(leaf_idx, bit32.lshift(1, j))
            table.insert(auth_path, tree_hash(public_seed, sibling_idx * 2^j, j))
        end

        local leaf = gen_leaf_wots(secret_seed, public_seed, leaf_idx)
        signature = signature .. leaf

        for _, node in ipairs(auth_path) do
            signature = signature .. node
        end

        R = f(R .. root)
        root = tree_hash(public_seed, leaf_idx * 2^(h//d), h//d)
    end

    return signature
end

function crypmethods.sphincs_verify(message: string, signature: string, public_seed: string): boolean
    local n = 32  -- Hash output length
    local w = 16  -- Winternitz parameter
    local h = 64  -- Height of the hypertree
    local d = 8   -- Number of layers in the hypertree

    local function wots_pk_from_sig(signature: string, message: string, public_seed: string, addr: number): string
        local l1 = math.ceil(8 * n / math.log(w, 2))
        local l2 = math.floor(math.log(l1 * (w - 1), 2) / math.log(w, 2)) + 1
        local l = l1 + l2

        local b = 0
        for i = 1, l1 do
            b = b + bit32.extract(string.unpack(">I4", message, (i-1)*4 + 1), 0, 8) * (w - 1)
        end

        local public_key = ""
        for i = 1, l do
            local sig_i = signature:sub((i-1)*n + 1, i*n)
            local a = i <= l1 and bit32.extract(string.unpack(">I4", message, (i-1)*4 + 1), 0, 8) or bit32.extract(b, (i-l1-1)*8, 8)
            public_key = public_key .. wots_chain(sig_i, w - 1 - a, w, public_seed)
        end

        return f(public_seed .. public_key .. string.pack(">I8", addr))
    end

    local r = signature:sub(1, n)
    local R = f(r .. public_seed .. message)
    signature = signature:sub(n + 1)

    local root = nil

    for i = 1, d do
        local leaf_idx = bit32.band(bit32.rshift(string.unpack(">I4", R, 1), h - (i-1)*h//d), 2^(h//d) - 1)
        local wots_sig = signature:sub(1, n * (2^(h//d) - 1))
        signature = signature:sub(n * (2^(h//d) - 1) + 1)

        local leaf = wots_pk_from_sig(wots_sig, R, public_seed, leaf_idx)

        local node = leaf
        for j = 0, h//d - 1 do
            local sibling = signature:sub(1, n)
            signature = signature:sub(n + 1)

            if bit32.band(bit32.rshift(leaf_idx, j), 1) == 0 then
                node = f(node .. sibling)
            else
                node = f(sibling .. node)
            end
        end

        if root == nil then
            root = node
        elseif not constant_time_compare(root, node) then
            return false
        end

        R = f(R .. root)
    end

    return constant_time_compare(root, tree_hash(public_seed, 0, h))
end

return crypmethods

--------------------------------------------------------------------------------------------------------

local function mod_inverse(a: number, m: number): number
	local m0, x0, x1 = m, 0, 1
	while a > 1 do
		local q = math.floor(a / m)
		m, a = a % m, m
		x0, x1 = x1 - q * x0, x0
	end
	return x1 < 0 and x1 + m0 or x1
end

local function random_vector(n: number, max: number): Vector
	local v = table.create(n)
	for i = 1, n do
		v[i] = math.random(0, max - 1)
	end
	return v
end

local function random_matrix(n: number, m: number, max: number): Matrix
	local mat = table.create(n)
	for i = 1, n do
		mat[i] = random_vector(m, max)
	end
	return mat
end

-- RING-LWE
local function polynomial_add(a: Vector, b: Vector, q: number): Vector
	local result = table.create(#a)
	for i = 1, #a do
		result[i] = (a[i] + b[i]) % q
	end
	return result
end

local function polynomial_multiply(a: Vector, b: Vector, q: number): Vector
	local n = #a
	local result = table.create(n, 0)
	for i = 1, n do
		for j = 1, n do
			local k = (i + j - 2) % n + 1
			result[k] = (result[k] + a[i] * b[j]) % q
		end
	end
	return result
end

function crypmethods.generate_key_pair(): (string, string)
    local secret_seed = secure_random(32)
    local public_seed = f(secret_seed .. "public")
    return secret_seed, public_seed
end

function crypmethods.generate_ring_lwe_keys(n: number, q: number): (Vector, {a: Vector, b: Vector})
	local a = random_vector(n, q)
	local s = random_vector(n, q)
	local e = random_vector(n, 3)  -- Small error vector
	local b = polynomial_add(polynomial_multiply(a, s, q), e, q)
	return s, {a = a, b = b}
end

function crypmethods.ring_lwe_encrypt(message: Vector, public_key: {a: Vector, b: Vector}, q: number): (Vector, Vector)
	local n = #public_key.a
	local r = random_vector(n, 3)  -- Small random vector
	local e1 = random_vector(n, 3)  -- Small error vector
	local e2 = random_vector(n, 3)  -- Small error vector
	local u = polynomial_add(polynomial_multiply(public_key.a, r, q), e1, q)
	local v = polynomial_add(polynomial_add(polynomial_multiply(public_key.b, r, q), e2, q), message, q)
	return u, v
end

function crypmethods.ring_lwe_decrypt(u: Vector, v: Vector, private_key: Vector, q: number): Vector
	local message = polynomial_add(v, polynomial_multiply(u, private_key, q), q)
	for i = 1, #message do
		message[i] = message[i] > q/2 and 0 or 1
	end
	return message
end

-- (McEliece)
local function generate_goppa_code(n: number, k: number, t: number): (Matrix, Matrix)
	-- Add goppa polynomial to this
	local g = random_matrix(k, n-k, 2)
	local h = random_matrix(n-k, n, 2)
	for i = 1, n-k do
		for j = 1, k do
			h[i][j] = g[j][i]
		end
		h[i][i+k] = 1
	end
	return g, h
end

function crypmethods.generate_mceliece_keys(n: number, k: number, t: number): ({G: Matrix, P: Matrix, S: Matrix}, Matrix)
	local g, h = generate_goppa_code(n, k, t)
	local s = random_matrix(k, k, 2)
	local p = random_matrix(n, n, 2)
	local g_pub = mtx.multiply(mtx.multiply(s, g), p)
	return {G = g, P = p, S = s}, g_pub
end

function crypmethods.mceliece_encrypt(message: Vector, public_key: Matrix): Vector
	local n = #public_key[1]
	local ciphertext = mtx.multiply({message}, public_key)[1]
	-- t random errors
	local t = math.floor((n - #message) / 2)
	for _ = 1, t do
		local pos = math.random(1, n)
		ciphertext[pos] = (ciphertext[pos] + 1) % 2
	end
	return ciphertext
end

function crypmethods.mceliece_decrypt(ciphertext: Vector, private_key: {G: Matrix, P: Matrix, S: Matrix}): Vector
	local n = #ciphertext
	local k = #private_key.S
	local permuted = table.create(n)
	for i = 1, n do
		for j = 1, n do
			if private_key.P[j][i] == 1 then
				permuted[j] = ciphertext[i]
				break
			end
		end
	end
	-- secret GOPPPA code :money_face: this is very simplified
	local decoded = table.create(k)
	for i = 1, k do
		decoded[i] = 0
		for j = 1, n do
			decoded[i] = (decoded[i] + permuted[j] * private_key.G[i][j]) % 2
		end
	end
	return mtx.multiply({decoded}, mtx.inverse(private_key.S, 2))[1]
end

local function f(m: string): string
	-- Please use the big ass table for this
end

local function h(m1: string, m2: string): string
	return f(m1 .. m2)
end

local function prf(seed: string, addr: number): string
	return f(seed .. string.pack(">I8", addr))
end

function crypmethods.sphincs_keygen(n: number): (string, string)
	local seed = string.rep("\0", n)
	for i = 1, n do
		seed = string.char(math.random(0, 255)) .. seed:sub(1, -2)
	end
	local public_seed = f(seed .. "public")
	local secret_seed = f(seed .. "secret")
	return secret_seed, public_seed
end

function crypmethods.sphincs_sign(message: string, secret_seed: string, public_seed: string): string
    local n = 32  -- Hash output length
    local w = 16  -- Winternitz parameter
    local h = 64  -- Height of the hypertree
    local d = 8   -- Number of layers in the hypertree

    local r = secure_random(n)
    local R = f(r .. public_seed .. message)

    local signature = r
    local root = tree_hash(public_seed, 0, h)

    for i = 1, d do
        local leaf_idx = bit32.band(bit32.rshift(string.unpack(">I4", R, 1), h - (i-1)*h//d), 2^(h//d) - 1)
        local auth_path = {}
        for j = 0, h//d - 1 do
            local sibling_idx = bit32.bxor(leaf_idx, bit32.lshift(1, j))
            table.insert(auth_path, tree_hash(public_seed, sibling_idx * 2^j, j))
        end

        local leaf = gen_leaf_wots(secret_seed, public_seed, leaf_idx)
        signature = signature .. leaf

        for _, node in ipairs(auth_path) do
            signature = signature .. node
        end

        R = f(R .. root)
        root = tree_hash(public_seed, leaf_idx * 2^(h//d), h//d)
    end

    return signature
end

function crypmethods.sphincs_verify(message: string, signature: string, public_seed: string): boolean
    local n = 32  -- Hash output length
    local w = 16  -- Winternitz parameter
    local h = 64  -- Height of the hypertree
    local d = 8   -- Number of layers in the hypertree

    local function wots_pk_from_sig(signature: string, message: string, public_seed: string, addr: number): string
        local l1 = math.ceil(8 * n / math.log(w, 2))
        local l2 = math.floor(math.log(l1 * (w - 1), 2) / math.log(w, 2)) + 1
        local l = l1 + l2

        local b = 0
        for i = 1, l1 do
            b = b + bit32.extract(string.unpack(">I4", message, (i-1)*4 + 1), 0, 8) * (w - 1)
        end

        local public_key = ""
        for i = 1, l do
            local sig_i = signature:sub((i-1)*n + 1, i*n)
            local a = i <= l1 and bit32.extract(string.unpack(">I4", message, (i-1)*4 + 1), 0, 8) or bit32.extract(b, (i-l1-1)*8, 8)
            public_key = public_key .. wots_chain(sig_i, w - 1 - a, w, public_seed)
        end

        return f(public_seed .. public_key .. string.pack(">I8", addr))
    end

    local r = signature:sub(1, n)
    local R = f(r .. public_seed .. message)
    signature = signature:sub(n + 1)

    local root = nil

    for i = 1, d do
        local leaf_idx = bit32.band(bit32.rshift(string.unpack(">I4", R, 1), h - (i-1)*h//d), 2^(h//d) - 1)
        local wots_sig = signature:sub(1, n * (2^(h//d) - 1))
        signature = signature:sub(n * (2^(h//d) - 1) + 1)

        local leaf = wots_pk_from_sig(wots_sig, R, public_seed, leaf_idx)

        local node = leaf
        for j = 0, h//d - 1 do
            local sibling = signature:sub(1, n)
            signature = signature:sub(n + 1)

            if bit32.band(bit32.rshift(leaf_idx, j), 1) == 0 then
                node = f(node .. sibling)
            else
                node = f(sibling .. node)
            end
        end

        if root == nil then
            root = node
        elseif not constant_time_compare(root, node) then
            return false
        end

        R = f(R .. root)
    end

    return constant_time_compare(root, tree_hash(public_seed, 0, h))
end

return crypmethods

local function point_double(P: EllipticCurvePoint, A: number, p: number): EllipticCurvePoint
	if P.x == 0 and P.y == 0 then
		return P
	end
	local m = (3 * P.x * P.x + A) * mod_inverse(2 * P.y, p) % p
	local x = (m * m - 2 * P.x) % p
	local y = (m * (P.x - x) - P.y) % p
	return {x = x, y = y}
end

local function point_add(P: EllipticCurvePoint, Q: EllipticCurvePoint, A: number, p: number): EllipticCurvePoint
	if P.x == 0 and P.y == 0 then
		return Q
	elseif Q.x == 0 and Q.y == 0 then
		return P
	elseif P.x == Q.x and P.y == Q.y then
		return point_double(P, A, p)
	elseif P.x == Q.x then
		return {x = 0, y = 0}
	end
	local m = (Q.y - P.y) * mod_inverse(Q.x - P.x, p) % p
	local x = (m * m - P.x - Q.x) % p
	local y = (m * (P.x - x) - P.y) % p
	return {x = x, y = y}
end

local function scalar_multiply(k: number, P: EllipticCurvePoint, A: number, p: number): EllipticCurvePoint
	local R = {x = 0, y = 0}
	local Q = P
	while k > 0 do
		if k % 2 == 1 then
			R = point_add(R, Q, A, p)
		end
		Q = point_double(Q, A, p)
		k = math.floor(k / 2)
	end
	return R
end

function crypmethods.sidh_keygen(p: number, l_A: number, e_A: number, l_B: number, e_B: number): ({number, EllipticCurvePoint}, {number, EllipticCurvePoint})
	local E = {a = 0, b = 1}  -- Starting curve y^2 = x^3 + x
	local x_P = 1
	local y_P = math.sqrt(x_P^3 + E.a*x_P + E.b) % p
	local P_A = {x = x_P, y = y_P}
	local Q_A = scalar_multiply(l_A^(e_A - 1), P_A, E.a, p)
	local sk_A = math.random(1, l_A^e_A - 1)
	local pk_A = scalar_multiply(sk_A, P_A, E.a, p)

	local x_P = 2
	local y_P = math.sqrt(x_P^3 + E.a*x_P + E.b) % p
	local P_B = {x = x_P, y = y_P}
	local Q_B = scalar_multiply(l_B^(e_B - 1), P_B, E.a, p)
	local sk_B = math.random(1, l_B^e_B - 1)
	local pk_B = scalar_multiply(sk_B, P_B, E.a, p)

	return {sk_A, P_A}, {sk_B, P_B}
end

function crypmethods.sidh_exchange(sk: number, P: EllipticCurvePoint, Q: EllipticCurvePoint, A: number, p: number): number
	local R = point_add(scalar_multiply(sk, P, A, p), Q, A, p)
	return R.x
end

local function secure_random(n: number): string
	local result = ""
	for _ = 1, n do
		result = result .. string.char(math.random(0, 255))
	end
	return result
end

local function constant_time_compare(a: string, b: string): boolean
	if #a ~= #b then
		return false
	end
	local result = 0
	for i = 1, #a do
		result = bit32.bor(result, bit32.bxor(string.byte(a, i), string.byte(b, i)))
	end
	return result == 0
end

-- implement a more secure version of SPHINCS+ with an improved Winternitz One-Time Signature (WOTS+)
local function wots_chain(x: string, i: number, w: number, public_seed: string): string
	for j = 0, i - 1 do
		x = f(public_seed .. x .. string.pack(">I2I2", w, j))
	end
	return x
end

local function wots_sign(message: string, secret_seed: string, public_seed: string, w: number): string
	local n = 32
	local l1 = math.ceil(8 * n / math.log(w, 2))
	local l2 = math.floor(math.log(l1 * (w - 1), 2) / math.log(w, 2)) + 1
	local l = l1 + l2

	local b = 0
	for i = 1, l1 do
		b = b + bit32.extract(message, (i-1)*4, 4) * (w - 1)
	end

	local signature = ""
	for i = 1, l do
		local sk_i = f(secret_seed .. string.pack(">I2", i))
		local a = i <= l1 and bit32.extract(message, (i-1)*4, 4) or bit32.extract(b, (i-l1-1)*4, 4)
		signature = signature .. wots_chain(sk_i, a, w, public_seed)
	end
	return signature
end

function crypmethods.sphincs_plus_sign(message: string, secret_seed: string, public_seed: string): string
    local n = 32  -- Hash output length
    local w = 16  -- Winternitz parameter
    local h = 64  -- Height of the hypertree
    local d = 8   -- Number of layers in the hypertree

    local function f_tree(seed: string, addr: number): string
        return f(seed .. string.pack(">I8", addr))
    end

    local function tree_hash(seed: string, leaf_idx: number, height: number): string
        local stack = {}
        for i = 0, 2^height - 1 do
            local node = f_tree(seed, leaf_idx + i)
            local height = 0
            while #stack > 0 and stack[#stack].height == height do
                local right = node
                local left = table.remove(stack).node
                node = f(left .. right)
                height = height + 1
            end
            table.insert(stack, {node = node, height = height})
        end
        return stack[1].node
    end

    local function gen_leaf_wots(secret_seed: string, public_seed: string, addr: number): string
        local leaf_secret = f(secret_seed .. string.pack(">I8", addr))
        return wots_sign(leaf_secret, leaf_secret, public_seed, w)
    end

    local r = secure_random(n)
    local R = f(r .. public_seed .. message)

    local signature = r
    local root = tree_hash(public_seed, 0, h)

    for i = 1, d do
        local leaf_idx = bit32.band(bit32.rshift(R, h - (i-1)*h/d), 2^(h/d) - 1)
        local auth_path = {}
        for j = 0, h/d - 1 do
            local sibling_idx = bit32.bxor(leaf_idx, bit32.lshift(1, j))
            table.insert(auth_path, tree_hash(public_seed, sibling_idx * 2^j, j))
        end

        local leaf = gen_leaf_wots(secret_seed, public_seed, leaf_idx)
        signature = signature .. leaf

        for _, node in ipairs(auth_path) do
            signature = signature .. node
        end

        R = f(R .. root)
        root = tree_hash(public_seed, leaf_idx * 2^(h/d), h/d)
    end

    return signature
end

function crypmethods.sphincs_plus_verify(message: string, signature: string, public_seed: string): boolean
    local n = 32  -- Hash output length
    local w = 16  -- Winternitz parameter
    local h = 64  -- Height of the hypertree
    local d = 8   -- Number of layers in the hypertree

    local function f_tree(seed: string, addr: number): string
        return f(seed .. string.pack(">I8", addr))
    end

    local function wots_pk_from_sig(signature: string, message: string, public_seed: string, addr: number): string
        local n = 32  -- Hash output length
        local l1 = math.ceil(8 * n / math.log(w, 2))
        local l2 = math.floor(math.log(l1 * (w - 1), 2) / math.log(w, 2)) + 1
        local l = l1 + l2

        local b = 0
        for i = 1, l1 do
            b = b + bit32.extract(message, (i-1)*4, 4) * (w - 1)
        end

        local public_key = ""
        for i = 1, l do
            local sig_i = signature:sub((i-1)*n + 1, i*n)
            local a = i <= l1 and bit32.extract(message, (i-1)*4, 4) or bit32.extract(b, (i-l1-1)*4, 4)
            public_key = public_key .. wots_chain(sig_i, w - 1 - a, w, public_seed)
        end

        return f(public_seed .. public_key .. string.pack(">I8", addr))
    end

    local r = signature:sub(1, n)
    local R = f(r .. public_seed .. message)
    signature = signature:sub(n + 1)

    local root = nil

    for i = 1, d do
        local leaf_idx = bit32.band(bit32.rshift(R, h - (i-1)*h/d), 2^(h/d) - 1)
        local wots_sig = signature:sub(1, n * (2^(h/d) - 1))
        signature = signature:sub(n * (2^(h/d) - 1) + 1)

        local leaf = wots_pk_from_sig(wots_sig, R, public_seed, leaf_idx)

        local node = leaf
        for j = 0, h/d - 1 do
            local sibling = signature:sub(1, n)
            signature = signature:sub(n + 1)

            if bit32.band(bit32.rshift(leaf_idx, j), 1) == 0 then
                node = f(node .. sibling)
            else
                node = f(sibling .. node)
            end
        end

        if root == nil then
            root = node
        elseif root ~= node then
            return false
        end

        R = f(R .. root)
    end

    return constant_time_compare(root, tree_hash(public_seed, 0, h))
end

-- secure version of SIDH
function crypmethods.sidh_secure_keygen(p: number, l_A: number, e_A: number, l_B: number, e_B: number): ({number, EllipticCurvePoint, EllipticCurvePoint}, {number, EllipticCurvePoint, EllipticCurvePoint})
    local E = {a = 0, b = 1}  -- Starting curve y^2 = x^3 + x
    local x_P = 1
    local y_P = math.sqrt(x_P^3 + E.a*x_P + E.b) % p
    local P_A = {x = x_P, y = y_P}
    local Q_A = scalar_multiply(l_A^(e_A - 1), P_A, E.a, p)
    local sk_A = math.random(1, l_A^e_A - 1)
    local pk_A = scalar_multiply(sk_A, P_A, E.a, p)
    local pk_A_Q = scalar_multiply(sk_A, Q_A, E.a, p)

    local x_P = 2
    local y_P = math.sqrt(x_P^3 + E.a*x_P + E.b) % p
    local P_B = {x = x_P, y = y_P}
    local Q_B = scalar_multiply(l_B^(e_B - 1), P_B, E.a, p)
    local sk_B = math.random(1, l_B^e_B - 1)
    local pk_B = scalar_multiply(sk_B, P_B, E.a, p)
    local pk_B_Q = scalar_multiply(sk_B, Q_B, E.a, p)

    return {sk_A, P_A, Q_A}, {sk_B, P_B, Q_B}, {pk_A, pk_A_Q}, {pk_B, pk_B_Q}
end

function crypmethods.sidh_secure_exchange(sk: number, P: EllipticCurvePoint, Q: EllipticCurvePoint, pk: EllipticCurvePoint, pk_Q: EllipticCurvePoint, A: number, p: number): number
    local R = point_add(scalar_multiply(sk, P, A, p), Q, A, p)
    local S = point_add(scalar_multiply(sk, pk, A, p), pk_Q, A, p)

    if (S.y^2 % p) ~= (S.x^3 + A*S.x + 1) % p then
        error("Invalid public key")
    end

    return R.x
end

function crypmethods.hybrid_encrypt(message: string, ring_lwe_pk: {a: Vector, b: Vector}, sidh_pk: EllipticCurvePoint, q: number, p: number): (Vector, Vector, EllipticCurvePoint, string)
    local n = #ring_lwe_pk.a
    local ring_lwe_message = {}
    for i = 1, n do
        ring_lwe_message[i] = string.byte(message, i) or 0
    end

    local u, v = crypmethods.ring_lwe_encrypt(ring_lwe_message, ring_lwe_pk, q)

    local sidh_sk, sidh_P, sidh_Q = crypmethods.sidh_secure_keygen(p, 2, 372, 3, 239)
    local shared_secret = crypmethods.sidh_secure_exchange(sidh_sk[1], sidh_pk, sidh_Q, sidh_P, sidh_Q, 0, p)

    local key = f(tostring(shared_secret))
    local encrypted_tail = ""
    for i = n + 1, #message do
        encrypted_tail = encrypted_tail .. string.char(bit32.bxor(string.byte(message, i), string.byte(key, (i-n-1) % 32 + 1)))
    end

    return u, v, sidh_P, encrypted_tail
end

function crypmethods.hybrid_decrypt(u: Vector, v: Vector, sidh_P: EllipticCurvePoint, encrypted_tail: string, ring_lwe_sk: Vector, sidh_sk: number, q: number, p: number): string
    local decrypted_head = crypmethods.ring_lwe_decrypt(u, v, ring_lwe_sk, q)
    local message = ""
    for i = 1, #decrypted_head do
        message = message .. string.char(decrypted_head[i])
    end

    local shared_secret = crypmethods.sidh_secure_exchange(sidh_sk, sidh_P, {x = 0, y = 0}, {x = 0, y = 0}, {x = 0, y = 0}, 0, p)
    local key = f(tostring(shared_secret))

    for i = 1, #encrypted_tail do
        message = message .. string.char(bit32.bxor(string.byte(encrypted_tail, i), string.byte(key, (i-1) % 32 + 1)))
    end

    return message
end

return crypmethods


--!native
--!optimize 2
--!strict

type complex = {r: number, i: number}
type qubit = {complex}
type qubits = {qubit}
type bits = {number}

local qb = require("qb")
local gate = require("gate")
local cpx = require("cpx")
local preset = require("preset")

local crypto = {}

local bb84 = {}

local function random_bit(): number
	return math.random(0, 1)
end

--[[

	how to use bb84:
	user 1:
	generate_random_bits -> generate_random_bases -> prepare_qubits
	user 2:
	generate_random_bits -> generate_random_bases -> prepare_qubits
	compare:
	compare_bases -> sift_key (user 1) -> sift_key (user 2)
	spy dtc:
	eavesdrop_detection
	privacy amp:
	privacy_amplification (user 1) -> privacy_amplification (user 2)

]]

function bb84.generate_random_bits(n: number): bits
	local bitstring = {}
	for i = 1, n do
		bitstring[i] = random_bit()
	end
	return bitstring
end

function bb84.generate_random_bases(n: number): bits
	local bases = {}
	for i = 1, n do
		bases[i] = random_bit() -- 0 for rectilinear (|0⟩, |1⟩), 1 for diagonal (|+⟩, |-⟩)
	end
	return bases
end

function bb84.prepare_qubits(bits: bits, bases: bits): qubits
	local qubits = {}
	for i = 1, #bits do
		local bit = bits[i]
		local basis = bases[i]
		local qubit
		if basis == 0 then
			-- rectilinear basis: |0⟩ or |1⟩
			qubit = bit == 0 and qb.new(cpx.new(1, 0), cpx.new(0, 0)) or qb.new(cpx.new(0, 0), cpx.new(1, 0))
		else
			-- diagonal basis: |+⟩ = (|0⟩ + |1⟩)/√2 or |−⟩ = (|0⟩ − |1⟩)/√2
			qubit = bit == 0 and qb.new(cpx.new(1 / math.sqrt(2), 0), cpx.new(1 / math.sqrt(2), 0))
				or qb.new(cpx.new(1 / math.sqrt(2), 0), cpx.new(-1 / math.sqrt(2), 0))
		end
		qubits[i] = qubit
	end
	return qubits
end

function bb84.measure_qubits(qubits: qubits, bases: bits): bits
	local measured_bits = {}
	for i = 1, #qubits do
		local qubit = qubits[i]
		local basis = bases[i]
		local result
		if basis == 0 then
			result = qb.measure(qubit)
		else
			qubit = gate.apply(qubit, preset.hadamard)
			result = qb.measure(qubit)
		end
		measured_bits[i] = result
	end
	return measured_bits
end

function bb84.compare_bases(bases_1: bits, bases_2: bits): bits
	local matching_indices = {}
	for i = 1, #bases_1 do
		if bases_1[i] == bases_2[i] then
			table.insert(matching_indices, i)
		end
	end
	return matching_indices
end

function bb84.sift_key(bits: bits, matching_indices: bits): bits
	local key = {}
	for _, index in ipairs(matching_indices) do
		table.insert(key, bits[index])
	end
	return key
end

function bb84.eavesdrop_detection(bits_1: bits, bits_2: bits, matching_indices: bits, sample_size: {
	sample: number?;
	tolerate: number?;
}?): boolean
	local a = sample_size or {
		sample = 0.1;
		tolerate = 0.2;
	}
	local samp = math.floor(#matching_indices * (a.sample or 0.1))
	local errors = 0

	for i = 1, samp do
		local index = matching_indices[i]
		if bits_1[index] ~= bits_2[index] then
			errors += 1
		end
	end

	return errors / samp < (a.tolerate or 0.2)
end

function bb84.privacy_amplification(key: bits): bits
	local new_key = {}
	for i = 1, #key, 2 do
		new_key[#new_key + 1] = bit32.bxor(bit32.lrotate(bit32.bxor(key[i], key[i + 1]) % 2, i),
			bit32.rshift(key[i] * math.random(1, #key), bit32.bxor(#new_key, os.time())))
	end
	return new_key
end

crypto.bb84 = bb84

local e91 = {}

function e91.entangle(n: number): {qubits}
	local entangled_pairs = {}
	for i = 1, n do
		-- epr pair (|00⟩ + |11⟩) / √2
		local qubits = qb.new_ex(2)
		qubits[1] = qb.new(cpx.new(1 / math.sqrt(2), 0), cpx.new(0, 0))
		qubits[2] = qb.new(cpx.new(0, 0), cpx.new(1 / math.sqrt(2), 0))
		table.insert(entangled_pairs, qubits)
	end
	return entangled_pairs
end

function e91.random_basis(n: number): {number}
	local bases = {}
	for i = 1, n do
		bases[i] = math.random(1, 3)
	end
	return bases
end

function e91.measure_qubit(qubit: qubit, basis: number): number
	if basis == 1 then
		return qb.measure(qubit)
	elseif basis == 2 then
		local rotated_qubit = gate.apply(qubit, preset.paulix)
		return qb.measure(rotated_qubit)
	elseif basis == 3 then
		local rotated_qubit = gate.apply(qubit, preset.pauliz)
		return qb.measure(rotated_qubit)
	end
	error("three basis only")
end

function e91.bell_test(results_1: {number}, results_2: {number}, threshold : number?): boolean
	local count = #results_1
	local violation_threshold = (threshold or 0.75) * count -- adjust based on bells inequality criteria
	local violations = 0

	for i = 1, count do
		if results_1[i] ~= results_2[i] then
			violations += 1
		end
	end

	-- since e91 is based on violation of the bells thing, if it is violated -> secure
	return violations / count > violation_threshold
end

function e91.sift_key(bits_1: {number}, bits_2: {number}, bases_1: {number}, bases_2: {number}): {number}
	local key = {}
	for i = 1, #bits_1 do
		if bases_1[i] == bases_2[i] then
			table.insert(key, bits_1[i])
		end
	end
	return key
end

crypto.e91 = e91

local b92 = {}

function b92.generatebits(n: number): bits
	local bitstring = table.create(n)
	for i = 1, n do
		bitstring[i] = math.random(0, 1)
	end
	return bitstring
end

function b92.preparequbit(bit: number): qubit
	local qubit = qb.new(cpx.new(1, 0), cpx.new(0, 0))
	if bit == 0 then
		qubit = gate.apply(qubit, preset.hadamard)
	else
		qubit = gate.apply(qubit, preset.paulix)
	end
	return qubit
end

function b92.preparequbits(bits: bits): qubits
	local prepared = table.create(#bits)
	for i, bit in ipairs(bits) do
		prepared[i] = b92.preparequbit(bit)
	end
	return prepared
end

function b92.measurequbit(qubit: qubit): number
	local basis = math.random(0, 1)
	if basis == 0 then
		return qb.measure(qubit)
	else
		qubit = gate.apply(qubit, preset.hadamard)
		return qb.measure(qubit)
	end
end

function b92.measurequbits(qubits: qubits): bits
	local results = table.create(#qubits)
	for i, qubit in ipairs(qubits) do
		results[i] = b92.measurequbit(qubit)
	end
	return results
end

function b92.siftkey(sendbits: bits, receivebits: bits): bits
	local key = {}
	for i = 1, #sendbits do
		if sendbits[i] == 0 and receivebits[i] == 1 then
			table.insert(key, 0)
		elseif sendbits[i] == 1 and receivebits[i] == 0 then
			table.insert(key, 1)
		end
	end
	return key
end

function b92.errorestimation(alicekey: bits, bobkey: bits, samplesize: number?): number
	local errors = 0
	local samples = math.min(samplesize or math.floor(#alicekey * 0.1), #alicekey)

	for i = 1, samples do
		if alicekey[i] ~= bobkey[i] then
			errors += 1
		end
	end

	return errors / samples
end

function b92.privacyamplification(key: bits, finallength: number?): bits
	local actuallength = finallength or math.floor(#key * 0.75)
	local newkey = table.create(actuallength)

	for i = 1, actuallength do
		local xorsum = 0
		for j = 1, #key do
			xorsum = bit32.bxor(xorsum, key[j] * ((j * 1103515245 + 12345) % 32768))
		end
		newkey[i] = bit32.band(xorsum, 1)
	end

	return newkey
end

crypto.b92 = b92

local sixstates = {}

local BASIS_Z = "Z"
local BASIS_X = "X"
local BASIS_Y = "Y"

function sixstates.generatebits(n: number): bits
	local bits = table.create(n)
	for i = 1, n do
		bits[i] = math.random(0, 1)
	end
	return bits
end

function sixstates.generatebases(n: number): {string}
	local bases = table.create(n)
	for i = 1, n do
		bases[i] = ({BASIS_Z, BASIS_X, BASIS_Y})[math.random(1, 3)]
	end
	return bases
end

function sixstates.preparequbits(bits: bits, bases: {string}): qubits
	local qubits = table.create(#bits)
	for i, bit in ipairs(bits) do
		local qubit = qb.new(cpx.new(1, 0), cpx.new(0, 0))
		if bit == 1 then
			qubit = gate.apply(qubit, preset.paulix)
		end
		if bases[i] == BASIS_X then
			qubit = gate.apply(qubit, preset.hadamard)
		elseif bases[i] == BASIS_Y then
			qubit = gate.apply(qubit, preset.hadamard)
			qubit = gate.apply(qubit, preset.phase)
		end
		qubits[i] = qubit
	end
	return qubits
end

function sixstates.measurequbits(qubits: qubits, bases: {string}): bits
	local results = table.create(#qubits)
	for i, qubit in ipairs(qubits) do
		local basis = bases[i]
		if basis == BASIS_X then
			qubit = gate.apply(qubit, preset.hadamard)
		elseif basis == BASIS_Y then
			qubit = gate.apply(qubit, gate.dagger(preset.phase))
			qubit = gate.apply(qubit, preset.hadamard)
		end
		results[i] = qb.measure(qubit)
	end
	return results
end

function sixstates.siftkey(bits: bits, basesa: {string}, basesb: {string}): bits
	local key = {}
	for i = 1, #bits do
		if basesa[i] == basesb[i] then
			table.insert(key, bits[i])
		end
	end
	return key
end

function sixstates.errorestimation(alicekey: bits, bobkey: bits, samplesize: number?): number
	local errors = 0
	local samples = math.min(samplesize or math.floor(#alicekey * 0.1), #alicekey)

	for i = 1, samples do
		if alicekey[i] ~= bobkey[i] then
			errors += 1
		end
	end

	return errors / samples
end

function sixstates.privacyamplification(key: bits, finallength: number?): bits
	local actuallength = finallength or math.floor(#key * 0.75)
	local newkey = table.create(actuallength)

	for i = 1, actuallength do
		local xorsum = 0
		for j = 1, #key do
			xorsum = bit32.bxor(xorsum, key[j] * ((j * 1103515245 + 12345) % 32768))
		end
		newkey[i] = bit32.band(xorsum, 1)
	end

	return newkey
end

crypto.sixstates = sixstates

local diqkd = {}

local function apply_epr_pair(): (qubit, qubit)
	local qubit_a = qb.new(cpx.new(1, 0), cpx.new(0, 0))
	local qubit_b = qb.new(cpx.new(1, 0), cpx.new(0, 0))
	qubit_a = gate.apply(qubit_a, preset.hadamard)

	local entangled_qubits = gate.apply_ex({qubit_a, qubit_b}, preset.cnot)
	return entangled_qubits[1], entangled_qubits[2]
end

function diqkd.share_entangled_qubits(n: number): {{qubit}}
	local p = {}
	for i = 1, n do
		local qubit_a, qubit_b = apply_epr_pair()
		table.insert(p, {qubit_a, qubit_b})
	end
	return p
end

function diqkd.measure_qubits(pairs: {{qubit}}, basis_1: {string}, basis_2: {string}): {{number}}
	local results = {}
	for i, pair in ipairs(pairs) do
		local qubit_a = pair[1]
		local qubit_b = pair[2]

		local a_basis = basis_1[i]
		local b_basis = basis_2[i]

		local a_result = diqkd.measure_in_basis(qubit_a, a_basis)
		local b_result = diqkd.measure_in_basis(qubit_b, b_basis)

		table.insert(results, {a_result, b_result})
	end
	return results
end

function diqkd.measure_in_basis(qubit: qubit, basis: string): number
	if basis == "X" then
		qubit = gate.apply(qubit, preset.hadamard)
	end
	return qb.measure(qubit)
end

function diqkd.compare_results(results: {{number}}): {boolean}
	local dtc = {}
	for _, result_pair in ipairs(results) do
		local result_1 = result_pair[1]
		local result_2 = result_pair[2]

		if result_1 ~= result_2 then
			table.insert(dtc, true)
		else
			table.insert(dtc, false)
		end
	end
	return dtc
end

crypto.diqkd = diqkd

local decoy = {}

function decoy.generate_states(intensity_signal: number, intensity_decoy: number, num_photons: number): qubits
	local states = {}
	for i = 1, num_photons do
		local intensity = math.random() < 0.5 and intensity_signal or intensity_decoy
		states[i] = {cpx.new(math.sqrt(intensity), 0), cpx.new(math.sqrt(1 - intensity), 0)}
	end
	return states
end

function decoy.send_and_receive(states: qubits, error_rate: number): qubits
	local received = {}
	for _, state in ipairs(states) do
		local error = math.random() < error_rate
		table.insert(received, error and {state[2], state[1]} or state)
	end
	return received
end

function decoy.analyze(states_signal: qubits, states_decoy: qubits, received_signal: qubits, received_decoy: qubits): (number, number)
	local error_rate_signal = decoy.calculate_error_rate(states_signal, received_signal)
	local error_rate_decoy = decoy.calculate_error_rate(states_decoy, received_decoy)

	if error_rate_decoy > error_rate_signal * 1.1 then
		print("potential eavesdropper detected due to increased decoy state error rate.")
	end

	return error_rate_signal, error_rate_decoy
end

function decoy.calculate_error_rate(sent: qubits, received: qubits): number
	local errors = 0
	for i = 1, #sent do
		if sent[i][1].r ~= received[i][1].r or sent[i][2].r ~= received[i][2].r then
			errors = errors + 1
		end
	end
	return errors / #sent
end

crypto.decoy = decoy

local cvqkd = {}
export type coherent_state = {alpha: complex, phase: number}

function cvqkd.prepare_state(amplitude: number, phase: number): coherent_state
	local alpha = {r = amplitude * math.cos(phase), i = amplitude * math.sin(phase)}
	return {alpha = alpha, phase = phase}
end

function cvqkd.transmit_state(state: coherent_state, noise_variance: number): coherent_state
	local noise_r = math.random() * noise_variance
	local noise_i = math.random() * noise_variance
	local noisy_alpha = {
		r = state.alpha.r + noise_r,
		i = state.alpha.i + noise_i
	}
	return {alpha = noisy_alpha, phase = state.phase}
end

function cvqkd.homodyne_measurement(state: coherent_state, basis: string): number
	if basis == "X" then
		return state.alpha.r
	elseif basis == "P" then
		return state.alpha.i
	else
		error("invalid basis, use 'X' for amplitude or 'P' for phase")
	end
end

function cvqkd.eavesdropper_simulation(state: coherent_state, eavesdropping_prob: number): coherent_state
	if math.random() < eavesdropping_prob then
		local additional_noise = 0.05
		local noise_r = math.random() * additional_noise
		local noise_i = math.random() * additional_noise
		local eavesdropped_alpha = {
			r = state.alpha.r + noise_r,
			i = state.alpha.i + noise_i
		}
		return {alpha = eavesdropped_alpha, phase = state.phase}
	else
		return state
	end
end

function cvqkd.reconciliation(key_1: {number}, key_2: {number}): {number}
	local reconciled_key = {}
	for i = 1, #key_1 do
		if key_1[i] == key_2[i] then
			table.insert(reconciled_key, key_1[i])
		else
		end
	end
	return reconciled_key
end

function cvqkd.privacy_amplification(key: {number}, eavesdropper_info: number): {number}
	local secure_key = {}
	local reduction_factor = 1 - eavesdropper_info
	for i = 1, math.floor(#key * reduction_factor) do
		secure_key[i] = key[i]
	end
	return secure_key
end

crypto.cvqkd = cvqkd

local charlie = {}
charlie.__index = charlie

local mdiqkd = {}
mdiqkd.__index = mdiqkd

function charlie.bell_state_measurement(state_a: qubit, state_b: qubit): (string, number)
	local outcome = math.random(1, 4)
	local bell_state = ""
	if outcome == 1 then
		bell_state = "Phi+"
	elseif outcome == 2 then
		bell_state = "Phi-"
	elseif outcome == 3 then
		bell_state = "Psi+"
	else
		bell_state = "Psi-"
	end
	return bell_state, outcome
end

mdiqkd.charlie = charlie

function mdiqkd.prepare_state(intensity: number): qubit
	local amplitude = math.sqrt(intensity)
	return {{r = amplitude, i = 0}, {r = 0, i = amplitude}}
end

function mdiqkd.transmit_to_charlie(state: qubit, noise: number): qubit
	local noise_r = math.random() * noise
	local noise_i = math.random() * noise
	local noisy_state = {
		{r = state[1].r + noise_r, i = state[1].i + noise_i},
		{r = state[2].r + noise_r, i = state[2].i + noise_i}
	}
	return noisy_state
end

function mdiqkd.sift_keys(results: {string}): {number}
	local sifted_key = {}
	for _, result in ipairs(results) do
		if result == "Phi+" or result == "Psi-" then
			table.insert(sifted_key, 0)
		else
			table.insert(sifted_key, 1)
		end
	end
	return sifted_key
end

function mdiqkd.error_correction(key_1: {number}, key_2: {number}): {number}
	local corrected_key = {}
	for i = 1, #key_1 do
		if key_1[i] == key_2[i] then
			table.insert(corrected_key, key_1[i])
		end
	end
	return corrected_key
end

function mdiqkd.privacy_amplification(key: {number}, eavesdropper_info: number): {number}
	local secure_key = {}
	local reduction_factor = 1 - eavesdropper_info
	for i = 1, math.floor(#key * reduction_factor) do
		table.insert(secure_key, key[i])
	end
	return secure_key
end

crypto.mdiqkd = mdiqkd

local qds = {}

type qds_qt_sig = {qubit}
type message = string

function qds.setup(num_qubits: number): qds_qt_sig
	local signature: qds_qt_sig = {}
	for _ = 1, num_qubits do
		table.insert(signature, qb.new(
			cpx.new(math.random() * 2 - 1, math.random() * 2 - 1),
			cpx.new(math.random() * 2 - 1, math.random() * 2 - 1))
		)
	end
	return signature
end

function qds.sign(message: message, signature: qds_qt_sig): {any}
	local measurement_basis = {}
	local signed_message = message .. "_signed"

	for _, qubit in ipairs(signature) do
		local measure_result = math.random(0, 1)
		table.insert(measurement_basis, measure_result)
	end

	return {signed_message, measurement_basis}
end
-- {message | {number}}
function qds.verify(received_signature: {any}, original_signature: qds_qt_sig, vrf_threshold: number?): boolean
	local signed_message, measurement_basis = received_signature[1], received_signature[2]
	local verification_threshold = vrf_threshold or 0.95
	local match_count = 0

	for i, qubit in ipairs(original_signature) do
		local expected_measurement = measurement_basis[i]::number
		local actual_measurement = math.random(0, 1)
		if expected_measurement == actual_measurement then
			match_count += 1
		end
	end

	local match_ratio = match_count / #original_signature
	local verification_pass = match_ratio >= verification_threshold

	if verification_pass then
		print("sig verified successfully for message:", signed_message, "; match ratio:", match_ratio)
	else
		print("sig verification failed for message:", signed_message, "; match ratio:", match_ratio)
	end

	return verification_pass
end

local adversary = {}

function adversary.intercept_and_measure(message: string, measurement_basis: {number}): {any}
	local tampered_message = message .. "_tampered"
	local tampered_basis = {}

	for _, result in ipairs(measurement_basis) do
		local tampered_result = (math.random() < 0.2) and (1 - result) or result
		table.insert(tampered_basis, tampered_result)
	end

	print("adversary tampered with the signature")
	return {tampered_message, tampered_basis}
end

qds.adversary = adversary

crypto.qds = qds

local qdsbs = {}

export type bell_pair = qubits
export type qdsbs_qt_sig = {bell_pair}

function qdsbs.setup(num_pairs: number): qdsbs_qt_sig
	local signature: qdsbs_qt_sig = {}
	for _ = 1, num_pairs do
		table.insert(signature, {
			qb.new(cpx.new(1 / math.sqrt(2), 0), cpx.new(0, 0));
			qb.new(cpx.new(0, 0), cpx.new(1 / math.sqrt(2), 0))
		})
	end
	return signature
end

function qdsbs.sign(message: message, signature: qdsbs_qt_sig): {any}
	local measurement_basis = {}
	local signed_message = message .. "_signed"

	for _, bell_pair in ipairs(signature) do
		local measure_result = math.random(0, 1)
		table.insert(measurement_basis, measure_result)
	end

	return {signed_message, measurement_basis}
end

function qdsbs.verify(received_signature: {any}, original_signature: qdsbs_qt_sig, vrf_threshold: number?): boolean
	local signed_message, measurement_basis = received_signature[1], received_signature[2]
	local verification_threshold = vrf_threshold or 0.95
	local match_count = 0

	for i, bell_pair in ipairs(original_signature) do
		local expected_measurement = measurement_basis[i]
		local actual_measurement = math.random(0, 1)
		if expected_measurement == actual_measurement then
			match_count += 1
		end
	end

	local match_ratio = match_count / #original_signature
	local verification_pass = match_ratio >= verification_threshold

	if verification_pass then
		print("bell state signature verified successfully for message:", signed_message, "; match ratio:", match_ratio)
	else
		print("bell state signature verification failed for message:", signed_message, "; match ratio:", match_ratio)
	end

	return verification_pass
end

qdsbs.adversary = adversary

function crypto.coin_flip(): boolean
	local coin = qb.new(cpx.new(1, 0), cpx.new(0, 0)) -- |0⟩ state
	coin = gate.apply(coin, preset.hadamard)
	local result = qb.measure(coin)
	return result == 0
end

function crypto.commit_bit(bit: boolean): (qubits, qubits)
	local commitment = qb.new_ex(2)
	local key = qb.new_ex(2)
	if bit then
		commitment[1] = gate.apply(commitment[1], preset.paulix)
	end
	commitment[2] = gate.apply(commitment[2], preset.hadamard)
	key[1] = gate.apply(key[1], preset.hadamard)
	key[2] = gate.apply(key[2], preset.hadamard)
	local entangled = gate.apply_ex(commitment, preset.cnot)
	entangled = gate.apply_ex(entangled, preset.cnot, {1, 3}) -- Control: commitment[1], Target: key[1]
	entangled = gate.apply_ex(entangled, preset.cnot, {2, 4}) -- Control: commitment[2], Target: key[2]
	return {entangled[1], entangled[2]}, {entangled[3], entangled[4]}
end

-- b
function crypto.verify_commitment(commitment: qubits, key: qubits): boolean
	local disentangled = gate.apply_ex(commitment, preset.cnot, {1, 3}) -- Control: commitment[1], Target: key[1]
	disentangled = gate.apply_ex(disentangled, preset.cnot, {2, 4}) -- Control: commitment[2], Target: key[2]
	local measured_commitment = {qb.measure(disentangled[1]), qb.measure(disentangled[2])}
	return measured_commitment[1] == measured_commitment[2]
end

function crypto.reveal_bit(key: qubits): boolean
	local measured_key = {qb.measure(key[1]), qb.measure(key[2])}
	return measured_key[1] ~= measured_key[2]
end

--[[
	TODO:
		DONE - decoy state protocol
		DONE - continuous var qkd
		DONE - measurement device independent qkd

		DONE - quantum signing
		DONE - qds using bell states

		DONE - quantum threshold secret sharing
		DONE - shamir's secret sharing

		DONE - 1-2 quantum oblivious transfer
		DONE - quantum secure multiparty computation

		DONE - coin flipping
		DONE - bit commitment

		- lattice-based cg
		- code-based cg
		- hash-based cg
		- multivariate polynomial cg
		- supersingular isogeny-based cg

		- qmacs
		- one-time pad

		- recycling protocol

		- zero-knewledge proof

		- fhe
		- homomorphic enc

		MAYBE:
		- crystals-kyber
		- falcon & dilithium
		- bike
		- sphincs

]]

return crypto


--!native
--!optimize 2
--!strict

type complex = {r: number, i: number}
type qubit = {complex}
type qubits = {qubit}
type gate = {{complex}}
type matrix = {{complex}}

local gate = require("gate")
local preset = require("preset")

local decay = {}

function decay.new(probabilities: {
	amplitude_damping: number?;
	phase_damping: number?;
	depolarizing: number?;
})
	local meta = {}
	meta.amplitude_damping = probabilities.amplitude_damping or 0.01
	meta.phase_damping = probabilities.phase_damping or 0.02
	meta.depolarizing = probabilities.depolarizing or 0.02
	setmetatable(meta, decay)

	return meta
end

function decay:apply_amplitude_damping(qubits: qubits): qubits
	local noisy_qubits = {}
	for i, qubit in ipairs(qubits) do
		if math.random() < self.amplitude_damping then
			noisy_qubits[i] = gate.apply(qubit, preset.identity)
		else
			noisy_qubits[i] = qubit
		end
	end
	return noisy_qubits
end

function decay:apply_phase_damping(qubits: qubits): qubits
	local noisy_qubits = {}
	for i, qubit in ipairs(qubits) do
		if math.random() < self.phase_damping then
			noisy_qubits[i] = gate.apply(qubit, preset.pauliz)
		else
			noisy_qubits[i] = qubit
		end
	end
	return noisy_qubits
end

function decay:apply_depolarizing(qubits: qubits): qubits
	local noisy_qubits = {}
	for i, qubit in ipairs(qubits) do
		if math.random() < self.depolarizing then
			local noise_type = math.random()
			if noise_type < 1/3 then
				noisy_qubits[i] = gate.apply(qubit, preset.paulix)
			elseif noise_type < 2/3 then
				noisy_qubits[i] = gate.apply(qubit, preset.pauliy)
			else
				noisy_qubits[i] = gate.apply(qubit, preset.pauliz)
			end
		else
			noisy_qubits[i] = qubit
		end
	end
	return noisy_qubits
end

function decay:apply_all(qubits: qubits): qubits
	qubits = self:apply_amplitude_damping(qubits)
	qubits = self:apply_phase_damping(qubits)
	qubits = self:apply_depolarizing(qubits)
	return qubits
end

return decay


--!native
--!optimize 2
--!strict

type complex = {r: number, i: number}
type qubit = {complex}
type qubits = {qubit}
type gate = {{complex}}
type matrix = {{complex}}

local gate = require("gate")
local qb = require("qb")
local preset = require("preset")
local cpx = require("cpx")

local ent = {}

local bell = {}

local phi = {}
local psi = {}

function phi.plus(qubits: qubits)
	qubits[1] = gate.apply(qubits[1], preset.hadamard)
	qubits = gate.apply_ex(qubits, preset.cnot)
end

function phi.minus(qubits: qubits)
	qubits[1] = gate.apply(qubits[1], preset.hadamard)
	qubits = gate.apply_ex(qubits, preset.cnot)
	qubits[2] = gate.apply(qubits[2], preset.pauliz)
end

function psi.plus(qubits: qubits)
	qubits[1] = gate.apply(qubits[1], preset.hadamard)
	qubits = gate.apply_ex(qubits, preset.cnot)
	qubits[2] = gate.apply(qubits[2], preset.paulix)
end

function psi.minus(qubits: qubits)
	qubits[1] = gate.apply(qubits[1], preset.hadamard)
	qubits = gate.apply_ex(qubits, preset.cnot)
	qubits[2] = gate.apply(qubits[2], preset.paulix)
	qubits[2] = gate.apply(qubits[2], preset.pauliz)
end

bell.phi = phi
bell.psi = psi

ent.bell = bell

function ent.ghz(n: number): qubits
	local qubits = qb.new_ex(n)
	qubits[1] = gate.apply(qubits[1], preset.hadamard)
	for i = 2, n do
		qubits = gate.apply_ex(qubits, gate.cnot(i - 1, i))
	end
	return qubits
end

function ent.w(n: number): qubits
    local qubits = qb.new_ex(n)

    -- Create superposition of states with exactly one qubit in state |1⟩
    -- First, set the system to |10...0⟩ state (first qubit in state |1⟩)
    qubits[1] = gate.apply(qubits[1], preset.paulix)
    -- Now create an equal superposition of all single-excitation states
    -- This requires creating superpositions and applying rotations
    -- For n=2, we need a simple operation to get |01⟩ + |10⟩
    if n == 2 then
        -- Apply a special rotation gate to get (|01⟩ + |10⟩)/√2
        local sqrt_half = 1 / math.sqrt(2)
        local rotation_gate = {
            {cpx.new(sqrt_half), cpx.new(sqrt_half)},
            {cpx.new(sqrt_half), cpx.new(-sqrt_half)}
        }
        qubits = gate.apply_ex(qubits, rotation_gate)
    else
        -- For n>2, we need more complex operations to create the W state.
		-- Made an entire gate for this, as for the n=2 case, it uses a simple rotation to create the Bell-like state (|01⟩ + |10⟩)/√2.
	    for i = 2, n do
            local cos_theta = math.sqrt((i-1)/i)
            local sin_theta = math.sqrt(1/i)
            local rotation_gate = gate.controlled_rotation(1, i, cos_theta, sin_theta)
            qubits = gate.apply_ex(qubits, rotation_gate)
        end
    end

    return qubits
end

-- Does Luau support that???
-- Does Luau support that???
--[[
--- @param qubits qubits
--- @return measured_qubits, measurement_probabilities
--]]

function ent.measure(qubits: qubits)
    local measured_qubits = {}
    local measurement_probabilities = {}
    for i, qubit in ipairs(qubits) do
        local probability_0, probability_1 = qb.measure(qubit)
        measured_qubits[i] = qb.new(probability_0 > probability_1)
        measurement_probabilities[i] = {probability_0, probability_1}
    end
    return measured_qubits, measurement_probabilities
end

function ent.cluster(n: number): qubits
    local qubits = qb.new_ex(n)
    qubits[1] = gate.apply(qubits[1], preset.hadamard)
    for i = 2, n do
        qubits = gate.apply_ex(qubits, gate.cnot(i - 1, i))
        qubits = gate.apply_ex(qubits, gate.cnot(i, i - 1))
    end
    return qubits
end

function ent.partial_trace(qubits: qubits, traced_indices: {number})
    if not qubits or #qubits == 0 then
        error("No qubits provided for partial trace")
    end

    if not traced_indices or #traced_indices == 0 then
        error("No indices provided for tracing out")
    end

    local trace_map = {}
    for _, index in ipairs(traced_indices) do
        if index < 1 or index > #qubits then
            error("Invalid qubit index: " .. tostring(index))
        end
        trace_map[index] = true
    end

    local total_qubits = #qubits
    local traced_qubits = #traced_indices
    local remaining_qubits = total_qubits - traced_qubits
    local kept_indices = {}

    for i = 1, total_qubits do
        if not trace_map[i] then
            table.insert(kept_indices, i)
        end
    end
    local full_density_matrix = qb.compute_density_matrix(qubits)

    local function compute_reduced_density()
        local reduced_density_matrix = {}
        local reduced_dim = 2^remaining_qubits
        for i = 1, reduced_dim do
            reduced_density_matrix[i] = {}
            for j = 1, reduced_dim do
                reduced_density_matrix[i][j] = cpx.new(0, 0)
            end
        end

        for basis_state_kept = 0, reduced_dim - 1 do
            for basis_state_traced = 0, 2^traced_qubits - 1 do
                local global_state = 0
                local kept_index = 1

                for qubit_index = 1, total_qubits do
                    if not trace_map[qubit_index] then
                        local bit = bit32.band(bit32.rshift(basis_state_kept, kept_index - 1), 1)
                        global_state = bit32.bor(global_state, bit32.lshift(bit, qubit_index - 1))
                        kept_index = kept_index + 1
                    else
                        local bit = bit32.band(bit32.rshift(basis_state_traced, qubit_index - #kept_indices - 1), 1)
                        global_state = bit32.bor(global_state, bit32.lshift(bit, qubit_index - 1))
                    end
                end
                reduced_density_matrix[basis_state_kept + 1][basis_state_kept + 1] =
                    cpx.add(reduced_density_matrix[basis_state_kept + 1][basis_state_kept + 1],
                           full_density_matrix[global_state + 1][global_state + 1])
            end
        end

        return reduced_density_matrix
    end

    local success, reduced_matrix = pcall(compute_reduced_density)
    if not success then
        error("Failed to compute partial trace: " .. tostring(reduced_matrix))
    end

    return reduced_matrix
end

function ent.print_density_matrix(matrix: matrix)
    for _, row in ipairs(matrix) do
        local formatted_row = {}
        for _, val in ipairs(row) do
            table.insert(formatted_row, string.format("%.4f", val.r))
        end
        print(table.concat(formatted_row, "\t"))
    end
end

return ent


--!native
--!optimize 2
--!strict

type complex = {r: number, i: number}
type qubit = {complex}
type qubits = {qubit}
type gate = {{complex}}
type matrix = {{complex}}

local cpx = require("cpx")
local gate = require("gate")
local mtx = require("mtx")
local qb = require("qb")
local preset = require("preset")

local err = {}

function err.shor(N: number): (number, number)
	local function mod_exp(base: number, exp: number, mod: number): number
		local result = 1
		base = base % mod
		while exp > 0 do
			if exp % 2 == 1 then
				result = (result * base) % mod
			end
			base = (base * base) % mod
			exp = math.floor(exp / 2)
		end
		return result
	end
	local function gcd(a: number, b: number): number
		while b ~= 0 do
			local temp = b
			b = a % b
			a = temp
		end
		return a
	end
	local function classical_shor(N: number): (number, number)
		if N % 2 == 0 then
			return 2, N / 2
		end

		local a = math.random(2, N - 1)
		local g = gcd(a, N)
		if g > 1 then
			return g, N / g -- i hope so.
		end

		return a, -1
	end
	local function continued_fraction(x: number, N: number): (number, number)
		local epsilon = 1e-6
		local a, b = 0, 1
		local prev_a, prev_b = 1, 0

		while math.abs(a / b - x) > epsilon do
			local q = math.floor(x)
			x = 1 / (x - q)
			local temp_a = a
			a = q * a + prev_a
			prev_a = temp_a

			local temp_b = b
			b = q * b + prev_b
			prev_b = temp_b
		end
		return a, b
	end
	local a, factor = classical_shor(N)
	if factor then
		return factor, N / factor
	end

	local qs = qb.new_ex(math.ceil(math.log(N) / math.log(2)))
	qs = gate.phase_estimation(a, N, qs)

	local period = qb.mmeasure(qs)
	local _, r = continued_fraction(period / (2^#qs), N)

	if r % 2 == 0 then
		local x1 = mod_exp(a, r // 2, N) - 1
		local x2 = mod_exp(a, r // 2, N) + 1
		local factor1 = gcd(x1, N)
		local factor2 = gcd(x2, N)

		if factor1 > 1 and factor2 > 1 then
			return factor1, factor2
		end
	end

	return err.shor(N)
end

local steane = {}

function steane.new(): qubits
	local qubits = {}
	for i = 1, 7 do
		qubits[i] = qb.new(cpx.new(1, 0), cpx.new(0, 0))
	end
	return qubits
end

function steane.encode(qubit: qubit): qubits
	local encoded = steane.new()

	encoded[1] = gate.apply(qubit, preset.paulix)
	encoded[2] = gate.apply(qubit, preset.paulix)
	encoded[3] = gate.apply(qubit, preset.hadamard)
	encoded[4] = gate.apply(qubit, preset.hadamard)
	encoded[5] = gate.apply(qubit, preset.paulix)
	encoded[6] = gate.apply(qubit, preset.paulix)
	encoded[7] = gate.apply(qubit, preset.hadamard)

	return encoded
end

function steane.measure_syndrome(qubits: qubits): matrix
	local synd = mtx.new({r = 3, c = 7})
	local H = mtx.new({r = 3, c = 7},
		cpx.new(1), cpx.new(1), cpx.new(1), cpx.new(1), cpx.new(), cpx.new(), cpx.new(),
		cpx.new(1), cpx.new(-1), cpx.new(1), cpx.new(-1), cpx.new(), cpx.new(), cpx.new(),
		cpx.new(1), cpx.new(1), cpx.new(-1), cpx.new(-1), cpx.new(1), cpx.new(1), cpx.new(-1)
	)
	for i = 1, 3 do
		for j = 1, 4 do
			synd[i][j] = cpx.new(0, 0)
			for k = 1, 7 do
				synd[i][j] = cpx.add(synd[i][j], cpx.mul(qubits[k][j], H[i][k]))
			end
		end
	end

	return synd
end

function steane.correct_errors(qubits: qubits, synd: matrix): qubits
	local ecr_mtx = mtx.identity(7)
	for i = 1, 7 do
		for j = 1, 3 do
			if synd[j][i] == cpx.new(1, 0) then
				qubits = gate.apply_ex(qubits, ecr_mtx)
			end
		end
	end

	return qubits
end

function steane.decode(qubits: qubits): qubit
	local decoded = qb.new(cpx.new(), cpx.new())
	decoded = gate.apply(decoded, preset.hadamard)
	decoded = gate.apply(decoded, preset.paulix)

	return decoded
end

err.steane = steane

local surface = {}

function surface.create_grid(size: number): qubits
	local qubits = {}
	for i = 1, size * size do
		qubits[i] = qb.new(cpx.new(1, 0), cpx.new(0, 0))
	end
	return qubits
end

function surface.stabilizers(size: number): (matrix, matrix)
	local x_stabilizers = {}
	local z_stabilizers = {}

	for i = 1, size do
		x_stabilizers[i] = {}
		z_stabilizers[i] = {}
		for j = 1, size do
			x_stabilizers[i][j] = cpx.new(0, 0)
			z_stabilizers[i][j] = cpx.new(0, 0)
		end
	end

	for i = 1, size do
		for j = 1, size do
			x_stabilizers[i][j] = cpx.new(1, 0)
		end
	end

	for i = 1, size - 1 do
		for j = 1, size - 1 do
			z_stabilizers[i][j] = cpx.new(1, 0)
		end
	end

	return x_stabilizers, z_stabilizers
end

function surface.measure_syndrome(qubits: qubits, x_stabilizers: matrix, z_stabilizers: matrix): ({complex}, {complex})
	local x_syndrome = {}
	local z_syndrome = {}

	for i = 1, #x_stabilizers do
		x_syndrome[i] = cpx.new()
		for j = 1, #x_stabilizers[i] do
			x_syndrome[i] = cpx.add(x_syndrome[i], cpx.mul(qubits[j][1], x_stabilizers[i][j]))
		end
	end

	for i = 1, #z_stabilizers do
		z_syndrome[i] = cpx.new()
		for j = 1, #z_stabilizers[i] do
			z_syndrome[i] = cpx.add(z_syndrome[i], cpx.mul(qubits[j][1], z_stabilizers[i][j]))
		end
	end

	return x_syndrome, z_syndrome
end

function surface.correct_errors(qubits: qubits, x_syndrome: {number}, z_syndrome: {number}): qubits
	for i = 1, #x_syndrome do
		if x_syndrome[i] == 1 then
			qubits[i] = gate.apply(qubits[i], preset.paulix)
		end
	end

	for i = 1, #z_syndrome do
		if z_syndrome[i] == 1 then
			qubits[i] = gate.apply(qubits[i], preset.pauliz)
		end
	end

	return qubits
end

err.surface = surface

local bitflip = {}

function bitflip.encode(qubit: qubit): qubits
	local encoded = {}
	for i = 1, 3 do
		encoded[i] = qb.new(cpx.new(qubit[1].r, qubit[1].i), cpx.new(0, 0))
	end
	return encoded
end

function bitflip.apply_error(encoded: qubits, qubit_index: number, error_type: string): qubits
	if error_type == "bitflip" then
		encoded[qubit_index] = gate.apply(encoded[qubit_index], preset.paulix)
	elseif error_type == "phaseflip" then
		encoded[qubit_index] = gate.apply(encoded[qubit_index], preset.pauliz)
	elseif error_type == "both" then
		encoded[qubit_index] = gate.apply(encoded[qubit_index], preset.paulix)
		encoded[qubit_index] = gate.apply(encoded[qubit_index], preset.pauliz)
	else
		error("Unknown error type")
	end
	return encoded
end

function bitflip.decode(encoded: qubits): qubit
	local counts = {0, 0}
	for i = 1, 3 do
		if encoded[i][1].r > 0.5 then
			counts[1] = counts[1] + 1
		else
			counts[2] = counts[2] + 1
		end
	end
	local state = counts[1] > counts[2] and cpx.new(1, 0) or cpx.new(0, 0)
	return qb.new(state, cpx.new(0, 0))
end

function bitflip.correct_errors(encoded: qubits): qubits
	local syndrome = {0, 0, 0}
	for i = 1, 3 do
		syndrome[i] = (encoded[i][1].r > 0.5 and 1 or 0)
	end

	local majority = (syndrome[1] + syndrome[2] + syndrome[3]) / 2
	for i = 1, 3 do
		if syndrome[i] ~= majority then
			encoded[i] = gate.apply(encoded[i], preset.paulix)
		end
	end

	return encoded
end

err.bitflip = bitflip

local signflip = {}

function signflip.encode(qubit: qubit): qubits
	local encoded = {}
	for i = 1, 3 do
		encoded[i] = qb.new(cpx.new(qubit[1].r, qubit[1].i), cpx.new(0, 0))
	end
	return encoded
end

function signflip.apply_error(encoded: qubits, qubit_index: number): qubits
	encoded[qubit_index] = gate.apply(encoded[qubit_index], preset.pauliz)
	return encoded
end

function signflip.decode(encoded: qubits): qubit
	local counts = {0, 0}
	for i = 1, 3 do
		if encoded[i][1].r > 0.5 then
			counts[1] = counts[1] + 1
		else
			counts[2] = counts[2] + 1
		end
	end
	local state = counts[1] > counts[2] and cpx.new(1, 0) or cpx.new(0, 0)
	return qb.new(state, cpx.new(0, 0))
end

function signflip.measure_syndrome(encoded: qubits): {number}
	local syndrome = {0, 0}
	for i = 1, 3 do
		if encoded[i][1].r < 0.5 then
			syndrome[i] = 1
		else
			syndrome[i] = 0
		end
	end
	return syndrome
end

function signflip.correct_errors(encoded: qubits, syndrome: {number}): qubits
	for i = 1, 3 do
		if syndrome[i] == 1 then
			encoded[i] = gate.apply(encoded[i], preset.pauliz)
		end
	end
	return encoded
end

err.signflip = signflip

local cat = {}

function cat.encode(qubit: qubit, n: number): qubits
	local encoded = {}
	for i = 1, n do
		encoded[i] = qb.new(cpx.new(qubit[1].r, qubit[1].i), cpx.new(qubit[2].r, qubit[2].i))
	end
	return encoded
end

function cat.apply_error(encoded: qubits, qubit_index: number, error_type: string): qubits
	if error_type == "bitflip" then
		encoded[qubit_index] = gate.apply(encoded[qubit_index], preset.paulix)
	elseif error_type == "phaseflip" then
		encoded[qubit_index] = gate.apply(encoded[qubit_index], preset.pauliz)
	else
		error("unknown error type")
	end
	return encoded
end

function cat.measure_syndrome(encoded: qubits): {number}
	local syndrome = {}
	for i = 1, #encoded - 1 do
		syndrome[i] = qb.mmeasure_ex(gate.apply_ex({encoded[i], encoded[i+1]}, preset.cnot))[1]
	end
	return syndrome
end

function cat.correct_errors(encoded: qubits, syndrome: {number}): qubits
	for i = 1, #syndrome do
		if syndrome[i] == 1 then
			encoded[i] = gate.apply(encoded[i], preset.paulix)
		end
	end
	return encoded
end

function cat.decode(encoded: qubits): qubit
	return encoded[1]
end

err.cat = cat

local binomial = {}

local function choose(n: number, k: number): number
	if k == 0 or k == n then
		return 1
	else
		return choose(n - 1, k - 1) + choose(n - 1, k)
	end
end

function binomial.encode(qubit: qubit, n: number): qubits
	local encoded = {}
	for i = 0, n do
		local amplitude = cpx.mul(cpx.new(math.sqrt(choose(n, i)), 0), qubit[1])
		encoded[i+1] = qb.new(amplitude, cpx.new(0, 0))
	end
	return encoded
end

function binomial.apply_error(encoded: qubits, error_type: string): qubits
	if error_type == "loss" then
		local index = math.random(1, #encoded)
		encoded[index] = qb.new(cpx.new(0, 0), cpx.new(0, 0))
	elseif error_type == "dephasing" then
		for i = 1, #encoded do
			encoded[i] = gate.apply(encoded[i], preset.pauliz)
		end
	else
		error("unknown error type")
	end
	return encoded
end

function binomial.measure_syndrome(encoded: qubits): number
	local total = 0
	for i = 1, #encoded do
	total = total + qb.measure(encoded[i]) * (i - 1)
	end
	return total
end

function binomial.correct_errors(encoded: qubits, syndrome: number): qubits
	local n = #encoded - 1
	if syndrome ~= n // 2 then
		local shift = n // 2 - syndrome
		local newencoded = {}
		for i = 1, #encoded do
			local newindex = i + shift
			if newindex >= 1 and newindex <= #encoded then
				newencoded[newindex] = encoded[i]
			end
		end
		encoded = newencoded
	end
	return encoded
end

function binomial.decode(encoded: qubits): qubit
	local n = #encoded - 1
	return encoded[n // 2 + 1]
end

err.binomial = binomial

-- repetition too
local repetition = {}

function repetition.encode(qubit: qubit, n: number): qubits
	local encoded = {}
	for i = 1, n do
		encoded[i] = qb.new(cpx.new(qubit[1].r, qubit[1].i), cpx.new(qubit[2].r, qubit[2].i))
	end
	return encoded
end

function repetition.apply_error(encoded: qubits, error_rate: number): qubits
	for i = 1, #encoded do
		if math.random() < error_rate then
			encoded[i] = gate.apply(encoded[i], preset.paulix)
		end
	end
	return encoded
end

function repetition.measure_syndrome(encoded: qubits): {number}
	local syndrome = {}
	for i = 1, #encoded - 1 do
		local parity = qb.mmeasure_ex(gate.apply_ex({encoded[i], encoded[i+1]}, preset.cnot))[1]
		syndrome[i] = parity
	end
	return syndrome
end

function repetition.correct_errors(encoded: qubits, syndrome: {number}): qubits
	local error_positions = {}
	for i = 1, #syndrome do
		if syndrome[i] == 1 then
			table.insert(error_positions, i)
		end
	end

	if #error_positions % 2 == 1 then
		table.insert(error_positions, #encoded)
	end

	for _, pos in ipairs(error_positions) do
		encoded[pos] = gate.apply(encoded[pos], preset.paulix)
	end

	return encoded
end

function repetition.decode(encoded: qubits): qubit
	local sum_x = cpx.new(0, 0)
	local sum_y = cpx.new(0, 0)
	for i = 1, #encoded do
		sum_x = cpx.add(sum_x, encoded[i][1])
		sum_y = cpx.add(sum_y, encoded[i][2])
	end
	return qb.new(cpx.div(sum_x, cpx.new(#encoded, 0)), cpx.div(sum_y, cpx.new(#encoded, 0)))
end

err.repetition = repetition

local toric = {}
toric.__index = toric

function toric.init(size: number?)
	local meta = {}
	meta.lattice = {}
	meta.size = size or 5
	setmetatable(meta, toric)

	for i = 1, meta.size do
		meta.lattice[i] = {}
		for j = 1, meta.size do
			meta.lattice[i][j] = qb.new(cpx.new(), cpx.new())
		end
	end

	return meta
end

function toric:apply_x_error(x: number, y: number)
	self.lattice[x][y] = gate.apply(self.lattice[x][y], preset.paulix)
end

function toric:apply_z_error(x: number, y: number)
	self.lattice[x][y] = gate.apply(self.lattice[x][y], preset.pauliz)
end

function toric:star_stabilizer(x: number, y: number)
	local qubits = {
		qb.new(cpx.new(x), cpx.new(y));
		qb.new(cpx.new((x - 1) % self.size + 1), cpx.new(y));
		qb.new(cpx.new((x + 1) % self.size + 1), cpx.new(y));
		qb.new(cpx.new(x), cpx.new((y - 1) % self.size + 1));
		qb.new(cpx.new(x), cpx.new((y + 1) % self.size + 1));
	}
	for _, pos in ipairs(qubits) do
		self.lattice[pos[1]][pos[2]] = gate.apply(self.lattice[pos[1]][pos[2]], preset.paulix)
	end
end

function toric:plaquette_stabilizer(x: number, y: number)
	local qubits = {
		qb.new(cpx.new(x), cpx.new(y));
		qb.new(cpx.new((x - 1) % self.size + 1), cpx.new(y));
		qb.new(cpx.new((x + 1) % self.size + 1), cpx.new(y));
		qb.new(cpx.new(x), cpx.new((y - 1) % self.size + 1));
		qb.new(cpx.new(x), cpx.new((y + 1) % self.size + 1));
	}
	for _, pos in ipairs(qubits) do
		self.lattice[pos[1]][pos[2]] = gate.apply(self.lattice[pos[1]][pos[2]], preset.pauliz)
	end
end

function toric:measure_star_stabilizer(x: number, y: number)
	return qb.measure(self.lattice[x][y]) == 1
end

function toric:measure_plaquette_stabilizer(x: number, y: number)
	return qb.measure(self.lattice[x][y]) == 1
end

function toric:syndrome_extraction()
	for i = 1, self.size do
		for j = 1, self.size do
			if self:measure_star_stabilizer(i, j) then
				print("x-type error at:", i, j)
			end
			if self:measure_plaquette_stabilizer(i, j) then
				print("z-type error at:", i, j)
			end
		end
	end
end

err.toric = toric

return err


--!native
--!optimize 2
--!strict
type complex = {r: number, i: number}
type qubit = {complex}
type qubits = {qubit}
type matrix = {{complex}}

export type gate = {{complex}}

local cpx = require("cpx")
local mtx = require("mtx")

local gate = {}

function gate.new(a11: complex, a12: complex, a21: complex, a22: complex): gate
	return {{a11, a12}, {a21, a22}}
end

function gate.new_ex(size: number, ...: complex): gate
	local matrix = {...}
	assert(#matrix == size^2, "matrix must be " .. size .. "x" .. size)
	local g = {}
	for i = 1, size do
		g[i] = {}
		for j = 1, size do
			g[i][j] = matrix[(i-1)*size + j]
		end
	end
	return g
end

function gate.apply(qubit: qubit, g: gate): qubit
	local alpha = cpx.add(
		cpx.mul(g[1][1], qubit[1]),
		cpx.mul(g[1][2], qubit[2])
	)
	local beta = cpx.add(
		cpx.mul(g[2][1], qubit[1]),
		cpx.mul(g[2][2], qubit[2])
	)
	return {alpha, beta}
end

function gate.apply_ex(qubits: qubits, g: gate, scope: {number}?): qubits
	local n = #g
	local new_qubits = {}
	for i = 1, n do
		new_qubits[i] = {cpx.new(0, 0), cpx.new(0, 0)}
		if scope then
			for _, v in ipairs(scope) do
				new_qubits[v][1] = cpx.add(new_qubits[v][1], cpx.mul(g[v][i], qubits[i][1]))
				new_qubits[v][2] = cpx.add(new_qubits[v][2], cpx.mul(g[v][i], qubits[i][2]))
			end
			break
		end
		for j = 1, n do
			new_qubits[i][1] = cpx.add(new_qubits[i][1], cpx.mul(g[i][j], qubits[j][1]))
			new_qubits[i][2] = cpx.add(new_qubits[i][2], cpx.mul(g[i][j], qubits[j][2]))
		end
	end
	return new_qubits
end

function gate.controlled_modular_exponentiation(a: number, N: number, n: number): gate
	local function mod_exp(base: number, exp: number, mod: number): number
		local result = 1
		base = base % mod
		while exp > 0 do
			if exp % 2 == 1 then
				result = (result * base) % mod
			end
			base = (base * base) % mod
			exp = math.floor(exp / 2)
		end
		return result
	end

	local g = mtx.identity(2^n)
	for x = 0, 2^n - 1 do
		local y = mod_exp(a, x, N)
		g[x + 1][y + 1] = cpx.new(1, 0)
		g[x + 1][x + 1] = cpx.new(0, 0)
	end
	return g
end

function gate.diffusion(qubits: qubits): qubits
	local n = #qubits
	local g = {}
	for i = 1, n do
		g[i] = {}
		for j = 1, n do
			if i == j then
				g[i][j] = cpx.new(2/n - 1, 0)
			else
				g[i][j] = cpx.new(2/n, 0)
			end
		end
	end
	return gate.apply_ex(qubits, g)
end

function gate.controlled_phase(angle: number): gate
	return gate.new_ex(4,
		cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(math.cos(angle), math.sin(angle))
	)
end

function gate.rx(angle: number): gate
	local half_angle = angle / 2
	local cos_half = math.cos(half_angle)
	local sin_half = math.sin(half_angle)
	return gate.new(
		cpx.new(cos_half, 0), cpx.new(0, -sin_half),
		cpx.new(0, sin_half), cpx.new(cos_half, 0)
	)
end

function gate.ry(angle: number): gate
	local half_angle = angle / 2
	local cos_half = math.cos(half_angle)
	local sin_half = math.sin(half_angle)
	return gate.new(
		cpx.new(cos_half, 0), cpx.new(-sin_half, 0),
		cpx.new(sin_half, 0), cpx.new(cos_half, 0)
	)
end

function gate.rz(angle: number): gate
	local half_angle = angle / 2
	return gate.new(
		cpx.new(math.cos(half_angle), -math.sin(half_angle)), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(math.cos(-half_angle), math.sin(half_angle))
	)
end

function gate.qft(n: number): gate
	local sz = 2^n
	local qft = mtx.new({r = sz, c = sz})

	for i = 1, sz do
		for j = 1, sz do
			qft[i][j] = cpx.new(0, 0)
		end
	end

	for i = 1, sz do
		for j = 1, sz do
			local angle = 2 * math.pi * (i - 1) * (j - 1) / sz
			qft[i][j] = cpx.exp(cpx.new(0, angle))
		end
	end

	local norm = 1 / math.sqrt(sz)
	for i = 1, sz do
		for j = 1, sz do
			qft[i][j] = cpx.mul(qft[i][j], cpx.new(norm, 0))
		end
	end

	return qft
end

function gate.inverse_qft(n: number): gate
	local qft_inv = mtx.identity(2^n)

	for i = n, 1, -1 do
		for j = i + 1, n do
			local angle = -2 * math.pi / (2^(j - i + 1))
			local phase_gate = gate.controlled_phase(angle)
			qft_inv = mtx.mul(qft_inv, phase_gate)
		end

		local expanded_hadamard = mtx.expand(gate.new(
			cpx.new(1/math.sqrt(2), 0), cpx.new(1/math.sqrt(2), 0),
			cpx.new(1/math.sqrt(2), 0), cpx.new(-1/math.sqrt(2), 0)
		), {r = n, c = n})
		qft_inv = mtx.mul(qft_inv, expanded_hadamard)
	end

	for i = 1, n // 2 do
		local expanded_swap = mtx.expand(gate.new_ex(4,
			cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0),
			cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0),
			cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0),
			cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0)
		), {r = n, c = n})
		qft_inv = mtx.mul(qft_inv, expanded_swap)
	end

	return qft_inv
end

--[[
function gate.new_ex(n: number): qubits
	local size = 2 ^ n
	local qubits = {}
	for i = 1, size do
		qubits[i] = {cpx.new(if i == 1 then 1 else 0, 0), cpx.new(0, 0)}
	end
	return qubits
end
]]--

function gate.phase_estimation(a: number, N: number, qubits: qubits): qubits
	local n = #qubits
	local qs = {}
	for i = 1, 2 ^ n do
		qs[i] = {cpx.new(if i == 1 then 1 else 0, 0), cpx.new(0, 0)}
	end

	for i = 1, n do
		qs = gate.apply_ex(qs, gate.new(
			cpx.new(1/math.sqrt(2), 0), cpx.new(1/math.sqrt(2), 0),
			cpx.new(1/math.sqrt(2), 0), cpx.new(-1/math.sqrt(2), 0)
		))
	end

	local cme_gate = gate.controlled_modular_exponentiation(a, N, n)
	qs = gate.apply_ex(qs, cme_gate)

	qs = gate.apply_ex(qs, gate.qft(n))

	return qs
end

function gate.cnot(control, target)
	local identity = {
		{cpx.new(1, 0), cpx.new(0, 0)},
		{cpx.new(0, 0), cpx.new(1, 0)}
	}
	local x_gate = {
		{cpx.new(0, 0), cpx.new(1, 0)},
		{cpx.new(1, 0), cpx.new(0, 0)}
	}
	local cnot = identity
	cnot[control][target] = x_gate[1][1]
	return cnot
end

function gate.cost_hamiltonian(qs: qubits, gamma: number): qubits
	for i = 1, #qs do
		local z_gate = gate.new(
			cpx.new(1, 0), cpx.new(0, 0),
			cpx.new(0, 0), cpx.exp(cpx.new(0, -gamma))
		)
		qs = gate.apply_ex(qs, z_gate)
	end
	return qs
end

function gate.mixer_hamiltonian(qs: qubits, beta: number): qubits
	for i = 1, #qs do
		local x_gate = gate.new(
			cpx.mul(cpx.exp(cpx.new(0, -beta)), cpx.new(0, 0)),
			cpx.mul(cpx.exp(cpx.new(0, beta)), cpx.new(1, 0)),
			cpx.new(1, 0), cpx.new(0, 0)
		)
		qs = gate.apply_ex(qs, x_gate)
	end
	return qs
end

function gate.dagger(g: gate): gate
	local n = #g
	local dagger = {}
	for i = 1, n do
		dagger[i] = {}
		for j = 1, n do
			dagger[i][j] = cpx.conj(g[j][i])
		end
	end
	return dagger
end

function gate.tensorproduct(g1: gate, g2: gate): gate
	local n1, n2 = #g1, #g2
	local result = {}
	for i = 1, n1 * n2 do
		result[i] = {}
		for j = 1, n1 * n2 do
			local i1, i2 = math.ceil(i / n2), (i - 1) % n2 + 1
			local j1, j2 = math.ceil(j / n2), (j - 1) % n2 + 1
			result[i][j] = cpx.mul(g1[i1][j1], g2[i2][j2])
		end
	end
	return result
end

function gate.controlledgate(g: gate): gate
	local n = #g
	local controlled = mtx.identity(2 * n)
	for i = 1, n do
		for j = 1, n do
			controlled[n + i][n + j] = g[i][j]
		end
	end
	return controlled
end

function gate.trace(g: gate): complex
	local tr = cpx.new(0, 0)
	for i = 1, #g do
		tr = cpx.add(tr, g[i][i])
	end
	return tr
end

function gate.fidelity(g1: gate, g2: gate): number
	local product = gate.apply_ex(g1, gate.dagger(g2))
	local tr = gate.trace(product)
	return math.abs(tr.r) / #g1
end

function gate.controlled_rotation(control: number, target: number, cos_theta: number, sin_theta: number): gate
    -- If control qubit is |0⟩, do nothing (identity)
    -- If control qubit is |1⟩, apply rotation by the specified angle
    return gate.new_ex(4,
        cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0),
        cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0),
        cpx.new(0, 0), cpx.new(0, 0), cpx.new(cos_theta, 0), cpx.new(-sin_theta, 0),
        cpx.new(0, 0), cpx.new(0, 0), cpx.new(sin_theta, 0), cpx.new(cos_theta, 0)
    )
end

function gate.decompose(g: gate): {
	gate | number
}
	assert(#g == 2 and #g[1] == 2, "decompose only supports 2x2 gates for now")

	local function atan2(y: number, x: number)
		if x > 0 then
			return math.atan(y / x)
		elseif x < 0 and y >= 0 then
			return math.atan(y / x) + math.pi
		elseif x < 0 and y < 0 then
			return math.atan(y / x) - math.pi
		elseif x == 0 and y > 0 then
			return math.pi / 2
		elseif x == 0 and y < 0 then
			return -math.pi / 2
		else
			return 0
		end
	end

	local det = cpx.sub(cpx.mul(g[1][1], g[2][2]), cpx.mul(g[1][2], g[2][1]))
	local phase = atan2(det.i, det.r) / 2
	local u = mtx.create_empty({r = 2, c = 2})
	for i = 1, 2 do
		for j = 1, 2 do
			u[i][j] = cpx.mul(g[i][j], cpx.exp(cpx.new(0, -phase)))
		end
	end

	local theta = 2 * math.acos(math.min(1, math.max(-1, u[1][1].r)))
	local phi = atan2(u[1][2].i, u[1][2].r)
	local lambda = atan2(u[2][1].i, u[2][1].r)

	return {
		gate.rz(phi),
		gate.ry(theta),
		gate.rz(lambda),
		phase
	}
end

function gate.toffoli(): gate
	local toffoli = mtx.identity(8)
	toffoli[7][7] = cpx.new(0, 0)
	toffoli[7][8] = cpx.new(1, 0)
	toffoli[8][7] = cpx.new(1, 0)
	toffoli[8][8] = cpx.new(0, 0)
	return toffoli
end

function gate.swap(): gate
	return {
		{cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0)},
		{cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0)},
		{cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0)},
		{cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0)}
	}
end

function gate.controlled_unitary(unitary: gate, control: number, target: number, exponent: number): gate
	local controlled_unitary = mtx.exp(unitary, exponent)

	local cu_gate = gate.new_ex(4,
		cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(0, 0), controlled_unitary[1][1], controlled_unitary[1][2],
		cpx.new(0, 0), cpx.new(0, 0), controlled_unitary[2][1], controlled_unitary[2][2]
	)

	return cu_gate
end

function gate.isunitary(g: gate)
	local dagger = gate.dagger(g)
	local product = gate.apply_ex(g, dagger)
	local identity = mtx.identity(#g)
	local epsilon = 1e-10

	for i = 1, #g do
		for j = 1, #g do
			if math.abs(product[i][j].r - identity[i][j].r) > epsilon or
			   math.abs(product[i][j].i - identity[i][j].i) > epsilon then
				return false
			end
		end
	end
	return true
end

return gate


--!native
--!optimize 2
--!strict

type kitaev_obj = {
	__index : kitaev_obj;
	init_chain: (chain_size: number?, delta: number?, mu: number?, t: number?, time_step: number?) -> kitaev_meta;
	apply_majorana: (self: kitaev_meta, i: number, mode: number) -> ();
	hamiltonian: (self: kitaev_meta) -> number;
	detect_majorana: (self: kitaev_meta) -> ();
	evolve_chain: (self: kitaev_meta, steps: number) -> ();
	braid_majorana: (self: kitaev_meta, i: number, j: number, mode_i: number, mode_j: number) -> ();
	error_detect: (self: kitaev_meta) -> number;
	error_correct: (self: kitaev_meta) -> ();
	measure_majorana: (self: kitaev_meta, i: number, mode: number) -> ();
	full_simulation: (self: kitaev_meta, steps: number) -> ();
	perform_braiding_sequence: (self: kitaev_meta, sequence: {{number}}) -> ();
	time_evolution: (self: kitaev_meta, steps: number, dt: number) -> ();
}

type kitaev_meta = typeof(setmetatable({} :: {
	lattice: {};
	chain_size: number;
	delta: number;
	mu: number;
	t: number;
	time_step: number;
}, {} :: kitaev_obj))

local kitaev: kitaev_obj = {} :: kitaev_obj
kitaev.__index = kitaev

-- chain_size: kitaevs chain size
-- delta: pairing param (controls strength of superconducting pairing)
-- mu: chemical potential
-- t: hopping param
-- time_step: simulation time step

function kitaev.init_chain(chain_size: number?, delta: number?, mu: number?, t: number?, time_step: number?)
	local meta = {}
	meta.lattice = {}
	meta.chain_size = chain_size or 10
	meta.delta = delta or 0.5
	meta.mu = mu or 0.5
	meta.t = t or 1
	meta.time_step = time_step or 0.1

	setmetatable(meta, kitaev)

	for i = 1, meta.chain_size do
		meta.lattice[i] = {majorana1 = 0, majorana2 = 0}
	end

	return meta
end

function kitaev:apply_majorana(i: number, mode: number)
	if mode == 1 then
		self.lattice[i].majorana1 = (self.lattice[i].majorana1 + 1) % 2
	else
		self.lattice[i].majorana2 = (self.lattice[i].majorana2 + 1) % 2
	end
end

function kitaev:hamiltonian(): number
	local energy = 0
	for i = 1, self.chain_size - 1 do
		-- hopping term (t * c_i† * c_{i+1} + h.c.)
		energy = energy - self.t * (self.lattice[i].majorana2 * self.lattice[i+1].majorana1)

		-- pairing term (Δ * c_i * c_{i+1} + h.c.)
		energy = energy + self.delta * (self.lattice[i].majorana1 * self.lattice[i+1].majorana2)
	end
	-- chemical potential term (μ * c_i† * c_i)
	for i = 1, self.chain_size do
		energy = energy - self.mu * (self.lattice[i].majorana1 * self.lattice[i].majorana2)
	end
	return energy
end

function kitaev:detect_majorana()
	local left_majorana = self.lattice[1].majorana1
	local right_majorana = self.lattice[self.chain_size].majorana2

	if left_majorana == 1 then
		print("unpaired majorana mode detected at the left end of the chain")
	end
	if right_majorana == 1 then
		print("unpaired majorana mode detected at the right end of the chain")
	end
end

function kitaev:evolve_chain(steps)
	for step = 1, steps do
		local energy = self:hamiltonian()
		print("energy at step " .. step .. ": " .. energy)
		self:detect_majorana()

		for i = 1, self.chain_size - 1 do
			self.lattice[i].majorana1 = (self.lattice[i].majorana1 + self.time_step * self.delta) % 2
			self.lattice[i].majorana2 = (self.lattice[i].majorana2 + self.time_step * self.t) % 2
		end
	end
end

function kitaev:braid_majorana(i: number, j: number, mode_i: number, mode_j: number)
	print("braiding majorana modes at positions " .. i .. " and " .. j)
	local temp
	if mode_i == 1 then
		temp = self.lattice[i].majorana1
		self.lattice[i].majorana1 = self.lattice[j].majorana1
		self.lattice[j].majorana1 = temp
	else
		temp = self.lattice[i].majorana2
		self.lattice[i].majorana2 = self.lattice[j].majorana2
		self.lattice[j].majorana2 = temp
	end
end

function kitaev:error_detect()
	local errors = 0
	for i = 1, self.chain_size do
		if self.lattice[i].majorana1 ~= 0 and self.lattice[i].majorana2 ~= 0 then
			print("error detected at site " .. i)
			errors = errors + 1
		end
	end
	return errors
end

function kitaev:error_correct()
	for i = 1, self.chain_size do
		self.lattice[i].majorana1 = 0
		self.lattice[i].majorana2 = 0
	end
end

function kitaev:measure_majorana(i: number, mode: number)
	if mode == 1 then
		return self.lattice[i].majorana1
	else
		return self.lattice[i].majorana2
	end
end

function kitaev:full_simulation(steps: number)
	for step = 1, steps do
		print("\nStep " .. step)
		self:evolve_chain(1)

		local errors = self:error_detect()
		if errors > 0 then
			self:error_correct()
		end

		if step % 5 == 0 then
			self:braid_majorana(1, self.chain_size, 1, 2)
		end
	end
end

function kitaev:perform_braiding_sequence(sequence: {{number}})
	for _, pair in ipairs(sequence) do
		local i, j, mode_i, mode_j = table.unpack(pair)
		self:braid_majorana(i, j, mode_i, mode_j)
	end
end

function kitaev:time_evolution(steps: number, dt: number)
	for step = 1, steps do
		local energy = self:hamiltonian()
		print("energy at step " .. step .. ": " .. energy)
		self:detect_majorana()

		for i = 1, self.chain_size - 1 do
			self.lattice[i].majorana1 = (self.lattice[i].majorana1 + dt * self.delta) % 2
			self.lattice[i].majorana2 = (self.lattice[i].majorana2 + dt * self.t) % 2
		end
	end
end

return kitaev


--!native
--!optimize 2
--!strict

type complex = {r: number, i: number}
type matrix = {{complex}}
type matrix_sz = {r: number, c: number}

local cpx = require("cpx")

local mtx = {}

function mtx.new(size: matrix_sz, ...: complex)
	local elements = {...}
	if #elements ~= size.r * size.c then
		error("number of elements doesnt match matrix size")
	end
	local matrix = {}
	for i = 1, size.r do
		matrix[i] = {}
		for j = 1, size.c do
			matrix[i][j] = elements[(i-1) * size.c + j]
		end
	end
	return matrix
end

function mtx.create_empty(size: matrix_sz)
	local matrix = {}
	for i = 1, size.r do
		matrix[i] = {}
		for j = 1, size.c do
			matrix[i][j] = cpx.new(0, 0)
		end
	end
	return matrix
end

function mtx.kronecker(a: matrix, b: matrix): matrix
	local ra, ca = #a, #a[1]
	local rb, cb = #b, #b[1]
	local result: matrix = {}

	for i = 1, ra * rb do
		result[i] = {}
		for j = 1, ca * cb do
			result[i][j] = cpx.new(0, 0)
		end
	end

	for i = 1, ra do
		for j = 1, ca do
			for m = 1, rb do
				for n = 1, cb do
					result[(i-1)*rb + m][(j-1)*cb + n] = cpx.mul(a[i][j], b[m][n])
				end
			end
		end
	end

	return result
end

function mtx.add(m1: matrix, m2: matrix): matrix
	local size1, size2 = #m1, #m2
	local cols1, cols2 = #m1[1], #m2[1]
	if size1 ~= size2 or cols1 ~= cols2 then
		error("matrix dimensions must match for addition")
	end
	local result = {}
	for i = 1, size1 do
		result[i] = {}
		for j = 1, cols1 do
			result[i][j] = cpx.add(m1[i][j], m2[i][j])
		end
	end
	return result
end

function mtx.sub(m1: matrix, m2: matrix): matrix
	local size1, size2 = #m1, #m2
	local cols1, cols2 = #m1[1], #m2[1]
	if size1 ~= size2 or cols1 ~= cols2 then
		error("matrix dimensions must match for subtraction")
	end
	local result = {}
	for i = 1, size1 do
		result[i] = {}
		for j = 1, cols1 do
			result[i][j] = cpx.sub(m1[i][j], m2[i][j])
		end
	end
	return result
end

function mtx.mul(m1: matrix, m2: matrix): matrix
	local rows1, cols1 = #m1, #m1[1]
	local rows2, cols2 = #m2, #m2[1]
	if cols1 ~= rows2 then
		error("number of columns of the first matrix must equal number of rows of the second matrix")
	end
	local result = {}
	for i = 1, rows1 do
		result[i] = {}
		for j = 1, cols2 do
			result[i][j] = cpx.new(0, 0)
			for k = 1, cols1 do
				result[i][j] = cpx.add(result[i][j], cpx.mul(m1[i][k], m2[k][j]))
			end
		end
	end
	return result
end

function mtx.mul_mv(matrix: matrix, vector: {complex}): {complex}
	local result = {}
	for i = 1, #matrix do
		result[i] = cpx.new(0, 0)
		for j = 1, #vector do
			result[i] = cpx.add(result[i], cpx.mul(matrix[i][j], vector[j]))
		end
	end
	return result
end

function mtx.scalar_mul(scalar: complex, m: matrix): matrix
	local rows, cols = #m, #m[1]
	local result = {}
	for i = 1, rows do
		result[i] = {}
		for j = 1, cols do
			result[i][j] = cpx.mul(scalar, m[i][j])
		end
	end
	return result
end

function mtx.identity(size: number): matrix
	local result = {}
	for i = 1, size do
		result[i] = {}
		for j = 1, size do
			if i == j then
				result[i][j] = cpx.new(1, 0)
			else
				result[i][j] = cpx.new(0, 0)
			end
		end
	end
	return result
end

function mtx.conj(m: matrix): matrix
	local result = {}
	for i = 1, #m do
		result[i] = {}
		for j = 1, #m[i] do
			local val = m[i][j]
			result[i][j] = cpx.new(val.r, -val.i)
		end
	end
	return result
end

function mtx.expand(m: matrix, new_sz: matrix_sz): matrix
	local result = {}

	for i = 1, new_sz.r do
		result[i] = {}
		for j = 1, new_sz.c do
			if i <= #m and j <= #m[1] then
				result[i][j] = m[i][j]
			else
				result[i][j] = cpx.new(0, 0)
			end
		end
	end

	return result
end

function mtx.exp(matrix: matrix, exponent: number): matrix
	local size = #matrix
	local result = mtx.identity(size)

	for _ = 1, exponent do
		result = mtx.mul(result, matrix)
	end

	return result
end

-- new funcs

function mtx.is_unitary(m: matrix): boolean
	local rows, cols = #m, #m[1]
	if rows ~= cols then
		return false
	end

	local m_conj_transpose = mtx.conj(mtx.transpose(m))
	local product = mtx.mul(m, m_conj_transpose)
	local identity = mtx.identity(rows)

	local epsilon = 1e-10
	for i = 1, rows do
		for j = 1, cols do
			if math.abs(product[i][j].r - identity[i][j].r) > epsilon or
				math.abs(product[i][j].i - identity[i][j].i) > epsilon then
				return false
			end
		end
	end

	return true
end

function mtx.transpose(m: matrix): matrix
	local rows, cols = #m, #m[1]
	local result = {}

	for j = 1, cols do
		result[j] = {}
		for i = 1, rows do
			result[j][i] = m[i][j]
		end
	end

	return result
end

function mtx.outer_product(v1: {complex}, v2: {complex}): matrix
	local result = {}
	for i = 1, #v1 do
		result[i] = {}
		for j = 1, #v2 do
			result[i][j] = cpx.mul(v1[i], cpx.conj(v2[j]))
		end
	end
	return result
end

return mtx


--!native
--!optimize 2
--!strict

type complex = {r: number, i: number}
type qubit = {complex}
type qubits = {qubit}
type gate = {{complex}}
type matrix = {{complex}}
export type noise_obj = {
	__index: noise_obj;
	new: ( probabilities: {
		bitflip: number?;
		phaseflip: number?;
		depolarizing: number?;
	}) -> noise_meta;
	apply_bitflip: (self: noise_meta, qubits: qubits) -> qubits;
	apply_phaseflip: (self: noise_meta, qubits: qubits) -> qubits;
	apply_depolarizing: (self: noise_meta, qubits: qubits) -> qubits;
	apply_all: (self: noise_meta, qubits: qubits) -> qubits;
}
export type noise_meta = typeof(setmetatable({} :: {
	bitflip: number;
	phaseflip: number;
	depolarizing: number;
}, {} :: noise_obj))

local gate = require("gate")
local preset = require("preset")

local noise: noise_obj = {} :: noise_obj
noise.__index = noise

function noise.new(probabilities: {
	bitflip: number?;
	phaseflip: number?;
	depolarizing: number?;
})
	local meta = {}
	meta.bitflip = probabilities.bitflip or 0.02
	meta.phaseflip = probabilities.phaseflip or 0.02
	meta.depolarizing = probabilities.depolarizing or 0.02
	setmetatable(meta, noise)

	return meta
end

function noise:apply_bitflip(qubits: qubits): qubits
	local noisy_qubits = {}
	for i, qubit in ipairs(qubits) do
		if math.random() < self.bitflip then
			noisy_qubits[i] = gate.apply(qubit, preset.paulix)
		else
			noisy_qubits[i] = qubit
		end
	end
	return noisy_qubits
end

function noise:apply_phaseflip(qubits: qubits): qubits
	local noisy_qubits = {}
	for i, qubit in ipairs(qubits) do
		if math.random() < self.phaseflip then
			noisy_qubits[i] = gate.apply(qubit, preset.pauliz)
		else
			noisy_qubits[i] = qubit
		end
	end
	return noisy_qubits
end

function noise:apply_depolarizing(qubits: qubits): qubits
	local noisy_qubits = {}
	for i, qubit in ipairs(qubits) do
		if math.random() < self.depolarizing then
			local noise_type = math.random()
			if noise_type < 1/3 then
				noisy_qubits[i] = gate.apply(qubit, preset.paulix)
			elseif noise_type < 2/3 then
				noisy_qubits[i] = gate.apply(qubit, preset.pauliy)
			else
				noisy_qubits[i] = gate.apply(qubit, preset.pauliz)
			end
		else
			noisy_qubits[i] = qubit
		end
	end
	return noisy_qubits
end

function noise:apply_all(qubits: qubits): qubits
	qubits = self:apply_bitflip(qubits)
	qubits = self:apply_phaseflip(qubits)
	qubits = self:apply_depolarizing(qubits)
	return qubits
end

return noise


--!native
--!optimize 2
--!strict

local cpx = require("cpx")
local gate = require("gate")

return {
	identity = gate.new(
		cpx.new(1, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(1, 0)
	);
	paulix = gate.new(
		cpx.new(0, 0), cpx.new(1, 0),
		cpx.new(1, 0), cpx.new(0, 0)
	);
	pauliy = gate.new(
		cpx.new(0, 0), cpx.new(0, -1),
		cpx.new(0, 1), cpx.new(0, 0)
	);
	pauliz = gate.new(
		cpx.new(1, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(-1, 0)
	);
	hadamard = gate.new(
		cpx.new(1/math.sqrt(2), 0), cpx.new(1/math.sqrt(2), 0),
		cpx.new(1/math.sqrt(2), 0), cpx.new(-1/math.sqrt(2), 0)
	);
	phase = gate.new(
		cpx.new(1, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(0, 1)
	);
	pi8 = gate.new(
		cpx.new(1, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(1/math.sqrt(2), 1/math.sqrt(2))
		-- you can represent e^(i*π/4) by using the eulers formula:
		-- e^(i*θ) = cos(θ) + i * sin(θ)
		-- in this case, θ = π/4
		-- e^(i*π/4) = cos(π/4) + i * sin(π/4)
		--           = 1/√(2) + 1/√(2) * i
		-- -> real part is 1/√(2), imaginary part is 1/√(2)
		-- (<REAL>+<IMG>i)
	);
	cnot = gate.new_ex(4,
		cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0),
		cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0)
	);
	swap = gate.new_ex(4,
		cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0)
	);
	ccnot = gate.new_ex(8,
		cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0), cpx.new(0, 0),
		cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0),
		cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(0, 0), cpx.new(1, 0), cpx.new(0, 0)
	);
}