import PIL
import torch
import numpy as np
from PIL import Image
from tqdm import tqdm
import torch.nn.functional as F
import torchvision.transforms as T
from diffusers import LMSDiscreteScheduler, DiffusionPipeline

# configurations
torch_device        = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
height, width       = 512, 512
guidance_scale      = 8
loss_scale          = 200
num_inference_steps = 50


model_path = "CompVis/stable-diffusion-v1-4"
sd_pipeline = DiffusionPipeline.from_pretrained(
    model_path,
    low_cpu_mem_usage = True,
    torch_dtype=torch.float32
).to(torch_device)


sd_pipeline.load_textual_inversion("sd-concepts-library/illustration-style")
sd_pipeline.load_textual_inversion("sd-concepts-library/line-art")
sd_pipeline.load_textual_inversion("sd-concepts-library/hitokomoru-style-nao")
sd_pipeline.load_textual_inversion("sd-concepts-library/style-of-marc-allante")
sd_pipeline.load_textual_inversion("sd-concepts-library/midjourney-style")
sd_pipeline.load_textual_inversion("sd-concepts-library/hanfu-anime-style")
sd_pipeline.load_textual_inversion("sd-concepts-library/birb-style")


styles_mapping = {
    "Illustration Style": '<illustration-style>', "Line Art":'<line-art>',
    "Hitokomoru Style":'<hitokomoru-style-nao>', "Marc Allante": '<Marc_Allante>',
    "Midjourney":'<midjourney-style>', "Hanfu Anime": '<hanfu-anime-style>',
    "Birb Style": '<birb-style>'
}

# Define seeds for all the styles
seed_list = [11, 56, 110, 65, 5, 29, 47]

# Loss Function based on Edge Detection
def edge_detection(image):
    channels = image.shape[1]

    # Define the kernels for Edge Detection
    ed_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0)
    ed_y = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0)

    # Replicate the Edge detection kernels for each channel
    ed_x = ed_x.repeat(channels, 1, 1, 1).to(image.device)
    ed_y = ed_y.repeat(channels, 1, 1, 1).to(image.device)

    # ed_x = ed_x.to(torch.float16)
    # ed_y = ed_y.to(torch.float16)

    # Convolve the image with the Edge detection kernels
    conv_ed_x = F.conv2d(image, ed_x, padding=1, groups=channels)
    conv_ed_y = F.conv2d(image, ed_y, padding=1, groups=channels)

    # Combine the x and y gradients after convolution
    ed_value = torch.sqrt(conv_ed_x**2 + conv_ed_y**2)

    return ed_value

def edge_loss(image):
    ed_value = edge_detection(image)
    ed_capped = (ed_value > 0.5).to(torch.float32)
    return F.mse_loss(ed_value, ed_capped)

def compute_loss(original_image, loss_type):

    if loss_type == 'blue':
        # blue loss
        # [:,2] -> all images in batch, only the blue channel
        error = torch.abs(original_image[:,2] - 0.9).mean()
    elif loss_type == 'edge':
        # edge loss
        error = edge_loss(original_image)        
    elif loss_type == 'contrast':
        # RGB to Gray loss
        transformed_image = T.functional.adjust_contrast(original_image, contrast_factor = 2)
        error = torch.abs(transformed_image - original_image).mean()
    elif loss_type == 'brightness':
        # brightnesss loss
        transformed_image = T.functional.adjust_brightness(original_image, brightness_factor = 2)
        error = torch.abs(transformed_image - original_image).mean()
    elif loss_type == 'sharpness':
        # sharpness loss
        transformed_image = T.functional.adjust_sharpness(original_image, sharpness_factor = 2)
        error = torch.abs(transformed_image - original_image).mean()
    elif loss_type == 'saturation':
        # saturation loss
        transformed_image = T.functional.adjust_saturation(original_image, saturation_factor = 10)
        error = torch.abs(transformed_image - original_image).mean()
    else:
        print("error. Loss not defined")

    return error


