needle = require("./needle")

function table.shallow_copy(t)
  local t2 = {}
  for k,v in pairs(t) do
    t2[k] = v
  end
  return t2
end

function multiple_path_backtrack_trace(matrix, seq1, seq2)
    local stack = {}
    local m=string.len(seq1)
    local n=string.len(seq2)
    local i=m
    local j=n
    table.insert(stack, 1, {i, j, {}})
    local trace = {}
    while #stack ~= 0 do
        local state = table.remove(stack, 1)
        table.insert(trace, #trace+1, state)
        local i=state[1]
        local j=state[2]
        local alignment = state[3]
        if (i > 0 and j > 0) then
            local nt1 = string.sub(seq1, i-1, i-1)
            local nt2 = string.sub(seq2, j-1, j-1)
            if (matrix[i][j] == matrix[i-1][j-1] + needle.sub(nt1, nt2)) then
                local new_alignment = table.shallow_copy(alignment)
                table.insert(new_alignment, 1, {nt1, nt2})
                table.insert(stack, 1, {i - 1, j - 1, new_alignment})
            end
            if (matrix[i][j] == matrix[i-1][j] + needle.gap_penalty) then
                local new_alignment = table.shallow_copy(alignment)
                table.insert(new_alignment, 1, {nt1, '-'})
                table.insert(stack, 1, {i-1, j, new_alignment})
            end
            if (matrix[i][j] == matrix[i][j-1] + needle.gap_penalty) then
                local new_alignment = table.shallow_copy(alignment)
                table.insert(new_alignment, 1, {'-', nt2})
                table.insert(stack, 1, {i, j-1, new_alignment})
            end
        end
        if (i > 0) then
            local nt1 = string.sub(seq1, i-1, i-1)
            local new_alignment = table.shallow_copy(alignment)
            table.insert(new_alignment, 1, {nt1, '-'})
            table.insert(stack, 1, {i-1, j, new_alignment})
        end
        if (j > 0) then
            local nt2 = string.sub(seq2, j-1, j-1)
            local new_alignment = table.shallow_copy(alignment)
            table.insert(new_alignment, 1, {'-', nt2})
            table.insert(stack, 1, {i, j-1, new_alignment})
        end
    end
    return trace
end

function repr_alignment(alignment)
    local repr = [[\begin{pmatrix}]]
    for i, vector in ipairs(alignment) do
        repr = repr .. vector[1]
        if i < #alignment then
            repr = repr .. " & "
        end
    end
    repr = repr .. [[\\]] .. " \n"
    for i, vector in ipairs(alignment) do
        repr = repr .. vector[2]
        if i < #alignment then
            repr = repr .. " & "
        end
    end
    repr = repr .. [[\end{pmatrix}]]
    return repr
end

function trace_repr(trace)
    local repr = ""
--     for stack_index, stack in ipairs(trace) do
--         repr = repr .. "iteration " .. stack_index .. " :" .. [[\\]]
        repr = repr .. [[\begin{tabular}{|c|} \\ \hline ]]
        for call_index, call in ipairs(trace) do
            local i = call[1]
            local j = call[2]
            local aligment = call[3]
            repr = repr .. [[ $\langle ]] .. i ..", " .. j .. ", " .. repr_alignment(alignment).. [[\rangle$ ]]
            repr = repr .. [[\\ \hline]]
        end
        repr = repr .. [[\end{tabular}]]
--     end
    return repr
end

function main()
    local seq1 = "ATCTGAT"
    local seq2 = "TGCATA"
    local matrix = needle.needle_matrix(seq1, seq2)
    local trace = multiple_path_backtrack_trace(matrix, seq1, seq2)
    print(#trace)
    print(trace_repr(trace))
end

main()