import glob
import os

import cv2
import matplotlib.pyplot as plt
import numpy as np
import scipy.stats as stats
from skimage.metrics import structural_similarity as ssim
from tqdm import tqdm

# use this: https://sketch.io

def calculate_rmse(imageA, imageB):
    err = np.sum((imageA - imageB) ** 2)
    err /= float(imageA.shape[0] * imageA.shape[1])
    return np.sqrt(err)

def compare_images(ground_truth_path, sample_paths):
    results = []
    gt_image = cv2.imread(ground_truth_path, cv2.IMREAD_GRAYSCALE)
    
    if gt_image is None:
        raise ValueError("Ground truth image could not be read. Please check the file path.")
    
    gt_image = gt_image.astype("float") / 255.0
    
    for path in tqdm(sample_paths):
        sample_image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        
        if sample_image is None:
            print(f"WARNING: Sample image at path {path} could not be read. Skipping this image.")
            continue
        
        sample_image = sample_image.astype("float") / 255.0
        
        rmse_value = calculate_rmse(gt_image, sample_image)
        ssim_value, _ = ssim(gt_image, sample_image, full=True, data_range=1)  # Corrected line
        
        diff_mask = cv2.absdiff(gt_image, sample_image)
        
        # plt.imshow(diff_mask * 255, cmap='gray')
        # plt.title(f'Difference Mask for {os.path.basename(path)}\nRMSE: {rmse_value:.5f} - SSIM: {ssim_value:.5f}')
        # plt.show()
        
        results.append({
            'path': path,
            'rmse': rmse_value,
            'ssim': ssim_value,
            'diff_mask': diff_mask
        })
    
    return results


ground_truth = 'ground_truth.png'
sample_images = glob.glob("samples/*.png")

results = compare_images(ground_truth, sample_images)

for res in results:
    print(f"Image: {res['path']} - RMSE: {res['rmse']} - SSIM: {res['ssim']}")

def calculate_confidence_interval(data, confidence_level=0.95):
    mean = np.mean(data)
    sem = stats.sem(data)
    df = len(data) - 1
    me = sem * stats.t.ppf((1 + confidence_level) / 2, df)
    return mean - me, mean + me

rmse_values = [res['rmse'] for res in results]
ssim_values = [res['ssim'] for res in results]

rmse_mean = np.mean(rmse_values)
rmse_median = np.median(rmse_values)
rmse_stdev = np.std(rmse_values, ddof=1)

ssim_mean = np.mean(ssim_values)
ssim_median = np.median(ssim_values)
ssim_stdev = np.std(ssim_values, ddof=1)

rmse_ci = calculate_confidence_interval(rmse_values)
ssim_ci = calculate_confidence_interval(ssim_values)

print(f"\nRMSE - Mean: {rmse_mean}, Median: {rmse_median}, Std Dev: {rmse_stdev}, 95% CI: {rmse_ci}")
print(f"SSIM - Mean: {ssim_mean}, Median: {ssim_median}, Std Dev: {ssim_stdev}, 95% CI: {ssim_ci}")

print(f"RMSE: {rmse_mean} ± {rmse_ci[1] - rmse_mean}")
print(f"SSIM: {ssim_mean} ± {ssim_ci[1] - ssim_mean}")

def save_average_diff_map(results, save_path='average_diff_map.png'):
    if not results:
        print("No results available to create an average diff map.")
        return
    
    avg_diff_map = None
    
    for res in results:
        if avg_diff_map is None:
            avg_diff_map = np.zeros_like(res['diff_mask'])
        
        avg_diff_map += res['diff_mask']
    
    avg_diff_map /= len(results)
    
    avg_diff_map = (avg_diff_map * 255).astype(np.uint8)
    
    cv2.imwrite(save_path, avg_diff_map)

# Usage
save_average_diff_map(results)