import gradio as gr
import torch
import torchvision.transforms as transforms
from sentence_transformers import SentenceTransformer, util
import json
import os
import matplotlib.pyplot as plt
import random
import torch
import torch.nn as nn
from generatorModel import Generator

noise_dim = 16
embed_dim = 384
embed_out_dim = 256
device = 'cpu'

generator = Generator(channels=3, embed_dim=embed_dim, noise_dim=noise_dim, embed_out_dim=embed_out_dim).to(device)


# Path to your .pth file
gen_weight = 'generator_20240421_3.pth'

# Load the weights
weights_gen = torch.load(gen_weight, map_location=torch.device(device))

# Apply the weights to your model
generator.load_state_dict(weights_gen)

model = SentenceTransformer('sentence-transformers/all-MiniLM-L12-v2')

def load_embedding(model):
    # Load your model and other components here
    with open(os.path.join("descriptions.json"), 'r') as file:
        dataset = json.load(file)
    
    classes = [e["text"] for e in dataset]
    embeddings_list = {cls: model.encode(cls, convert_to_tensor=True) for cls in classes}
    return embeddings_list

def generate_image(caption):
    embeddings = load_embedding(model)
    noise_dim = 16
    results = [(util.pytorch_cos_sim(model.encode(caption, convert_to_tensor=True), embeddings[cls]).item(), cls) for cls in classes]
    sorted_results = sorted(results, key=lambda x: x[0], reverse=True)[:5]

    threshold = 0.40
    coeff = 0.89

    if sorted_results[0][0] <= threshold:
        caption = sorted_results[0][1]
        results = [(util.pytorch_cos_sim(model.encode(caption, convert_to_tensor=True), embeddings[cls]).item(), cls) for cls in classes]
        sorted_results = sorted(results, key=lambda x: x[0], reverse=True)[:5]

    if sorted_results[0][0] >= 0.99:
        coeff = 0.85

    last_score = sorted_results[0][0]
    filtered_results = []
    for score, cls in sorted_results:
        if score >= last_score * coeff:
            filtered_results.append((score, cls))
            last_score = score
        else:
            break

    items = [cls for score, cls in filtered_results]
    probabilities = [score for score, cls in filtered_results]
    sampled_item = random.choices(items, weights=probabilities, k=1)[0]

    noise = torch.randn(1, noise_dim, 1, 1, device=device)  # Adjust noise_dim if different
    fake_images = generator(noise, embeddings[sampled_item].unsqueeze(0).unsqueeze(0))
    img = fake_images.squeeze(0).permute(1, 2, 0).cpu().detach().numpy()
    img = (img - img.min()) / (img.max() - img.min())

    return img

iface = gr.Interface(fn=generate_image,
                     inputs=gr.Textbox(lines=2, placeholder="Enter Caption Here..."),
                     outputs=gr.Image(type="numpy"),
                     title="CUHK Shenzhen Building Text-to-Image Generation",
                     description="Enter a caption of some specific building in CUHK-Shenzhen to generate an image..")

iface.launch()