import numpy as np
import os
from scipy.io import loadmat
import cv2
from menpo.shape.pointcloud import PointCloud
from menpo.transform import ThinPlateSplines
import menpo.io as mio
import matplotlib.pyplot as plt
from scipy.ndimage import zoom
from glob import glob
from deformation_functions import *

'''********* bounding box and image loading functions *********'''


def center_margin_bb(bb, img_bounds, margin=0.25):
    bb_size = ([bb[0, 2] - bb[0, 0], bb[0, 3] - bb[0, 1]])
    margins = (np.max(bb_size) * (1 + margin) - bb_size) / 2

    bb_new = np.zeros_like(bb)
    bb_new[0, 0] = np.maximum(bb[0, 0] - margins[0], 0)
    bb_new[0, 2] = np.minimum(bb[0, 2] + margins[0], img_bounds[1])
    bb_new[0, 1] = np.maximum(bb[0, 1] - margins[1], 0)
    bb_new[0, 3] = np.minimum(bb[0, 3] + margins[1], img_bounds[0])
    return bb_new


def load_bb_files(bb_file_dirs):
    bb_files_dict = {}
    for bb_file in bb_file_dirs:
        bb_mat = loadmat(bb_file)['bounding_boxes']
        num_imgs = np.max(bb_mat.shape)
        for i in range(num_imgs):
            name = bb_mat[0][i][0][0][0][0]
            bb_init = bb_mat[0][i][0][0][1] - 1  # matlab indicies
            bb_gt = bb_mat[0][i][0][0][2] - 1  # matlab indicies
            if str(name) in bb_files_dict.keys():
                print str(name), 'already loaded from: ', bb_file
            bb_files_dict[str(name)] = (bb_init, bb_gt)
    return bb_files_dict


def load_bb_dictionary(bb_dir, mode, test_data='full'):
    if mode == 'TRAIN':
        bb_dirs = \
            ['bounding_boxes_afw.mat', 'bounding_boxes_helen_trainset.mat', 'bounding_boxes_lfpw_trainset.mat']
    else:
        if test_data == 'common':
            bb_dirs = \
                ['bounding_boxes_helen_testset.mat', 'bounding_boxes_lfpw_testset.mat']
        elif test_data == 'challenging':
            bb_dirs = ['bounding_boxes_ibug.mat']
        elif test_data == 'full':
            bb_dirs = \
                ['bounding_boxes_ibug.mat', 'bounding_boxes_helen_testset.mat', 'bounding_boxes_lfpw_testset.mat']
        elif test_data == 'training':
            bb_dirs = \
                ['bounding_boxes_afw.mat', 'bounding_boxes_helen_trainset.mat', 'bounding_boxes_lfpw_trainset.mat']
        else:
            bb_dirs=None

    if mode == 'TEST' and test_data not in ['full', 'challenging', 'common', 'training']:
        bb_files_dict = None
    else:
        bb_dirs = [os.path.join(bb_dir, dataset) for dataset in bb_dirs]
        bb_files_dict = load_bb_files(bb_dirs)

    return bb_files_dict


def crop_to_face_image(img, bb_dictionary=None, gt=True, margin=0.25, image_size=256):
    name = img.path.name
    img_bounds = img.bounds()[1]

    if bb_dictionary is None:
        bb_menpo = img.landmarks['PTS'].bounding_box().points
        bb = np.array([[bb_menpo[0, 1], bb_menpo[0, 0], bb_menpo[2, 1], bb_menpo[2, 0]]])
    else:
        if gt:
            bb = bb_dictionary[name][1]  # ground truth
        else:
            bb = bb_dictionary[name][0]  # init from face detector

    bb = center_margin_bb(bb, img_bounds, margin=margin)

    bb_pointcloud = PointCloud(np.array([[bb[0, 1], bb[0, 0]],
                                         [bb[0, 3], bb[0, 0]],
                                         [bb[0, 3], bb[0, 2]],
                                         [bb[0, 1], bb[0, 2]]]))

    face_crop = img.crop_to_pointcloud(bb_pointcloud).resize([image_size, image_size])

    return face_crop


