import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms as tfms
import torchvision.models as models

from PIL import Image
import numpy as np
from diffusers import LMSDiscreteScheduler, DiffusionPipeline

import random
import os
import subprocess

from matplotlib import pyplot as plt
from pathlib import Path
from torch import autocast
from tqdm.auto import tqdm

# Set device
torch_device = "cuda" if torch.cuda.is_available() else "cpu"

# Load a pre-trained VGG model (you can use other models as well)
vgg_model = models.vgg16(pretrained=True).features
vgg_model = vgg_model.to(torch_device)

# Create a new model that extracts features from the chosen layers
feature_extractor = nn.Sequential()
for name, layer in vgg_model._modules.items():
    if name == '0':  # Stop at the 0th layer
        break
    feature_extractor.add_module(name, layer)
feature_extractor = feature_extractor.to(torch_device)

pretrained_model_name_or_path = "segmind/tiny-sd"
pipe = DiffusionPipeline.from_pretrained(
    pretrained_model_name_or_path,
    torch_dtype=torch.float32
).to(torch_device)

# The noise scheduler
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)

concept_dict={'anime_bg_v2':('sd-concepts-library/anime-background-style-v2','<anime-background-style-v2>',31),
              'birb':('sd-concepts-library/birb-style','<birb-style>',32),
              'depthmap':('sd-concepts-library/depthmap','<depthmap>',33),
              'gta5_artwork':('sd-concepts-library/gta5-artwork','<gta5_artwork>',34),
              'midjourney':('sd-concepts-library/midjourney-style','<midjourney-style>',35),
              'beetlejuice':('sd-concepts-library/beetlejuice-cartoon-style','<beetlejuice-cartoon>',36)}

cache_style_list = []

def transform_pattern_image(pattern_image):
    preprocess = tfms.Compose([
        tfms.Resize((320, 320)),
        tfms.ToTensor(),
    ])
    tfms_pattern_image = preprocess(pattern_image).unsqueeze(0)
    return tfms_pattern_image

def load_required_style(style):
    for concept, value in concept_dict.items():
        if style in concept:
            concept_key = value[1]
            concept_seed = value[2]
            if style not in cache_style_list:
                pipe.load_textual_inversion(value[0])
                cache_style_list.append(style)
            break
    return concept_key, concept_seed

def pil_to_latent(input_im):
    # Single image -> single latent in a batch (so size 1, 4, 64, 64)
    with torch.no_grad():
        latent = pipe.vae.encode(tfms.ToTensor()(input_im).unsqueeze(0).to(torch_device)*2-1) # Note scaling
    return 0.18215 * latent.latent_dist.sample()  # [1, 4, 64, 64]

def latents_to_pil(latents):
    # bath of latents -> list of images
    latents = (1 / 0.18215) * latents
    with torch.no_grad():
        image = pipe.vae.decode(latents).sample
    image = (image / 2 + 0.5).clamp(0, 1)
    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
    images = (image * 255).round().astype("uint8")
    pil_images = [Image.fromarray(image) for image in images]
    return pil_images

def perceptual_loss(images, pattern):
    """
    This function calculates the perceptual loss between the output image and the target image.

    Parameters:
    """
    criterion = nn.MSELoss()
    mse_loss = criterion(images, pattern)
    return mse_loss

