import os
from PIL import Image
import numpy as np

# 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_dir = r"G:\Users\Chipistil\Desktop\High-throughput-T2NP-Memristor-based-Dual-clock-edge-sampling-TRNG\Dual-Image Encryption Verification\02_Fidelity Analysis\01_mean_squared_error_MSE"  
output_file = os.path.join(output_dir, "MSE_calculation_results.txt")  

os.makedirs(output_dir, exist_ok=True)

try:
    print("Reading images...")
    original_image = Image.open(original_image_path).convert('RGB')
    encrypted_image = Image.open(encrypted_image_path).convert('RGB')
    
    print("Resizing images to 512x512...")
    original_image = original_image.resize((512, 512))
    encrypted_image = encrypted_image.resize((512, 512))
    
    original_image = np.array(original_image)
    encrypted_image = np.array(encrypted_image)
    
    W, H = 512, 512
    print(f"Image size: {W}x{H}")
    
    total_pixels = W * H
    print(f"Total pixels: {total_pixels}")
    
    print("Calculating pixel differences...")
    diff = original_image.astype(np.float64) - encrypted_image.astype(np.float64)
    
    print("Calculating squared differences...")
    squared_diff = diff ** 2
    
    print("Calculating squared sums for each channel...")
    
    squared_diff_b = squared_diff[:, :, 0]
    squared_diff_g = squared_diff[:, :, 1]
    squared_diff_r = squared_diff[:, :, 2]
    
    sum_squared_diff_b = np.sum(squared_diff_b)
    sum_squared_diff_g = np.sum(squared_diff_g)
    sum_squared_diff_r = np.sum(squared_diff_r)
    
    non_zero_count_b = np.count_nonzero(squared_diff_b)
    non_zero_count_g = np.count_nonzero(squared_diff_g)
    non_zero_count_r = np.count_nonzero(squared_diff_r)
    
    print("Calculating MSE for each channel...")
    mse_b = sum_squared_diff_b / total_pixels
    mse_g = sum_squared_diff_g / total_pixels
    mse_r = sum_squared_diff_r / total_pixels
    
    total_sum_squared_diff = sum_squared_diff_b + sum_squared_diff_g + sum_squared_diff_r
    mse_overall = total_sum_squared_diff / (3 * total_pixels)
    
    output_content = []
    output_content.append("===== Mean Square Error (MSE) Detailed Calculation Process =====")
    output_content.append("")
    output_content.append(f"Original image path: {original_image_path}")
    output_content.append(f"Encrypted image path: {encrypted_image_path}")
    output_content.append("")
    output_content.append("===== Calculation Parameters =====")
    output_content.append(f"Image width (W): {W}")
    output_content.append(f"Image height (H): {H}")
    output_content.append(f"Total pixels (W×H): {total_pixels}")
    output_content.append("")
    output_content.append("===== Detailed Calculation Steps =====")
    output_content.append("Calculation formula: MSE(P, E) = (1/(WH)) * ΣΣ(P(i,j) - E(i,j))²")
    output_content.append("Where P is the original image and E is the encrypted image")
    output_content.append("")
    output_content.append("Step 1: Read and preprocess images")
    output_content.append(f"   - Original image shape: {original_image.shape}")
    output_content.append(f"   - Encrypted image shape: {encrypted_image.shape}")
    output_content.append("")
    output_content.append("Step 2: Calculate pixel difference matrix")
    output_content.append(f"   - Pixel difference range: [{np.min(diff):.2f}, {np.max(diff):.2f}]")
    output_content.append("")
    output_content.append("Step 3: Calculate squared differences")
    output_content.append(f"   - Minimum squared difference: {np.min(squared_diff):.2f}")
    output_content.append(f"   - Maximum squared difference: {np.max(squared_diff):.2f}")
    output_content.append("")
    output_content.append("Step 4: Calculate squared sums for each channel")
    output_content.append("   Calculation process: Squared sum = ΣΣ(P(i,j) - E(i,j))², i.e., iterate through all pixels and accumulate the squared difference of each pixel")
    output_content.append("   Blue channel:")
    output_content.append(f"     - Non-zero difference pixel count: {non_zero_count_b:,}")
    output_content.append(f"     - Blue channel squared sum: {sum_squared_diff_b:,.2f}")
    output_content.append("   Green channel:")
    output_content.append(f"     - Non-zero difference pixel count: {non_zero_count_g:,}")
    output_content.append(f"     - Green channel squared sum: {sum_squared_diff_g:,.2f}")
    output_content.append("   Red channel:")
    output_content.append(f"     - Non-zero difference pixel count: {non_zero_count_r:,}")
    output_content.append(f"     - Red channel squared sum: {sum_squared_diff_r:,.2f}")
    output_content.append("   Total:")
    output_content.append(f"     - Total squared sum for all channels: {total_sum_squared_diff:,.2f}")
    output_content.append(f"     - Total non-zero difference pixel count: {non_zero_count_b + non_zero_count_g + non_zero_count_r:,}")
    output_content.append("")
    output_content.append("Step 5: Calculate MSE for each channel")
    output_content.append(f"   - Blue channel MSE = {sum_squared_diff_b:,.2f} / {total_pixels:,} = {mse_b:.6f}")
    output_content.append(f"   - Green channel MSE = {sum_squared_diff_g:,.2f} / {total_pixels:,} = {mse_g:.6f}")
    output_content.append(f"   - Red channel MSE = {sum_squared_diff_r:,.2f} / {total_pixels:,} = {mse_r:.6f}")
    output_content.append("")
    output_content.append("Step 6: Calculate overall MSE")
    output_content.append(f"   - Overall MSE = {total_sum_squared_diff:,.2f} / ({3 * total_pixels:,}) = {mse_overall:.6f}")
    output_content.append("")
    output_content.append("===== Final Calculation Results =====")
    output_content.append(f"Blue channel MSE: {mse_b:.6f}")
    output_content.append(f"Green channel MSE: {mse_g:.6f}")
    output_content.append(f"Red channel MSE: {mse_r:.6f}")
    output_content.append(f"Overall MSE: {mse_overall:.6f}")
    output_content.append("")
    output_content.append("===== Calculation Completed =====")
    
    for line in output_content:
        print(line)
    
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write('\n'.join(output_content))
    
    print(f"\nCalculation results saved to: {output_file}")
    
    print("\nVerification information:")
    print(f"Original image shape: {original_image.shape}")
    print(f"Encrypted image shape: {encrypted_image.shape}")
    print(f"Pixel difference range: [{np.min(diff):.2f}, {np.max(diff):.2f}]")
    
except Exception as e:
    print(f"Error during calculation: {str(e)}")
    error_output = ["Error during calculation:", str(e)]
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write('\n'.join(error_output))