def get_examples():
   examples = [
      ['A bird sitting on a tree', 'Midjourney', 'edge', 5],
      ['Cats fighting on the road', 'Marc Allante', 'brightness', 65],
      ['A mouse with the head of a puppy', 'Hitokomoru Style', 'contrast', 110],
      ['A woman with a smiling face in front of an Italian Pizza', 'Hanfu Anime', 'brightness', 29],
      ['A campfire (oil on canvas)', 'Birb Style', 'blue', 47],
   ]
   return(examples)


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


def show_image(prompt, concept, guidance_type):

  for idx, sd in enumerate(styles_mapping.keys()):
    if(sd == concept):
      break
  seed =  seed_list[idx]
  prompt = f"{prompt} in the style of {styles_mapping[sd]}"
  styled_image_without_loss = latents_to_pil(generate_image(seed, prompt, guidance_type, loss_flag=False))
  styled_image_with_loss = latents_to_pil(generate_image(seed, prompt, guidance_type, loss_flag=True))
  return([styled_image_without_loss, styled_image_with_loss])


def generate_image(seed, prompt, loss_type, loss_flag=False):

    generator           = torch.manual_seed(seed)
    batch_size          = 1

    # scheduler
    scheduler    = LMSDiscreteScheduler(beta_start = 0.00085, beta_end = 0.012, beta_schedule = "scaled_linear", num_train_timesteps = 1000)
    scheduler.set_timesteps(num_inference_steps)
    scheduler.timesteps = scheduler.timesteps.to(torch.float32)

    # text embeddings of the prompt
    text_input = sd_pipeline.tokenizer(prompt, padding='max_length', max_length = sd_pipeline.tokenizer.model_max_length, truncation= True, return_tensors="pt")
    input_ids = text_input.input_ids.to(torch_device)

    with torch.no_grad():
        text_embeddings = sd_pipeline.text_encoder(text_input.input_ids.to(torch_device))[0]

    max_length = text_input.input_ids.shape[-1]
    uncond_input = sd_pipeline.tokenizer(
          [""] * batch_size, padding="max_length", max_length= max_length, return_tensors="pt"
    )

    with torch.no_grad():
        uncond_embeddings = sd_pipeline.text_encoder(uncond_input.input_ids.to(torch_device))[0]

    text_embeddings = torch.cat([uncond_embeddings,text_embeddings]) # shape: 2,77,768

    # random latent
    latents = torch.randn(
        (batch_size, sd_pipeline.unet.config.in_channels, height// 8, width //8),
        generator = generator,
    ) .to(torch.float32)


    latents = latents.to(torch_device)
    latents = latents * scheduler.init_noise_sigma

    for i, t in tqdm(enumerate(scheduler.timesteps), total = len(scheduler.timesteps)):

        latent_model_input = torch.cat([latents] * 2)
        sigma = scheduler.sigmas[i]
        latent_model_input = scheduler.scale_model_input(latent_model_input, t)

        with torch.no_grad():
            noise_pred = sd_pipeline.unet(latent_model_input.to(torch.float32), t, encoder_hidden_states=text_embeddings)["sample"]

        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

        if loss_flag and i%5 == 0:

            latents = latents.detach().requires_grad_()
            # the following line alone does not work, it requires change to reduce step only once
            # hence commenting it out
            #latents_x0 = scheduler.step(noise_pred,t, latents).pred_original_sample 
            latents_x0 = latents - sigma * noise_pred

            # use vae to decode the image
            denoised_images = sd_pipeline.vae.decode((1/ 0.18215) * latents_x0).sample / 2 + 0.5 # range(0,1)

            loss = compute_loss(denoised_images, loss_type) * loss_scale
            #loss = loss.to(torch.float16)
            print(f"{i} loss {loss}")

            cond_grad = torch.autograd.grad(loss, latents)[0]
            latents = latents.detach() - cond_grad * sigma**2

        latents = scheduler.step(noise_pred,t, latents).prev_sample

    return latents