import os
import sys
import time
import datetime
import json
from PIL import Image, ImageFilter
import cv2 as cv
import numpy as np
from . import main, utils
from .. import basemodels
from flask import current_app as app
from .paths import get_paths
celery = basemodels.celery
from celery.utils.log import get_task_logger
celery_logger = get_task_logger(__name__)
import imutils


# init image processors
sys.path.append('/dulldream/src/')
from .img_proc_config import ImgProcConfig
from image_processors.mask_rcnn import MaskRCNN
from image_processors.pix2pix import Pix2Pix
from utils import imx
from utils import fiox


# initialize image processor
img_proc_config = ImgProcConfig()
p2p = Pix2Pix(img_proc_config.p2p_ckpts_dir,epoch=img_proc_config.p2p_epoch)
p2p_objects = Pix2Pix(img_proc_config.p2p_bg_ckpts_dir,epoch=img_proc_config.p2p_epoch)

mask_rcnn = MaskRCNN(img_proc_config.mask_rcnn_class_config,
    model_path=img_proc_config.mask_rcnn_model)


@celery.task(bind=True)
def task_dull(self, uuid_name, agree, mask_rcnn_result):
  """Process image and update during"""
  celery_logger.debug('process_image_task, uuid: {}'.format(uuid_name))
  
  upload_dir, render_dir, json_dir, upload_uri, render_uri = get_paths(agree)

  files = []
  im = Image.open(os.path.join(upload_dir, uuid_name + '.jpg')).convert('RGB')
  #im_np = cv.cvtColor(imx.ensure_np(im),cv.COLOR_RGB2BGR)
  im_np = imx.ensure_np(im)
  im_np = im_np[:,:,::-1]
  im = im.resize((256,256))
  im_np_256 = imutils.resize(im_np,width=256)

  # Add original
  fpath = os.path.join(render_dir, uuid_name + '_orig.jpg')
  im.save(fpath, 'JPEG', quality=95)
  files.append({
      'title': 'Original',
      'fn': render_uri + uuid_name + '_orig.jpg'
  })
  
  if mask_rcnn_result['valid']:
    # -----------------------------------------------
    # Segment image (processed in views)
    # seems to be an error with async celery processor?
    # -----------------------------------------------

    # parse mrcnn data
    im_mask = cv.imread(mask_rcnn_result['fp_im_mask'])
    seg_mask = cv.imread(mask_rcnn_result['fp_seg_mask'])
    #score = mask_rcnn_result['score']
    #name = mask_rcnn_result['name']
    #color = mask_rcnn_result['color']
    files.append({
        'title': 'Semantic Segmentation',
        'fn': render_uri + uuid_name + '_seg_mask.jpg'
    })
    files.append({
        'title': 'Semantic Segmentation Isolate',
        'fn': render_uri + uuid_name + '_im_mask.jpg'
    })

    
    # -----------------------------------------------
    # run rag generator
    # -----------------------------------------------

    self.update_state(
    state = 'PROCESSING',
      meta = {
        'percent': 0.50,
        'message': 'Applying Region Adjacency Graph',
        'uuid': uuid_name
    })

    # save the regions adjancency graph
    im_rag = imx.create_rag_mean(im_mask,compactness=30,n_segments=128)
    fpath = os.path.join(render_dir, uuid_name + '_rgraph.jpg')
    imx.save_np_as_pil(fpath,im_rag,quality=95)
    files.append({
        'title': 'Region Adjacency Graph',
        'fn': render_uri + uuid_name + '_rgraph.jpg'
    })


    # -----------------------------------------------
    # generate p2p fake
    # -----------------------------------------------

    self.update_state(
      state = 'PROCESSING',
      meta = {
        'percent': 0.75,
        'message': 'Running generative adversarial network...',
        'uuid': uuid_name
    })

    
    # convert segmentation to mask
    seg_mask_gray = cv.cvtColor(seg_mask,cv.COLOR_BGR2GRAY)
    seg_mask_gray[seg_mask_gray > 1] = 255

    # find best P2P fit
    ims_p2p = []
    match_amts = []
    iters = 15
    for i in range(0,iters):
      im_p2p = p2p.create_p2p(im_rag)
      ims_p2p.append(im_p2p)
      im_p2p_mask = cv.cvtColor(im_p2p,cv.COLOR_RGB2GRAY)
      im_p2p_mask[im_p2p_mask > 1] = 255
      # find where masks intersect
      matches = np.bitwise_and(im_p2p_mask,seg_mask_gray)
      amt = len(np.where(matches == 255)[0])
      match_amts.append(amt)
      self.update_state(
      state = 'PROCESSING',
      meta = {
        'percent': 0.75,
        'message': 'Generating ({}/{})'.format(i,iters),
        'uuid': uuid_name
      })
      
    best_idx = np.argmax(match_amts)
    im_p2p = ims_p2p[best_idx]

    fpath = os.path.join(render_dir, uuid_name + '_gan.jpg')
    imx.save_np_as_pil(fpath,im_p2p,quality=95)
    files.append({
        'title': 'Generative Adversarial Network',
        'fn': render_uri + uuid_name + '_gan.jpg'
    })


    # -----------------------------------------------
    # generate p2p fake
    # -----------------------------------------------

    # announce to user
    self.update_state(
      state = 'PROCESSING',
      meta = {
        'percent': 0.90,
        'message': 'Compositing images...',
        'uuid': uuid_name
    })


    # apply masked cloning
    im_p2p_gray = cv.cvtColor(im_p2p,cv.COLOR_BGR2GRAY)
    im_clone_mask = np.zeros_like(im_p2p_gray,dtype=np.uint8)
    im_clone_mask[im_p2p_gray > 1] = 255


    # apply smoothed copy+paste clone
    im_blur_mask = np.zeros(im_np_256.shape[:2],dtype=np.float64)
    im_blur_mask[im_p2p_gray > 1] = 1.0
    im_blur_mask = np.array([im_blur_mask,im_blur_mask,im_blur_mask]).transpose((1,2,0))
    
    # erode mask to remove black border
    kernel = np.ones((3,3),np.uint8)
    im_blur_mask = cv.erode(im_blur_mask,kernel,iterations = 3)

    # feather mask
    feather_amt = (3,3)
    im_blur_mask = (cv.GaussianBlur(im_blur_mask,feather_amt, 0) > 0) * 1.0 #?
    im_blur_mask = cv.GaussianBlur(im_blur_mask,feather_amt, 0)
    im_blur_mask = np.clip(im_blur_mask,0.0,1.0)

    # mask p2p fg --> photo bg
    im_dull = im_np_256.astype(np.float64) * (1.0 - im_blur_mask) + im_p2p.astype(np.float64) * im_blur_mask
    im_dull = im_dull.astype(np.uint8)
    

  else:
    print('No person. Apply background P2P')
    celery_logger.debug('No person. Apply background P2P, uuid: {}'.format(uuid_name))
    im_bg_blur = cv.GaussianBlur(im_np_256,(31,31),0)
    im_bg_rag = imx.create_rag_mean(im_bg_blur,compactness=30,n_segments=64)

    # apply gan
    im_dull = p2p_objects.create_p2p(im_bg_rag)

  # resize back to full 512px
  im_dull_512 = imutils.resize(im_dull,width=512)

  # save dulldream image
  fpath = os.path.join(render_dir, uuid_name + '_dull.jpg')
  imx.save_np_as_pil(fpath,im_dull_512,quality=95)
  files.append({
      'title': 'Your DullDream',
      'fn': render_uri + uuid_name + '_dull.jpg'
  })
  

  # -----------------------------------------------
  # Write data to disk
  # -----------------------------------------------

  data = {
      'uuid': uuid_name,
      'date': str(datetime.datetime.now()),
      'files': files
  }

  json_path = os.path.join(json_dir, uuid_name + '.json')
  with open(json_path, 'w') as json_file:
      json.dump(data, json_file)
 
  return {
      'percent': 100,
      'state': 'complete',
      'uuid': uuid_name
  }




