"""model.py"""

import numpy as np
import scipy
import math
import numbers
from PIL import Image
import torch
import imageio

def latent_traversal_1d_multi_dim(model,
                                  latent_vector,
                                  device,
                                  dimensions=None,
                                  values=None,
                                  transpose=False):
  """Creates latent traversals for a latent vector along multiple dimensions.
  Creates a 2d grid image where each grid image is generated by passing a
  modified version of latent_vector to the generator_fn. In each column, a
  fixed dimension of latent_vector is modified. In each row, the value in the
  modified dimension is replaced by a fixed value.
  Args:
    generator_fn: Function that computes (fixed size) images from latent
      representation. It should accept a single Numpy array argument of the same
      shape as latent_vector and return a Numpy array of images where the first
      dimension corresponds to the different vectors in latent_vectors.
    latent_vector: 1d Numpy array with the base latent vector to be used.
    dimensions: 1d Numpy array with the indices of the dimensions that should be
      modified. If an integer is passed, the dimensions 0, 1, ...,
      (dimensions - 1) are modified. If None is passed, all dimensions of
      latent_vector are modified.
    values: 1d Numpy array with the latent space values that should be used for
      modifications. If an integer is passed, a linear grid between -1 and 1
      with that many points is constructed. If None is passed, a default grid is
      used (whose specific design is not guaranteed).
    transpose: Boolean which indicates whether rows and columns of the 2d grid
      should be transposed.
  Returns:
    Numpy array with image.
  """
  if latent_vector.ndim != 1:
    raise ValueError("Latent vector needs to be 1-dimensional.")

  if dimensions is None:
    # Default case, use all available dimensions.
    dimensions = np.arange(latent_vector.shape[0])
  elif isinstance(dimensions, numbers.Integral):
    # Check that there are enough dimensions in latent_vector.
    if dimensions > latent_vector.shape[0]:
      raise ValueError("The number of dimensions of latent_vector is less than"
                       " the number of dimensions requested in the arguments.")
    if dimensions < 1:
      raise ValueError("The number of dimensions has to be at least 1.")
    dimensions = np.arange(dimensions)
  if dimensions.ndim != 1:
    raise ValueError("Dimensions vector needs to be 1-dimensional.")

  if values is None:
    # Default grid of values.
    values = np.linspace(-1., 1., num=11)
  elif isinstance(values, numbers.Integral):
    if values <= 1:
      raise ValueError("If an int is passed for values, it has to be >1.")
    values = np.linspace(-1., 1., num=values)
  if values.ndim != 1:
    raise ValueError("Values vector needs to be 1-dimensional.")

  # We iteratively generate the rows/columns for each dimension as different
  # Numpy arrays. We do not preallocate a single final Numpy array as this code
  # is not performance critical and as it reduces code complexity.
  num_values = len(values)
  row_or_columns = []
  for dimension in dimensions:
    # Creates num_values copy of the latent_vector along the first axis.
    latent_traversal_vectors = np.tile(latent_vector, [num_values, 1])
    # Intervenes in the latent space.
    latent_traversal_vectors[:, dimension] = values
    latent_traversal_vectors = torch.from_numpy(latent_traversal_vectors).to(device)
    # Generate the batch of images
    images = model.decode(latent_traversal_vectors)
    images = images.cpu().data.numpy()
    images = np.transpose(images, (0, 3, 2, 1))
    # Adds images as a row or column depending whether transpose is True.
    axis = (1 if transpose else 0)
    row_or_columns.append(np.concatenate(images, axis))
  axis = (0 if transpose else 1)
  return np.concatenate(row_or_columns, axis)


def grid_save_images(images, image_path):
  """Saves images in list of [0,1]-valued np.arrays on a grid.
  Args:
    images: List of Numpy arrays of shape (height, width, {1,3}) with values in
      [0, 1].
    image_path: String with path to output image.
  """
  side_length = int(math.floor(math.sqrt(len(images))))
  image_rows = [
      np.concatenate(
          images[side_length * i:side_length * i + side_length], axis=0)
      for i in range(side_length)
  ]
  tiled_image = np.concatenate(image_rows, axis=1)
  print(tiled_image.shape)
  save_image(tiled_image, image_path)

def save_image(image, image_path):
  """Saves an image in the [0,1]-valued Numpy array to image_path.
  Args:
    image: Numpy array of shape (height, width, {1,3}) with values in [0, 1].
    image_path: String with path to output image.
  """
  # Copy the single channel if we are provided a grayscale image.
  if image.shape[2] == 1:
    image = np.repeat(image, 3, axis=2)
  image = np.ascontiguousarray(image)
  image *= 255.
  image = image.astype("uint8")
  img = Image.fromarray(image, mode="RGB")
  img.save(image_path)

