# AUTOGENERATED! DO NOT EDIT! File to edit: ../62_Generators_and_Losses.ipynb.

# %% auto 0
__all__ = ['PixelGenerator', 'MSELossToTarget', 'ImStackGenerator', 'calc_vgg_features', 'ContentLossToTarget',
           'OTStyleLossToTarget', 'GramStyleLossToTarget', 'VincentStyleLossToTarget', 'SirenGenerator', 'DotGenerator',
           'MakeCutouts', 'spherical_dist_loss', 'set_requires_grad', 'clip_loss_embeddings', 'CLIPLossToTargets',
           'optimise']

# %% ../62_Generators_and_Losses.ipynb 2
from .utils import *

# %% ../62_Generators_and_Losses.ipynb 7
import torch
from torch import nn
import torch.nn.functional as F
from torchvision import transforms as T
from matplotlib import pyplot as plt
from imstack.core import ImStack
from siren_pytorch import Sine, Siren, SirenNet, SirenWrapper
import open_clip
from PIL import Image
import numpy as np
from tqdm.notebook import tqdm
import torchvision.models as models

# %% ../62_Generators_and_Losses.ipynb 10
class PixelGenerator(nn.Module):
    """A tensor of shape [1, 3, size, size] to represent an image"""
    def __init__(self, size=256, init_image=None):
        super(PixelGenerator, self).__init__()
        self.im_array = nn.Parameter(torch.rand(1, 3, size, size))
        if init_image:
            self.im_array = nn.Parameter(pil_to_tensor(init_image.convert('RGB').resize((size, size))))

    def forward(self):
        return self.im_array

# %% ../62_Generators_and_Losses.ipynb 16
class MSELossToTarget(nn.Module):
    """ MSE between input and target, resizing if needed"""
    def __init__(self, target, size=256):
        super(MSELossToTarget, self).__init__()
        self.resize = T.Resize(size)
        self.target = self.resize(target) # resize target image to size

    def forward(self, input):
        input = self.resize(input) # set size (assumes square images)
        squared_error = (self.target - input)**2
        return squared_error.mean() # MSE

# %% ../62_Generators_and_Losses.ipynb 28
class ImStackGenerator(nn.Module):
    """An imstack to represent the image"""
    def __init__(self, size=256, n_layers=4, base_size=16, 
               layer_decay=0.7, init_image=None, scale=2):
        super(ImStackGenerator, self).__init__()
        self.imstack = ImStack(n_layers=n_layers, out_size=size, 
                               base_size=base_size, init_image=init_image, 
                               decay=layer_decay, scale=scale)

    def parameters(self): # How to access the learnable parameters
        return self.imstack.layers

    def forward(self):
        return self.imstack()

# %% ../62_Generators_and_Losses.ipynb 36
# Extracting features from an image using this pretrained model:
def calc_vgg_features(vgg16, imgs, use_layers=[1, 6, 11, 18, 25]):
  mean = torch.tensor([0.485, 0.456, 0.406])[:,None,None].to(device)
  std = torch.tensor([0.229, 0.224, 0.225])[:,None,None].to(device)
  x = (imgs-mean) / std
  b, c, h, w = x.shape
  features = [x.reshape(b, c, h*w)] # This reshape is for convenience later
  for i, layer in enumerate(vgg16[:max(use_layers)+1]):
    x = layer(x)
    if i in use_layers:
      b, c, h, w = x.shape
      features.append(x.reshape(b, c, h*w))
  return features

# %% ../62_Generators_and_Losses.ipynb 41
class ContentLossToTarget(nn.Module):
    """ Perceptual loss between input and target, resizing if needed, based on vgg16"""
    def __init__(self, target, vgg16=None, size=128, content_layers = [14, 19]):
        super(ContentLossToTarget, self).__init__()
        self.resize = T.Resize(size)
        self.target = self.resize(target) # resize target image to size
        self.content_layers = content_layers
        self.vgg16 = vgg16
        if self.vgg16 == None:
            vgg16 = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).to(device)
            vgg16.eval()
            self.vgg16 = vgg16.features
        
        with torch.no_grad():
            self.target_features = calc_vgg_features(self.vgg16, self.target, use_layers = self.content_layers)

    def forward(self, input):
        input = self.resize(input) # set size (assumes square images)
        input_features = calc_vgg_features(self.vgg16, input, use_layers = self.content_layers)
        l = 0
        # Run through all features and take l1 loss (mean error) between them
        for im_features, target_features in zip(input_features, self.target_features):
            l += nn.L1Loss()(im_features, target_features)
        return l/len(input_features)

