# -*- coding: utf-8 -*-

import numpy as np
import time
import hashlib
import glob
import os
import progressbar
import cv2

from .pysixd_stuff import transform
from .pysixd_stuff import view_sampler
from .utils import lazy_property


class Dataset(object):

    def __init__(self, dataset_path, **kw):

        self.shape = (int(kw['h']), int(kw['w']), int(kw['c']))
        self.noof_training_imgs = int(kw['noof_training_imgs'])
        self.dataset_path = dataset_path

        self.bg_img_paths = glob.glob(kw['background_images_glob'])
        self.noof_bg_imgs = min(int(kw['noof_bg_imgs']), len(self.bg_img_paths))

        self._kw = kw
        # self._aug = eval(self._kw['code'])

        self.train_x = np.empty( (self.noof_training_imgs,) + self.shape, dtype=np.uint8 )
        self.mask_x = np.empty( (self.noof_training_imgs,) + self.shape[:2], dtype= bool)
        self.noof_obj_pixels = np.empty( (self.noof_training_imgs,), dtype= bool)
        self.train_y = np.empty( (self.noof_training_imgs,) + self.shape, dtype=np.uint8 )
        self.bg_imgs = np.empty( (self.noof_bg_imgs,) + self.shape, dtype=np.uint8 )
        if np.float(eval(self._kw['realistic_occlusion'])):
            self.random_syn_masks


    @lazy_property
    def viewsphere_for_embedding(self):
        kw = self._kw
        num_cyclo = int(kw['num_cyclo'])
        azimuth_range = (0, 2 * np.pi)
        elev_range = (-0.5 * np.pi, 0.5 * np.pi)
        views, _ = view_sampler.sample_views(
            int(kw['min_n_views']),
            float(kw['radius']),
            azimuth_range,
            elev_range
        )
        Rs = np.empty( (len(views)*num_cyclo, 3, 3) )
        i = 0
        for view in views:
            for cyclo in np.linspace(0, 2.*np.pi, num_cyclo):
                rot_z = np.array([[np.cos(-cyclo), -np.sin(-cyclo), 0], [np.sin(-cyclo), np.cos(-cyclo), 0], [0, 0, 1]])
                Rs[i,:,:] = rot_z.dot(view['R'])
                i += 1
        return Rs

    @lazy_property
    def renderer(self):
        from auto_pose.meshrenderer import meshrenderer, meshrenderer_phong
        if self._kw['model'] == 'cad':
            renderer = meshrenderer.Renderer(
               [self._kw['model_path']],
               int(self._kw['antialiasing']),
               self.dataset_path,
               float(self._kw['vertex_scale'])
            )
        elif self._kw['model'] == 'reconst':
            renderer = meshrenderer_phong.Renderer(
               [self._kw['model_path']],
               int(self._kw['antialiasing']),
               self.dataset_path,
               float(self._kw['vertex_scale'])
            )
        else:
            'Error: neither cad nor reconst in model path!'
            exit()
        return renderer

    def get_training_images(self, dataset_path, args):
        current_config_hash = hashlib.md5((str(args.items('Dataset')+args.items('Paths'))).encode('utf-8')).hexdigest()
        current_file_name = os.path.join(dataset_path, current_config_hash + '.npz')

        if os.path.exists(current_file_name):
            training_data = np.load(current_file_name)
            self.train_x = training_data['train_x'].astype(np.uint8)
            self.mask_x = training_data['mask_x']
            self.train_y = training_data['train_y'].astype(np.uint8)
        else:
            self.render_training_images()
            np.savez(current_file_name, train_x = self.train_x, mask_x = self.mask_x, train_y = self.train_y)
        self.noof_obj_pixels = np.count_nonzero(self.mask_x==0,axis=(1,2))
        print('loaded %s training images' % len(self.train_x))

    def get_sprite_training_images(self, train_args):

        dataset_path= train_args.get('Paths','MODEL_PATH')
        dataset_zip = np.load(dataset_path)

        # print('Keys in the dataset:', dataset_zip.keys())
        imgs = dataset_zip['imgs']
        latents_values = dataset_zip['latents_values']
        latents_classes = dataset_zip['latents_classes']
        metadata = dataset_zip['metadata'][()]

        latents_sizes = metadata['latents_sizes']
        latents_bases = np.concatenate((latents_sizes[::-1].cumprod()[::-1][1:],
                                        np.array([1,])))

        latents_classes_heart = latents_classes[:245760]
        latents_classes_heart_rot = latents_classes_heart.copy()

        latents_classes_heart_rot[:, 0] = 0
        latents_classes_heart_rot[:, 1] = 0
        latents_classes_heart_rot[:, 2] = 5
        latents_classes_heart_rot[:, 4] = 16
        latents_classes_heart_rot[:, 5] = 16

        def latent_to_index(latents):
          return np.dot(latents, latents_bases).astype(int)

        indices_sampled = latent_to_index(latents_classes_heart_rot)
        imgs_sampled_rot = imgs[indices_sampled]
        indices_sampled = latent_to_index(latents_classes_heart)
        imgs_sampled_all = imgs[indices_sampled]

        self.train_x = np.expand_dims(imgs_sampled_all, 3)*255
        self.train_y = np.expand_dims(imgs_sampled_rot, 3)*255


    # def get_embedding_images(self, dataset_path, args):

    #     current_config_hash = hashlib.md5(str(args.items('Embedding') + args.items('Dataset')+args.items('Paths'))).hexdigest()
    #     current_file_name = os.path.join(dataset_path, current_config_hash + '.npz')

    #     if os.path.exists(current_file_name):
    #         embedding_data = np.load(current_file_name)
    #         self.embedding_data = embedding_data.astype(np.uint8)
    #     else:
    #         self.render_embedding_images()
    #         np.savez(current_file_name, train_x = self.train_x, mask_x = self.mask_x, train_y = self.train_y)
    #     print 'loaded %s training images' % len(self.train_x)

    def load_bg_images(self, dataset_path):
        current_config_hash = hashlib.md5((str(self.shape) + str(self.noof_bg_imgs) + str(self._kw['background_images_glob'])).encode('utf-8')).hexdigest()
        current_file_name = os.path.join(dataset_path, current_config_hash +'.npy')
        if os.path.exists(current_file_name):
            self.bg_imgs = np.load(current_file_name)
        else:
            file_list = self.bg_img_paths[:self.noof_bg_imgs]
            from random import shuffle
            shuffle(file_list)


            for j,fname in enumerate(file_list):
                print('loading bg img %s/%s' % (j,self.noof_bg_imgs))
                bgr = cv2.imread(fname)
                H,W = bgr.shape[:2]
                y_anchor = int(np.random.rand() * (H-self.shape[0]))
                x_anchor = int(np.random.rand() * (W-self.shape[1]))
                # bgr = cv2.resize(bgr, self.shape[:2])
                bgr = bgr[y_anchor:y_anchor+self.shape[0],x_anchor:x_anchor+self.shape[1],:]
                if bgr.shape[0]!=self.shape[0] or bgr.shape[1]!=self.shape[1]:
                    continue
                if self.shape[2] == 1:
                    bgr = cv2.cvtColor(np.uint8(bgr), cv2.COLOR_BGR2GRAY)[:,:,np.newaxis]
                self.bg_imgs[j] = bgr
            np.save(current_file_name,self.bg_imgs)



        print('loaded %s bg images' % self.noof_bg_imgs)


    def render_rot(self, R, t=None ,downSample = 1):
        kw = self._kw
        h, w = self.shape[:2]
        radius = float(kw['radius'])
        render_dims = eval(kw['render_dims'])
        K = eval(kw['k'])
        K = np.array(K).reshape(3,3)
        K[:2,:] = K[:2,:] / downSample

        clip_near = float(kw['clip_near'])
        clip_far = float(kw['clip_far'])
        pad_factor = float(kw['pad_factor'])

        t = np.array([0, 0, float(kw['radius'])])

        bgr_y, depth_y = self.renderer.render(
            obj_id=0,
            W=render_dims[0]/downSample,
            H=render_dims[1]/downSample,
            K=K.copy(),
            R=R,
            t=t,
            near=clip_near,
            far=clip_far,
            random_light=False
        )

        ys, xs = np.nonzero(depth_y > 0)
        obj_bb = view_sampler.calc_2d_bbox(xs, ys, render_dims)
        x, y, w, h = np.array(obj_bb).astype(np.int32)

        size = int(np.maximum(h, w) * pad_factor)

        left = int(np.maximum(x+w/2-size/2, 0))
        right = int(np.minimum(x+w/2+size/2, bgr_y.shape[1]))
        top = int(np.maximum(y+h/2-size/2, 0))
        bottom = int(np.minimum(y+h/2+size/2, bgr_y.shape[0]))

        bgr_y = bgr_y[top:bottom, left:right]
        return cv2.resize(bgr_y, self.shape[:2])


    def render_training_images(self):
        kw = self._kw
        H, W = int(kw['h']), int(kw['w'])
        render_dims = eval(kw['render_dims'])
        K = eval(kw['k'])
        K = np.array(K).reshape(3,3)
        clip_near = float(kw['clip_near'])
        clip_far = float(kw['clip_far'])
        pad_factor = float(kw['pad_factor'])
        max_rel_offset = float(kw['max_rel_offset'])
        t = np.array([0, 0, float(kw['radius'])])

        widgets = ['Training: ', progressbar.Percentage(),
             ' ', progressbar.Bar(),
             ' ', progressbar.Counter(), ' / %s' % self.noof_training_imgs,
             ' ', progressbar.ETA(), ' ']
        bar = progressbar.ProgressBar(maxval=self.noof_training_imgs,widgets=widgets)
        bar.start()

        for i in np.arange(self.noof_training_imgs):
            bar.update(i)

            # print '%s/%s' % (i,self.noof_training_imgs)
            # start_time = time.time()
            R = transform.random_rotation_matrix()[:3,:3]
            bgr_x, depth_x = self.renderer.render(
                obj_id=0,
                W=render_dims[0],
                H=render_dims[1],
                K=K.copy(),
                R=R,
                t=t,
                near=clip_near,
                far=clip_far,
                random_light=True
            )
            bgr_y, depth_y = self.renderer.render(
                obj_id=0,
                W=render_dims[0],
                H=render_dims[1],
                K=K.copy(),
                R=R,
                t=t,
                near=clip_near,
                far=clip_far,
                random_light=False
            )
            # render_time = time.time() - start_time
            # cv2.imshow('bgr_x',bgr_x)
            # cv2.imshow('bgr_y',bgr_y)
            # cv2.waitKey(0)

            ys, xs = np.nonzero(depth_x > 0)

            try:
                obj_bb = view_sampler.calc_2d_bbox(xs, ys, render_dims)
            except ValueError as e:
                print('Object in Rendering not visible. Have you scaled the vertices to mm?')
                break


            x, y, w, h = obj_bb

            rand_trans_x = np.random.uniform(-max_rel_offset, max_rel_offset) * w
            rand_trans_y = np.random.uniform(-max_rel_offset, max_rel_offset) * h

            obj_bb_off = obj_bb + np.array([rand_trans_x,rand_trans_y,0,0])

            bgr_x = self.extract_square_patch(bgr_x, obj_bb_off, pad_factor,resize=(W,H),interpolation = cv2.INTER_NEAREST)
            depth_x = self.extract_square_patch(depth_x, obj_bb_off, pad_factor,resize=(W,H),interpolation = cv2.INTER_NEAREST)
            mask_x = depth_x == 0.


            ys, xs = np.nonzero(depth_y > 0)
            obj_bb = view_sampler.calc_2d_bbox(xs, ys, render_dims)

            bgr_y = self.extract_square_patch(bgr_y, obj_bb, pad_factor,resize=(W,H),interpolation = cv2.INTER_NEAREST)

            if self.shape[2] == 1:
                bgr_x = cv2.cvtColor(np.uint8(bgr_x), cv2.COLOR_BGR2GRAY)[:,:,np.newaxis]
                bgr_y = cv2.cvtColor(np.uint8(bgr_y), cv2.COLOR_BGR2GRAY)[:,:,np.newaxis]

            self.train_x[i] = bgr_x.astype(np.uint8)
            self.mask_x[i] = mask_x
            self.train_y[i] = bgr_y.astype(np.uint8)

            #print 'rendertime ', render_time, 'processing ', time.time() - start_time
        bar.finish()

    def render_embedding_image_batch(self, start, end):
        kw = self._kw
        h, w = self.shape[:2]
        azimuth_range = (0, 2 * np.pi)
        elev_range = (-0.5 * np.pi, 0.5 * np.pi)
        radius = float(kw['radius'])
        render_dims = eval(kw['render_dims'])
        K = eval(kw['k'])
        K = np.array(K).reshape(3,3)

        clip_near = float(kw['clip_near'])
        clip_far = float(kw['clip_far'])
        pad_factor = float(kw['pad_factor'])

        t = np.array([0, 0, float(kw['radius'])])
        batch = np.empty( (end-start,)+ self.shape)
        obj_bbs = np.empty( (end-start,)+ (4,))

        for i, R in enumerate(self.viewsphere_for_embedding[start:end]):
            bgr_y, depth_y = self.renderer.render(
                obj_id=0,
                W=render_dims[0],
                H=render_dims[1],
                K=K.copy(),
                R=R,
                t=t,
                near=clip_near,
                far=clip_far,
                random_light=False
            )
            # cv2.imshow('depth',depth_y)
            # cv2.imshow('bgr',bgr_y)
            # print depth_y.max()
            # cv2.waitKey(0)
            ys, xs = np.nonzero(depth_y > 0)
            obj_bb = view_sampler.calc_2d_bbox(xs, ys, render_dims)

            obj_bbs[i] = obj_bb

            resized_bgr_y = self.extract_square_patch(bgr_y, obj_bb, pad_factor,resize=self.shape[:2],interpolation = cv2.INTER_NEAREST)

            if self.shape[2] == 1:
                resized_bgr_y = cv2.cvtColor(resized_bgr_y, cv2.COLOR_BGR2GRAY)[:,:,np.newaxis]
            batch[i] = resized_bgr_y / 255.
        return (batch, obj_bbs)

    def extract_square_patch(self, scene_img, bb_xywh, pad_factor,resize=(128,128),interpolation=cv2.INTER_NEAREST,black_borders=False):

        x, y, w, h = np.array(bb_xywh).astype(np.int32)
        size = int(np.maximum(h, w) * pad_factor)

        left = int(np.maximum(x+w/2-size/2, 0))
        right = int(np.minimum(x+w/2+size/2, scene_img.shape[1]))
        top = int(np.maximum(y+h/2-size/2, 0))
        bottom = int(np.minimum(y+h/2+size/2, scene_img.shape[0]))

        scene_crop = scene_img[top:bottom, left:right].copy()

        if black_borders:
            scene_crop[:(y-top),:] = 0
            scene_crop[(y+h-top):,:] = 0
            scene_crop[:,:(x-left)] = 0
            scene_crop[:,(x+w-left):] = 0

        scene_crop = cv2.resize(scene_crop, resize, interpolation = interpolation)
        return scene_crop

    @property
    def embedding_size(self):
        return len(self.viewsphere_for_embedding)


    @lazy_property
    def _aug(self):
        from imgaug.augmenters import Sequential,SomeOf,OneOf,Sometimes,WithColorspace,WithChannels, \
            Noop,Lambda,AssertLambda,AssertShape,Scale,CropAndPad, \
            Pad,Crop,Fliplr,Flipud,Superpixels,ChangeColorspace, PerspectiveTransform, \
            Grayscale,GaussianBlur,AverageBlur,MedianBlur,Convolve, \
            Sharpen,Emboss,EdgeDetect,DirectedEdgeDetect,Add,AddElementwise, \
            AdditiveGaussianNoise,Multiply,MultiplyElementwise,Dropout, \
            CoarseDropout,Invert,ContrastNormalization,Affine,PiecewiseAffine, \
            ElasticTransformation
        return eval(self._kw['code'])

    @lazy_property
    def _aug_occl(self):
        from imgaug.augmenters import Sequential,SomeOf,OneOf,Sometimes,WithColorspace,WithChannels, \
            Noop,Lambda,AssertLambda,AssertShape,Scale,CropAndPad, \
            Pad,Crop,Fliplr,Flipud,Superpixels,ChangeColorspace, PerspectiveTransform, \
            Grayscale,GaussianBlur,AverageBlur,MedianBlur,Convolve, \
            Sharpen,Emboss,EdgeDetect,DirectedEdgeDetect,Add,AddElementwise, \
            AdditiveGaussianNoise,Multiply,MultiplyElementwise,Dropout, \
            CoarseDropout,Invert,ContrastNormalization,Affine,PiecewiseAffine, \
            ElasticTransformation
        return Sequential([Sometimes(0.7, CoarseDropout( p=0.4, size_percent=0.01) )])


    @lazy_property
    def random_syn_masks(self):
        import bitarray
        workspace_path = os.environ.get('AE_WORKSPACE_PATH')

        random_syn_masks = bitarray.bitarray()
        with open(os.path.join(workspace_path,'random_tless_masks/arbitrary_syn_masks_1000.bin'), 'r') as fh:
            random_syn_masks.fromfile(fh)
        occlusion_masks = np.fromstring(random_syn_masks.unpack(), dtype=np.bool)
        occlusion_masks = occlusion_masks.reshape(-1,224,224,1).astype(np.float32)
        print(occlusion_masks.shape)

        occlusion_masks = np.array([cv2.resize(mask,(self.shape[0],self.shape[1]), interpolation = cv2.INTER_NEAREST) for mask in occlusion_masks])
        return occlusion_masks


    def augment_occlusion_mask(self, masks, verbose=False, min_trans = 0.2, max_trans=0.7, max_occl = 0.25,min_occl = 0.0):


        new_masks = np.zeros_like(masks,dtype=np.bool)
        occl_masks_batch = self.random_syn_masks[np.random.choice(len(self.random_syn_masks),len(masks))]
        for idx,mask in enumerate(masks):
            occl_mask = occl_masks_batch[idx]
            while True:
                trans_x = int(np.random.choice([-1,1])*(np.random.rand()*(max_trans-min_trans) + min_trans)*occl_mask.shape[0])
                trans_y = int(np.random.choice([-1,1])*(np.random.rand()*(max_trans-min_trans) + min_trans)*occl_mask.shape[1])
                M = np.float32([[1,0,trans_x],[0,1,trans_y]])

                transl_occl_mask = cv2.warpAffine(occl_mask,M,(occl_mask.shape[0],occl_mask.shape[1]))

                overlap_matrix = np.invert(mask.astype(np.bool)) * transl_occl_mask.astype(np.bool)
                overlap = len(overlap_matrix[overlap_matrix==True])/float(len(mask[mask==0]))

                if overlap < max_occl and overlap > min_occl:
                    new_masks[idx,...] = np.logical_xor(mask.astype(np.bool), overlap_matrix)
                    if verbose:
                        print('overlap is ', overlap)
                    break

        return new_masks
    def augment_squares(self,masks,rand_idcs,max_occl=0.25):
        new_masks = np.invert(masks)

        idcs = np.arange(len(masks))
        while len(idcs) > 0:
            new_masks[idcs] = self._aug_occl.augment_images(np.invert(masks[idcs]))
            new_noof_obj_pixels = np.count_nonzero(new_masks,axis=(1,2))
            idcs = np.where(new_noof_obj_pixels/self.noof_obj_pixels[rand_idcs].astype(np.float32) < 1-max_occl)[0]
            print(idcs)
        return np.invert(new_masks)

    def batch(self, batch_size):

        # batch_x = np.empty( (batch_size,) + self.shape, dtype=np.uint8 )
        # batch_y = np.empty( (batch_size,) + self.shape, dtype=np.uint8 )

        rand_idcs = np.random.choice(self.noof_training_imgs, batch_size, replace=False)

        assert self.noof_bg_imgs > 0

        rand_idcs_bg = np.random.choice(self.noof_bg_imgs, batch_size, replace=False)

        batch_x, masks, batch_y = self.train_x[rand_idcs], self.mask_x[rand_idcs], self.train_y[rand_idcs]
        rand_vocs = self.bg_imgs[rand_idcs_bg]

        if eval(self._kw['realistic_occlusion']):
            masks = self.augment_occlusion_mask(masks.copy(),max_occl=np.float(self._kw['realistic_occlusion']))

        if eval(self._kw['square_occlusion']):
            masks = self.augment_squares(masks.copy(),rand_idcs,max_occl=np.float(self._kw['square_occlusion']))

        batch_x[masks] = rand_vocs[masks]

        # random in-plane rotation, not necessary
        # for i in range(batch_size):
        #   rot_angle= np.random.rand()*360
        #   cent = int(self.shape[0]/2)
        #   M = cv2.getRotationMatrix2D((cent,cent),rot_angle,1)
        #   batch_x[i] = cv2.warpAffine(batch_x[i],M,self.shape[:2])[:,:,np.newaxis]
        #   batch_y[i] = cv2.warpAffine(batch_y[i],M,self.shape[:2])[:,:,np.newaxis]


        #needs uint8
        batch_x = self._aug.augment_images(batch_x)

        #slow...
        batch_x = batch_x / 255.
        batch_y = batch_y / 255.


        return (batch_x, batch_y)
