import random
import gradio as gr
from ngram_score import ngram_score


def get_keydict(key_dict, current_key, original_alphabet):
    for i in range(len(current_key)):
        key_dict[current_key[i]] = original_alphabet[i]  # 这是将当前的密文进行映射
    return key_dict


def exchange(mydict, message):
    message = list(message)
    cnt = 0
    for i in message:
        if i in mydict:
            message[cnt] = mydict[i]
        cnt = cnt + 1
    return "".join(message)


def solve_cipher(ciphertext):
    S_new = ciphertext.replace(" ", "")
    S_new = S_new.replace(",", "")
    S_new = S_new.replace("-", "")
    S_new = S_new.replace(".", "")
    # 参数初始化
    m_message = S_new.upper()  # 这是全部改变为大写的密文
    current_key = list('ABCDEFGHIJKLMNOPQRSTUVWXYZ')  # 这是当前的密文
    original_alphabet = list('ABCDEFGHIJKLMNOPQRSTUVWXYZ')
    key_dict = dict()  # 这是一个字典，用来将字母映射到上面字母表上去
    fitness = ngram_score('english_quadgrams.txt')

    last_score = -2 ** 31
    current_max_score = -2 ** 31
    generation = 0  # generation就是迭代的数量

    best_plaintext = ""

    while generation < 10:
        # 上面是迭代最高次数，一般10以内就能出结果
        generation = generation + 1

        # 随机改变顺序
        random.shuffle(current_key)
        key_dict = get_keydict(key_dict, current_key, original_alphabet)  # 获得明密文映射
        last_score = fitness.score(exchange(key_dict, m_message))  # 计算适应度

        count = 0
        while count < 1000:
            a = random.randint(0, 25)
            b = random.randint(0, 25)
            # 随机交换并进行比较
            child_current_key = current_key[:]
            child_current_key[a], child_current_key[b] = child_current_key[b], child_current_key[a]

            child_key_dict = dict()
            child_key_dict = get_keydict(child_key_dict, child_current_key, original_alphabet)
            score = fitness.score(exchange(child_key_dict, m_message))
            # 说明新的key_dict更高效
            if score > last_score:
                last_score = score
                current_key = child_current_key
                count = 0
            count = count + 1

        # 输出结果
        if last_score > current_max_score:
            current_max_score = last_score
            maxkey = current_key
            key_dict = get_keydict(key_dict, current_key, original_alphabet)
            best_plaintext = exchange(key_dict, ciphertext.upper()).lower()

    return best_plaintext

iface = gr.Interface(
    fn=solve_cipher,
    inputs=gr.Textbox(lines=10, placeholder="Enter ciphertext here..."),
    outputs="text",
    title="Substitution Cipher Solver",
    description="Enter the ciphertext and see the decrypted plaintext."
)

if __name__ == "__main__":
    iface.launch()