@celery.task(bind=True)
def blur_task(self, uuid_name, agree, extra):
    """Process image and update during"""
    celery_logger.debug('process_image_task, uuid: {}'.format(uuid_name))

    upload_dir, render_dir, json_dir, upload_uri, render_uri = get_paths(agree)

    files = []

    im = Image.open(os.path.join(upload_dir, uuid_name + '.jpg')).convert('RGB')
    im = im.resize((256,256))
    files.append({
        'title': 'Original image',
        'fn': upload_uri + uuid_name + '.jpg'
    })

    self.update_state(
        state = 'PROCESSING',
        meta = {
            'percent': 0.25,
            'message': 'Applying blur',
            'uuid': uuid_name
        })

    im_np = utils.ensure_np(im)
    im_blur = cv.blur(im_np, (5,5), 1.0)
    im_blur_pil = utils.ensure_pil(im_blur)

    fn = uuid_name + '_blur.jpg'
    fpath = os.path.join(render_dir, fn)
    im_blur_pil.save(fpath, 'JPEG', quality=95)

    files.append({
        'title': 'Blurred image',
        'fn': render_uri + uuid_name + '_blur.jpg'
    })

    time.sleep(3)

    self.update_state(
        state = 'PROCESSING',
        meta = {
            'percent': 0.50,
            'message': 'Sleeping for some reason',
            'uuid': uuid_name
        })
    time.sleep(2)

    self.update_state(
        state = 'PROCESSING',
        meta = {
            'percent': 0.75,
            'message': 'Sleeping some more',
            'uuid': uuid_name
        })
    time.sleep(2)

    data = {
        'uuid': uuid_name,
        'date': str(datetime.datetime.now()),
        'files': files
    }

    json_path = os.path.join(json_dir, uuid_name + '.json')
    with open(json_path, 'w') as json_file:
        json.dump(data, json_file)

    celery_logger.debug('ok')
        
    return {
        'percent': 100,
        'state': 'complete',
        'uuid': uuid_name,
    }

