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 == 0:
    return img

  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

last_im = None
def mix_next_image(opt, data_opt, rpc_client, im, i, sequence, sequence_i):
  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)
  meta = { 'i': i, 'sequence_i': sequence_i, 'sequence_len': len(sequence) }
  if data_opt.sequence and len(sequence):
    sequence_path = sequence[sequence_i]
    if sequence_i >= len(sequence)-1:
      print('(((( sequence looped ))))')
      sequence_i = 1
    else:
      sequence_i += 1

  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)

  if data_opt.recursive and last_im is not None:
    if data_opt.sequence and len(sequence):
      A_img = Image.open(sequence_path).convert('RGB')
      A_im = np.asarray(A_img)
      frac_a = data_opt.recursive_frac
      frac_b = data_opt.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:
      frac_a = data_opt.recursive_frac
      frac_b = 1.0 - frac_a
      array_a = np.multiply(last_im.astype('float64'), frac_a)
      array_b = np.multiply(im.astype('float64'), frac_b)
      next_im = np.add(array_a, array_b).astype('uint8')

    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):
    A_img = Image.open(sequence_path).convert('RGB')
    A_im = np.asarray(A_img)
    frac_b = 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)
      frac_b *= 1.0 - t
    frac_c = 1.0 - frac_b
    array_b = np.multiply(A_im.astype('float64'), frac_b)
    array_c = np.multiply(im.astype('float64'), frac_c)
    array_bc = np.add(array_b, array_c)
    next_im = array_bc.astype('uint8')

  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