import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"


import os
import tempfile
import numpy as np
import cv2
import gradio as gr
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess_input
from tensorflow.keras.preprocessing import image
from skimage.metrics import structural_similarity as ssim
from PIL import Image

# Disable GPU for TensorFlow
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

class ImageCharacterClassifier:
    def __init__(self, similarity_threshold=0.5):
        self.model = ResNet50(weights='imagenet', include_top=False, pooling='avg')
        self.similarity_threshold = similarity_threshold

    def load_and_preprocess_image(self, img):
        # Convert image to array and preprocess it
        img = img.convert('RGB')
        img_array = np.array(img)
        img_array = cv2.resize(img_array, (224, 224))  # Ensure correct size
        img_array = np.expand_dims(img_array, axis=0)
        img_array = preprocess_input(img_array)
        return img_array

    def extract_features(self, img):
        preprocessed_img = self.load_and_preprocess_image(img)
        features = self.model.predict(preprocessed_img)
        return features

    def calculate_ssim(self, img1, img2):
        img1_gray = cv2.cvtColor(img1, cv2.COLOR_RGB2GRAY)
        img2_gray = cv2.cvtColor(img2, cv2.COLOR_RGB2GRAY)
        img2_gray = cv2.resize(img2_gray, (img1_gray.shape[1], img1_gray.shape[0]))
        return ssim(img1_gray, img2_gray)

def process_images(reference_image, comparison_images, similarity_threshold):
    try:
        if reference_image is None:
            return "Please upload a reference image.", []
        if not comparison_images:
            return "Please upload comparison images.", []

        classifier = ImageCharacterClassifier(similarity_threshold)

        # Convert reference image to NumPy array
        ref_image = Image.fromarray(reference_image)
        ref_features = classifier.extract_features(ref_image)

        results = []
        html_output = "<h3>Comparison Results:</h3>"

        for comp_image in comparison_images:
            try:
                # Read image file as PIL Image
                comp_pil = Image.open(comp_image)
                comp_pil = comp_pil.convert("RGB")

                # Convert to NumPy format for SSIM
                comp_array = np.array(comp_pil)

                # Calculate SSIM score
                ssim_score = classifier.calculate_ssim(reference_image, comp_array)

                # Extract features
                comp_features = classifier.extract_features(comp_pil)
                max_feature_diff = np.max(np.abs(ref_features - comp_features))
                is_similar = max_feature_diff < 6.0

                status_text = "SIMILAR" if is_similar else "NOT SIMILAR"
                status_color = "green" if is_similar else "red"

                html_output += f"<p style='color:{status_color};'>{comp_image.name}: {status_text}</p>"
                results.append(comp_array)

            except Exception as e:
                html_output += f"<p style='color:red;'>Error processing {comp_image.name}: {str(e)}</p>"

        return html_output, results

    except Exception as e:
        return f"<p style='color:red;'>Error: {str(e)}</p>", []

def create_interface():
    with gr.Blocks() as interface:
        gr.Markdown("# Image Similarity Classifier")
        gr.Markdown("Upload a reference image and multiple comparison images.")

        with gr.Row():
            with gr.Column():
                reference_input = gr.Image(label="Reference Image", type="numpy")
                comparison_input = gr.Files(label="Comparison Images", type="filepath")
                threshold_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, step=0.05, label="Similarity Threshold")
                submit_button = gr.Button("Compare Images")

            with gr.Column():
                output_html = gr.HTML(label="Results")
                output_gallery = gr.Gallery(label="Processed Images", columns=3)

        submit_button.click(
            fn=process_images,
            inputs=[reference_input, comparison_input, threshold_slider],
            outputs=[output_html, output_gallery]
        )

    return interface

if __name__ == "__main__":
    interface = create_interface()
    interface.launch(share=True)