# ------------------------------------------------------------------------------
# Util functions to visualize images.
# ------------------------------------------------------------------------------

import numpy as np
import cv2 as cv
from PIL import Image

def split(x):
  assert type(x) == int
  t = int(np.floor(np.sqrt(x)))
  for a in range(t, 0, -1):
    if x % a == 0:
      return a, x / a

def grid_transform(x):
  n, c, h, w = x.shape
  a, b = split(n)
  x = np.transpose(x, [0, 2, 3, 1])
  x = np.reshape(x, [int(a), int(b), int(h), int(w), int(c)])
  x = np.transpose(x, [0, 2, 1, 3, 4])
  x = np.reshape(x, [int(a * h), int(b * w), int(c)])
  if x.shape[2] == 1:
    x = np.squeeze(x, axis=2)
  return x

def seq_transform(x):
  n, c, h, w = x.shape
  x = np.transpose(x, [2, 0, 3, 1])
  x = np.reshape(x, [h, n * w, c])
  return x

# Converts image pixels from range [-1, 1] to [0, 255].
def data2img(data):
  rescaled = np.divide(data + 1.0, 2.0) * 255.
  rescaled = np.clip(rescaled, 0, 255)
  return np.rint(rescaled).astype('uint8')

def data2pil(data):
  return Image.fromarray(data2img(data), mode='RGB')

def interleave(a, b):
  res = np.empty([a.shape[0] + b.shape[0]] + list(a.shape[1:]), dtype=a.dtype)
  res[0::2] = a
  res[1::2] = b
  return res

def save_image(filepath, img):
  pilimg = Image.fromarray(img)
  pilimg.save(filepath)

def imread(filename):
  img = cv.imread(filename, cv.IMREAD_UNCHANGED)
  if img is not None:
    if len(img.shape) > 2:
      img = img[...,::-1]
  return img

def imconvert_float32(im):
  im = np.float32(im)
  im = (im / 255) * 2.0 - 1
  return im

def load_image(opt_fp_in, opt_dims=128):
  target_im = imread(opt_fp_in)
  w = target_im.shape[1]
  h = target_im.shape[0]
  if w <= h:
    scale = opt_dims / w
  else:
    scale = opt_dims / h
  target_im = cv.resize(target_im,(0,0), fx=scale, fy=scale)
  w = target_im.shape[1]
  h = target_im.shape[0]

  x0 = 0
  x1 = opt_dims
  y0 = 0
  y1 = opt_dims
  if w > opt_dims:
    x0 += int((w - opt_dims) / 2)
    x1 += x0
  if h > opt_dims:
    y0 += int((h - opt_dims) / 2)
    y1 += y0
  phi_target = imconvert_float32(target_im)
  phi_target = phi_target[y0:y1,x0:x1]
  if phi_target.shape[2] == 4:
    phi_target = phi_target[:,:,1:4]
  b = np.dsplit(phi_target, 3)
  phi_target = np.stack(b).reshape((3,opt_dims, opt_dims))
  #print(phi_target.shape)
  #phi_target = np.expand_dims(phi_target, 0)
  #phi_target = np.reshape(3, opt_dims, opt_dims)
  return phi_target