import pydiffvg
import torch
import torchvision
from PIL import Image
import numpy as np

# Use GPU if available
pydiffvg.set_use_gpu(torch.cuda.is_available())

def inv_exp(a,x,xpow=1):
    return pow(a,pow(1.-x,xpow))

import math
import numbers
import torch
from torch import nn
from torch.nn import functional as F

import visdom

class GaussianSmoothing(nn.Module):
    """
    Apply gaussian smoothing on a
    1d, 2d or 3d tensor. Filtering is performed seperately for each channel
    in the input using a depthwise convolution.
    Arguments:
        channels (int, sequence): Number of channels of the input tensors. Output will
            have this number of channels as well.
        kernel_size (int, sequence): Size of the gaussian kernel.
        sigma (float, sequence): Standard deviation of the gaussian kernel.
        dim (int, optional): The number of dimensions of the data.
            Default value is 2 (spatial).
    """
    def __init__(self, channels, kernel_size, sigma, dim=2):
        super(GaussianSmoothing, self).__init__()
        if isinstance(kernel_size, numbers.Number):
            kernel_size = [kernel_size] * dim
        if isinstance(sigma, numbers.Number):
            sigma = [sigma] * dim

        # The gaussian kernel is the product of the
        # gaussian function of each dimension.
        kernel = 1
        meshgrids = torch.meshgrid(
            [
                torch.arange(size, dtype=torch.float32)
                for size in kernel_size
            ]
        )
        for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
            mean = (size - 1) / 2
            kernel *= 1 / (std * math.sqrt(2 * math.pi)) * \
                      torch.exp(-((mgrid - mean) / std) ** 2 / 2)

        # Make sure sum of values in gaussian kernel equals 1.
        kernel = kernel / torch.sum(kernel)

        # Reshape to depthwise convolutional weight
        kernel = kernel.view(1, 1, *kernel.size())
        kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))

        self.register_buffer('weight', kernel)
        self.groups = channels

        if dim == 1:
            self.conv = F.conv1d
        elif dim == 2:
            self.conv = F.conv2d
        elif dim == 3:
            self.conv = F.conv3d
        else:
            raise RuntimeError(
                'Only 1, 2 and 3 dimensions are supported. Received {}.'.format(dim)
            )

    def forward(self, input):
        """
        Apply gaussian filter to input.
        Arguments:
            input (torch.Tensor): Input to apply gaussian filter on.
        Returns:
            filtered (torch.Tensor): Filtered output.
        """
        return self.conv(input, weight=self.weight, groups=self.groups)

vis=visdom.Visdom(port=8080)

smoothing = GaussianSmoothing(4, 5, 1)

settings=pydiffvg.SvgOptimizationSettings()
settings.global_override(["optimize_color"],False)
settings.global_override(["optimize_alpha"],False)
settings.global_override(["gradients","optimize_color"],False)
settings.global_override(["gradients","optimize_alpha"],False)
settings.global_override(["gradients","optimize_stops"],False)
settings.global_override(["gradients","optimize_location"],False)
settings.global_override(["optimizer"],"Adam")
settings.global_override(["paths","optimize_points"],False)
settings.global_override(["transforms","transform_lr"],1e-2)
settings.undefault("linearGradient3152")
settings.retrieve("linearGradient3152")[0]["transforms"]["optimize_transforms"]=False

#optim=pydiffvg.OptimizableSvg("note_small.svg",settings,verbose=True)
optim=pydiffvg.OptimizableSvg("heart_green.svg",settings,verbose=True)

#img=torchvision.transforms.ToTensor()(Image.open("note_transformed.png")).permute(1,2,0)
img=torchvision.transforms.ToTensor()(Image.open("heart_green_90.png")).permute(1,2,0)

name="heart_green_90"

pydiffvg.imwrite(img.cpu(), 'results/simple_transform_svg/target.png')
target = img.clone().detach().requires_grad_(False)

img=optim.render()
pydiffvg.imwrite(img.cpu(), 'results/simple_transform_svg/init.png')

def smooth(input, kernel):
    input=torch.nn.functional.pad(input.permute(2,0,1).unsqueeze(0), (2, 2, 2, 2), mode='reflect')
    output=kernel(input)
    return output

def printimg(optim):
    img=optim.render()
    comp = img.clone().detach()
    bg = torch.tensor([[[1., 1., 1.]]])
    comprgb = comp[:, :, 0:3]
    compalpha = comp[:, :, 3].unsqueeze(2)
    comp = comprgb * compalpha \
           + bg * (1 - compalpha)
    return comp