def augment_face_image(img, image_size=256, crop_size=248, angle_range=30, flip=True):

    # taken from MDM
    jaw_indices = np.arange(0, 17)
    lbrow_indices = np.arange(17, 22)
    rbrow_indices = np.arange(22, 27)
    upper_nose_indices = np.arange(27, 31)
    lower_nose_indices = np.arange(31, 36)
    leye_indices = np.arange(36, 42)
    reye_indices = np.arange(42, 48)
    outer_mouth_indices = np.arange(48, 60)
    inner_mouth_indices = np.arange(60, 68)

    mirrored_parts_68 = np.hstack([
        jaw_indices[::-1], rbrow_indices[::-1], lbrow_indices[::-1],
        upper_nose_indices, lower_nose_indices[::-1],
        np.roll(reye_indices[::-1], 4), np.roll(leye_indices[::-1], 4),
        np.roll(outer_mouth_indices[::-1], 7),
        np.roll(inner_mouth_indices[::-1], 5)
    ])

    def mirror_landmarks_68(lms, im_size):
        return PointCloud(abs(np.array([0, im_size[1]]) - lms.as_vector(
        ).reshape(-1, 2))[mirrored_parts_68])

    def mirror_image(im):
        im = im.copy()
        im.pixels = im.pixels[..., ::-1].copy()

        for group in im.landmarks:
            lms = im.landmarks[group]
            if lms.points.shape[0] == 68:
                im.landmarks[group] = mirror_landmarks_68(lms, im.shape)

        return im

    lim = image_size - crop_size
    min_crop_inds = np.random.randint(0, lim, 2)
    max_crop_inds = min_crop_inds + crop_size
    flip_rand = np.random.random() > 0.5
    rot_angle = 2 * angle_range * np.random.random_sample() - angle_range

    if flip and flip_rand:
        rand_crop = img.crop(min_crop_inds, max_crop_inds)
        rand_crop = mirror_image(rand_crop)
        rand_crop = rand_crop.rotate_ccw_about_centre(rot_angle).resize([image_size, image_size])

    else:
        rand_crop = img.crop(min_crop_inds, max_crop_inds). \
            rotate_ccw_about_centre(rot_angle).resize([image_size, image_size])

    return rand_crop


def load_menpo_image_list(img_dir, mode, bb_dictionary=None, image_size=256, margin=0.25, bb_type='gt',
                          test_data='full', augment=True):
    def crop_to_face_image_gt(img, bb_dictionary=bb_dictionary, margin=margin, image_size=image_size):
        return crop_to_face_image(img, bb_dictionary, gt=True, margin=margin, image_size=image_size)

    def crop_to_face_image_init(img, bb_dictionary=bb_dictionary, margin=margin, image_size=image_size):
        return crop_to_face_image(img, bb_dictionary, gt=False, margin=margin, image_size=image_size)

    if mode is 'TRAIN':
        img_set_dir = os.path.join(img_dir, 'training_set')

    else:
        img_set_dir = os.path.join(img_dir, test_data + '_set')

    image_menpo_list = mio.import_images(img_set_dir, verbose=True)

    if bb_type is 'gt':
        face_crop_image_list = image_menpo_list.map(crop_to_face_image_gt)
    else:
        face_crop_image_list = image_menpo_list.map(crop_to_face_image_init)

    if mode is 'TRAIN' and augment:
        out_image_list = face_crop_image_list.map(augment_face_image)
    else:
        out_image_list = face_crop_image_list

    return out_image_list


def augment_menpo_img_ns(img, img_dir_ns, p_ns=0):
    img = img.copy()
    texture_aug = p_ns > 0.5
    if texture_aug:
        ns_augs = glob(os.path.join(img_dir_ns, img.path.name.split('.')[0] + '*'))
        num_augs = len(ns_augs)
        if num_augs > 1:
            ns_ind = np.random.randint(1, num_augs)
            ns_aug = mio.import_image(ns_augs[ns_ind])
            ns_pixels = ns_aug.pixels
            img.pixels = ns_pixels
    return img


