import os
import numpy as np
from PIL import Image

# 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\03_structural_similarity_index_measure_SSIM"
output_file = os.path.join(output_dir, "SSIM_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))
    
    Q = np.array(original_image, dtype=np.float64)
    R = np.array(encrypted_image, dtype=np.float64)
    
    W, H, channels = Q.shape
    print(f"Image size: {W}x{H}x{channels}")
    
    C1 = 6.5025
    C2 = 58.5225
    print(f"SSIM parameters: C1={C1}, C2={C2}")
    
    output_content = []
    output_content.append("===== Structural Similarity Index Measure (SSIM) 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"Number of channels: {channels}")
    output_content.append(f"C1: {C1}")
    output_content.append(f"C2: {C2}")
    output_content.append("")
    output_content.append("===== SSIM Calculation Formula =====")
    output_content.append("SSIM(Q, R) = [(2μQμR + C1)(2σQR + C2)] / [(μQ² + μR² + C1)(σQ² + σR² + C2)]")
    output_content.append("Where:")
    output_content.append("μQ = Mean pixel value of original image")
    output_content.append("μR = Mean pixel value of encrypted image")
    output_content.append("σQR = Covariance between original and encrypted images")
    output_content.append("σQ² = Variance of original image")
    output_content.append("σR² = Variance of encrypted image")
    output_content.append("")
    
    ssim_values = []
    channel_names = ['Blue Channel', 'Green Channel', 'Red Channel']
    
    for c in range(channels):
        output_content.append(f"===== {channel_names[c]} SSIM Calculation =====")
        
        Q_channel = Q[:, :, c]
        R_channel = R[:, :, c]
        
        mu_Q = np.mean(Q_channel)
        mu_R = np.mean(R_channel)
        output_content.append(f"Step 1: Calculate mean values")
        output_content.append(f"   - μQ ({channel_names[c]} original image mean) = {mu_Q:.6f}")
        output_content.append(f"   - μR ({channel_names[c]} encrypted image mean) = {mu_R:.6f}")
        
        sigma_Q_sq = np.mean((Q_channel - mu_Q) ** 2)
        sigma_R_sq = np.mean((R_channel - mu_R) ** 2)
        output_content.append(f"Step 2: Calculate variance")
        output_content.append(f"   - σQ² ({channel_names[c]} original image variance) = {sigma_Q_sq:.6f}")
        output_content.append(f"   - σR² ({channel_names[c]} encrypted image variance) = {sigma_R_sq:.6f}")
        
        sigma_QR = np.mean((Q_channel - mu_Q) * (R_channel - mu_R))
        output_content.append(f"Step 3: Calculate covariance")
        output_content.append(f"   - σQR ({channel_names[c]} covariance) = {sigma_QR:.6f}")
        
        numerator = (2 * mu_Q * mu_R + C1) * (2 * sigma_QR + C2)
        denominator = (mu_Q ** 2 + mu_R ** 2 + C1) * (sigma_Q_sq + sigma_R_sq + C2)
        output_content.append(f"Step 4: Calculate SSIM numerator and denominator")
        output_content.append(f"   - Numerator = (2 * {mu_Q:.6f} * {mu_R:.6f} + {C1}) * (2 * {sigma_QR:.6f} + {C2}) = {numerator:.6f}")
        output_content.append(f"   - Denominator = ({mu_Q:.6f}² + {mu_R:.6f}² + {C1}) * ({sigma_Q_sq:.6f} + {sigma_R_sq:.6f} + {C2}) = {denominator:.6f}")
        
        ssim = numerator / denominator
        ssim_values.append(ssim)
        output_content.append(f"Step 5: Calculate SSIM")
        output_content.append(f"   - SSIM = {numerator:.6f} / {denominator:.6f} = {ssim:.6f}")
        output_content.append("")
    
    overall_ssim = np.mean(ssim_values)
    
    output_content.append("===== SSIM Calculation Results for Each Channel =====")
    for i in range(channels):
        output_content.append(f"{channel_names[i]} SSIM: {ssim_values[i]:.6f}")
    output_content.append(f"Overall SSIM (channel average): {overall_ssim:.6f}")
    output_content.append("")
    
    output_content.append("===== SSIM Detailed Calculation Verification =====")
    
    for i in range(channels):
        Q_channel = Q[:, :, i]
        R_channel = R[:, :, i]
        mu_Q = np.mean(Q_channel)
        mu_R = np.mean(R_channel)
        output_content.append(f"{channel_names[i]}:")
        output_content.append(f"   - μQ = Sum of all pixel values / Total number of pixels = {mu_Q:.6f}")
        output_content.append(f"   - μR = Sum of all pixel values / Total number of pixels = {mu_R:.6f}")
    output_content.append("")
    
    for i in range(channels):
        Q_channel = Q[:, :, i]
        R_channel = R[:, :, i]
        mu_Q = np.mean(Q_channel)
        mu_R = np.mean(R_channel)
        sigma_Q_sq = np.mean((Q_channel - mu_Q) ** 2)
        sigma_R_sq = np.mean((R_channel - mu_R) ** 2)
        sigma_QR = np.mean((Q_channel - mu_Q) * (R_channel - mu_R))
        
        output_content.append(f"{channel_names[i]}:")
        output_content.append(f"   - σQ² = Σ(Q(i,j) - μQ)² / WH = {sigma_Q_sq:.6f}")
        output_content.append(f"   - σR² = Σ(R(i,j) - μR)² / WH = {sigma_R_sq:.6f}")
        output_content.append(f"   - σQR = Σ(Q(i,j) - μQ)(R(i,j) - μR) / WH = {sigma_QR:.6f}")
    output_content.append("")
    
    output_content.append("===== SSIM Final Calculation Verification =====")
    for i in range(channels):
        Q_channel = Q[:, :, i]
        R_channel = R[:, :, i]
        mu_Q = np.mean(Q_channel)
        mu_R = np.mean(R_channel)
        sigma_Q_sq = np.mean((Q_channel - mu_Q) ** 2)
        sigma_R_sq = np.mean((R_channel - mu_R) ** 2)
        sigma_QR = np.mean((Q_channel - mu_Q) * (R_channel - mu_R))
        
        numerator = (2 * mu_Q * mu_R + C1) * (2 * sigma_QR + C2)
        denominator = (mu_Q ** 2 + mu_R ** 2 + C1) * (sigma_Q_sq + sigma_R_sq + C2)
        ssim = numerator / denominator
        
        output_content.append(f"{channel_names[i]} SSIM calculation verification:")
        output_content.append(f"   SSIM = [(2μQμR + C1)(2σQR + C2)] / [(μQ² + μR² + C1)(σQ² + σR² + C2)]")
        output_content.append(f"   = [(2*{mu_Q:.2f}*{mu_R:.2f} + {C1})(2*{sigma_QR:.2f} + {C2})] / [({mu_Q:.2f}² + {mu_R:.2f}² + {C1})({sigma_Q_sq:.2f} + {sigma_R_sq:.2f} + {C2})]")
        output_content.append(f"   = {ssim:.6f}")
    output_content.append(f"Overall SSIM = Average of all channel SSIM values = {overall_ssim:.6f}")
    output_content.append("")
    
    output_content.append("===== SSIM Explanation =====")
    output_content.append("SSIM (Structural Similarity Index Measure) is a metric for evaluating image similarity.")
    output_content.append("SSIM values range from [-1, 1], where:")
    output_content.append("- SSIM = 1 indicates that two images are completely identical")
    output_content.append("- SSIM = 0 indicates that two images have completely different structures")
    output_content.append("- Higher SSIM values indicate better structural similarity between two images")
    output_content.append("")
    output_content.append(f"In this calculation, parameters C1={C1} and C2={C2} are used, which are set based on pixel range [0, 255].")
    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: {Q.shape}")
    print(f"Encrypted image shape: {R.shape}")
    print(f"Original image pixel value range: [{np.min(Q):.2f}, {np.max(Q):.2f}]")
    print(f"Encrypted image pixel value range: [{np.min(R):.2f}, {np.max(R):.2f}]")
    
    print("\nSSIM results for each channel:")
    for i in range(channels):
        print(f"{channel_names[i]} SSIM: {ssim_values[i]:.6f}")
    print(f"Overall SSIM: {overall_ssim:.6f}")
    
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))