# %% ../62_Generators_and_Losses.ipynb 46
class OTStyleLossToTarget(nn.Module):
  """ Optimal Transport Loss for style comparison"""
  def __init__(self, target, vgg16=None, size=128, style_layers = [1, 6, 11, 18, 25], scale_factor=1e-5):
    super(OTStyleLossToTarget, self).__init__()
    self.resize = T.Resize(size)
    self.target = self.resize(target) # resize target image to size
    self.style_layers = style_layers
    self.scale_factor = scale_factor # Defaults tend to be very large, we scale to make them easier to work with
    self.vgg16 = vgg16
    if self.vgg16 == None:
        vgg16 = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).to(device)
        vgg16.eval()
        self.vgg16 = vgg16.features
    with torch.no_grad():
      self.target_features = calc_vgg_features(self.vgg16, self.target, use_layers = self.style_layers)

  def project_sort(self, x, proj):
    return torch.einsum('bcn,cp->bpn', x, proj).sort()[0]

  def ot_loss(self, source, target, proj_n=32):
    ch, n = source.shape[-2:]
    projs = F.normalize(torch.randn(ch, proj_n).to(device), dim=0)
    source_proj = self.project_sort(source, projs)
    target_proj = self.project_sort(target, projs)
    target_interp = F.interpolate(target_proj, n, mode='nearest')
    return (source_proj-target_interp).square().sum()

  def forward(self, input):
    input = self.resize(input) # set size (assumes square images)
    input_features = calc_vgg_features(self.vgg16, input, use_layers = self.style_layers)
    l = 0
    # Run through all features and take l1 loss (mean error) between them
    return sum(self.ot_loss(x, y) for x, y in zip(input_features, self.target_features)) * self.scale_factor

# %% ../62_Generators_and_Losses.ipynb 48
class GramStyleLossToTarget(nn.Module):
  """ Gram matrix based style loss"""
  def __init__(self, target, vgg16=None, size=128, style_layers = [1, 6, 11, 18, 25], scale_factor=0.1):
    super(GramStyleLossToTarget, self).__init__()
    self.resize = T.Resize(size)
    self.target = self.resize(target) # resize target image to size
    self.style_layers = style_layers
    self.scale_factor = scale_factor # Defaults tend to be very large, we scale to make them easier to work with
    self.vgg16 = vgg16
    if self.vgg16 == None:
        vgg16 = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).to(device)
        vgg16.eval()
        self.vgg16 = vgg16.features
    with torch.no_grad():
      self.target_features = calc_vgg_features(self.vgg16, self.target, use_layers = self.style_layers)

  def gram(self, x):
    b,c,hxw = x.size();
    h, w = int(hxw**0.5), int(hxw**0.5)
    x = x.view(b*c, -1);
    return torch.mm(x, x.t())/(h*w) # /(h*w) normalizes by size to avoid over-focusing on early layers


  def forward(self, input):
    input = self.resize(input) # set size (assumes square images)
    input_features = calc_vgg_features(self.vgg16, input, use_layers = self.style_layers)
    loss = 0
    for im_features, target_features in zip(input_features, self.target_features):
      im_gram = self.gram(im_features)
      target_gram = self.gram(target_features)
      loss += F.mse_loss(im_gram, target_gram);
    return (loss/len(input_features)) * self.scale_factor