def augment_menpo_img_geom(img, p_geom=0):
    img = img.copy()
    if p_geom > 0.5:
        lms_geom_warp=deform_face_geometric_style(img.landmarks['PTS'].points.copy(),p_scale=p_geom,p_shift=p_geom)
        img = warp_face_image_tps(img,PointCloud(lms_geom_warp))
    return img


def warp_face_image_tps(img,new_shape):
    tps = ThinPlateSplines(new_shape, img.landmarks['PTS'])
    img_warp=img.warp_to_shape(img.shape,tps)
    img_warp.landmarks['PTS']=new_shape
    return img_warp


def load_menpo_image_list_artistic_aug(
        img_dir, train_crop_dir, img_dir_ns, mode, bb_dictionary=None, image_size=256, margin=0.25,
        bb_type='gt', test_data='full',augment_basic=True, augment_texture=False, p_texture=0,
        augment_geom=False, p_geom=0):

    def crop_to_face_image_gt(img):
        return crop_to_face_image(img, bb_dictionary, gt=True, margin=margin, image_size=image_size)

    def crop_to_face_image_init(img):
        return crop_to_face_image(img, bb_dictionary, gt=False, margin=margin, image_size=image_size)

    def augment_menpo_img_ns_rand(img):
        return augment_menpo_img_ns(img, img_dir_ns, p_ns=1. * (np.random.rand() <= p_texture))

    def augment_menpo_img_geom_rand(img):
        return augment_menpo_img_geom(img, p_geom=1. * (np.random.rand() <= p_geom))

    if mode is 'TRAIN':
        img_set_dir = os.path.join(img_dir, train_crop_dir)
        out_image_list = mio.import_images(img_set_dir, verbose=True)

        if augment_texture:
            out_image_list = out_image_list.map(augment_menpo_img_ns_rand)
        if augment_geom:
            out_image_list = out_image_list.map(augment_menpo_img_geom_rand)
        if augment_basic:
            out_image_list = out_image_list.map(augment_face_image)

    else:
        img_set_dir = os.path.join(img_dir, test_data + '_set')
        out_image_list = mio.import_images(img_set_dir, verbose=True)
        if test_data in ['full', 'challenging', 'common', 'training', 'test']:
            if bb_type is 'gt':
                out_image_list = out_image_list.map(crop_to_face_image_gt)
            elif bb_type is 'init':
                out_image_list = out_image_list.map(crop_to_face_image_init)

    return out_image_list


def reload_img_menpo_list_artistic_aug_train(
        img_dir, train_crop_dir, img_dir_ns, mode, train_inds, image_size=256,
        augment_basic=True, augment_texture=False, p_texture=0, augment_geom=False, p_geom=0):

    img_menpo_list = load_menpo_image_list_artistic_aug(
        img_dir=img_dir, train_crop_dir=train_crop_dir, img_dir_ns=img_dir_ns, mode=mode,image_size=image_size,
        augment_basic=augment_basic, augment_texture=augment_texture, p_texture=p_texture, augment_geom=augment_geom,
        p_geom=p_geom)

    img_menpo_list_train = img_menpo_list[train_inds]

    return img_menpo_list_train


'''********* heat-maps and image loading functions *********'''


# look for: ECT-FaceAlignment/caffe/src/caffe/layers/data_heatmap.cpp
def gaussian(x, y, x0, y0, sigma=6):
    return 1./(np.sqrt(2*np.pi)*sigma) * np.exp(-0.5 * ((x-x0)**2 + (y-y0)**2) / sigma**2)


def create_heat_maps(landmarks, num_landmarks=68, image_size=256, sigma=6):

    x, y = np.mgrid[0:image_size, 0:image_size]

    maps = np.zeros((image_size, image_size, num_landmarks))

    for i in range(num_landmarks):
        out = gaussian(x, y, landmarks[i,0], landmarks[i,1], sigma=sigma)
        maps[:, :, i] = (8./3)*sigma*out  # copied from ECT

    return maps


