import os
import sys
sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../live-cortex/rpc/'))
from options.test_options import TestOptions
from options.dataset_options import DatasetOptions
from data import CreateRecursiveDataLoader
from models import create_model
from util.util import mkdirs
from shutil import copyfile, rmtree
from PIL import Image, ImageOps
import numpy as np
import cv2
from datetime import datetime
import re
import gevent

from img_ops import process_image, mix_next_image
from listener import Listener, read_sequence

module_name = 'pix2pix'

def load_opt():
  opt_parser = TestOptions()
  opt = opt_parser.parse()
  data_opt_parser = DatasetOptions()
  data_opt = data_opt_parser.parse(opt.unknown)
  opt.nThreads = 1   # test code only supports nThreads = 1
  opt.batchSize = 1  # test code only supports batchSize = 1
  opt.serial_batches = True  # no shuffle
  opt.no_flip = True  # no flip

  data_opt.tag = get_tag(opt, data_opt)
  opt.render_dir = opt.results_dir + opt.name + "/" + data_opt.tag + "/"
  return opt, data_opt, data_opt_parser

def get_tag(opt, data_opt):
  if data_opt.tag == '':
    d = datetime.now()
    tag = data_opt.tag = "{}_{}_{}".format(
      opt.name,
      'live',
      d.strftime('%Y%m%d%H%M')
    )
  else:
    tag = data_opt.tag
  return tag

def create_render_dir(opt):
  print("create render_dir: {}".format(opt.render_dir))
  if os.path.exists(opt.render_dir):
      rmtree(opt.render_dir)
  mkdirs(opt.render_dir)

def load_first_frame(opt, data_opt, i=0):
  start_img_path = os.path.join(opt.render_dir, "frame_{:05}.png".format(i))
  if data_opt.just_copy:
    copyfile(opt.start_img, start_img_path)
    A_img = None
    A_im = None
    A_offset = 0
  else:
    print("preload {}".format(opt.start_img))
    A_img = Image.open(opt.start_img).convert('RGB')
    A_im = np.asarray(A_img)
    A = process_image(opt, data_opt, A_im)
    cv2.imwrite(start_img_path, A)

  numz = re.findall(r'\d+', os.path.basename(opt.start_img))
  # print(numz)
  if len(numz) > 0:
    A_offset = int(numz[0])
    # print(A_offset)
    if A_offset:
      print(">> starting offset: {}".format(A_offset))
      A_dir = opt.start_img.replace(numz[0], "{:05d}")
      print(A_dir)
    else:
      print("Sequence not found")
  return A_offset, A_im, A_dir

def process_live_input(opt, data_opt, rpc_client):
  print(">>> Process live input")
  if data_opt.processing:
    print("Already processing...")
  data_opt.processing = True
  data_loader = CreateRecursiveDataLoader(opt)
  dataset = data_loader.load_data()

  create_render_dir(opt)
  sequence = read_sequence(data_opt.sequence_name, opt.module_name)
  print("Got sequence {}, {} images".format(data_opt.sequence, len(sequence)))
  if len(sequence) == 0:
    print("Got empty sequence...")
    data_opt.processing = False
    rpc_client.send_status('processing', False)
    return
  print("First image: {}".format(sequence[0]))

  rpc_client.send_status('processing', True)

  start_img_path = os.path.join(opt.render_dir, "frame_{:05d}.png".format(0))
  copyfile(sequence[0], start_img_path)

  model = create_model(opt)

  sequence_i = 1

  print("generating...")
  for i, data in enumerate(data_loader):
    if i >= opt.how_many:
      print("generated {} images, exiting".format(i))
      break

    if data_opt.load_checkpoint is True:
      model.save_dir = os.path.join(opt.checkpoints_dir, opt.module_name, data_opt.checkpoint_name)
      model.load_network(model.netG, 'G', data_opt.epoch)
      data_opt.load_checkpoint = False
    if data_opt.load_sequence is True:
      data_opt.load_sequence = False
      new_sequence = read_sequence(data_opt.sequence_name, opt.module_name)
      if len(new_sequence) != 0:
        print("Got sequence {}, {} images, first: {}".format(data_opt.sequence_name, len(new_sequence), new_sequence[0]))
        sequence = new_sequence
        sequence_i = 1
      else:
        print("Sequence not found")
    if data_opt.seek_to != 1:
      if data_opt.seek_to > 0 and data_opt.seek_to < len(sequence):
        sequence_i = data_opt.seek_to
      data_opt.seek_to = 1

    model.set_input(data)
    model.test()
    visuals = model.get_current_visuals()
    img_path = model.get_image_paths()

    sequence_i = mix_next_image(opt, data_opt, rpc_client, visuals['fake_B'], i, sequence, sequence_i)

    if data_opt.pause:
      data_opt.pause = False
      break
    gevent.sleep(data_opt.frame_delay)

  data_opt.processing = False
  rpc_client.send_status('processing', False)

if __name__ == '__main__':
  opt, data_opt, data_opt_parser = load_opt()
  listener = Listener(opt, data_opt, data_opt_parser, process_live_input)
  listener.connect()