# %% ../62_Generators_and_Losses.ipynb 50
class VincentStyleLossToTarget(nn.Module):
  """ Attempting to copy  """
  def __init__(self, target, vgg16=None, size=128, style_layers = [1, 6, 11, 18, 25], scale_factor=1e-5):
    super(VincentStyleLossToTarget, self).__init__()
    self.resize = T.Resize(size)
    self.target = self.resize(target) # resize target image to size
    self.style_layers = style_layers
    self.scale_factor = scale_factor # Defaults tend to be very large, we scale to make them easier to work with
    self.vgg16 = vgg16
    if self.vgg16 == None:
        vgg16 = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).to(device)
        vgg16.eval()
        self.vgg16 = vgg16.features
    with torch.no_grad():
      self.target_features = calc_vgg_features(self.vgg16, self.target, use_layers = self.style_layers)

  def calc_2_moments(self, x):
    b, c, n = x.shape
    mu = x.mean(dim=-1, keepdim=True) # b, c, 1
    cov = torch.matmul(x-mu, torch.transpose(x-mu, -1, -2))
    return mu, cov

  def matrix_diag(self, diagonal):
    # Try to copy tf.diag
    N = diagonal.shape[-1]
    shape = diagonal.shape[:-1] + (N, N)
    device, dtype = diagonal.device, diagonal.dtype
    result = torch.zeros(shape, dtype=dtype, device=device)
    indices = torch.arange(result.numel(), device=device).reshape(shape)
    indices = indices.diagonal(dim1=-2, dim2=-1)
    result.view(-1)[indices] = diagonal
    return result

  def l2wass_dist(self, mean_stl, cov_stl, mean_synth, cov_synth):
        
    # Calculate tr_cov and root_cov from mean_stl and cov_stl
    eigvals,eigvects = torch.linalg.eigh(cov_stl) # eig returns complex tensors, I think eigh matches tf self_adjoint_eig
    eigroot_mat = self.matrix_diag(torch.sqrt(eigvals.clip(0)))
    root_cov_stl = torch.matmul(torch.matmul(eigvects, eigroot_mat),torch.transpose(eigvects, -1, -2))
    tr_cov_stl = torch.sum(eigvals.clip(0), dim=1, keepdim=True)

    # eigvals_synth,eigvects_synth = torch.linalg.eigh(cov_synth) # Replaced withtorch.linalg.eigvalsh
    tr_cov_synth = torch.sum(torch.linalg.eigvalsh(cov_synth).clip(0), dim=1, keepdim=True)
    mean_diff_squared = torch.mean((mean_synth - mean_stl)**2)
    cov_prod = torch.matmul(torch.matmul(root_cov_stl,cov_synth),root_cov_stl)
    var_overlap = torch.sum(torch.sqrt(torch.linalg.eigvalsh(cov_prod).clip(0.1)), dim=1, keepdim=True) # .clip(0) meant errors getting eigvals 
    dist = mean_diff_squared+tr_cov_stl+tr_cov_synth-2*var_overlap
    return dist

  def forward(self, input):
    input = self.resize(input) # set size (assumes square images)
    input_features = calc_vgg_features(self.vgg16, input, use_layers = self.style_layers) # get features
    l = 0
    for x, y in zip(input_features, self.target_features):
      mean_synth, cov_synth = self.calc_2_moments(x) # input mean and cov
      mean_stl, cov_stl = self.calc_2_moments(y) # target mean and cov
      l += self.l2wass_dist(mean_stl, cov_stl, mean_synth, cov_synth) 
    return l * self.scale_factor

# %% ../62_Generators_and_Losses.ipynb 53
class SirenGenerator(nn.Module):
  """An SIREN network to represent the image"""
  def __init__(self, size=256, dim_hidden = 64, 
               num_layers = 5, w0_initial = 30.): 
    super(SirenGenerator, self).__init__()
    self.net = SirenNet(
        dim_in = 2,                        # input dimension, ex. 2d coordinates
        dim_hidden = dim_hidden,           # hidden dimension
        dim_out = 3,                       # output dimension, ex. rgb value
        num_layers = num_layers,           # number of layers
        final_activation = nn.Sigmoid(),   # activation of final layer (nn.Identity() for direct output)
        w0_initial = w0_initial            # different signals may require different omega_0 in the first layer - this is a hyperparameter
    )
    
    self.wrapper = SirenWrapper(
        self.net,
        image_width = size,
        image_height = size
    )
  
  def parameters(self): # How to access the learnable parameters
    return self.net.parameters()

  def forward(self):
    return self.wrapper()