def padded_grid(images, num_rows=None, padding_px=10, value=None):
  """Creates a grid with padding in between images."""
  num_images = len(images)
  if num_rows is None:
    num_rows = best_num_rows(num_images)

  # Computes how many empty images we need to add.
  num_cols = int(np.ceil(float(num_images) / num_rows))
  num_missing = num_rows * num_cols - num_images

  # Add the empty images at the end.
  all_images = images + [np.ones_like(images[0])] * num_missing

  # Create the final grid.
  rows = [padded_stack(all_images[i * num_cols:(i + 1) * num_cols], padding_px,
                       1, value=value) for i in range(num_rows)]
  return padded_stack(rows, padding_px, axis=0, value=value)


def padded_stack(images, padding_px=10, axis=0, value=None):
  """Stacks images along axis with padding in between images."""
  padding_arr = padding_array(images[0], padding_px, axis, value=value)
  new_images = [images[0]]
  for image in images[1:]:
    new_images.append(padding_arr)
    new_images.append(image)
  return np.concatenate(new_images, axis=axis)


def padding_array(image, padding_px, axis, value=None):
  """Creates padding image of proper shape to pad image along the axis."""
  shape = list(image.shape)
  shape[axis] = padding_px
  if value is None:
    return np.ones(shape, dtype=image.dtype)
  else:
    assert len(value) == shape[-1]
    shape[-1] = 1
    return np.tile(value, shape)


def best_num_rows(num_elements, max_ratio=4):
  """Automatically selects a smart number of rows."""
  best_remainder = num_elements
  best_i = None
  i = int(np.sqrt(num_elements))
  while True:
    if num_elements > max_ratio * i * i:
      return best_i
    remainder = (i - num_elements % i) % i
    if remainder == 0:
      return i
    if remainder < best_remainder:
      best_remainder = remainder
      best_i = i
    i -= 1


def pad_around(image, padding_px=10, axis=None, value=None):
  """Adds a padding around each image."""
  # If axis is None, pad both the first and the second axis.
  if axis is None:
    image = pad_around(image, padding_px, axis=0, value=value)
    axis = 1
  padding_arr = padding_array(image, padding_px, axis, value=value)
  return np.concatenate([padding_arr, image, padding_arr], axis=axis)


def add_below(image, padding_px=10, value=None):
  """Adds a footer below."""
  if len(image.shape) == 2:
    image = np.expand_dims(image, -1)
  if image.shape[2] == 1:
    image = np.repeat(image, 3, 2)
  if image.shape[2] != 3:
    raise ValueError("Could not convert image to have three channels.")
  with open("disentanglement.png", "rb") as f:
    footer = np.array(Image.open(f).convert("RGB")) * 1.0 / 255.
  missing_px = image.shape[1] - footer.shape[1]
  if missing_px < 0:
    return image
  if missing_px > 0:
    padding_arr = padding_array(footer, missing_px, axis=1, value=value)
    footer = np.concatenate([padding_arr, footer], axis=1)
  return padded_stack([image, footer], padding_px, axis=0, value=value)


def save_animation(list_of_animated_images, image_path, fps):
  full_size_images = []
  for single_images in zip(*list_of_animated_images):
    full_size_images.append(
        pad_around(add_below(padded_grid(list(single_images)))))
  imageio.mimwrite(image_path, full_size_images, fps=fps)


def cycle_factor(starting_index, num_indices, num_frames):
  """Cycles through the state space in a single cycle."""
  grid = np.linspace(starting_index, starting_index + 2*num_indices,
                     num=num_frames, endpoint=False)
  grid = np.array(np.ceil(grid), dtype=np.int64)
  grid -= np.maximum(0, 2*grid - 2*num_indices + 1)
  grid += np.maximum(0, -2*grid - 1)
  return grid


def cycle_gaussian(starting_value, num_frames, loc=0., scale=1.):
  """Cycles through the quantiles of a Gaussian in a single cycle."""
  starting_prob = scipy.stats.norm.cdf(starting_value, loc=loc, scale=scale)
  grid = np.linspace(starting_prob, starting_prob + 2.,
                     num=num_frames, endpoint=False)
  grid -= np.maximum(0, 2*grid - 2)
  grid += np.maximum(0, -2*grid)
  grid = np.minimum(grid, 0.999)
  grid = np.maximum(grid, 0.001)
  return np.array([scipy.stats.norm.ppf(i, loc=loc, scale=scale) for i in grid])