import os
import shutil

import cv2
import numpy as np
from skimage.metrics import structural_similarity as ssim

def load_images(folder):
    """Load all PNG images from a folder, categorized by prefix."""
    true_images = {}
    false_images = {}
    for filename in os.listdir(folder):
        if filename.endswith('.png'):
            filepath = os.path.join(folder, filename)
            image = cv2.imread(filepath, cv2.IMREAD_UNCHANGED)  # Load image with alpha channel
            if filename.startswith("True_"):
                true_images[filename] = image
            elif filename.startswith("False_"):
                false_images[filename] = image
    return true_images, false_images

def resize_image(image, size=(125, 50)):
    """Resize an image to a standard size."""
    return cv2.resize(image, size, interpolation=cv2.INTER_AREA)

def calculate_ssim(image1, image2):
    """Calculate SSIM between two images."""
    if image1.shape[2] == 4:  # Handle alpha channel
        image1 = cv2.cvtColor(image1, cv2.COLOR_BGRA2BGR)
    if image2.shape[2] == 4:  # Handle alpha channel
        image2 = cv2.cvtColor(image2, cv2.COLOR_BGRA2BGR)
    gray1 = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
    score, _ = ssim(gray1, gray2, full=True)
    return score

def calculate_color_similarity(image1, image2):
    """Compare color histograms of two images."""
    if image1.shape[2] == 4:  # Handle alpha channel
        image1 = cv2.cvtColor(image1, cv2.COLOR_BGRA2BGR)
    if image2.shape[2] == 4:  # Handle alpha channel
        image2 = cv2.cvtColor(image2, cv2.COLOR_BGRA2BGR)
    hist1 = cv2.calcHist([image1], [0, 1, 2], None, [8, 8, 8], [0, 256, 0, 256, 0, 256])
    hist2 = cv2.calcHist([image2], [0, 1, 2], None, [8, 8, 8], [0, 256, 0, 256, 0, 256])
    hist1 = cv2.normalize(hist1, hist1).flatten()
    hist2 = cv2.normalize(hist2, hist2).flatten()
    return cv2.compareHist(hist1, hist2, cv2.HISTCMP_CORREL)

def calculate_combined_similarity(image1, image2):
    """Combine SSIM and color histogram similarity."""
    # ssim_score = calculate_ssim(image1, image2)
    color_score = calculate_color_similarity(image1, image2)
    # Weight SSIM and color similarity equally (adjust weights as needed)
    return 1 * color_score

def find_most_similar_images(true_images, false_images):
    """Find the most similar image for each 'True_' image among 'False_' images."""
    results = {}
    for true_name, true_image in true_images.items():
        max_similarity = -1
        most_similar_image = None
        for false_name, false_image in false_images.items():
            # resized_true = resize_image(true_image)
            # resized_false = resize_image(false_image)

            if true_name.split("_")[2] == false_name.split("_")[2]:
                continue

            similarity = calculate_combined_similarity(true_image, false_image)
            if similarity > max_similarity:
                max_similarity = similarity
                most_similar_image = false_name
        results[true_name] = (most_similar_image, max_similarity)
        os.makedirs("output14\\"+true_name+"_"+most_similar_image,exist_ok=True)
        shutil.copyfile('output13\\'+true_name, 'output14\\'+true_name+"_"+most_similar_image+"\\"+true_name)
        shutil.copyfile('output13\\' + most_similar_image, 'output14\\' +true_name+"_"+most_similar_image+"\\" + most_similar_image)

        # shutil.copyfile('output12\\'+ "segment_"+true_name.split("_")[2]+".png", 'output14\\'+true_name+"_"+most_similar_image+"\\segment_"+true_name.split("_")[2]+".png")
        # shutil.copyfile('output12\\' +"segment_"+most_similar_image.split("_")[2]+".png", 'output14\\'+true_name+"_"+most_similar_image+"\\segment_"+most_similar_image.split("_")[2]+".png")

        shutil.copyfile('output10\\'+ "segment_"+true_name.split("_")[2]+".png", 'output14\\'+true_name+"_"+most_similar_image+"\\segment_"+true_name.split("_")[2]+".png")
        shutil.copyfile('output10\\' +"segment_"+most_similar_image.split("_")[2]+".png", 'output14\\'+true_name+"_"+most_similar_image+"\\segment_"+most_similar_image.split("_")[2]+".png")


    return results

def main():
    folder = 'output13'  # Replace with your folder path
    true_images, false_images = load_images(folder)
    results = find_most_similar_images(true_images, false_images)
    for true_image, (most_similar, similarity) in results.items():
        print(f"True Image: {true_image}, Most Similar False Image: {most_similar}, Similarity: {similarity:.4f}")


if __name__ == '__main__':
    main()