# %% ../62_Generators_and_Losses.ipynb 60
class DotGenerator(nn.Module):
  def __init__(self, n_points=100, size=256, device='cpu'):
    super(DotGenerator, self).__init__()

    # Set some attributes
    w, h = size, size # You can take both as args for non-square images
    self.n_points = n_points
    self.w = w
    self.h = h
    self.sig = nn.Sigmoid()
    self.device=device

    # A grid of x, y locations (from 0 to 1):
    gridy, gridx = torch.meshgrid(1-torch.arange(h)/h, torch.arange(w)/w, indexing='ij')
    self.grid = torch.stack([gridx, gridy]).to(device) # shape: (2, w, h)

    # The parameters for the points:
    self.coords = nn.Parameter(torch.rand((2, n_points))) # X and Y positions
    self.sizes = nn.Parameter(torch.rand(n_points)*10) # << Tweak for bigger or smaller dots
    self.falloffs = nn.Parameter((torch.rand(n_points)*100+0.2*100)) # << Tweak 0.2 and 100 to change look
    self.colors = nn.Parameter(torch.rand(3, n_points)) # RGB colours
    self.intensities = nn.Parameter(torch.rand(n_points)*3) # << Change 3 to change max brightness

    # The 'image' we'll draw to - RGB, same size as grid (w x h)
    self.canvas = torch.zeros_like(self.grid[0].expand(3, -1, -1)).to(device)

  def to_device(self, device):
    self.to(device)
    self.grid = self.grid.to(device)
    

  def forward(self):
    self.canvas = torch.zeros_like(self.canvas).to(self.device) # Zero out our canvas
    for i in range(self.n_points): # Iterating over points (not ideal!!)
      # Point attributes:
      r, g, b = self.colors[:,i]
      intensity = self.intensities[i]
      px, py = self.coords[:,i]
      size = self.sizes[i]
      falloff = self.falloffs[i]

      # Distance from each pixel to this point:
      dx, dy = torch.abs(self.grid[0]-px), torch.abs(self.grid[1]-py)
      dist = (dx**2 + dy**2)**0.5

      # Draw the point with intensity dropping as distance increases
      sig_dist = 1-self.sig(dist*falloff - size)
      self.canvas[0] += r*sig_dist*intensity
      self.canvas[1] += g*sig_dist*intensity
      self.canvas[2] += b*sig_dist*intensity

    return self.canvas.unsqueeze(0)/(0.5*self.n_points**0.5) # Scale the final result (this is approximate) 

# %% ../62_Generators_and_Losses.ipynb 65
class MakeCutouts(nn.Module):
    def __init__(self, cut_size, cut_power=1.0): # TODO add augs
        super().__init__()

        self.cut_size = cut_size
        self.cut_power = cut_power

    def forward(self, pixel_values, num_cutouts):
        sideY, sideX = pixel_values.shape[2:4]
        max_size = min(sideX, sideY)
        min_size = min(sideX, sideY, self.cut_size)
        cutouts = []
        for _ in range(num_cutouts):
            size = int(torch.rand([]) ** self.cut_power * (max_size - min_size) + min_size)
            offsetx = torch.randint(0, sideX - size + 1, ())
            offsety = torch.randint(0, sideY - size + 1, ())
            cutout = pixel_values[:, :, offsety : offsety + size, offsetx : offsetx + size]
            cutouts.append(F.adaptive_avg_pool2d(cutout, self.cut_size))
        return torch.cat(cutouts)


def spherical_dist_loss(x, y):
    x = F.normalize(x, dim=-1)
    y = F.normalize(y, dim=-1)
    return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)


def set_requires_grad(model, value):
    for param in model.parameters():
        param.requires_grad = value