def load_data(img_list, batch_inds, image_size=256, c_dim=3, num_landmarks=68 , sigma=6, scale='255',
              save_landmarks=False, primary=False):

    num_inputs = len(batch_inds)
    batch_menpo_images = img_list[batch_inds]

    images = np.zeros([num_inputs, image_size, image_size, c_dim]).astype('float32')
    maps_small = np.zeros([num_inputs, image_size/4, image_size/4, num_landmarks]).astype('float32')

    if primary:
        maps = None
    else:
        maps = np.zeros([num_inputs, image_size, image_size, num_landmarks]).astype('float32')

    if save_landmarks:
        landmarks = np.zeros([num_inputs, num_landmarks, 2]).astype('float32')
    else:
        landmarks = None

    for ind, img in enumerate(batch_menpo_images):

        images[ind, :, :, :] = np.rollaxis(img.pixels, 0, 3)

        if primary:
            lms = img.resize([image_size/4,image_size/4]).landmarks['PTS'].points
            maps_small[ind, :, :, :] = create_heat_maps(lms, num_landmarks, image_size/4, sigma)
        else:
            lms = img.landmarks['PTS'].points
            maps[ind, :, :, :] = create_heat_maps(lms, num_landmarks, image_size, sigma)
            maps_small[ind, :, :, :]=zoom(maps[ind, :, :, :],(0.25,0.25,1))

        if save_landmarks:
            landmarks[ind, :, :] = lms

    if scale is '255':
        images *= 255  # SAME AS ECT?
    elif scale is '0':
        images = 2 * images - 1

    return images, maps, maps_small, landmarks


def heat_maps_to_image(maps, landmarks=None, image_size=256, num_landmarks=68):

    if landmarks is None:
        landmarks = heat_maps_to_landmarks(maps, image_size=image_size, num_landmarks=num_landmarks)

    x, y = np.mgrid[0:image_size, 0:image_size]

    pixel_dist = np.sqrt(
        np.square(np.expand_dims(x, 2) - landmarks[:, 0]) + np.square(np.expand_dims(y, 2) - landmarks[:, 1]))

    nn_landmark = np.argmin(pixel_dist, 2)

    map_image = maps[x, y, nn_landmark]
    map_image = (map_image-map_image.min())/(map_image.max()-map_image.min())  # normalize for visualization

    return map_image


def heat_maps_to_landmarks(maps, image_size=256, num_landmarks=68):

    landmarks = np.zeros((num_landmarks,2)).astype('float32')

    for m_ind in range(num_landmarks):
        landmarks[m_ind, :] = np.unravel_index(maps[:, :, m_ind].argmax(), (image_size, image_size))

    return landmarks


def batch_heat_maps_to_landmarks(batch_maps, batch_size, image_size=256, num_landmarks=68):
    batch_landmarks = np.zeros((batch_size,num_landmarks, 2)).astype('float32')
    for i in range(batch_size):
        batch_landmarks[i,:,:]=heat_maps_to_landmarks(
            batch_maps[i,:,:,:], image_size=image_size, num_landmarks=num_landmarks)

    return batch_landmarks


def print_training_params_to_file(init_locals):
    del init_locals['self']
    with open(os.path.join(init_locals['save_log_path'], 'Training_Parameters.txt'), 'w') as f:
        f.write('Training Parameters:\n\n')
        for key, value in init_locals.items():
            f.write('* %s: %s\n' % (key, value))


def create_img_with_landmarks(image, landmarks, image_size=256, num_landmarks=68, scale='255', circle_size=2):
    image = image.reshape(image_size, image_size, -1)

    if scale is '0':
        image = 127.5 * (image + 1)
    elif scale is '1':
        image *= 255

    landmarks = landmarks.reshape(num_landmarks, 2)
    landmarks = np.clip(landmarks, 0, image_size)

    for (y, x) in landmarks.astype('int'):
        cv2.circle(image, (x, y), circle_size, (255, 0, 0), -1)

    return image


