st

import random
from matplotlib import pyplot as plt
from PIL import Image


def str_to_bitlst(strr):
    """ Zwraca listę bitów w binarnym zapisie napisu"""
    res = []
    for i in strr:
        res.extend(list(map(int, bin(ord(i))[2:].zfill(8))))
    return res


def bitgen(lst):
    for i in range(len(lst) // 2):
        yield 2*lst[2 * i] + lst[2 * i + 1]


def encrypt(text, image_name):
    XOR = [0, 1, 2, 3]
    binary_text = str_to_bitlst(text)

    image = Image.open(image_name)
    image = image.convert('RGB')
    width, height = image.size
    image1 = image.copy()
    image2 = image.copy()

    generator = bitgen(binary_text)

    breaks = 0
    for i in range(height):
        for j in range(width):
            pixel_id = (i, j)
            image_pixel = list(image.getpixel(pixel_id))
            image1_pixel = image_pixel.copy()
            image2_pixel = image_pixel.copy()
            for k in range(3):
                try:
                    bits_to_encrypt = next(generator)
                    xor = random.choice(XOR)
                    tmp = image_pixel[k] ^ (image_pixel[k] & 3)
                    image1_pixel[k] = tmp ^ xor
                    image2_pixel[k] = tmp ^ bits_to_encrypt ^ xor
                except StopIteration:
                    breaks = 1
                if breaks:
                    break
            image1.putpixel(pixel_id, tuple(image1_pixel))
            image2.putpixel(pixel_id, tuple(image2_pixel))
            if breaks:
                break
        if breaks:
            break

    image1.save(image_name[:-4] + '1' + image_name[-4:])
    image2.save(image_name[:-4] + '2' + image_name[-4:])


def decrypt(image_name1, image_name2):
    image1 = Image.open(image_name1)
    image2 = Image.open(image_name2)
    width, height = image1.size
    decrypted_bits = []

    for i in range(height):
        for j in range(width):
            image1_pixel = image1.getpixel((i, j))
            image2_pixel = image2.getpixel((i, j))
            for k, l in zip(image1_pixel, image2_pixel):
                tmp = k ^ l
                decrypted_bits.append(bin(tmp)[2:].zfill(2))

    result = []
    for i in range((len(decrypted_bits) + 3) // 4):
        result.append(chr(int("".join(decrypted_bits[4 * i: 4 * (i + 1)]), 2)))

    return "".join(result).strip()


def show_images(*image_names):
    images = [Image.open(i) for i in image_names]
    fig, ax = plt.subplots(1, len(image_names))

    for idx, i in enumerate(images):
        ax[idx].imshow(i)
        ax[idx].axis('off')

    plt.show()


def main():
    image_name = input('Podaj nazwę obrazka w którym ma zostać ukryty tekst: ')
    text = input('Podaj napis, który ma zostać ukryty: ')
    encrypt(text, image_name)
    show_images(image_name, image_name[:-4] + '1' + image_name[-4:],
                image_name[:-4] + '2' + image_name[-4:])


main()