import numpy as np
from PIL import Image
import os
import random

# 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"
encrypted_image_path = r"G:\Users\Chipistil\Desktop\High-throughput-T2NP-Memristor-based-Dual-clock-edge-sampling-TRNG\AHU Black Swan_encrypted_image.png"
output_base_path = r"G:\Users\Chipistil\Desktop\High-throughput-T2NP-Memristor-based-Dual-clock-edge-sampling-TRNG\Dual-Image Encryption Verification\03_Robustness Testing"

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):
    img = Image.open(image_path)
    img_array = np.array(img)
    
    original_key = read_key(key_path)
    key_length = len(original_key)
    key = original_key
    
    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 = 0
    
    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
    
    encrypted_img = Image.fromarray(encrypted_array.astype('uint8'))
    encrypted_img.save(output_encrypted_path, 'PNG')
    print(f"Encrypted image saved to: {output_encrypted_path}")
    
    return encrypted_array

def add_salt_and_pepper_noise(image_array, noise_percentage):
    noisy_array = np.copy(image_array)
    
    height, width, channels = noisy_array.shape
    
    total_pixels = height * width
    num_noisy_pixels = int(total_pixels * noise_percentage / 100)
    
    if noise_percentage <= 0:
        return noisy_array
    
    for _ in range(num_noisy_pixels):
        i = random.randint(0, height - 1)
        j = random.randint(0, width - 1)
        
        if random.random() < 0.5:
            noisy_array[i, j] = [255, 255, 255]
        else:
            noisy_array[i, j] = [0, 0, 0]
    
    return noisy_array

def decrypt_image(encrypted_array, key_path, output_decrypted_path):
    original_key = read_key(key_path)
    key_length = len(original_key)
    key = original_key
    
    height, width, channels = encrypted_array.shape
    
    total_pixels = height * width * channels
    required_bits = total_pixels * 8
    
    if key_length < required_bits:
        key = key * (required_bits // key_length + 1)
    
    decrypted_array = np.copy(encrypted_array)
    bit_index = 0
    
    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}")

if __name__ == "__main__":
    os.makedirs(output_base_path, exist_ok=True)
    
    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:
        encrypted_array = encrypt_image(image_path, key_path, encrypted_image_path)
        
        noise_percentages = [0, 6.25, 12.5, 25, 50]
        
        for percentage in noise_percentages:
            noisy_array = add_salt_and_pepper_noise(encrypted_array, percentage)
            
            noisy_encrypted_path = os.path.join(output_base_path, f"encrypted_image_{percentage}pct_noise.png")
            noisy_encrypted_img = Image.fromarray(noisy_array.astype('uint8'))
            noisy_encrypted_img.save(noisy_encrypted_path, 'PNG')
            print(f"Noisy encrypted image saved to: {noisy_encrypted_path}")
            
            decrypted_path = os.path.join(output_base_path, f"decrypted_image_{percentage}pct_noise.png")
            decrypt_image(noisy_array, key_path, decrypted_path)
            
        print("\nAll processing completed!")