import cv2
import random
import numpy as np

import spiga.data.loaders.dl_config as dl_cfg
import spiga.data.loaders.dataloader as dl
import spiga.data.visualize.plotting as plot


def inspect_parser():
    import argparse
    pars = argparse.ArgumentParser(description='Data augmentation and dataset visualization. '
                                               'Press Q to quit,'
                                               'N to visualize the next image'
                                               ' and any other key to visualize the next default data.')
    pars.add_argument('database', type=str,
                      choices=['wflw', '300wpublic', '300wprivate', 'cofw68', 'merlrav'], help='Database name')
    pars.add_argument('-a', '--anns', type=str, default='train', help='Annotation type: test, train or valid')
    pars.add_argument('-np', '--nopose', action='store_false', default=True, help='Avoid pose generation')
    pars.add_argument('-c', '--clean', action='store_true', help='Process without data augmentation for train')
    pars.add_argument('--shape', nargs='+', type=int, default=[256, 256], help='Image cropped shape (W,H)')
    pars.add_argument('--img', nargs='+', type=int, default=None, help='Select specific image ids')
    return pars.parse_args()


class DatasetInspector:

    def __init__(self, database, anns_type, data_aug=True, pose=True, image_shape=(256,256)):

        data_config = dl_cfg.AlignConfig(database, anns_type)
        data_config.image_size = image_shape
        data_config.ftmap_size = image_shape
        data_config.generate_pose = pose

        if not data_aug:
            data_config.aug_names = []

        self.data_config = data_config
        dataloader, dataset = dl.get_dataloader(1, data_config, debug=True)
        self.dataset = dataset
        self.dataloader = dataloader
        self.colors_dft = {'lnd': (plot.GREEN, plot.RED), 'pose': (plot.BLUE, plot.GREEN, plot.RED)}

    def show_dataset(self, ids_list=None):

        if ids_list is None:
            ids = self.get_idx(shuffle=self.data_config.shuffle)
        else:
            ids = ids_list

        for img_id in ids:
            data_dict = self.dataset[img_id]
            crop_imgs, full_img = self.plot_features(data_dict)

            # Plot crop
            if 'merge' in crop_imgs.keys():
                crop = crop_imgs['merge']
            else:
                crop = crop_imgs['lnd']
            cv2.imshow('crop', crop)

            # Plot full
            cv2.imshow('image', full_img['lnd'])

            key = cv2.waitKey()
            if key == ord('q'):
                break

    def plot_features(self, data_dict, colors=None):

        # Init variables
        crop_imgs = {}
        full_imgs = {}
        if colors is None:
            colors = self.colors_dft

        # Cropped image
        image = data_dict['image']
        landmarks = data_dict['landmarks']
        visible = data_dict['visible']
        if np.any(np.isnan(visible)):
            visible = None
        mask = data_dict['mask_ldm']

        # Full image
        if 'image_ori' in data_dict.keys():
            image_ori = data_dict['image_ori']
        else:
            image_ori = cv2.imread(data_dict['imgpath'])
        landmarks_ori = data_dict['landmarks_ori']
        visible_ori = data_dict['visible_ori']
        if np.any(np.isnan(visible_ori)):
            visible_ori = None
        mask_ori = data_dict['mask_ldm_ori']

        # Plot landmarks
        crop_imgs['lnd'] = self._plot_lnd(image, landmarks, visible, mask, colors=colors['lnd'])
        full_imgs['lnd'] = self._plot_lnd(image_ori, landmarks_ori, visible_ori, mask_ori, colors=colors['lnd'])

        if self.data_config.generate_pose:
            rot, trl, cam_matrix = self._extract_pose(data_dict)

            # Plot pose
            crop_imgs['pose'] = plot.draw_pose(image, rot, trl, cam_matrix, euler=True, colors=colors['pose'])

            # Plot merge features
            crop_imgs['merge'] = plot.draw_pose(crop_imgs['lnd'], rot, trl, cam_matrix, euler=True, colors=colors['pose'])

        return crop_imgs, full_imgs

    def get_idx(self, shuffle=False):
        ids = list(range(len(self.dataset)))
        if shuffle:
            random.shuffle(ids)
        return ids

    def reload_dataset(self, data_config=None):
        if data_config is None:
            data_config = self.data_config
        dataloader, dataset = dl.get_dataloader(1, data_config, debug=True)
        self.dataset = dataset
        self.dataloader = dataloader

    def _extract_pose(self, data_dict):
        # Rotation and translation matrix
        pose = data_dict['pose']
        rot = pose[:3]
        trl = pose[3:]

        # Camera matrix
        cam_matrix = data_dict['cam_matrix']

        # Check for ground truth anns
        if 'headpose_ori' in data_dict.keys():
            if len(self.data_config.aug_names) == 0:
                print('Image headpose generated by ground truth data')
                pose_ori = data_dict['headpose_ori']
                rot = pose_ori

        return rot, trl, cam_matrix

    def _plot_lnd(self, image, landmarks, visible, mask, max_shape_thr=720, colors=None):

        if colors is None:
            colors = self.colors_dft['lnd']

        # Full image plots
        W, H, C = image.shape

        # Original image resize if need it
        if W > max_shape_thr or H > max_shape_thr:
            max_shape = max(W, H)
            scale_factor = max_shape_thr / max_shape
            resize_shape = (int(H * scale_factor), int(W * scale_factor))
            image_out = plot.draw_landmarks(image, landmarks, visible=visible, mask=mask,
                                            thick_scale=1 / scale_factor, colors=colors)
            image_out = cv2.resize(image_out, resize_shape)
        else:
            image_out = plot.draw_landmarks(image, landmarks, visible=visible, mask=mask, colors=colors)

        return image_out


if __name__ == '__main__':
    args = inspect_parser()
    data_aug = True
    database = args.database
    anns_type = args.anns
    pose = args.nopose
    select_img = args.img
    if args.clean:
        data_aug = False

    if len(args.shape) != 2:
        raise ValueError('--shape requires two values: width and height. Ej: --shape 256 256')
    else:
        img_shape = tuple(args.shape)

    visualizer = DatasetInspector(database, anns_type, data_aug=data_aug, pose=pose, image_shape=img_shape)
    visualizer.show_dataset(ids_list=select_img)