#2. Paillier cryptosystem

import random
from math import gcd, lcm
from sympy import isprime, public

ALPHABET = " ABCDEFGHIJKLMNOPQRSTUVWXYZ"
ALPHABET_MAP = {char: idx for idx, char in enumerate(ALPHABET)}
REVERSE_ALPHABET_MAP = {idx: char for idx, char in enumerate(ALPHABET)}

def text_to_numeric(plaintext):
    return [ALPHABET_MAP[char] for char in plaintext if char in ALPHABET_MAP]

def numeric_to_text(numeric_plaintext):
    return "".join([REVERSE_ALPHABET_MAP[idx] for idx in numeric_plaintext])


def generate_paillier_keypair(bit_length=512):
    p = generate_large_prime(bit_length)
    q = generate_large_prime(bit_length)
    n = p * q
    n_squared = n * n

    lambda_n = lcm(p - 1, q - 1)

    mmi = 0

    while True:
        g = random.randint(1, n_squared-1)
        if g // n == g / n:
            continue
        mmi = pow(L(pow(g, lambda_n, n_squared), n),-1,n)
        break

    public_key = (n, g)
    private_key = (lambda_n, mmi)
    return public_key, private_key

def L(x, n):
    return (x - 1) // n

def generate_large_prime(bit_length):
    while True:
        candidate = random.getrandbits(bit_length)
        if isprime(candidate):
            return candidate
    
def paillier_encrypt(plaintext, public_key):
    n, g = public_key
    r = random.randint(1, n-1)
    while gcd(r, n) != 1:
        # Extremly unlikely to happen
        r = random.randint(1, n-1)
    ciphertext = []
    for int_letter in text_to_numeric(plaintext):
        c = (pow(g, int_letter, n*n) * pow(r, n, n*n)) % (n*n)
        ciphertext.append(c)
    return ciphertext


def paillier_decrypt(ciphertext, private_key, public_key):
    lambda_n, mmi = private_key
    n, g = public_key
    plaintext = []
    for c in ciphertext:
        l = L(pow(c, lambda_n, n*n), n)
        m = (l * mmi) % n
        plaintext.append(m)
    return numeric_to_text(plaintext)
    

if __name__ == "__main__":
    public_key, private_key = generate_paillier_keypair()
    plaintext = "HELLO WORLD"
    print("Plaintext:", plaintext)
    ciphertext = paillier_encrypt(plaintext, public_key)
    print("Ciphertext:", ciphertext)
    decrypted = paillier_decrypt(ciphertext, private_key, public_key)
    print("Decrypted:", decrypted)