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

# Function to calculate SSIM between two images
def calculate_similarity(img1, img2):
    if len(img1.shape) == 2:
        img1 = cv2.cvtColor(img1, cv2.COLOR_GRAY2RGB)
    if len(img2.shape) == 2:
        img2 = cv2.cvtColor(img2, cv2.COLOR_GRAY2RGB)
    return ssim(img1, img2, win_size=3)


# Function to compute similarity scores for all images
def compute_similarity(target_image, image_list):
    scores = []
    target_image_resized = cv2.resize(target_image, (target_image.shape[1], target_image.shape[0]))
    for image in image_list:
        image_resized = cv2.resize(image, (target_image.shape[1], target_image.shape[0]))
        similarity_score = calculate_similarity(target_image_resized, image_resized)
        scores.append(similarity_score)
    return scores


# Function to handle the Gradio interface
def image_similarity(target_image, image):
    target_image = target_image.astype(np.uint8)
    image = image.astype(np.uint8)
    scores = compute_similarity(target_image, [image])
    result = (format_image(image), f"Score: {scores[0]:.4f}")
    return [result]

# Function to format the image representation as a string
def format_image(image):
    formatted_image = ""
    for row in image:
        formatted_row = ", ".join([f"[{', '.join(map(str, pixel))}]" for pixel in row])
        formatted_image += f"[{formatted_row}], "
    return formatted_image[:-2]

# Prepare Gradio interface
iface = gr.Interface(
    fn=image_similarity,
    inputs=[
        gr.inputs.Image(type="numpy", label="Target Image"),
        gr.inputs.Image(type="numpy", label="Image")
    ],
    outputs="text",
    title="Image Similarity Calculator",
    description="Upload a target image and another image. Get the similarity score."
)

# Launch the interface
iface.launch()