def merge_images_landmarks_maps(images, maps, image_size=256, num_landmarks=68, num_samples=9, scale='255',
                                circle_size=2):
    images = images[:num_samples]
    if maps.shape[1] is not image_size:
        images = zoom(images, (1, 0.25, 0.25, 1))
        image_size /= 4
    cmap = plt.get_cmap('jet')

    row = int(np.sqrt(num_samples))
    merged = np.zeros([row * image_size, row * image_size * 2, 3])

    for idx, img in enumerate(images):
        i = idx // row
        j = idx % row

        img_lamdmarks = heat_maps_to_landmarks(maps[idx, :, :, :], image_size=image_size, num_landmarks=num_landmarks)
        map_image = heat_maps_to_image(maps[idx, :, :, :], img_lamdmarks, image_size=image_size,
                                       num_landmarks=num_landmarks)

        rgba_map_image = cmap(map_image)
        map_image = np.delete(rgba_map_image, 3, 2) * 255

        img = create_img_with_landmarks(img, img_lamdmarks, image_size, num_landmarks, scale=scale,
                                        circle_size=circle_size)

        merged[i * image_size:(i + 1) * image_size, (j * 2) * image_size:(j * 2 + 1) * image_size, :] = img
        merged[i * image_size:(i + 1) * image_size, (j * 2 + 1) * image_size:(j * 2 + 2) * image_size, :] = map_image

    return merged


def merge_compare_maps(maps_small, maps, image_size=64, num_landmarks=68, num_samples=9):

    maps_small = maps_small[:num_samples]
    maps = maps[:num_samples]

    if maps_small.shape[1] is not image_size:
        image_size = maps_small.shape[1]

    if maps.shape[1] is not maps_small.shape[1]:
        maps_rescale = zoom(maps, (1, 0.25, 0.25, 1))
    else:
        maps_rescale = maps

    cmap = plt.get_cmap('jet')

    row = int(np.sqrt(num_samples))
    merged = np.zeros([row * image_size, row * image_size * 2, 3])

    for idx, map_small in enumerate(maps_small):
        i = idx // row
        j = idx % row

        map_image_small = heat_maps_to_image(map_small, image_size=image_size, num_landmarks=num_landmarks)
        map_image = heat_maps_to_image(maps_rescale[idx, :, :, :], image_size=image_size, num_landmarks=num_landmarks)

        rgba_map_image = cmap(map_image)
        map_image = np.delete(rgba_map_image, 3, 2) * 255

        rgba_map_image_small = cmap(map_image_small)
        map_image_small = np.delete(rgba_map_image_small, 3, 2) * 255

        merged[i * image_size:(i + 1) * image_size, (j * 2) * image_size:(j * 2 + 1) * image_size, :] = map_image_small
        merged[i * image_size:(i + 1) * image_size, (j * 2 + 1) * image_size:(j * 2 + 2) * image_size, :] = map_image

    return merged


def normalize_map(map_in):
    return (map_in - map_in.min()) / (map_in.max() - map_in.min())


def map_to_rgb(map_gray):
    cmap = plt.get_cmap('jet')
    rgba_map_image = cmap(map_gray)
    map_rgb = np.delete(rgba_map_image, 3, 2) * 255
    return map_rgb


def load_art_data(img_list, batch_inds, image_size=256, c_dim=3, scale='255'):

    num_inputs = len(batch_inds)
    batch_menpo_images = img_list[batch_inds]

    images = np.zeros([num_inputs, image_size, image_size, c_dim]).astype('float32')

    for ind, img in enumerate(batch_menpo_images):
        images[ind, :, :, :] = np.rollaxis(img.pixels, 0, 3)

    if scale is '255':
        images *= 255  # SAME AS ECT?
    elif scale is '0':
        images = 2 * images - 1

    return images


