from PIL import Image
import requests
from io import BytesIO
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
import numpy as np
import os
import plotly.graph_objects as go
from plotly.subplots import make_subplots

url = 'https://upload.wikimedia.org/wikipedia/commons/thumb/8/8c/Glockenkelter_Stetten_i.R.jpg/330px-Glockenkelter_Stetten_i.R.jpg'
response = requests.get(url)
img = Image.open(BytesIO(response.content))

gray_img = img.convert('L')

#########################################

gray_arr = np.array(gray_img)
bytes_data = gray_arr.tobytes()
shape = gray_arr.shape
dtype = gray_arr.dtype

vigenere_key = bytearray([b * 2 for b in b"Wikipedia is great!"])
vigenere_data = b""
for i in range(len(bytes_data)):
    key_value = vigenere_key[i % len(vigenere_key)]
    text_value = bytes_data[i]
    value = (text_value + key_value) % 256
    vigenere_data += value.to_bytes(1, byteorder='big')

aes_key = os.urandom(32)
aes_cipher = Cipher(algorithms.AES(aes_key), modes.ECB())
aes_encryptor = aes_cipher.encryptor()
aes_padder = padding.PKCS7(128).padder()
aes_padded_data = aes_padder.update(bytes_data) + aes_padder.finalize()
aes_data = aes_encryptor.update(aes_padded_data) + aes_encryptor.finalize()

images = {}
images["Original"] = gray_arr
images["Caesar cipher"] = gray_arr + 60
images["Vigenère cipher"] = np.frombuffer(vigenere_data, dtype=dtype).reshape(shape)
images["AES EBC"] = np.frombuffer(aes_data[:np.prod(shape)], dtype=dtype).reshape(shape)

#########################################

img_height, img_width = list(images.values())[0].shape

horizontal_spacing_px = 10
vertical_spacing_px = 35
fig_width = img_width * 2 + horizontal_spacing_px + 50
fig_height = img_height * 2 + vertical_spacing_px + 95

fig = make_subplots(rows=2, cols=2, horizontal_spacing=horizontal_spacing_px / fig_width, vertical_spacing=vertical_spacing_px / fig_height, subplot_titles=list(images.keys()))

for i in range(2):
    for j in range(2):
        rgb_array = np.stack([list(images.values())[i*2+j]] * 3, axis=-1)
        fig.add_trace(
            go.Image(z=rgb_array),
            row=i + 1, col=j + 1
        )

fig.update_layout(
    margin = dict(l=25, r=25, t=70, b=25),
    height = fig_height,
    width = fig_width,
    title = "Multiple encryption methods applied on image",
    title_x = 0.5
)

for i in range(1, 5):
    fig.layout[f'xaxis{i}'].visible = False
    fig.layout[f'yaxis{i}'].visible = False
    fig.layout[f'yaxis{i}'].scaleanchor = f'x{i}'

fig.show()