#!/usr/bin/python
# _*_ coding:utf-8 _*_

#alphabeta
letter_list = 'ABCDEFGHIKLMNOPQRSTUVWXYZ'

#key_list
ciphertable = ['','','','','']

#get the decrypt or encrypt mode
def get_mode():
    print("Do you want to Decrypt or Encrypt")
    mode = input().lower()
    if mode in "decrypt d encrypt e".split():
        return mode
    else:
        print("Either decrypt or d or encrypt e")


#remove duplicate word
def Remove_dulicates(key):
    key = key.upper()
    new_key = ''

    for ch in key:
        if ch == "J":
            ch = 'I'
        if ch in new_key:
            continue
        else:
            new_key += ch
    return new_key


#build_Key_List
def creatPasswordList(key):
    key = Remove_dulicates(key)  # remove duplicate word
    key = key.replace(' ', '')  # remove space of key

    for ch in letter_list:
        if ch not in key:
            key += ch

    j = 0
    for i in range(len(key)):
        ciphertable[j] += key[i]
        if (i+1)%5 == 0:
            j+=1

# get the position of plaintext in ciphertable
def Get_MatrixIndex(ch):
    for i in range(len(ciphertable)):
        for j in range(len(ciphertable)):
            if ch == ciphertable[i][j]:
                return i, j  # i is column and j is row

#Encrypt
def Encrypt(plaintext,ciphertable):
    ciphertext = ''
    plaintext = plaintext.replace(' ','')

    if (len(plaintext)% 2) != 0: #if the length of the text is odd we need to add an "Z" to the end of the sentence
        plaintext += 'Z'


    i = 0
    while i < len(plaintext):
        if plaintext[i].isalpha():
            j = i+1
            while j<len(plaintext):
                if plaintext[j].isalpha():
                    if plaintext[i] == 'J':
                        x = Get_MatrixIndex('I')
                    else:
                        x = Get_MatrixIndex( plaintext[i] . upper() )
                    if plaintext[j].upper() == 'J':
                        y = Get_MatrixIndex ('I')
                    else:
                        y = Get_MatrixIndex(plaintext[j].upper())

                    if x[0] == y[0]:
                        ciphertext += ciphertable[x[0]][(x[1]+1)%5]+ciphertable[y[0]][(y[1]+1)%5]
                    elif x[1] == y[1]:
                        ciphertext += ciphertable[(x[0]+1)%5][x[1]]+ciphertable[(y[0]+1)%5][y[1]]
                    else:
                        ciphertext += ciphertable[x[0]][y[1]]+ciphertable[y[0]][x[1]]
                    break
                j+= 1
            i=j+1
            continue
        else:
            ciphertext += plaintext[i]
        i += 1

    return ciphertext


#Decrypt
def Decrypt(ciphertext,ciphertable):
    plaintext = ''

    if (len(ciphertext) % 2) != 0:  # if the length of the text is odd we need to add an "Z" to the end of the sentence
        ciphertext += 'Z'

    i = 0
    while i < len(ciphertext):
        if ciphertext[i].isalpha():
            j = i + 1
            while j < len(ciphertext):
                if ciphertext[j].isalpha ():
                    if ciphertext[i] == 'J':
                        x = Get_MatrixIndex ( 'I' )
                    else:
                        x = Get_MatrixIndex(ciphertext[i].upper())
                    if ciphertext[j].upper() == 'J':
                        y = Get_MatrixIndex('I')
                    else:
                        y = Get_MatrixIndex(ciphertext[j].upper () )

                    if x[0] == y[0]:
                        plaintext += ciphertable[x[0]][(x[1] - 1) % 5] + ciphertable[y[0]][(y[1] - 1) % 5]
                    elif x[1] == y[1]:
                        plaintext += ciphertable[(x[0] - 1) % 5][x[1]] + ciphertable[(y[0] - 1) % 5][y[1]]
                    else:
                        plaintext += ciphertable[x[0]][y[1]] + ciphertable[y[0]][x[1]]
                    break;
                j += 1
            i = j + 1
            continue
        else:
            plaintext += ciphertext[i]
        i += 1

    return plaintext



mode = get_mode()

print('Please enter the key：')
key = input()

creatPasswordList(key)

if mode in 'encrypt e'.split():
    print("Please enter your plaintext:")
    plaintext = input()
    print("Ciphertext is:",Encrypt(plaintext, ciphertable))
elif mode in "decrypt d".split():
    print("Please enter your ciphertext:")
    ciphertext = input()
    print("Plaintext is :",Decrypt(ciphertext, ciphertable))



