from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
#import mxnet as mx
#from mxnet import ndarray as nd
import argparse
import cv2
import pickle
import numpy as np
import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'common'))
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'RetinaFace'))
import face_align
from retinaface import RetinaFace

def to_rgb(img):
    w, h = img.shape
    ret = np.empty((w, h, 3), dtype=np.uint8)
    ret[:, :, 0] = ret[:, :, 1] = ret[:, :, 2] = img
    return ret


def IOU(Reframe,GTframe):
  x1 = Reframe[0];
  y1 = Reframe[1];
  width1 = Reframe[2]-Reframe[0];
  height1 = Reframe[3]-Reframe[1];

  x2 = GTframe[0]
  y2 = GTframe[1]
  width2 = GTframe[2]-GTframe[0]
  height2 = GTframe[3]-GTframe[1]

  endx = max(x1+width1,x2+width2)
  startx = min(x1,x2)
  width = width1+width2-(endx-startx)

  endy = max(y1+height1,y2+height2)
  starty = min(y1,y2)
  height = height1+height2-(endy-starty)

  if width <=0 or height <= 0:
    ratio = 0
  else:
    Area = width*height
    Area1 = width1*height1
    Area2 = width2*height2
    ratio = Area*1./(Area1+Area2-Area)
  return ratio

parser = argparse.ArgumentParser(description='Package eval images')
# general
parser.add_argument('--data-dir', default='', help='')
parser.add_argument('--image-size', type=int, default=112, help='')
parser.add_argument('--gpu', type=int, default=0, help='')
parser.add_argument('--det-prefix', type=str, default='./model/R50', help='')
parser.add_argument('--output', default='./', help='path to save.')
parser.add_argument('--align-mode', default='arcface', help='align mode.')
args = parser.parse_args()

gpu_id = args.gpu

detector = RetinaFace(args.det_prefix, 0, gpu_id, network='net3')
target_size = 400
max_size = 800

def get_norm_crop(image_path):
  im = cv2.imread(image_path)
  im_shape = im.shape
  im_size_min = np.min(im_shape[0:2])
  im_size_max = np.max(im_shape[0:2])
  im_scale = float(target_size) / float(im_size_min)
  # prevent bigger axis from being more than max_size:
  if np.round(im_scale * im_size_max) > max_size:
      im_scale = float(max_size) / float(im_size_max)
  bbox, landmark = detector.detect(im, threshold=0.5, scales=[im_scale])
  #print(im.shape, bbox.shape, landmark.shape)
  if bbox.shape[0]==0:
    bbox, landmark = detector.detect(im, threshold=0.05, scales=[im_scale*0.75, im_scale, im_scale*2.0])
    print('refine', im.shape, bbox.shape, landmark.shape)
  nrof_faces = bbox.shape[0]
  if nrof_faces>0:
    det = bbox[:,0:4]
    img_size = np.asarray(im.shape)[0:2]
    bindex = 0
    if nrof_faces>1:
      bounding_box_size = (det[:,2]-det[:,0])*(det[:,3]-det[:,1])
      img_center = img_size / 2
      offsets = np.vstack([ (det[:,0]+det[:,2])/2-img_center[1], (det[:,1]+det[:,3])/2-img_center[0] ])
      offset_dist_squared = np.sum(np.power(offsets,2.0),0)
      bindex = np.argmax(bounding_box_size-offset_dist_squared*2.0) # some extra weight on the centering
    #_bbox = bounding_boxes[bindex, 0:4]
    _landmark = landmark[bindex]
    warped = face_align.norm_crop(im, landmark = _landmark, image_size=args.image_size, mode=args.align_mode)
    return warped
  else:
    return None


bins = []
issame_list = []
pp = 0
for line in open(os.path.join(args.data_dir, 'pairs_label.txt'), 'r'):
  pp+=1
  if pp%100==0:
    print('processing', pp)
  line = line.strip().split()
  assert len(line)==3
  path1 = os.path.join(args.data_dir, line[0])
  path2 = os.path.join(args.data_dir, line[1])
  im1 = get_norm_crop(path1)
  im2 = get_norm_crop(path2)
  issame = True
  if line[2]=='0':
    issame = False
  issame_list.append(issame)
  for im in [im1, im2]:
    _, s = cv2.imencode('.jpg', im)
    bins.append(s)

with open(args.output, 'wb') as f:
  pickle.dump((bins, issame_list), f, protocol=pickle.HIGHEST_PROTOCOL)