def clip_loss_embeddings(image_embed, target_embeds, weights=None):
  input_normed = F.normalize(image_embed.unsqueeze(1), dim=-1)
  embed_normed = F.normalize(target_embeds.unsqueeze(0), dim=-1)
  dists = input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2) # Squared Great Circle Distance
  return dists.mean()

# %% ../62_Generators_and_Losses.ipynb 72
class CLIPLossToTargets(nn.Module):
  """ Gram matrix based style loss"""
  def __init__(self, text_prompts=[], image_prompts=[], clip_model=None, n_cuts=16, device=None):
    super(CLIPLossToTargets, self).__init__()
    
    if device==None:
        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    self.device=device
    
    # Set up CLIP model
    self.clip_model = clip_model
    if self.clip_model == None:
        clip_model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32-quickgelu', pretrained='laion400m_e32')
        self.clip_model = clip_model
    self.preprocess = T.Compose([
        T.Resize(size=224, max_size=None, antialias=None),
        T.CenterCrop(size=(224, 224)),
        T.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))
    ])
    self.clip_model.to(device)
    self.clip_model.eval();
    
    # Get target embeddings
    self.target_embeds = []
    with torch.no_grad():
      for text_prompt in text_prompts:
        tokenized_text = open_clip.tokenize([text_prompt]).to(device)
        self.target_embeds.append(clip_model.encode_text(tokenized_text))
      for image_prompt in image_prompts:
        image_embed = clip_model.encode_image(self.preprocess(image_prompt))
        self.target_embeds.append(image_embed)

    self.target_embeds = torch.cat(self.target_embeds) # All the target embeddings

    self.n_cuts = n_cuts
    if self.n_cuts > 1:
      self.make_cutouts = MakeCutouts(224)
    
    
      

  def forward(self, input):
    if self.n_cuts > 1:
      input = self.make_cutouts(input, num_cutouts=self.n_cuts)
    input_embed = self.clip_model.encode_image(self.preprocess(input))
    return clip_loss_embeddings(input_embed, self.target_embeds)

# %% ../62_Generators_and_Losses.ipynb 75
import os
from fastprogress.fastprogress import master_bar, progress_bar

def optimise(gen, loss_functions, n_steps=100, optimizer=None, loss_weights=None, save_previews=True, display_every=10, save_every=10):
  
  # The optimizer (defaults to Adam)
  if optimizer == None:
    optimizer = torch.optim.Adam(gen.parameters(), lr=0.01, weight_decay=1e-6)

  # The loss weights
  if loss_weights == None:
    loss_weights = [1 for _ in loss_functions]

  # Make an empty folder to save steps
  os.system('rm -rf steps/')
  os.system('mkdir steps/')

  losses = [] # Keep track of our losses

  # Some fancier progress bar stuff
  mb = master_bar(range(1))
  mb.names=['loss']
  mb.graph_fig, axs = plt.subplots(1, 2, figsize=(15, 5)) # For custom display
  mb.graph_ax = axs[0]
  mb.img_ax = axs[1]
  mb.graph_out = display(mb.graph_fig, display_id=True)

  for p in mb:
    for i in progress_bar(range(n_steps), parent=mb):
      # Get the generator output
      output = gen()

      # Calculate our loss
      loss = 0

      for idx, loss_fn in enumerate(loss_functions):
        loss += loss_fn(output) * loss_weights[idx]

      # Store the loss for later
      losses.append(loss.detach().item())

      # View every 5 iterations
      if i % display_every == 0:
        # Show progress and loss plot
        mb.update_graph([[range(len(losses)), losses]])
        # Update image:
        img = tensor_to_pil(output)
        mb.img_ax.imshow(img)
        mb.graph_out.update(mb.graph_fig)
      

      # Save the output occasionally
      if i % save_every == 0:
        tensor_to_pil(output).save(f'steps/{i//save_every:05}.jpeg')

      # Backpropagate the loss and use it to update the parameters
      optimizer.zero_grad() # Reset everything related to gradient calculations
      loss.backward() # This does all the gradient calculations
      optimizer.step() # The optimizer does the update
