import numpy as np
from PIL import Image
import os
import math
from pathlib import Path

# Path settings
original_image_path = r"G:\Users\Chipistil\Desktop\High-throughput-T2NP-Memristor-based-Dual-clock-edge-sampling-TRNG\AHU Black Swan.png"
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_file_path = r"G:\Users\Chipistil\Desktop\High-throughput-T2NP-Memristor-based-Dual-clock-edge-sampling-TRNG\Dual-Image Encryption Verification\01_Statistical Analysis\01_information_entropy.txt"

def calculate_entropy(image_path):
    if not Path(image_path).exists():
        print(f"Error: File {image_path} does not exist")
        return None
    
    try:
        img = Image.open(image_path)
        img_array = np.array(img)
        
        height, width, channels = img_array.shape
        total_pixels = height * width
        
        r = img_array[:, :, 0]
        g = img_array[:, :, 1]
        b = img_array[:, :, 2]
        
        def channel_entropy_detail(channel_data, channel_name):
            hist, bin_edges = np.histogram(channel_data.flatten(), bins=256, range=[0, 256])
            probabilities = hist / total_pixels
            non_zero_indices = probabilities > 0
            non_zero_pixels = np.where(non_zero_indices)[0]
            non_zero_probs = probabilities[non_zero_indices]
            contributions = -non_zero_probs * np.log2(non_zero_probs)
            entropy = np.sum(contributions)
            
            return {
                'entropy': entropy,
                'histogram': hist,
                'probabilities': probabilities,
                'non_zero_pixels': non_zero_pixels,
                'non_zero_probs': non_zero_probs,
                'contributions': contributions,
                'total_pixels': total_pixels,
                'image_dimensions': (width, height)
            }
        
        red_data = channel_entropy_detail(r, 'red')
        green_data = channel_entropy_detail(g, 'green')
        blue_data = channel_entropy_detail(b, 'blue')
        
        total_entropy = (red_data['entropy'] + green_data['entropy'] + blue_data['entropy']) / 3
        
        return {
            'red': red_data,
            'green': green_data,
            'blue': blue_data,
            'total_entropy': total_entropy,
            'image_dimensions': (width, height),
            'total_pixels': total_pixels
        }
        
    except Exception as e:
        print(f"Error processing image: {e}")
        return None

