import numpy as np
from PIL import Image
import os

# Path settings
image_path = r"G:\Users\Chipistil\Desktop\High-throughput-T2NP-Memristor-based-Dual-clock-edge-sampling-TRNG\AHU Black Swan.png"
key_path = r"G:\Users\Chipistil\Desktop\High-throughput-T2NP-Memristor-based-Dual-clock-edge-sampling-TRNG\key_AHU Black Swan_spike.txt"
output_encrypted_path = r"G:\Users\Chipistil\Desktop\High-throughput-T2NP-Memristor-based-Dual-clock-edge-sampling-TRNG\True Random Numbers Verification\01_XOR_encryption_and_decryption\AHU Black Swan_encrypted_image.png"
output_decrypted_path = r"G:\Users\Chipistil\Desktop\High-throughput-T2NP-Memristor-based-Dual-clock-edge-sampling-TRNG\True Random Numbers Verification\01_XOR_encryption_and_decryption\AHU Black Swan_decrypted_image.png"
start_bit = 0

def read_key(key_path):
    with open(key_path, 'r') as f:
        key = f.read().strip()
    key = ''.join([c for c in key if c in '01'])
    return key

def encrypt_image(image_path, key_path, output_encrypted_path, output_decrypted_path, start_bit=0):
    img = Image.open(image_path)
    img_array = np.array(img)
    
    original_key = read_key(key_path)
    key_length = len(original_key)
    key = original_key
    
    if start_bit < 0:
        start_bit = 0
    elif start_bit >= key_length:
        start_bit = key_length - 1
    
    height, width, channels = img_array.shape
    
    total_pixels = height * width * channels
    required_bits = total_pixels * 8
    
    if key_length < required_bits:
        key = key * (required_bits // key_length + 1)
    
    encrypted_array = np.copy(img_array)
    bit_index = start_bit
    
    for i in range(height):
        for j in range(width):
            for c in range(channels):
                pixel_value = img_array[i, j, c]
                
                pixel_bits = format(pixel_value, '08b')
                
                encrypted_bits = []
                for bit_pos in range(8):
                    if bit_index < len(key):
                        encrypted_bit = str(int(pixel_bits[bit_pos]) ^ int(key[bit_index]))
                        encrypted_bits.append(encrypted_bit)
                        bit_index += 1
                    else:
                        encrypted_bits.append(pixel_bits[bit_pos])
                
                encrypted_value = int(''.join(encrypted_bits), 2)
                encrypted_array[i, j, c] = encrypted_value
    
    actual_bits_used = bit_index - start_bit
    bits_used_from_original = min(actual_bits_used, key_length - start_bit)
    bits_remaining = max(0, key_length - (start_bit + bits_used_from_original))
    bits_used = bits_used_from_original
    
    encrypted_img = Image.fromarray(encrypted_array.astype('uint8'))
    encrypted_img.save(output_encrypted_path, 'PNG')
    print(f"Encrypted image saved to: {output_encrypted_path}")
    
    decrypted_array = np.copy(encrypted_array)
    bit_index = start_bit
    
    for i in range(height):
        for j in range(width):
            for c in range(channels):
                encrypted_value = encrypted_array[i, j, c]
                
                encrypted_bits = format(encrypted_value, '08b')
                
                decrypted_bits = []
                for bit_pos in range(8):
                    if bit_index < len(key):
                        decrypted_bit = str(int(encrypted_bits[bit_pos]) ^ int(key[bit_index]))
                        decrypted_bits.append(decrypted_bit)
                        bit_index += 1
                    else:
                        decrypted_bits.append(encrypted_bits[bit_pos])
                
                decrypted_value = int(''.join(decrypted_bits), 2)
                decrypted_array[i, j, c] = decrypted_value
    
    decrypted_img = Image.fromarray(decrypted_array.astype('uint8'))
    decrypted_img.save(output_decrypted_path, 'PNG')
    print(f"Decrypted image saved to: {output_decrypted_path}")
    
    print(f"\nKey statistics:")
    print(f"- Original key length: {key_length} bits")
    print(f"- Start bit position: {start_bit}")
    print(f"- Bits used for encryption: {bits_used} bits")
    print(f"- Remaining unused key bits: {bits_remaining} bits")
    print(f"- Total bits required for encryption: {required_bits} bits")
    
    if key_length < required_bits:
        cycles = required_bits // key_length
        remaining_in_last_cycle = required_bits % key_length
        print(f"- Key reused {cycles} times, {remaining_in_last_cycle} bits used in last cycle")

if __name__ == "__main__":
    if not os.path.exists(image_path):
        print(f"Error: Image file not found {image_path}")
    elif not os.path.exists(key_path):
        print(f"Error: Key file not found {key_path}")
    else:
        encrypt_image(image_path, key_path, output_encrypted_path, output_decrypted_path, start_bit)