from io import BytesIO
import os
import json

import requests
from PIL import Image
from imagehash import average_hash, phash, dhash, whash
import numpy as np
import gradio as gr

from similarity.similarity import Similarity


def preprocess_image(url):
    response = requests.get(url)
    img = Image.open(BytesIO(response.content))
    img = img.convert('RGB')
    img = img.resize((224, 224))  # adjust as per your model requirements
    img_array = np.array(img) / 255.0
    return img_array


def calculate_hashes(img1_array, img2_array):
    hash1 = average_hash(Image.fromarray((img1_array * 255).astype(np.uint8)))
    hash2 = average_hash(Image.fromarray((img2_array * 255).astype(np.uint8)))
    average_similarity = 1 - (hash1 - hash2) / len(hash1.hash) ** 2

    hash1 = phash(Image.fromarray((img1_array * 255).astype(np.uint8)))
    hash2 = phash(Image.fromarray((img2_array * 255).astype(np.uint8)))
    perceptual_similarity = 1 - (hash1 - hash2) / len(hash1.hash) ** 2

    hash1 = dhash(Image.fromarray((img1_array * 255).astype(np.uint8)))
    hash2 = dhash(Image.fromarray((img2_array * 255).astype(np.uint8)))
    difference_similarity = 1 - (hash1 - hash2) / len(hash1.hash) ** 2

    hash1 = whash(Image.fromarray((img1_array * 255).astype(np.uint8)))
    hash2 = whash(Image.fromarray((img2_array * 255).astype(np.uint8)))
    wavelet_similarity = 1 - (hash1 - hash2) / len(hash1.hash) ** 2

    return {
        'average_similarity': float(average_similarity),  # ensures the output will be printable by json.dumps
        'perceptual_similarity': float(perceptual_similarity),
        'difference_similarity': float(difference_similarity),
        'wavelet_similarity': float(wavelet_similarity)
    }


def cosine_similarity_bit(url1, url2):
    # slimmed down version of https://huggingface.co/spaces/AnnasBlackHat/Image-Similarity/tree/main
    urls = [url1, url2]
    similarity = Similarity()
    model = similarity.get_model()
    similarity_score = similarity.check_similarity(urls, model)
    return float(similarity_score)


def main(url1, url2):
    img1_array = preprocess_image(url1)
    img2_array = preprocess_image(url2)
    scores = calculate_hashes(img1_array, img2_array)
    scores['cosine_bit'] = cosine_similarity_bit(url1, url2)
    return json.dumps(scores, indent=4)


demo = gr.Interface(fn=main, inputs=['text', 'text'], outputs='text')
demo.launch()