import torch
import torch.optim as optim
import torch.nn.functional as F
from torchvision import transforms, models
from torch.autograd import Variable

from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
from skimage.transform import resize

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

def im_convert(tensor):
  
    image = tensor.to("cpu").clone().detach()
    image = image.numpy().squeeze()
    image = image.transpose(1, 2, 0) # B G R
    # undo normalizations
    image = image * np.array((0.229, 0.224, 0.225)) + np.array((0.485, 0.456, 0.406)) 
    image = image.clip(0, 1)

    return image

def load_image(path, max_size = 400, shape = None):
    ''' Load and transform image and make sure the image <= 400
    pixels in the x-y dims. '''

    image = Image.open(path).convert('RGB')

    size = max_size if (max(image.size) > max_size) else max(image.size)

    if shape is not None:
        size = shape


    in_transform = transforms.Compose([
                        transforms.Resize(size),
                        transforms.ToTensor(),
                        transforms.Normalize((0.485, 0.456, 0.406),
                                            (0.229, 0.224, 0.225))
    ])

    # discard the transparent, alpha channel(:3) and add the batch dimension
    image = in_transform(image)[:3, :, :].unsqueeze(0)

    return image

def image_to_tensor(x):
    """
    Transforms np.array to torch.Tensor
        (W, H)       -> (1, 1, W, H)
        (W, H, C)    -> (1, C, W, H)
        (B, W, H, C) -> (B, C, W, H)
    """
    if x.ndim == 2:
        return torch.Tensor(x).unsqueeze(0).unsqueeze(0)
    if x.ndim == 3:
        return torch.Tensor(x.transpose(2, 0, 1)).unsqueeze(0)
    if x.ndim == 4:
        return torch.Tensor(x.transpose(0, 3, 1, 2))
    raise RuntimeError("np.array's ndim is out of range 2, 3 or 4.")
    
    
def extract_masks(segment):
    """
    Extracts the segmentation masks from the segmentated image.
    Allowed colors are:
        blue, green, black, white, red,
        yellow, grey, light_blue, purple.
    """
    extracted_colors = []

    # BLUE
    mask_r = segment[..., 0] < 0.1
    mask_g = segment[..., 1] < 0.1
    mask_b = segment[..., 2] > 0.9
    mask = mask_r & mask_g & mask_b
    extracted_colors.append(mask)

    # GREEN
    mask_r = segment[..., 0] < 0.1
    mask_g = segment[..., 1] > 0.9
    mask_b = segment[..., 2] < 0.1
    mask = mask_r & mask_g & mask_b
    extracted_colors.append(mask)

    # BLACK
    mask_r = segment[..., 0] < 0.1
    mask_g = segment[..., 1] < 0.1
    mask_b = segment[..., 2] < 0.1
    mask = mask_r & mask_g & mask_b
    extracted_colors.append(mask)

    # WHITE
    mask_r = segment[..., 0] > 0.9
    mask_g = segment[..., 1] > 0.9
    mask_b = segment[..., 2] > 0.9
    mask = mask_r & mask_g & mask_b
    extracted_colors.append(mask)

    # RED
    mask_r = segment[..., 0] > 0.9
    mask_g = segment[..., 1] < 0.1
    mask_b = segment[..., 2] < 0.1
    mask = mask_r & mask_g & mask_b
    extracted_colors.append(mask)

    # YELLOW
    mask_r = segment[..., 0] > 0.9
    mask_g = segment[..., 1] > 0.9
    mask_b = segment[..., 2] < 0.1
    mask = mask_r & mask_g & mask_b
    extracted_colors.append(mask)

    # GREY
    mask_r = (segment[..., 0] > 0.4) & (segment[..., 0] < 0.6)
    mask_g = (segment[..., 1] > 0.4) & (segment[..., 1] < 0.6)
    mask_b = (segment[..., 2] > 0.4) & (segment[..., 2] < 0.6)
    mask = mask_r & mask_g & mask_b
    extracted_colors.append(mask)

    # LIGHT_BLUE
    mask_r = segment[..., 0] < 0.1
    mask_g = segment[..., 1] > 0.9
    mask_b = segment[..., 2] > 0.9
    mask = mask_r & mask_g & mask_b
    extracted_colors.append(mask)

    # PURPLE
    mask_r = segment[..., 0] > 0.9
    mask_g = segment[..., 1] < 0.1
    mask_b = segment[..., 2] > 0.9
    mask = mask_r & mask_g & mask_b
    extracted_colors.append(mask)

    return extracted_colors


def get_all_masks(path):
    """
    Returns the segmentation masks from the segmentated image.
    """
    image = Image.open(path)
    np_image = np.array(image, dtype=np.float) / 255
    return extract_masks(np_image)


def is_nonzero(mask, thrs=0.01):
    """
    Checks segmentation mask is dense.
    """
    return np.sum(mask) / mask.size > thrs


def get_masks(path_style, path_content):
    """
    Returns the meaningful segmentation masks.
    Avoides "orphan semantic labels" problem.
    """
    masks_style = get_all_masks(path_style)
    masks_content = get_all_masks(path_content)

    non_zero_masks = [
        is_nonzero(mask_c) and is_nonzero(mask_s)
        for mask_c, mask_s in zip(masks_content, masks_style)
    ]

    masks_style = [mask for mask, cond in zip(masks_style, non_zero_masks) if cond]
    masks_content = [mask for mask, cond in zip(masks_content, non_zero_masks) if cond]

    return masks_style, masks_content


def resize_masks(masks_style, masks_content, size):
    """
    Resizes masks to given size.
    """
    resize_mask = lambda mask: resize(mask, size, mode="reflect")

    masks_style = [resize_mask(mask) for mask in masks_style]
    masks_content = [resize_mask(mask) for mask in masks_content]

    return masks_style, masks_content


