import numpy as np
import gmpy2


def _minor(matrix: np.array, i, j):
    return matrix[
        np.array(list(range(i))+list(range(i+1, matrix.shape[0])))[:, np.newaxis],
        np.array(list(range(j))+list(range(j+1, matrix.shape[1])))
    ]


def _adjoint(matrix: np.array):
    minors = np.zeros_like(matrix)
    for i in range(matrix.shape[0]):
        for j in range(matrix.shape[1]):
            minors[i][j] = round(np.linalg.det(_minor(matrix, i, j))) * (-1) ** (i + j)

    return minors.T


def _matrix_modulo_invert(matrix: np.array, modulo):
    inverse_det = int(gmpy2.invert(int(round(np.linalg.det(matrix))), modulo))
    return inverse_det * _adjoint(matrix) % modulo


class Hill:
    def __init__(self, key: str):
        if len(key) != 4:
            raise ValueError("Wrong key length")

        key = key.upper()
        self.encrypt_key = np.array([ord(x) - ord("A") for x in key])
        self.encrypt_key = self.encrypt_key.reshape((2, 2))
        self.decrypt_key = _matrix_modulo_invert(self.encrypt_key, 26)

    def _encrypt_with_key(self, message: str, key: np.array) -> str:
        message = message.upper()
        ciphertext = []
        for i in range(0, len(message), key.shape[0]):
            block = message[i:i + key.shape[0]]
            block_vector = np.vstack(
                np.array([ord(x) - ord("A") for x in block])
            )

            ciphertext_block_vector = np.matmul(
                key, block_vector
            ) % 26

            ciphertext_block = [
                chr(x + ord("A")) for x in ciphertext_block_vector
            ]
            ciphertext.extend(ciphertext_block)

        return "".join(ciphertext)

    def encrypt(self, message: str) -> str:
        return self._encrypt_with_key(message, self.encrypt_key)

    def decrypt(self, message: str) -> str:
        return self._encrypt_with_key(message, self.decrypt_key)


def main():
    hill = Hill("GHJL")
    encrypt = hill.encrypt("ZOHF")
    print(encrypt)
    print(hill.decrypt(encrypt))


if __name__ == "__main__":
    main()