#Generating image with the modified embeddings with pattern loss guidance and saving the images to steps/{concept} folder
def generate_with_embs_pattern_loss(prompt, concept_seed, tfm_pattern_image, num_inf_steps):
    height = 320                        # default height of Stable Diffusion
    width = 320                         # default width of Stable Diffusion
    num_inference_steps = num_inf_steps # Number of denoising steps
    guidance_scale = 8                  # Scale for classifier-free guidance
    generator = torch.manual_seed(concept_seed) # Seed generator to create the inital latent noise
    batch_size = 1
    pattern_loss_scale = 20

    text_input = pipe.tokenizer(prompt, padding="max_length", max_length=pipe.tokenizer.model_max_length, truncation=True, return_tensors="pt")
    input_ids = text_input.input_ids.to(torch_device)
    with torch.no_grad():
        text_embeddings = pipe.text_encoder(text_input.input_ids.to(torch_device))[0]

    max_length = text_input.input_ids.shape[-1]
    uncond_input = pipe.tokenizer([""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt")
    with torch.no_grad():
        uncond_embeddings = pipe.text_encoder(uncond_input.input_ids.to(torch_device))[0]
    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

    # Prep Scheduler
    scheduler.set_timesteps(num_inference_steps)

    # Prep latents
    latents = torch.randn((batch_size, pipe.unet.in_channels, height // 8, width // 8),
                           generator=generator,)
    latents = latents.to(torch_device)
    latents = latents * scheduler.init_noise_sigma

    # Loop
    for i, t in tqdm(enumerate(scheduler.timesteps)):
        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
        latent_model_input = torch.cat([latents] * 2)
        sigma = scheduler.sigmas[i]
        latent_model_input = scheduler.scale_model_input(latent_model_input, t)

        # predict the noise residual
        with torch.no_grad():
            noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]

        # perform CFG (Classifier Free Guidance)
        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
        #### ADDITIONAL GUIDANCE ###
        if (i%3 == 0):
            # Requires grad on the latents
            latents = latents.detach().requires_grad_()

            # Get the predicted x0:
            latents_x0 = latents - sigma * noise_pred
            # latents_x0 = scheduler.step(noise_pred, t, latents).pred_original_sample

            # Decode to image space
            denoised_images = pipe.vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5 # range (0, 1)
            # Calculate loss
            denoised_images_extr = feature_extractor(denoised_images)
            reference_img_extr = feature_extractor(tfm_pattern_image)
            loss = perceptual_loss(denoised_images_extr, reference_img_extr) * pattern_loss_scale
            # Get gradient
            cond_grad = torch.autograd.grad(loss, latents)[0]

            # Modify the latents based on this gradient
            latents = latents.detach() - cond_grad * sigma**2

        # Now step with scheduler. compute the previous noisy sample x_t -> x_t-1
        latents = scheduler.step(noise_pred, t, latents).prev_sample

    return latents

def generate_image(prompt, pattern_image, style, num_inf_steps):
    tfm_pattern_image = transform_pattern_image(pattern_image)  # Transform the pattern image to be fed to feature extractor
    tfm_pattern_image = tfm_pattern_image.to(torch_device)
    if style == "no-style":
        concept_seed = 40
        main_prompt = str(prompt)
    else:
        concept_key, concept_seed = load_required_style(style)
        main_prompt = f"{str(prompt)} in the style of {concept_key}"
    latents = generate_with_embs_pattern_loss(main_prompt, concept_seed, tfm_pattern_image, num_inf_steps)
    generated_image = latents_to_pil(latents)[0]
    return generated_image

def gradio_fn(prompt, pattern_image, style, num_inf_steps):
    output_pil_image = generate_image(prompt, pattern_image, style, num_inf_steps)
    return output_pil_image

demo = gr.Interface(fn=gradio_fn,
                    inputs=[gr.Textbox(info="Example prompt: 'A toddler gazing at sky'"),
                            gr.Image(type="pil", height=224, width=224),
                            gr.Radio(["anime","birb","depthmap","gta5","midjourney","beetlejuice","no-style"], label="Style",
                                     info="Choose the style in which image to be made"),
                            gr.Slider(50, 100, value=50, label="Num_inference_steps", info="Choose between 50 & 100")],
                    outputs=gr.Image(height=320, width=320),
                    title="ImageAlchemy using Stable Diffusion",
                    description="- Stable Diffusion model that generates single image to fit \
                                  (a) given text prompt (b) given reference image and (c) selected style.")

demo.launch(share=True)