def comp_loss_and_grad(img, tgt, it, sz):
    dif=img-tgt

    loss=dif.pow(2).mean()

    dif=dif.detach()

    cdif=dif.clone().abs()
    cdif[:,:,3]=1.

    resdif=torch.nn.functional.interpolate(cdif.permute(2,0,1).unsqueeze(0),sz,mode='bilinear').squeeze().permute(1,2,0).abs()
    pydiffvg.imwrite(resdif[:,:,0:4], 'results/simple_transform_svg/dif_{:04}.png'.format(it))

    dif=dif.numpy()
    padded=np.pad(dif,[(1,1),(1,1),(0,0)],mode='edge')
    #print(padded[:-2,:,:].shape)
    grad_x=(padded[:-2,:,:]-padded[2:,:,:])[:,1:-1,:]
    grad_y=(padded[:,:-2,:]-padded[:,2:,:])[1:-1,:,:]

    resshape=dif.shape
    resshape=(resshape[0],resshape[1],2)
    res=np.zeros(resshape)

    for x in range(resshape[0]):
        for y in range(resshape[1]):
            A=np.concatenate((grad_x[x,y,:][:,np.newaxis],grad_y[x,y,:][:,np.newaxis]),axis=1)
            b=-dif[x,y,:]
            v=np.linalg.lstsq(np.dot(A.T,A),np.dot(A.T,b))
            res[x,y,:]=v[0]

    return loss, res

import colorsys
def print_gradimg(gradimg,it,shape=None):
    out=torch.zeros((gradimg.shape[0],gradimg.shape[1],3),requires_grad=False,dtype=torch.float32)
    for x in range(gradimg.shape[0]):
        for y in range(gradimg.shape[1]):
            h=math.atan2(gradimg[x,y,1],gradimg[x,y,0])
            s=math.tanh(np.linalg.norm(gradimg[x,y,:]))
            v=1.
            vec=(gradimg[x,y,:].clip(min=-1,max=1)/2)+.5
            #out[x,y,:]=torch.tensor(colorsys.hsv_to_rgb(h,s,v),dtype=torch.float32)
            out[x,y,:]=torch.tensor([vec[0],vec[1],0])

    if shape is not None:
        out=torch.nn.functional.interpolate(out.permute(2,0,1).unsqueeze(0),shape,mode='bilinear').squeeze().permute(1,2,0)
    pydiffvg.imwrite(out.cpu(), 'results/simple_transform_svg/grad_{:04}.png'.format(it))

# Run 150 Adam iterations.
for t in range(1000):
    print('iteration:', t)
    optim.zero_grad()
    with open('results/simple_transform_svg/viter_{:04}.svg'.format(t),"w") as f:
        f.write(optim.write_xml())
    scale=inv_exp(1/16,math.pow(t/1000,1),0.5)
    #print(scale)
    #img = optim.render(seed=t+1,scale=scale)
    img = optim.render(seed=t + 1, scale=None)
    vis.line(torch.tensor([img.shape[0]]), X=torch.tensor([t]), win=name + " size", update="append",
             opts={"title": name + " size"})
    #print(img.shape)
    #img = optim.render(seed=t + 1)

    ptgt=target.permute(2,0,1).unsqueeze(0)
    sz=img.shape[0:2]
    restgt=torch.nn.functional.interpolate(ptgt,size=sz,mode='bilinear').squeeze().permute(1,2,0)

    # Compute the loss function. Here it is L2.
    #loss = (smooth(img,smoothing) - smooth(restgt,smoothing)).pow(2).mean()
    #loss = (img - restgt).pow(2).mean()
    #loss=(img-target).pow(2).mean()
    loss,gradimg=comp_loss_and_grad(img, restgt,t,target.shape[0:2])
    print_gradimg(gradimg,t,target.shape[0:2])
    print('loss:', loss.item())
    vis.line(loss.unsqueeze(0), X=torch.tensor([t]), win=name+" loss", update="append",
             opts={"title": name + " loss"})

    # Backpropagate the gradients.
    loss.backward()

    # Take a gradient descent step.
    optim.step()

    # Save the intermediate render.
    comp=printimg(optim)
    pydiffvg.imwrite(comp.cpu(), 'results/simple_transform_svg/iter_{:04}.png'.format(t))


# Render the final result.

img = optim.render()
# Save the images and differences.
pydiffvg.imwrite(img.cpu(), 'results/simple_transform_svg/final.png')
with open('results/simple_transform_svg/final.svg', "w") as f:
    f.write(optim.write_xml())

# Convert the intermediate renderings to a video.
from subprocess import call
call(["ffmpeg", "-framerate", "24", "-i",
    "results/simple_transform_svg/iter_%04d.png", "-vb", "20M",
    "results/simple_transform_svg/out.mp4"])

call(["ffmpeg", "-framerate", "24", "-i",
    "results/simple_transform_svg/grad_%04d.png", "-vb", "20M",
    "results/simple_transform_svg/out_grad.mp4"])