import os
import numpy as np
import cv2
import math
from PIL import Image, ImageOps
from skimage.transform import resize
from scipy.misc import imresize

def clamp(n,a,b):
  return max(a, min(n, b))

def lerp(n,a,b):
  return (b-a)*n+a

def process_image(opt, data_opt, im):
  img = im[:, :, ::-1].copy()
  processed = False

  if data_opt.process_frac == 1.0:
    return img

  if data_opt.transform != False:
    (h, w) = img.shape[:2]
    center = (w / 2, h / 2)
    M = cv2.getRotationMatrix2D(center, data_opt.rotate, data_opt.scale)
    img = cv2.warpAffine(img, M, (w, h))
  if data_opt.hsl is True:
    processed = True
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv)
    cv2.add(h, data_opt.hue, h)
    cv2.add(s, data_opt.saturation, s)
    cv2.add(v, data_opt.luminosity, v)
    hslimg = cv2.merge((h,s,v))
    img = cv2.cvtColor(hslimg, cv2.COLOR_HSV2BGR)
  if data_opt.clahe is True:
    processed = True
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=data_opt.clip_limit, tileGridSize=(8,8))
    l = clahe.apply(l)
    limg = cv2.merge((l,a,b))
    img = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
  if data_opt.posterize is True:
    processed = True
    img = cv2.pyrMeanShiftFiltering(img, data_opt.spatial_window, data_opt.color_window)
  if data_opt.grayscale is True:
    processed = True
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  if data_opt.blur is True:
    processed = True
    img = cv2.GaussianBlur(img, (data_opt.blur_radius, data_opt.blur_radius), data_opt.blur_sigma)
  if data_opt.canny is True:
    processed = True
    img = cv2.Canny(img, data_opt.canny_lo, data_opt.canny_hi)
    img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

  if processed is False:
    return img

  src_img = im[:, :, ::-1].copy()
  frac_a = data_opt.process_frac
  frac_b = 1.0 - frac_a
  array_a = np.multiply(src_img.astype('float64'), frac_a)
  array_b = np.multiply(img.astype('float64'), frac_b)
  img = np.add(array_a, array_b).astype('uint8')
  return img

def mix(frac, a, b):
  frac_a = frac
  frac_b = 1.0 - frac_a
  array_a = np.multiply(a.astype('float64'), frac_a)
  array_b = np.multiply(b.astype('float64'), frac_b)
  return np.add(array_a, array_b).astype('uint8')

last_im = None
def mix_next_image(opt, data_opt, rpc_client, im, sequence, i=0, sequence_i=0, skip_i=0):
  global last_im

  if (i % 100) == 0:
    print('%04d: process image...' % (i))

  last_path = opt.render_dir + "frame_{:05d}.png".format(i)
  tmp_path = opt.render_dir + "frame_{:05d}_tmp.png".format(i+1)
  next_path = opt.render_dir + "frame_{:05d}.png".format(i+1)
  current_path = opt.render_dir + "ren_{:05d}.png".format(i+1)

  # skip a certain number of frames
  if data_opt.sequence_skip == 0:
    skip_i = 0
  else:
    skip_i += 1
    if skip_i >= data_opt.sequence_skip:
        skip_i = 0
        sequence_i += data_opt.sequence_step * (data_opt.sequence_skip - 1)
  
  # get the sequence path for the current frame
  sequence_path = sequence[math.floor(sequence_i)]

  if skip_i == 0 and data_opt.sequence_playing and data_opt.sequence_step != 0:
    sequence_i += data_opt.sequence_step

  if sequence_i >= len(sequence):
    print('(((( sequence looped ))))')
    sequence_i = sequence_i % len(sequence)
  if sequence_i < 0:
    print('(((( sequence looped ))))')
    sequence_i += len(sequence)

  meta = {
    'i': i,
    'sequence_i': math.floor(sequence_i),
    'skip_i': math.floor(skip_i),
    'sequence_len': len(sequence)
  }

  if data_opt.store_a is not True:
    os.remove(last_path)

  if data_opt.store_b is True:
    image_pil = Image.fromarray(im, mode='RGB')
    image_pil.save(tmp_path)
    os.rename(tmp_path, current_path)

  recursive_frac = data_opt.recursive_frac
  sequence_frac = data_opt.sequence_frac
  
  if data_opt.sequence_skip > 0 and data_opt.fade_sequence > 1:
    sequence_frac *= (1.0 - (skip_i * data_opt.fade_sequence) / (data_opt.sequence_skip - 1))
    sequence_frac = max(0, min(sequence_frac, 1))

  if data_opt.recursive and last_im is not None:
    if data_opt.sequence and len(sequence) and sequence_frac > 0:
      A_img = Image.open(sequence_path).convert('RGB')
      A_im = np.asarray(A_img)
      frac_a = recursive_frac
      frac_b = sequence_frac
      frac_sum = frac_a + frac_b
      if frac_sum > 1.0:
        frac_a = frac_a / frac_sum
        frac_b = frac_b / frac_sum
      if data_opt.transition:
        t = lerp(math.sin(i / data_opt.transition_period * math.pi * 2.0 ) / 2.0 + 0.5, data_opt.transition_min, data_opt.transition_max)
        frac_a *= 1.0 - t
        frac_b *= 1.0 - t
      frac_c = 1.0 - frac_a - frac_b
      array_a = np.multiply(last_im.astype('float64'), frac_a)
      array_b = np.multiply(A_im.astype('float64'), frac_b)
      array_c = np.multiply(im.astype('float64'), frac_c)
      array_ab = np.add(array_a, array_b)
      array_abc = np.add(array_ab, array_c)
      next_im = array_abc.astype('uint8')

    else:
      next_im = mix(data_opt.recursive_frac, last_im, im)

    if data_opt.recurse_roll != 0:
      last_im = np.roll(im, data_opt.recurse_roll, axis=data_opt.recurse_roll_axis)
    else:
      last_im = next_im.copy().astype('uint8')

  elif data_opt.sequence and len(sequence) and sequence_frac > 0:
    A_img = Image.open(sequence_path).convert('RGB')
    A_im = np.asarray(A_img)
    sequence_frac = data_opt.sequence_frac
    if data_opt.transition:
      t = lerp(math.sin(i / data_opt.transition_period * math.pi * 2.0 ) / 2.0 + 0.5, data_opt.transition_min, data_opt.transition_max)
      sequence_frac *= 1.0 - t
    next_im = mix(sequence_frac, A_im, im)

  else:
    last_im = im.copy().astype('uint8')
    next_im = im

  next_img = process_image(opt, data_opt, next_im)

  img_to_send = None
  if data_opt.send_image == 'a':
    rgb_im = cv2.cvtColor(next_img, cv2.COLOR_BGR2RGB)
    img_to_send = Image.fromarray(rgb_im)
  if data_opt.send_image == 'b':
    img_to_send = Image.fromarray(im, mode='RGB')
  if data_opt.send_image == 'sequence':
    img_to_send = A_img
  if data_opt.send_image == 'recursive':
    img_to_send = Image.fromarray(next_im)

  if img_to_send is not None:
    if data_opt.resize_before_sending:
      img_to_send.resize((256, 256), Image.BICUBIC)
    rpc_client.send_pil_image("frame_{:05d}.png".format(i+1), meta, img_to_send, data_opt.output_format)

  cv2.imwrite(tmp_path, next_img)
  os.rename(tmp_path, next_path)

  if (i % 20) == 0:
    print("created {}".format(next_path))

  return sequence_i, skip_i