#!/usr/bin/env lua
local automaton = require "automaton"
local generate = require "generate"
local reading = require "reading"
local format = require "format"
local init = automaton(generate.trange(1, 25))


local function add_node (graph, id)
	local node = graph[id]
	if node then
		return node
	end
	node = {
		id = id,
		outcoming = {},
		incoming = {},
		paths = {},
	}
	graph[id] = node
	return node
end


local function add_edges (graph, state, kind)
	for dst, src in ipairs(state) do
		local a = add_node(graph, src)
		local b = add_node(graph, dst)
		a.outcoming[kind] = b
		b.incoming[kind] = a
	end
end


local graph = {}
add_edges(graph, init:left(), "left")
add_edges(graph, init:right(), "right")
add_edges(graph, init:pivot(), "centre")
add_edges(graph, init:up(), "up")
add_edges(graph, init:down(), "down")
local map = {
	centre = 0,
	up = 1,
	right = 2,
	down = 3,
	left = 4,
}


local function trigram (path)
	return reading.trigram_to_value(
		map[path[1]],
		map[path[2]],
		map[path[3]])
end


local function shallow_copy (obj)
	local new = {}
	for k, v in ipairs(obj) do
		new[k] = v
	end
	return new
end


local function generate_paths_from (start, length)
	local queue = {{node=start, length=0, path={}}}
	while #queue > 0 do
		local entry = table.remove(queue, 1)
		if entry.length == length then
			local paths = entry.node.paths[start.id] or {}
			table.insert(paths, entry.path)
			entry.node.paths[start.id] = paths
		end
		if entry.length < length then
			for kind, node in pairs(entry.node.incoming) do
				local path = shallow_copy(entry.path)
				table.insert(path, 1, kind)
				table.insert(queue, {node=node, length=entry.length + 1, path=path})
			end
		end
	end
end


local function error_no_path_between (start, dest)
	if not start.paths[dest.id] then
		io.stderr:write(string.format("cannot move from %d to %d\n", start.id, dest.id))
	end
end


generate_paths_from(graph[3], 3)
generate_paths_from(graph[19], 3)
for _, node in ipairs(graph) do
	for target, paths in pairs(node.paths) do
		for _, path in ipairs(paths) do
			-- print(string.format("%d,%d,%d,%s,%s,%s", node.id, target, trigram(path), path[1], path[2], path[3]))
		end
	end
	error_no_path_between(node, graph[3])
	error_no_path_between(node, graph[19])
end


local function stupid_invert (obj)
	local lookup = {}
	for k, v in pairs(obj) do
		lookup[v] = k
	end
	return lookup
end


local function encode (graph, state, message)
	state = automaton(state)
	local lookup = stupid_invert(state)
	local values = {}
	for i=1, #message do
		local node = graph[lookup[message:sub(i, i)]]
		if not node then
			error(message:sub(i, i))
		end
		local output = 3
		if (i - 1) % 3 == 1 then
			output = 19
		end
		local path = node.paths[output][math.random(1, #node.paths[output])]
		table.insert(values, trigram(path))
		state = state:apply(map[path[1]]):apply(map[path[2]]):apply(map[path[3]])
		lookup = stupid_invert(state)
	end
	return values
end


for i=1, 5 do
	print(format.csv(encode(graph, "abcdefghiwklm opqrstuvxyz", "we cam try this method to see what kimd of potemial ciphertext we will get out of it")))
	print(format.csv(encode(graph, "abcdefghiwklm opqrstuvxyz", "amd them try it with amother plaimtext iust to make sure this actually works as expected with lomger texts too")))
end