def save_entropy_results(original_entropy, encrypted_entropy, output_path):
    output_dir = os.path.dirname(output_path)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write("Image Encryption Information Entropy Analysis Results (Detailed Calculation Process)\n")
        f.write("="*60 + "\n\n")
        
        f.write("[Calculation Principle]\n")
        f.write("Information entropy formula: H = -Σ p(i) * log₂(p(i))\n")
        f.write("Where p(i) is the probability of intensity level i appearing in the image, i ranges from 0 to 255\n")
        f.write("p(i) = Number of pixels with value i / Total number of pixels\n")
        f.write("p(i) = Number of pixels with value i / Total number of pixels\n\n")
        
        f.write("[Original Image Information]\n")
        width, height = original_entropy['image_dimensions']
        total_pixels = original_entropy['total_pixels']
        f.write(f"Image size: {width}x{height}\n")
        f.write(f"Total pixels: {total_pixels}\n\n")
        
        f.write("[Original Image Information Entropy Detailed Calculation Process]\n")
        f.write("="*60 + "\n\n")
        
        f.write("1. Original Image Red Channel Entropy Calculation\n")
        f.write("="*60 + "\n")
        red_data = original_entropy['red']
        f.write(f"Red channel total pixels: {red_data['total_pixels']}\n")
        f.write("Calculation steps:\n")
        f.write("a. Count the occurrence of each pixel value (0-255)\n")
        
        f.write("   Sample pixel value statistics: \n")
        sample_indices = list(range(0, 256, 32))
        for i in sample_indices:
            count = red_data['histogram'][i]
            prob = red_data['probabilities'][i]
            f.write(f"   Pixel value {i}: Appeared {count} times, probability = {prob:.6f}\n")
        
        non_zero_indices = np.nonzero(red_data['histogram'])[0]
        if len(non_zero_indices) > 0:
            max_count_idx = non_zero_indices[np.argmax(red_data['histogram'][non_zero_indices])]
            min_count_idx = non_zero_indices[np.argmin(red_data['histogram'][non_zero_indices])]
            f.write(f"   Pixel value with most occurrences: {max_count_idx} (appeared {red_data['histogram'][max_count_idx]} times)\n")
            f.write(f"   Pixel value with least occurrences: {min_count_idx} (appeared {red_data['histogram'][min_count_idx]} times)\n")
        
        f.write("\nb. Calculate information entropy contribution for each non-zero probability pixel value\n")
        f.write("   Information entropy contribution formula: -p(i) * log₂(p(i))\n")
        
        if len(red_data['non_zero_pixels']) > 0:
            show_count = min(10, len(red_data['non_zero_pixels']))
            f.write(f"   First {show_count} non-zero pixel value contribution examples: \n")
            for i in range(show_count):
                pixel_val = red_data['non_zero_pixels'][i]
                prob = red_data['non_zero_probs'][i]
                contrib = red_data['contributions'][i]
                f.write(f"   Pixel value {pixel_val}: p={prob:.6f}, -p*log2(p)={contrib:.6f}\n")
            
            if len(red_data['non_zero_pixels']) > show_count:
                f.write(f"   ...\n")
                f.write(f"   Last {show_count} non-zero pixel value contribution examples: \n")
                for i in range(-show_count, 0):
                    pixel_val = red_data['non_zero_pixels'][i]
                    prob = red_data['non_zero_probs'][i]
                    contrib = red_data['contributions'][i]
                    f.write(f"   Pixel value {pixel_val}: p={prob:.6f}, -p*log2(p)={contrib:.6f}\n")
        
        f.write("\nc. Calculate total information entropy\n")
        f.write(f"   Red channel information entropy = Σ[-p(i) * log₂(p(i))] = {red_data['entropy']:.6f} bits\n")
        f.write("\n")
        
        f.write("2. Original Image Green Channel Entropy Calculation\n")
        f.write("="*60 + "\n")
        green_data = original_entropy['green']
        f.write(f"Green channel information entropy = {green_data['entropy']:.6f} bits\n\n")
        
        f.write("3. Original Image Blue Channel Entropy Calculation\n")
        f.write("="*60 + "\n")
        blue_data = original_entropy['blue']
        f.write(f"Blue channel information entropy = {blue_data['entropy']:.6f} bits\n\n")
        
        f.write("4. Original Image Total Information Entropy Calculation\n")
        f.write("="*60 + "\n")
        f.write(f"Total information entropy = (Red channel entropy + Green channel entropy + Blue channel entropy) / 3\n")
        f.write(f"Total information entropy = ({red_data['entropy']:.6f} + {green_data['entropy']:.6f} + {blue_data['entropy']:.6f}) / 3 = {original_entropy['total_entropy']:.6f} bits\n\n")
        
        f.write("[Encrypted Image Information Entropy Detailed Calculation Process]\n")
        f.write("="*60 + "\n\n")
        
        f.write("1. Encrypted Image Red Channel Entropy Calculation\n")
        f.write("="*60 + "\n")
        enc_red_data = encrypted_entropy['red']
        f.write(f"Red channel information entropy = {enc_red_data['entropy']:.6f} bits\n\n")
        
        f.write("2. Encrypted Image Green Channel Entropy Calculation\n")
        f.write("="*60 + "\n")
        enc_green_data = encrypted_entropy['green']
        f.write(f"Green channel information entropy = {enc_green_data['entropy']:.6f} bits\n\n")
        
        f.write("3. Encrypted Image Blue Channel Entropy Calculation\n")
        f.write("="*60 + "\n")
        enc_blue_data = encrypted_entropy['blue']
        f.write(f"Blue channel information entropy = {enc_blue_data['entropy']:.6f} bits\n\n")
        
        f.write("4. Encrypted Image Total Information Entropy Calculation\n")
        f.write("="*60 + "\n")
        f.write(f"Total information entropy = (Red channel entropy + Green channel entropy + Blue channel entropy) / 3\n")
        f.write(f"Total information entropy = ({enc_red_data['entropy']:.6f} + {enc_green_data['entropy']:.6f} + {enc_blue_data['entropy']:.6f}) / 3 = {encrypted_entropy['total_entropy']:.6f} bits\n\n")
        
        f.write("[Entropy Change Detailed Calculation]\n")
        f.write("="*60 + "\n")
        
        red_diff = enc_red_data['entropy'] - red_data['entropy']
        f.write(f"Red channel entropy change = Encrypted entropy - Original entropy = {enc_red_data['entropy']:.6f} - {red_data['entropy']:.6f} = {red_diff:.6f} bits\n")
        
        green_diff = enc_green_data['entropy'] - green_data['entropy']
        f.write(f"Green channel entropy change = Encrypted entropy - Original entropy = {enc_green_data['entropy']:.6f} - {green_data['entropy']:.6f} = {green_diff:.6f} bits\n")
        
        blue_diff = enc_blue_data['entropy'] - blue_data['entropy']
        f.write(f"Blue channel entropy change = Encrypted entropy - Original entropy = {enc_blue_data['entropy']:.6f} - {blue_data['entropy']:.6f} = {blue_diff:.6f} bits\n")
        
        total_diff = encrypted_entropy['total_entropy'] - original_entropy['total_entropy']
        f.write(f"Total information entropy change = Encrypted total entropy - Original total entropy = {encrypted_entropy['total_entropy']:.6f} - {original_entropy['total_entropy']:.6f} = {total_diff:.6f} bits\n\n")
        
        f.write("[Analysis]\n")
        f.write("="*60 + "\n")
        if encrypted_entropy['total_entropy'] > original_entropy['total_entropy']:
            f.write("The information entropy of the encrypted image increased, indicating that the pixel distribution is more uniform and random, with good encryption effect.\n")
            f.write(f"Entropy increased by {total_diff:.6f} bits, relative increase ratio: {total_diff/original_entropy['total_entropy']*100:.2f}%\n")
        elif encrypted_entropy['total_entropy'] < original_entropy['total_entropy']:
            f.write("The information entropy of the encrypted image decreased, which may be due to issues with the encryption algorithm or key selection.\n")
            f.write(f"Entropy decreased by {-total_diff:.6f} bits, relative decrease ratio: {-total_diff/original_entropy['total_entropy']*100:.2f}%\n")
        else:
            f.write("The information entropy of the image remained unchanged before and after encryption.\n")
        
        f.write("\n[Information Entropy Reference]\n")
        f.write("="*60 + "\n")
        max_theoretical_entropy = 8
        f.write(f"For 8-bit images (0-255), the theoretical maximum information entropy is {max_theoretical_entropy} bits/pixel\n")
        f.write(f"Ratio of original image total entropy to theoretical maximum: {original_entropy['total_entropy']/max_theoretical_entropy*100:.2f}%\n")
        f.write(f"Ratio of encrypted image total entropy to theoretical maximum: {encrypted_entropy['total_entropy']/max_theoretical_entropy*100:.2f}%\n")
    
    print(f"Information entropy analysis results saved to: {output_path}")

def main():
    print("Starting image information entropy calculation...")
    
    print(f"Calculating original image entropy: {original_image_path}")
    original_entropy = calculate_entropy(original_image_path)
    if not original_entropy:
        return
    
    print(f"Calculating encrypted image entropy: {encrypted_image_path}")
    encrypted_entropy = calculate_entropy(encrypted_image_path)
    if not encrypted_entropy:
        return
    
    save_entropy_results(original_entropy, encrypted_entropy, output_file_path)
    
    print("Information entropy calculation completed!")


if __name__ == "__main__":
    main()