from PIL import Image
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, models

IMG_MEAN = [0.485, 0.456, 0.406]
IMG_STD = [0.229, 0.224, 0.225]
CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
CLIP_STD = [0.26862954, 0.26130258, 0.27577711]
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

LR = 5e-4
CROP_SIZE = 128
RESIZE = 224
NUM_EPOCHS = 200
NUM_CROPS = 64
PATCH_THRESHOLD = 0.7
L_TV = 2e-3
L_PATCH = 9000
L_DIR = 500
L_CONTENT = 150
IMG_SIZE = 512
# copied from openai clip
IMAGENET_TEMPLATES = [
    "a bad photo of a {}.",
    "a photo of many {}.",
    "a sculpture of a {}.",
    "a photo of the hard to see {}.",
    "a low resolution photo of the {}.",
    "a rendering of a {}.",
    "graffiti of a {}.",
    "a bad photo of the {}.",
    "a cropped photo of the {}.",
    "a tattoo of a {}.",
    "the embroidered {}.",
    "a photo of a hard to see {}.",
    "a bright photo of a {}.",
    "a photo of a clean {}.",
    "a photo of a dirty {}.",
    "a dark photo of the {}.",
    "a drawing of a {}.",
    "a photo of my {}.",
    "the plastic {}.",
    "a photo of the cool {}.",
    "a close-up photo of a {}.",
    "a black and white photo of the {}.",
    "a painting of the {}.",
    "a painting of a {}.",
    "a pixelated photo of the {}.",
    "a sculpture of the {}.",
    "a bright photo of the {}.",
    "a cropped photo of a {}.",
    "a plastic {}.",
    "a photo of the dirty {}.",
    "a jpeg corrupted photo of a {}.",
    "a blurry photo of the {}.",
    "a photo of the {}.",
    "a good photo of the {}.",
    "a rendering of the {}.",
    "a {} in a video game.",
    "a photo of one {}.",
    "a doodle of a {}.",
    "a close-up photo of the {}.",
    "a photo of a {}.",
    "the origami {}.",
    "the {} in a video game.",
    "a sketch of a {}.",
    "a doodle of the {}.",
    "a origami {}.",
    "a low resolution photo of a {}.",
    "the toy {}.",
    "a rendition of the {}.",
    "a photo of the clean {}.",
    "a photo of a large {}.",
    "a rendition of a {}.",
    "a photo of a nice {}.",
    "a photo of a weird {}.",
    "a blurry photo of a {}.",
    "a cartoon {}.",
    "art of a {}.",
    "a sketch of the {}.",
    "a embroidered {}.",
    "a pixelated photo of a {}.",
    "itap of the {}.",
    "a jpeg corrupted photo of the {}.",
    "a good photo of a {}.",
    "a plushie {}.",
    "a photo of the nice {}.",
    "a photo of the small {}.",
    "a photo of the weird {}.",
    "the cartoon {}.",
    "art of the {}.",
    "a drawing of the {}.",
    "a photo of the large {}.",
    "a black and white photo of a {}.",
    "the plushie {}.",
    "a dark photo of a {}.",
    "itap of a {}.",
    "graffiti of the {}.",
    "a toy {}.",
    "itap of my {}.",
    "a photo of a cool {}.",
    "a photo of a small {}.",
    "a tattoo of the {}.",
]

def get_mean(mean_dist):
    mean = torch.tensor(mean_dist).to(DEVICE)
    return mean.view(1, -1, 1, 1)

def get_std(std_dist):
    std = torch.tensor(std_dist).to(DEVICE)
    return std.view(1, -1, 1, 1)

def normalize(data):
    mean = get_mean(IMG_MEAN)
    std = get_std(IMG_STD)

    norm_data = (data - mean) / std
    return norm_data

def clip_normalize(data):
    resized = nn.functional.interpolate(data, size=RESIZE, mode='bicubic')
    mean = get_mean(CLIP_MEAN)
    std = get_std(CLIP_STD)
    norm_data = (resized - mean) / std
    return norm_data

def load_image(image):
    image = image.resize((IMG_SIZE, IMG_SIZE))
    transform = transforms.Compose([transforms.ToTensor()])
    return transform(image)[:3, :, :].unsqueeze(0)

def get_features(image, vgg19):
    # uses vgg19 model to extract content features
    layers = {'0': 'conv1_1',
              '5': 'conv2_1',
              '10': 'conv3_1',
              '19': 'conv4_1',
              '21': 'conv4_2',
              '28': 'conv5_1',
              '31': 'conv5_2'
              }
    features = {}
    x = image
    for name, layer in vgg19._modules.items():
        x = layer(x)
        if name in layers:
            features[layers[name]] = x

    return features

def prompt_ensemble(prompt):
    return [template.format(prompt) for template in IMAGENET_TEMPLATES]

def get_image_prior_losses(target):
    diff1 = target[:, :, :, :-1] - target[:, :, :, 1:]
    diff2 = target[:, :, :-1, :] - target[:, :, 1:, :]
    diff3 = target[:, :, 1:, :-1] - target[:, :, :-1, 1:]
    diff4 = target[:, :, :-1, :-1] - target[:, :, 1:, 1:]

    loss_var_l2 = torch.norm(diff1) + torch.norm(diff2) + torch.norm(diff3) + torch.norm(diff4)
    
    return loss_var_l2