def merge_images_landmarks_maps_gt(images, maps, maps_gt, image_size=256, num_landmarks=68, num_samples=9, scale='255',
                                   circle_size=2, test_data='full', fast=False):
    images = images[:num_samples]
    if maps.shape[1] is not image_size:
        images = zoom(images, (1, 0.25, 0.25, 1))
        image_size /= 4
    if maps_gt.shape[1] is not image_size:
        maps_gt = zoom(maps_gt, (1, 0.25, 0.25, 1))

    cmap = plt.get_cmap('jet')

    row = int(np.sqrt(num_samples))
    merged = np.zeros([row * image_size, row * image_size * 3, 3])

    if fast:
        maps_gt_images = np.amax(maps_gt, 3)
        maps_images = np.amax(maps, 3)

    for idx, img in enumerate(images):
        i = idx // row
        j = idx % row

        img_landmarks = heat_maps_to_landmarks(maps[idx, :, :, :], image_size=image_size, num_landmarks=num_landmarks)

        if fast:
            map_image = maps_images[idx]
        else:
            map_image = heat_maps_to_image(maps[idx, :, :, :], img_landmarks, image_size=image_size,
                                           num_landmarks=num_landmarks)
        rgba_map_image = cmap(map_image)
        map_image = np.delete(rgba_map_image, 3, 2) * 255

        if test_data not in ['full', 'challenging', 'common', 'training']:
            map_gt_image = map_image.copy()
        else:
            if fast:
                map_gt_image = maps_gt_images[idx]
            else:
                map_gt_image = heat_maps_to_image(maps_gt[idx, :, :, :], image_size=image_size, num_landmarks=num_landmarks)
            rgba_map_gt_image = cmap(map_gt_image)
            map_gt_image = np.delete(rgba_map_gt_image, 3, 2) * 255

        img = create_img_with_landmarks(img, img_landmarks, image_size, num_landmarks, scale=scale,
                                        circle_size=circle_size)

        merged[i * image_size:(i + 1) * image_size, (j * 3) * image_size:(j * 3 + 1) * image_size, :] = img
        merged[i * image_size:(i + 1) * image_size, (j * 3 + 1) * image_size:(j * 3 + 2) * image_size, :] = map_image
        merged[i * image_size:(i + 1) * image_size, (j * 3 + 2) * image_size:(j * 3 + 3) * image_size, :] = map_gt_image

    return merged


def map_comapre_channels(images,maps1, maps2, image_size=64, num_landmarks=68, scale='255',test_data='full'):
    map1 = maps1[0]
    map2 = maps2[0]
    image = images[0]

    if image.shape[0] is not image_size:
        image = zoom(image, (0.25, 0.25, 1))
    if scale is '1':
        image *= 255
    elif scale is '0':
        image = 127.5 * (image + 1)

    row = np.ceil(np.sqrt(num_landmarks)).astype(np.int64)
    merged = np.zeros([row * image_size, row * image_size * 2, 3])

    for idx in range(num_landmarks):
        i = idx // row
        j = idx % row
        channel_map = map_to_rgb(normalize_map(map1[:, :, idx]))
        if test_data not in ['full', 'challenging', 'common', 'training']:
            channel_map2=channel_map.copy()
        else:
            channel_map2 = map_to_rgb(normalize_map(map2[:, :, idx]))

        merged[i * image_size:(i + 1) * image_size, (j * 2) * image_size:(j * 2 + 1) * image_size, :] = channel_map
        merged[i * image_size:(i + 1) * image_size, (j * 2 + 1) * image_size:(j * 2 + 2) * image_size, :] = channel_map2

    i = (idx + 1) // row
    j = (idx + 1) % row
    merged[i * image_size:(i + 1) * image_size, (j * 2) * image_size:(j * 2 + 1) * image_size, :] = image

    return merged


def train_val_shuffle_inds_per_epoch(valid_inds, train_inds, train_iter, batch_size, log_path, save_log=True):
    np.random.seed(0)
    num_train_images = len(train_inds)
    num_epochs = int(np.ceil((1. * train_iter) / (1. * num_train_images / batch_size)))+1
    epoch_inds_shuffle = np.zeros((num_epochs, num_train_images)).astype(int)
    img_inds = np.arange(num_train_images)
    for i in range(num_epochs):
        np.random.shuffle(img_inds)
        epoch_inds_shuffle[i, :] = img_inds

    if save_log:
        with open(os.path.join(log_path, "train_val_shuffle_inds.csv"), "wb") as f:
            if valid_inds is not None:
                f.write(b'valid inds\n')
                np.savetxt(f, valid_inds.reshape(1, -1), fmt='%i', delimiter=",")
            f.write(b'train inds\n')
            np.savetxt(f, train_inds.reshape(1, -1), fmt='%i', delimiter=",")
            f.write(b'shuffle inds\n')
            np.savetxt(f, epoch_inds_shuffle, fmt='%i', delimiter=",")

    return epoch_inds_shuffle