# Шифр Уитстона
from collections import OrderedDict
from textwrap import wrap


def input_keys():
    keys = input("Введите ключ: ").lower().replace("q", "").split(" ")
    if len(keys) == 1 or len(keys) == 2:
        return keys
    else:
        print("Ошибка ввода.")
        input_keys()


def input_text():
    text = wrap(input("Введите текст: ").replace(" ", ""), 2)
    if len(text) == 0:
        print("Ошибка ввода.")
        input_text()
    if len(text[-1]) == 1:
        text[-1] = "{}x".format(text[-1])
    return text


def print_matrix(matrix):
    for i in matrix:
        for j in i:
            print(j, end="\t")
        print()


def gen_key_matrix(key):
    abc = "abcdefghijklmnoprstuvwxyz"
    abc = "{}{}".format(key, abc)
    abc = list(OrderedDict.fromkeys("".join(abc)))
    key_matrix = [[abc[5 * i + j] for j in range(5)] for i in range(5)]
    return key_matrix


def find_a_b(matrix1, matrix2, bigram):
    for i in matrix1:
        if bigram[0] in i:
            row1 = i
    for j in matrix2:
        if bigram[1] in j:
            row2 = j
    a_j = row1.index(bigram[0])
    b_j = row2.index(bigram[1])
    a_i = matrix1.index(row1)
    b_i = matrix2.index(row2)
    return [a_i, a_j, b_i, b_j]


def main():
    keys = input_keys()
    if len(keys) == 1:
        key1, key2 = keys[0], ""
    else:
        key1, key2 = keys
    bigrams = input_text()
    print(bigrams)
    matrix1 = gen_key_matrix(key1)
    matrix2 = gen_key_matrix(key2)
    print("===== Ключ 1 =====")
    print_matrix(matrix1)
    print("===== Ключ 2 =====")
    print_matrix(matrix2)
    result = []
    for b in bigrams:
        a_i, a_j, b_i, b_j = find_a_b(matrix1, matrix2, b)
        if a_j == b_j:
            new_a_i, new_a_j, new_b_i, new_b_j = b_i, b_j, a_i, a_j
        else:
            new_a_i, new_a_j, new_b_i, new_b_j = a_i, b_j, b_i, a_j
        result.append("{}{}".format(matrix1[new_a_i][new_a_j], matrix2[new_b_i][new_b_j]))
    print("{} - Исходные данные\n{} - Результат\n".format(" ".join(bigrams), " ".join(result)))


if __name__ == "__main__":
    main()