def masks_to_tensor(masks_style, masks_content):
    """
    Transforms masks to torch.Tensor from np.array.
    """
    masks_style = [image_to_tensor(mask) for mask in masks_style]
    masks_content = [image_to_tensor(mask) for mask in masks_content]

    return masks_style, masks_content


def masks_loader(path_style, path_content, size):
    """
    Loads masks.
    """
    style_masks, content_masks = get_masks(path_style, path_content)
    style_masks, content_masks = resize_masks(style_masks, content_masks, size)
    style_masks, content_masks = masks_to_tensor(style_masks, content_masks)

    return style_masks, content_masks

def get_features(image, model, layers=None):
    """ Run an image forward through a model and get the features for 
          a set of layers. Default layers are for VGGNet matching Gatys et al (2016)
    """

    if layers is None:
        layers = {
            '0':'conv1_1',
            '5':'conv2_1',
            '10':'conv3_1',
            '19':'conv4_1',
            '21':'conv4_2',
            '28':'conv5_1'
        }

    features = {}
    x = image
    # model._modules is a dictionary holding each module in the model
    for name, layer in model._modules.items():
        x = layer(x)
          
        if name in layers:
            features[layers[name]] = x

    return features

def gram_matrix(tensor):
    """ Calculate the Gram Matrix of a given tensor 
        Gram Matrix: https://en.wikipedia.org/wiki/Gramian_matrix
    """

    # get the batch_size, depth, height, and width of the Tensor
    _, d, h, w = tensor.size()

    # reshape so we're multiplying the features for each channel
    tensor = tensor.view(d, h * w)

    # calculate the gram matrix
    gram = torch.mm(tensor, tensor.t())

    return gram

def PhotorealismRegularization(image, target):
    L = compute_laplacian(image)
    grad = L.dot(target.reshape(-1, 3))
    loss = np.sum(target.reshape(3, -1).dot(grad))
    return torch.tensor(loss), 2. * grad.reshape(*target.shape)

import io

style_path = io.BytesIO(uploaded['tar4.png'])
content_path = io.BytesIO(uploaded['in4.png'])

style_mask_path = io.BytesIO(uploaded['seg_tar4.png'])
content_mask_path = io.BytesIO(uploaded['seg_in4.png'])

vgg = models.vgg19(pretrained = True).features

for param in vgg.parameters():
    param.requires_grad_(False)

vgg.to(device)

# load content and style
content = load_image(content_path).to(device)
# resize style to match content
style = load_image(style_path, shape = content.shape[-2:]).to(device)

# load the masks
style_masks, content_masks = masks_loader(
    style_mask_path,
    content_mask_path,
    content.shape[-2:])

# get content and style features
content_features = get_features(content, vgg)
style_features = get_features(style, vgg)

# get the gram matrices
style_grams = {layer : gram_matrix(style_features[layer]) for layer in style_features}

# get the target image
target = content.clone().requires_grad_(True).to(device)

style_weights = {'conv1_1': 0.2,
                 'conv2_1': 0.2,
                 'conv3_1': 0.2,
                 'conv4_1': 0.2,
                 'conv5_1': 0.2}

content_weight = 1  # Alpha
style_weight = 1e2  # Gamma
photorealismRegularization_weight = 1e4 # lambda

# constants
content_layer = 'conv4_2'

# for displaying the target image, intermittently
show_every = 400

# iteration hyperparameters
optimizer = optim.Adam([target], lr=0.003)
steps = 2000  # decide how many iterations to update your image (5000)

for ii in range(1, steps+1):
    
    # get the features from your target image
    target_features = get_features(target, vgg)
                                   
    # content loss
    content_loss = torch.mean((target_features[content_layer] - content_features[content_layer])**2)
    
    # photorealistic regularization loss
    reg_loss, reg_grad = PhotorealismRegularization(content.detach().to("cpu").numpy().transpose(0, 2, 3, 1).squeeze().clip(0, 1), target.detach().to("cpu").numpy().transpose(0, 2, 3, 1).squeeze().clip(0, 1))
    
    reg_grad_tensor = image_to_tensor(reg_grad).to(device)
    
    #target.grad += reg_grad_tensor
    
    # augmented style loss
    #style_loss = augmented_style_loss(target, vgg, content_masks, style_masks)
    
    layers = {
        '0':'conv1_1',
        '5':'conv2_1',
        '10':'conv3_1',
        '19':'conv4_1',
        '28':'conv5_1'
    }

    style_loss = 0
    x = target.clone().to(device)
    
    # model._modules is a dictionary holding each module in the model
    for name, layer in vgg._modules.items():
        x = layer(x)
        if isinstance(layer, torch.nn.MaxPool2d):
            style_masks = [layer(mask) for mask in style_masks]
            content_masks = [layer(mask) for mask in content_masks]
          
        if name in layers:
            target_gram = [gram_matrix(x * Variable(mask.to(device), requires_grad = True)) for mask in content_masks]
            style_gram = [gram_matrix(x * Variable(mask.to(device), requires_grad = True)) for mask in style_masks]
            
            layer_loss = style_weights[layers[name]] * sum(F.mse_loss(t_gram, s_gram) for t_gram, s_gram in zip(target_gram, style_gram))
            
            target_feature = target_features[layers[name]]
            _, d, h, w = target_feature.shape
            
            layer_loss = layer_loss / (d * h * w)
            style_loss = style_loss + layer_loss

    
    # total loss
    total_loss = content_weight * content_loss + style_weight * style_loss + photorealismRegularization_weight * reg_loss

    # update target image
    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()

    # display intermediate images and print the loss
    if ii % show_every == 0:
        print('Total loss: ', total_loss.item())
        plt.imshow(im_convert(target))
        plt.show()