import sys
import numpy as np

def cipher_encryption(plain, key):

    # Jeżeli długość wiadomości jest nieparzysta, dodaję zero na końcu
    len_chk = 0
    if len(plain) % 2 != 0:
        plain += "0"
        len_chk = 1
        
    # Zamiana wiadomości na macierz
    row = 2
    col = int(len(plain)/2)
    msg2d = np.zeros((row, col), dtype=int)
    
    itr1 = 0
    itr2 = 0
    for i in range(len(plain)):
        if i%2 == 0:
            msg2d[0][itr1]= int(ord(plain[i]) - 65)
            itr1 += 1
        else:
            msg2d[1][itr2] = int(ord(plain[i]) - 65)
            itr2 += 1
    
    
    # Zamiana klucza na macierz 2x2
    key2d = np.zeros((2,2), dtype=int)
    itr3 = 0
    for i in range(2):
        for j in range(2):
            key2d[i][j] = ord(key[itr3]) - 65
            itr3 += 1
            
    print(key2d)
            
    # Sprawdzanie poprawności klucza
    # Obliczanie wyznacznika
    deter = key2d[0][0] * key2d[1][1] - key2d[0][1] * key2d[1][0]
    deter = deter % 26
    
    # Szukanie odwrotności macierzy
    for i in range(26):
        temp_inv = deter * i
        if temp_inv % 26 == 1:
            mul_inv = i
            break
        else:
            continue
    
    if mul_inv == -1:
        print("Klucz jest niepoprawny")
        sys.exit()
    
    encryp_text = ""
    itr_count = int(len(plain)/2)
    if len_chk == 0:
        for i in range(itr_count):
            temp1 = msg2d[0][i] * key2d[0][0] + msg2d[1][i] * key2d[0][1]
            encryp_text += chr((temp1 % 26) + 65)
            temp2 = msg2d[0][i] * key2d[1][0] + msg2d[1][i] * key2d[1][1]
            encryp_text += chr((temp2 % 26) + 65)
    else:
        for i in range(itr_count-1):
            temp1 = msg2d[0][i] * key2d[0][0] + msg2d[1][i] * key2d[0][1]
            encryp_text += chr((temp1 % 26) + 65)
            temp2 = msg2d[0][i] * key2d[1][0] + msg2d[1][i] * key2d[1][1]
            encryp_text += chr((temp2 % 26) + 65)
    
    print("Zaszyfrowana wiadomość: {}".format(encryp_text))
    return encryp_text
    
def cipher_decryption(cipher, key):

    # Jeżeli długość wiadomości jest nieparzysta, dodaję zero na końcu
    len_chk = 0
    if len(cipher) % 2 != 0:
        cipher += "0"
        len_chk = 1
        
    # Zamiana wiadomości na macierz
    row = 2
    col = int(len(cipher)/2)
    msg2d = np.zeros((row, col), dtype=int)
    
    itr1 = 0
    itr2 = 0
    for i in range(len(cipher)):
        if i%2 == 0:
            msg2d[0][itr1]= int(ord(cipher[i]) - 65)
            itr1 += 1
        else:
            msg2d[1][itr2] = int(ord(cipher[i]) - 65)
            itr2 += 1

    
    # Zamiana klucza na macierz 2x2
    key2d = np.zeros((2,2), dtype=int)
    itr3 = 0
    for i in range(2):
        for j in range(2):
            key2d[i][j] = ord(key[itr3]) - 65
            itr3 += 1

    # Obliczanie wyznacznika
    deter = key2d[0][0] * key2d[1][1] - key2d[0][1] * key2d[1][0]
    deter = deter % 26

    # Szukanie odwrotności macierzy
    for i in range(26):
        temp_inv = deter * i
        if temp_inv % 26 == 1:
            mul_inv = i
            break
        else:
            continue
    
    # Macierz dołączona
    # Zamiana wartości
    key2d[0][0], key2d[1][1] = key2d[1][1], key2d[0][0]
    
    # Zamiana znaków
    key2d[0][1] *= -1
    key2d[1][0] *= -1
    
    key2d[0][1] = key2d[0][1] % 26
    key2d[1][0] = key2d[1][0] % 26
    
    # Mnożenie odwrotności przez macierz dołączoną
    for i in range(2):
        for j in range(2):
            key2d[i][j] *= mul_inv
    
    # Modulo
    for i in range(2):
        for j in range(2):
            key2d[i][j] = key2d[i][j] % 26
            
    # Odszyfrowywanie
    decryp_text = ""
    itr_count = int(len(cipher)/2)
    if len_chk == 0:
        for i in range(itr_count):
            temp1 = msg2d[0][i] * key2d[0][0] + msg2d[1][i] * key2d[0][1]
            decryp_text += chr((temp1 % 26) + 65)
            temp2 = msg2d[0][i] * key2d[1][0] + msg2d[1][i] * key2d[1][1]
            decryp_text += chr((temp2 % 26) + 65)
        # for
    else:
        for i in range(itr_count-1):
            temp1 = msg2d[0][i] * key2d[0][0] + msg2d[1][i] * key2d[0][1]
            decryp_text += chr((temp1 % 26) + 65)
            temp2 = msg2d[0][i] * key2d[1][0] + msg2d[1][i] * key2d[1][1]
            decryp_text += chr((temp2 % 26) + 65)
        # for
    # if else
    
    print("Odszyfrowana wiadomość: {}".format(decryp_text))
    
plaintext = "Tajna wiadomosc"
plaintext = plaintext.upper().replace(" ","")
key = "hill"
key = key.upper().replace(" ","")
ciphertext = cipher_encryption(plaintext, key)
cipher_decryption(ciphertext, key)
# Zaszyfrowana wiadomość: DBLIUIEKDFOAMM
# Odszyfrowana wiadomość: TAJNAWIADOMOSC