@celery.task(bind=True)
def sleep_task(self, uuid_name):
    celery_logger.debug('sleep_task'.format(uuid_name))
    msgs = [ 
        {'msg':'Uploaded OK','time':.1},
        {'msg':'Segmenting Image...','time':2},
        {'msg':'Found: Person, Horse','time':1},
        {'msg':'Creating Pix2Pix','time':2}
    ]
    for i,m in enumerate(msgs):
        percent = int(float(i)/float(len(msgs))*100.0)
        self.update_state(
            state = 'PROCESSING',
            meta = {
                'percent': percent,
                'message': m['msg'],
                'uuid': uuid_name
            })
        celery_logger.debug(m['msg'])
        time.sleep(m['time'])

    return {
        'percent': 100,
        'state': 'complete',
        'uuid': uuid_name
    }

def make_task_json():
    dropdown = {}
    for k,v in task_lookup.items():
        if 'active' not in v or v['active'] is not False:
            is_default = 'default' in v and v['default'] is True
            task = {
                'name': k,
                'title': v['title'],
                'selected': is_default,
            }
            dropdown[k] = task
    return json.dumps(dropdown)

# Add all valid tasks to this lookup.
# Set 'active': False to disable a task
# Set 'default': True to define the default task

task_lookup = {
    'sleep': {
        'title': 'Sleep Test',
        'task': sleep_task,
        'active': False
    },
    'blur': {
        'title': 'Blur',
        'task': blur_task,
        'active': False
    },
    'task_dull': {
        'title': 'DullDream V2',
        'task': task_dull,
        'active': True,
        'default': True
    }
}