from abc import ABCMeta, abstractmethod

import cPickle as pkl
import logging
import matplotlib.pyplot as plt
import numpy as np
import os
import sys
from time import time

from sklearn.mixture import GaussianMixture

import autolab_core.utils as utils
from autolab_core import Point
from perception import BinaryImage, ColorImage, DepthImage, RgbdImage, SegmentationImage, CameraIntrinsics

from gqcnn import Grasp2D, SuctionPoint2D, ImageGraspSamplerFactory, GQCNN, InputDataMode, GraspQualityFunctionFactory, GQCnnQualityFunction, NoValidGraspsException
from gqcnn import Visualizer as vis

FIGSIZE = 16
SEED = 5234709

class RgbdImageState(object):
    """ State to encapsulate RGB-D images.

    Attributes
    ----------
    rgbd_im : :obj:`perception.RgbdImage`
        an RGB-D image to plan grasps on
    camera_intr : :obj:`perception.CameraIntrinsics`
        intrinsics of the RGB-D camera
    segmask : :obj:`perception.BinaryImage`
        segmentation mask for the image
    obj_segmask : :obj:`perception.SegmentationImage`
        segmentation mask for the different objects in the image
    full_observed : :obj:`object`
        representation of the fully observed state
    """
    def __init__(self, rgbd_im, camera_intr,
                 segmask=None,
                 obj_segmask=None,
                 fully_observed=None):
        self.rgbd_im = rgbd_im
        self.camera_intr = camera_intr
        self.segmask = segmask
        self.obj_segmask = obj_segmask
        self.fully_observed = fully_observed
        logging.debug(self.camera_intr)

    def save(self, save_dir):
        if not os.path.exists(save_dir):
            os.mkdir(save_dir)
        color_image_filename = os.path.join(save_dir, 'color.png')
        depth_image_filename = os.path.join(save_dir, 'depth.npy')
        camera_intr_filename = os.path.join(save_dir, 'camera.intr')
        segmask_filename = os.path.join(save_dir, 'segmask.npy')
        obj_segmask_filename = os.path.join(save_dir, 'obj_segmask.npy')
        state_filename = os.path.join(save_dir, 'state.pkl')
        self.rgbd_im.color.save(color_image_filename)
        self.rgbd_im.depth.save(depth_image_filename)
        self.camera_intr.save(camera_intr_filename)
        if self.segmask is not None:
            self.segmask.save(segmask_filename)
        if self.obj_segmask is not None:
            self.obj_segmask.save(obj_segmask_filename)
        if self.fully_observed is not None:
            pkl.dump(self.fully_observed, open(state_filename, 'wb'))

    @staticmethod
    def load(save_dir):
        if not os.path.exists(save_dir):
            raise ValueError('Directory %s does not exist!' %(save_dir))
        color_image_filename = os.path.join(save_dir, 'color.png')
        depth_image_filename = os.path.join(save_dir, 'depth.npy')
        camera_intr_filename = os.path.join(save_dir, 'camera.intr')
        segmask_filename = os.path.join(save_dir, 'segmask.npy')
        obj_segmask_filename = os.path.join(save_dir, 'obj_segmask.npy')
        state_filename = os.path.join(save_dir, 'state.pkl')
        camera_intr = CameraIntrinsics.load(camera_intr_filename)
        color = ColorImage.open(color_image_filename, frame=camera_intr.frame)
        depth = DepthImage.open(depth_image_filename, frame=camera_intr.frame)
        segmask = None
        if os.path.exists(segmask_filename):
            segmask = BinaryImage.open(segmask_filename, frame=camera_intr.frame)
        obj_segmask = None
        if os.path.exists(obj_segmask_filename):
            obj_segmask = SegmentationImage.open(obj_segmask_filename, frame=camera_intr.frame)
        fully_observed = None
        if os.path.exists(state_filename):
            fully_observed = pkl.load(open(state_filename, 'rb'))
        return RgbdImageState(RgbdImage.from_color_and_depth(color, depth),
                              camera_intr,
                              segmask=segmask,
                              obj_segmask=obj_segmask,
                              fully_observed=fully_observed)

class ParallelJawGrasp(object):
    """ Action to encapsulate parallel jaw grasps.
    """
    def __init__(self, grasp, q_value, image):
        self.grasp = grasp
        self.q_value = q_value
        self.image = image

    def save(self, save_dir):
        if not os.path.exists(save_dir):
            os.mkdir(save_dir)
        grasp_filename = os.path.join(save_dir, 'grasp.pkl')
        q_value_filename = os.path.join(save_dir, 'pred_robustness.pkl')
        image_filename = os.path.join(save_dir, 'tf_image.npy')
        pkl.dump(self.grasp, open(grasp_filename, 'wb'))
        pkl.dump(self.q_value, open(q_value_filename, 'wb'))
        self.image.save(image_filename)

    @staticmethod
    def load(save_dir):
        if not os.path.exists(save_dir):
            raise ValueError('Directory %s does not exist!' %(save_dir))
        grasp_filename = os.path.join(save_dir, 'grasp.pkl')
        q_value_filename = os.path.join(save_dir, 'pred_robustness.pkl')
        image_filename = os.path.join(save_dir, 'tf_image.npy')
        grasp = pkl.load(open(grasp_filename, 'rb'))
        q_value = pkl.load(open(q_value_filename, 'rb'))
        image = DepthImage.open(image_filename)
        return ParallelJawGrasp(grasp, q_value, image)

class GraspAction(object):
    """ Action to encapsulate parallel jaw grasps.
    """
    def __init__(self, grasp, q_value, image):
        self.grasp = grasp
        self.q_value = q_value
        self.image = image

    def save(self, save_dir):
        if not os.path.exists(save_dir):
            os.mkdir(save_dir)
        grasp_filename = os.path.join(save_dir, 'grasp.pkl')
        q_value_filename = os.path.join(save_dir, 'pred_robustness.pkl')
        image_filename = os.path.join(save_dir, 'tf_image.npy')
        pkl.dump(self.grasp, open(grasp_filename, 'wb'))
        pkl.dump(self.q_value, open(q_value_filename, 'wb'))
        self.image.save(image_filename)

class Policy(object):
    """ Abstract policy class. """
    __metaclass__ = ABCMeta

    def __call__(self, state):
        return self.action(state)

    @abstractmethod
    def action(self, state):
        """ Returns an action for a given state.
        """
        pass

class GraspingPolicy(Policy):
    """ Policy for robust grasping with Grasp Quality Convolutional Neural Networks (GQ-CNN).
    Attributes
    ----------
    config : dict
        dictionary of parameters for the policy

    Notes
    -----
    Required configuration parameters are specified in Other Parameters

    Other Parameters
    ----------------
    sampling : dict
        dictionary of parameters for grasp sampling, see gqcnn/image_grasp_sampler.py
    gqcnn_model : str
        string path to a trained GQ-CNN model see gqcnn/neural_networks.py
    """
    def __init__(self, config):
        # store parameters
        self._config = config
        self._gripper_width = np.inf
        if 'gripper_width' in config.keys():
            self._gripper_width = config['gripper_width']

        # init grasp sampler
        self._sampling_config = config['sampling']
        self._sampling_config['gripper_width'] = self._gripper_width
        sampler_type = self._sampling_config['type']
        self._grasp_sampler = ImageGraspSamplerFactory.sampler(sampler_type,
                                                               self._sampling_config)

        # init grasp quality function
        self._metric_config = config['metric']
        metric_type = self._metric_config['type']
        self._grasp_quality_fn = GraspQualityFunctionFactory.quality_function(metric_type,
                                                                              self._metric_config)

    @property
    def config(self):
        """ Returns the policy parameters. """
        return self._config

    @property
    def grasp_sampler(self):
        """ Returns the grasp sampler. """
        return self._grasp_sampler

    @property
    def grasp_quality_fn(self):
        """ Returns the grasp sampler. """
        return self._grasp_quality_fn

    @property
    def gqcnn(self):
        """ Returns the GQ-CNN. """
        return self._gqcnn

    @abstractmethod
    def action(self, state):
        """ Returns an action for a given state.
        """
        pass

class UniformRandomGraspingPolicy(GraspingPolicy):
    """ Returns a grasp uniformly at random. """
    def __init__(self, config):
        GraspingPolicy.__init__(self, config)
        self._num_grasp_samples = 1

    def action(self, state):
        """ Plans the grasp with the highest probability of success on
        the given RGB-D image.

        Attributes
        ----------
        state : :obj:`RgbdImageState`
            image to plan grasps on

        Returns
        -------
        :obj:`GraspAction`
            grasp to execute
        """
        # check valid input
        if not isinstance(state, RgbdImageState):
            raise ValueError('Must provide an RGB-D image state.')

        # parse state
        rgbd_im = state.rgbd_im
        camera_intr = state.camera_intr
        segmask = state.segmask

        # sample grasps
        grasps = self._grasp_sampler.sample(rgbd_im, camera_intr,
                                            self._num_grasp_samples,
                                            segmask=segmask,
                                            visualize=self.config['vis']['grasp_sampling'],
                                            seed=None)
        num_grasps = len(grasps)
        if num_grasps == 0:
            logging.warning('No valid grasps could be found')
            raise NoValidGraspsException()

        # set grasp
        grasp = grasps[0]

        # form tensors
        return GraspAction(grasp, 0.0, state.rgbd_im.depth)

class RobustGraspingPolicy(GraspingPolicy):
    """ Samples a set of grasp candidates in image space,
    ranks the grasps by the predicted probability of success from a GQ-CNN,
    and returns the grasp with the highest probability of success.

    Notes
    -----
    Required configuration parameters are specified in Other Parameters

    Other Parameters
    ----------------
    num_grasp_samples : int
        number of grasps to sample
    gripper_width : float, optional
        width of the gripper in meters
    logging_dir : str, optional
        directory in which to save the sampled grasps and input images
    """
    def __init__(self, config, filters=None):
        GraspingPolicy.__init__(self, config)
        self._parse_config()
        self._filters = filters

    def _parse_config(self):
        """ Parses the parameters of the policy. """
        self._num_grasp_samples = self.config['sampling']['num_grasp_samples']
        self._max_grasps_filter = 1
        if 'max_grasps_filter' in self.config.keys():
            self._max_grasps_filter = self.config['max_grasps_filter']
        self._gripper_width = np.inf
        if 'gripper_width' in self.config.keys():
            self._gripper_width = self.config['gripper_width']
        self._logging_dir = None
        if 'logging_dir' in self.config.keys():
            self._logging_dir = self.config['logging_dir']

    def select(self, grasps, q_value):
        """ Selects the grasp with the highest probability of success.
        Can override for alternate policies (e.g. epsilon greedy).
        """
        # sort grasps
        num_grasps = len(grasps)
        grasps_and_predictions = zip(np.arange(num_grasps), q_value)
        grasps_and_predictions.sort(key = lambda x : x[1], reverse=True)

        # return top grasps
        if self._filters is None:
            return grasps_and_predictions[0][0]

        # filter grasps
        logging.info('Filtering grasps')
        i = 0
        while i < self._max_grasps_filter and i < len(grasps_and_predictions):
            index = grasps_and_predictions[i][0]
            grasp = grasps[index]
            valid = True
            for filter_name, is_valid in self._filters.iteritems():
                valid = is_valid(grasp)
                logging.debug('Grasp {} filter {} valid: {}'.format(i, filter_name, valid))
                if not valid:
                    valid = False
                    break
            if valid:
                return index
            i += 1
        raise NoValidGraspsException('No grasps satisfied filters')

    def action(self, state):
        """ Plans the grasp with the highest probability of success on
        the given RGB-D image.

        Attributes
        ----------
        state : :obj:`RgbdImageState`
            image to plan grasps on

        Returns
        -------
        :obj:`GraspAction`
            grasp to execute
        """
        # check valid input
        if not isinstance(state, RgbdImageState):
            raise ValueError('Must provide an RGB-D image state.')

        # parse state
        rgbd_im = state.rgbd_im
        camera_intr = state.camera_intr
        segmask = state.segmask

        # sample grasps
        grasps = self._grasp_sampler.sample(rgbd_im, camera_intr,
                                            self._num_grasp_samples,
                                            segmask=segmask,
                                            visualize=self.config['vis']['grasp_sampling'],
                                            seed=None)
        num_grasps = len(grasps)
        if num_grasps == 0:
            logging.warning('No valid grasps could be found')
            raise NoValidGraspsException()

        # save if specified
        if self._logging_dir is not None:
            if not os.path.exists(self._logging_dir):
                raise ValueError('Logging directory %s does not exist' %(self._logging_dir))

            # create output dir
            test_case_id = utils.gen_experiment_id()
            test_case_dir = os.path.join(self._logging_dir, 'test_case_%s' %(test_case_id))

            # re-sample if the test case dir exists
            while os.path.exists(test_case_dir):
                test_case_id = utils.gen_experiment_id()
                test_case_dir = os.path.join(self._logging_dir, 'test_case_%s' %(test_case_id))

            # create the directory and save
            os.mkdir(test_case_dir)
            candidate_actions_filename = os.path.join(test_case_dir, 'actions.pkl')
            pkl.dump(grasps, open(candidate_actions_filename, 'wb'))
            image_state_filename = os.path.join(test_case_dir, 'state.pkl')
            pkl.dump(state, open(image_state_filename, 'wb'))

        # compute grasp quality
        compute_start = time()
        q_values = self._grasp_quality_fn(state, grasps, params=self._config)
        logging.debug('Grasp evaluation took %.3f sec' %(time()-compute_start))

        if self.config['vis']['grasp_candidates']:
            # display each grasp on the original image, colored by predicted success
            norm_q_values = (q_values - np.min(q_values)) / (np.max(q_values) - np.min(q_values))
            vis.figure(size=(FIGSIZE,FIGSIZE))
            vis.imshow(rgbd_im.depth,
                       vmin=self.config['vis']['vmin'],
                       vmax=self.config['vis']['vmax'])
            for grasp, q in zip(grasps, norm_q_values):
                vis.grasp(grasp, scale=1.0,
                          grasp_center_size=10,
                          grasp_center_thickness=2.5,
                          jaw_width=2.5,
                          show_center=False,
                          show_axis=True,
                          color=plt.cm.RdYlGn(q))
            vis.title('Sampled grasps')
            vis.show()

        # select grasp
        index = self.select(grasps, q_values)
        grasp = grasps[index]
        q_value = q_values[index]
        if self.config['vis']['grasp_plan']:
            vis.figure()
            vis.imshow(rgbd_im.depth)#,
                      # vmin=self.config['vis']['vmin'],
                      # vmax=self.config['vis']['vmax'])
            vis.grasp(grasp, scale=2.0, show_axis=True)
            vis.title('Best Grasp: d=%.3f, q=%.3f' %(grasp.depth, q_value))
            vis.show()

        return GraspAction(grasp, q_value, state.rgbd_im.depth)

class CrossEntropyRobustGraspingPolicy(GraspingPolicy):
    """ Optimizes a set of grasp candidates in image space using the
    cross entropy method:
    (1) sample an initial set of candidates
    (2) sort the candidates
    (3) fit a GMM to the top P%
    (4) re-sample grasps from the distribution
    (5) repeat steps 2-4 for K iters
    (6) return the best candidate from the final sample set

    Parameters
    ----------
    filters : :obj:`dict` mapping names to functions
        list of functions to apply to filter invalid grasps

    Notes
    -----
    Required configuration parameters are specified in Other Parameters

    Other Parameters
    ----------------
    num_seed_samples : int
        number of candidate to sample in the initial set
    num_gmm_samples : int
        number of candidates to sample on each resampling from the GMMs
    num_iters : int
        number of sample-and-refit iterations of CEM
    gmm_refit_p : float
        top p-% of grasps used for refitting
    gmm_component_frac : float
        percentage of the elite set size used to determine number of GMM components
    gmm_reg_covar : float
        regularization parameters for GMM covariance matrix, enforces diversity of fitted distributions
    deterministic : bool, optional
        whether to set the random seed to enforce deterministic behavior
    gripper_width : float, optional
        width of the gripper in meters
    """
    def __init__(self, config, filters=None):
        GraspingPolicy.__init__(self, config)
        self._parse_config()
        self._filters = filters

    def _parse_config(self):
        """ Parses the parameters of the policy. """
        # cross entropy method parameters
        self._num_seed_samples = self.config['num_seed_samples']
        self._num_gmm_samples = self.config['num_gmm_samples']
        self._num_iters = self.config['num_iters']
        self._gmm_refit_p = self.config['gmm_refit_p']
        self._gmm_component_frac = self.config['gmm_component_frac']
        self._gmm_reg_covar = self.config['gmm_reg_covar']

        self._max_grasps_filter = 1
        if 'max_grasps_filter' in self.config.keys():
            self._max_grasps_filter = self.config['max_grasps_filter']

        self._max_resamples_per_iteration = 100
        if 'max_resamples_per_iteration' in self.config.keys():
            self._max_resamples_per_iteration = self.config['max_resamples_per_iteration']

        self._max_approach_angle = np.inf
        if 'max_approach_angle' in self.config.keys():
            self._max_approach_angle = np.deg2rad(self.config['max_approach_angle'])

        # gripper parameters
        self._seed = None
        if self.config['deterministic']:
            self._seed = SEED
        self._gripper_width = np.inf
        if 'gripper_width' in self.config.keys():
            self._gripper_width = self.config['gripper_width']

        # optional, logging dir
        self._logging_dir = None
        if 'logging_dir' in self.config.keys():
            self._logging_dir = self.config['logging_dir']
            if not os.path.exists(self._logging_dir):
                os.mkdir(self._logging_dir)

    def select(self, grasps, q_values):
        """ Selects the grasp with the highest probability of success.
        Can override for alternate policies (e.g. epsilon greedy).
        """
        # sort
        logging.info('Sorting grasps')
        num_grasps = len(grasps)
        if num_grasps == 0:
            raise NoValidGraspsException('Zero grasps')
        grasps_and_predictions = zip(np.arange(num_grasps), q_values)
        grasps_and_predictions.sort(key = lambda x : x[1], reverse=True)

        # return top grasps
        if self._filters is None:
            return grasps_and_predictions[0][0]

        # filter grasps
        logging.info('Filtering grasps')
        i = 0
        while i < self._max_grasps_filter and i < len(grasps_and_predictions):
            index = grasps_and_predictions[i][0]
            grasp = grasps[index]
            valid = True
            for filter_name, is_valid in self._filters.iteritems():
                valid = is_valid(grasp)
                logging.info('Grasp {} filter {} valid: {}'.format(i, filter_name, valid))
                if not valid:
                    valid = False
                    break
            if valid:
                return index
            i += 1
        raise NoValidGraspsException('No grasps satisfied filters')

    def action(self, state):
        """ Plans the grasp with the highest probability of success on
        the given RGB-D image.

        Attributes
        ----------
        state : :obj:`RgbdImageState`
            image to plan grasps on

        Returns
        -------
        :obj:`GraspAction`
            grasp to execute
        """
        # check valid input
        if not isinstance(state, RgbdImageState):
            raise ValueError('Must provide an RGB-D image state.')

        if self._logging_dir is not None:
            policy_id = utils.gen_experiment_id()
            policy_dir = os.path.join(self._logging_dir, 'policy_output_%s' %(policy_id))
            while os.path.exists(policy_dir):
                policy_id = utils.gen_experiment_id()
                policy_dir = os.path.join(self._logging_dir, 'policy_output_%s' %(policy_id))
            os.mkdir(policy_dir)
            state_dir = os.path.join(policy_dir, 'state')
            state.save(state_dir)

        # parse state
        seed_set_start = time()
        rgbd_im = state.rgbd_im
        depth_im = rgbd_im.depth
        camera_intr = state.camera_intr
        segmask = state.segmask
        point_cloud_im = camera_intr.deproject_to_image(depth_im)
        normal_cloud_im = point_cloud_im.normal_cloud_im()

        # sample grasps
        grasps = self._grasp_sampler.sample(rgbd_im, camera_intr,
                                            self._num_seed_samples,
                                            segmask=segmask,
                                            visualize=self.config['vis']['grasp_sampling'],
                                            seed=self._seed)
        num_grasps = len(grasps)
        if num_grasps == 0:
            logging.warning('No valid grasps could be found')
            raise NoValidGraspsException()

        grasp_type = 'parallel_jaw'
        if isinstance(grasps[0], SuctionPoint2D):
            grasp_type = 'suction'

        logging.info('Sampled %d grasps' %(len(grasps)))
        logging.info('Computing the seed set took %.3f sec' %(time() - seed_set_start))

        # iteratively refit and sample
        for j in range(self._num_iters):
            logging.info('CEM iter %d' %(j))

            # predict grasps
            predict_start = time()
            q_values = self._grasp_quality_fn(state, grasps, params=self._config)
            logging.info('Prediction took %.3f sec' %(time()-predict_start))

            # sort grasps
            resample_start = time()
            q_values_and_indices = zip(q_values, np.arange(num_grasps))
            q_values_and_indices.sort(key = lambda x : x[0], reverse=True)

            if self.config['vis']['grasp_candidates']:
                # display each grasp on the original image, colored by predicted success
                norm_q_values = q_values #(q_values - np.min(q_values)) / (np.max(q_values) - np.min(q_values))
                vis.figure(size=(FIGSIZE,FIGSIZE))
                vis.imshow(rgbd_im.depth)#,
                           #vmin=self.config['vis']['vmin'],
                           #vmax=self.config['vis']['vmax'])
                for grasp, q in zip(grasps, norm_q_values):
                    vis.grasp(grasp, scale=2.0,
                              jaw_width=2.0,
                              show_center=False,
                              show_axis=True,
                              color=plt.cm.RdYlGn(q))
                vis.title('Sampled grasps iter %d' %(j))
                filename = None
                if self._logging_dir is not None:
                    filename = os.path.join(self._logging_dir, 'cem_iter_%d.png' %(j))
                vis.show(filename)

            # fit elite set
            elite_start = time()
            num_refit = max(int(np.ceil(self._gmm_refit_p * num_grasps)), 1)
            elite_q_values = [i[0] for i in q_values_and_indices[:num_refit]]
            elite_grasp_indices = [i[1] for i in q_values_and_indices[:num_refit]]
            elite_grasps = [grasps[i] for i in elite_grasp_indices]
            elite_grasp_arr = np.array([g.feature_vec for g in elite_grasps])

            if self.config['vis']['elite_grasps']:
                # display each grasp on the original image, colored by predicted success
                norm_q_values = (elite_q_values - np.min(elite_q_values)) / (np.max(elite_q_values) - np.min(elite_q_values))
                vis.figure(size=(FIGSIZE,FIGSIZE))
                vis.imshow(rgbd_im.depth)#,
                           #vmin=self.config['vis']['vmin'],
                           #vmax=self.config['vis']['vmax'])
                for grasp, q in zip(elite_grasps, norm_q_values):
                    vis.grasp(grasp, scale=1.5, show_center=False, show_axis=True,
                              color=plt.cm.RdYlGn(q))
                vis.title('Elite grasps iter %d' %(j))
                filename = None
                if self._logging_dir is not None:
                    filename = os.path.join(self._logging_dir, 'elite_set_iter_%d.png' %(j))
                vis.show(filename)

            # normalize elite set
            elite_grasp_mean = np.mean(elite_grasp_arr, axis=0)
            elite_grasp_std = np.std(elite_grasp_arr, axis=0)
            elite_grasp_std[elite_grasp_std == 0] = 1.0
            elite_grasp_arr = (elite_grasp_arr - elite_grasp_mean) / elite_grasp_std
            logging.info('Elite set computation took %.3f sec' %(time()-elite_start))

            # fit a GMM to the top samples
            num_components = max(int(np.ceil(self._gmm_component_frac * num_refit)), 1)
            uniform_weights = (1.0 / num_components) * np.ones(num_components)
            gmm = GaussianMixture(n_components=num_components,
                                  weights_init=uniform_weights,
                                  reg_covar=self._gmm_reg_covar)
            train_start = time()
            gmm.fit(elite_grasp_arr)
            logging.info('GMM fitting with %d components took %.3f sec' %(num_components, time()-train_start))

            # sample the next grasps
            grasps = []
            loop_start = time()
            num_tries = 0
            while len(grasps) < self._num_gmm_samples and num_tries < self._max_resamples_per_iteration:
                # sample from GMM
                sample_start = time()
                grasp_vecs, _ = gmm.sample(n_samples=self._num_gmm_samples)
                grasp_vecs = elite_grasp_std * grasp_vecs + elite_grasp_mean
                logging.info('GMM sampling took %.3f sec' %(time()-sample_start))

                # convert features to grasps and store if in segmask
                for grasp_vec in grasp_vecs:
                    feature_start = time()
                    if grasp_type == 'parallel_jaw':
                        # form grasp object
                        grasp = Grasp2D.from_feature_vec(grasp_vec,
                                                         width=self._gripper_width,
                                                         camera_intr=camera_intr)
                    elif grasp_type == 'suction':
                        # read depth and approach axis
                        u = int(min(max(grasp_vec[1], 0), depth_im.height-1))
                        v = int(min(max(grasp_vec[0], 0), depth_im.width-1))
                        grasp_depth = depth_im[u, v]

                        # approach_axis
                        grasp_axis = -normal_cloud_im[u, v]

                        # form grasp object
                        grasp = SuctionPoint2D.from_feature_vec(grasp_vec,
                                                                camera_intr=camera_intr,
                                                                depth=grasp_depth,
                                                                axis=grasp_axis)
                    logging.debug('Feature vec took %.5f sec' %(time()-feature_start))

                    bounds_start = time()
                    # check in bounds
                    if state.segmask is None or \
                        (grasp.center.y >= 0 and grasp.center.y < state.segmask.height and \
                         grasp.center.x >= 0 and grasp.center.x < state.segmask.width and \
                         np.any(state.segmask[int(grasp.center.y), int(grasp.center.x)] != 0) and \
                         grasp.approach_angle < self._max_approach_angle):

                        # check validity according to filters
                        grasps.append(grasp)
                    logging.debug('Bounds took %.5f sec' %(time()-bounds_start))
                    num_tries += 1

            # check num grasps
            num_grasps = len(grasps)
            if num_grasps == 0:
                logging.warning('No valid grasps could be found')
                raise NoValidGraspsException()
            logging.info('Resample loop took %.3f sec' %(time()-loop_start))
            logging.info('Resampling took %.3f sec' %(time()-resample_start))

        # predict final set of grasps
        predict_start = time()
        q_values = self._grasp_quality_fn(state, grasps, params=self._config)
        logging.info('Final prediction took %.3f sec' %(time()-predict_start))

        if self.config['vis']['grasp_candidates']:
            # display each grasp on the original image, colored by predicted success
            norm_q_values = q_values #(q_values - np.min(q_values)) / (np.max(q_values) - np.min(q_values))
            vis.figure(size=(FIGSIZE,FIGSIZE))
            vis.imshow(rgbd_im.depth)#,
                       #vmin=self.config['vis']['vmin'],
                       #vmax=self.config['vis']['vmax'])
            for grasp, q in zip(grasps, norm_q_values):
                vis.grasp(grasp, scale=2.0,
                          jaw_width=2.0,
                          show_center=False,
                          show_axis=True,
                          color=plt.cm.RdYlGn(q))
            vis.title('Final sampled grasps')
            filename = None
            if self._logging_dir is not None:
                filename = os.path.join(self._logging_dir, 'final_grasps.png')
            vis.show(filename)

        # select grasp
        index = self.select(grasps, q_values)
        grasp = grasps[index]
        q_value = q_values[index]
        if self.config['vis']['grasp_plan']:
            vis.figure()
            vis.imshow(rgbd_im.depth)#,
                       #vmin=self.config['vis']['vmin'],
                       #vmax=self.config['vis']['vmax'])
            vis.grasp(grasp, scale=5.0, show_center=False, show_axis=True, jaw_width=1.0, grasp_axis_width=0.2)
            vis.title('Best Grasp: d=%.3f, q=%.3f' %(grasp.depth,
                                                     q_value))
            filename = None
            if self._logging_dir is not None:
                filename = os.path.join(self._logging_dir, 'planned_grasp.png')
            vis.show(filename)

        # form return image
        image = state.rgbd_im.depth
        if isinstance(self._grasp_quality_fn, GQCnnQualityFunction):
            image_arr, _ = self._grasp_quality_fn.grasps_to_tensors([grasp], state)
            image = DepthImage(image_arr[0,...],
                               frame=state.rgbd_im.frame)

        # return action
        action = GraspAction(grasp, q_value, image)

        if self._logging_dir is not None:
            action_dir = os.path.join(policy_dir, 'action')
            action.save(action_dir)

        return action

class QFunctionRobustGraspingPolicy(CrossEntropyRobustGraspingPolicy):
    """ Optimizes a set of antipodal grasp candidates in image space using the
    cross entropy method with a GQ-CNN that estimates the Q-function
    for use in Q-learning.

    Notes
    -----
    Required configuration parameters are specified in Other Parameters

    Other Parameters
    ----------------
    reinit_pc1 : bool
        whether or not to reinitialize the pc1 layer of the GQ-CNN
    reinit_fc3: bool
        whether or not to reinitialize the fc3 layer of the GQ-CNN
    reinit_fc4: bool
        whether or not to reinitialize the fc4 layer of the GQ-CNN
    reinit_fc5: bool
        whether or not to reinitialize the fc5 layer of the GQ-CNN
    num_seed_samples : int
        number of candidate to sample in the initial set
    num_gmm_samples : int
        number of candidates to sample on each resampling from the GMMs
    num_iters : int
        number of sample-and-refit iterations of CEM
    gmm_refit_p : float
        top p-% of grasps used for refitting
    gmm_component_frac : float
        percentage of the elite set size used to determine number of GMM components
    gmm_reg_covar : float
        regularization parameters for GMM covariance matrix, enforces diversity of fitted distributions
    deterministic : bool, optional
        whether to set the random seed to enforce deterministic behavior
    gripper_width : float, optional
        width of the gripper in meters
    """
    def __init__(self, config):
        CrossEntropyRobustGraspingPolicy.__init__(self, config)
        QFunctionRobustGraspingPolicy._parse_config(self)
        self._setup_gqcnn()

    def _parse_config(self):
        """ Parses the parameters of the policy. """
        self._reinit_pc1 = self.config['reinit_pc1']
        self._reinit_fc3 = self.config['reinit_fc3']
        self._reinit_fc4 = self.config['reinit_fc4']
        self._reinit_fc5 = self.config['reinit_fc5']

    def _setup_gqcnn(self):
        """ Sets up the GQ-CNN. """
        # close existing session (from superclass initializer)
        self.gqcnn.close_session()

        # check valid output size
        if self.gqcnn.fc5_out_size != 1 and not self._reinit_fc5:
            raise ValueError('Q function must return scalar values')

        # reinitialize layers
        if self._reinit_fc5:
            self.gqcnn.fc5_out_size = 1

        # TODO: implement reinitialization of pc0
        self.gqcnn.reinitialize_layers(self._reinit_fc3,
                                       self._reinit_fc4,
                                       self._reinit_fc5)
        self.gqcnn.initialize_network(add_softmax=False)

class EpsilonGreedyQFunctionRobustGraspingPolicy(QFunctionRobustGraspingPolicy):
    """ Optimizes a set of antipodal grasp candidates in image space
    using the cross entropy method with a GQ-CNN that estimates the
    Q-function for use in Q-learning, and chooses a random antipodal
    grasp with probability epsilon.

    Notes
    -----
    Required configuration parameters are specified in Other Parameters

    Other Parameters
    ----------------
    epsilon : float
    """
    def __init__(self, config):
        QFunctionRobustGraspingPolicy.__init__(self, config)
        self._parse_config()

    def _parse_config(self):
        """ Parses the parameters of the policy. """
        self._epsilon = self.config['epsilon']

    @property
    def epsilon(self):
        return self._epsilon

    @epsilon.setter
    def epsilon(self, val):
        self._epsilon = val

    def greedy_action(self, state):
        """ Plans the grasp with the highest probability of success on
        the given RGB-D image.

        Attributes
        ----------
        state : :obj:`RgbdImageState`
            image to plan grasps on

        Returns
        -------
        :obj:`GraspAction`
            grasp to execute
        """
        return CrossEntropyRobustGraspingPolicy.action(self, state)

    def action(self, state):
        """ Plans the grasp with the highest probability of success on
        the given RGB-D image.

        Attributes
        ----------
        state : :obj:`RgbdImageState`
            image to plan grasps on

        Returns
        -------
        :obj:`GraspAction`
            grasp to execute
        """
        # take the greedy action with prob 1 - epsilon
        if np.random.rand() > self.epsilon:
            logging.debug('Taking greedy action')
            return CrossEntropyRobustGraspingPolicy.action(self, state)

        # otherwise take a random action
        logging.debug('Taking random action')

        # check valid input
        if not isinstance(state, RgbdImageState):
            raise ValueError('Must provide an RGB-D image state.')

        # parse state
        rgbd_im = state.rgbd_im
        camera_intr = state.camera_intr
        segmask = state.segmask

        # sample random antipodal grasps
        grasps = self._grasp_sampler.sample(rgbd_im, camera_intr,
                                            self._num_seed_samples,
                                            segmask=segmask,
                                            visualize=self.config['vis']['grasp_sampling'],
                                            seed=self._seed)

        num_grasps = len(grasps)
        if num_grasps == 0:
            logging.warning('No valid grasps could be found')
            raise NoValidGraspsException()

        # choose a grasp uniformly at random
        grasp_ind = np.random.choice(num_grasps, size=1)[0]
        grasp = grasps[grasp_ind]
        depth = grasp.depth

        # create transformed image
        image_tensor, pose_tensor = self.grasps_to_tensors([grasp], state)
        image = DepthImage(image_tensor[0,...])

        # predict prob success
        output_arr = self.gqcnn.predict(image_tensor, pose_tensor)
        q_value = output_arr[0,-1]

        # visualize planned grasp
        if self.config['vis']['grasp_plan']:
            scale_factor = float(self.gqcnn.im_width) / float(self._crop_width)
            scaled_camera_intr = camera_intr.resize(scale_factor)
            vis_grasp = Grasp2D(Point(image.center), 0.0, depth,
                                width=self._gripper_width,
                                camera_intr=scaled_camera_intr)
            vis.figure()
            vis.imshow(image)
            vis.grasp(vis_grasp, scale=1.5, show_center=False, show_axis=True)
            vis.title('Best Grasp: d=%.3f, q=%.3f' %(depth, q_value))
            vis.show()

        # return action
        return GraspAction(grasp, q_value, image)


import time
import threading
import urllib.parse
import subprocess
import json
import re

import fastparquet
from fastparquet import core

from webhdfs import WebHdfsFile

import pyarrow.parquet
import numpy

reader_cache = dict()
reader_cache_lock = threading.Lock()

def get_reader(path: str, storage_ip=None):
    reader_cache_lock.acquire()
    if path not in reader_cache.keys():
        reader_cache[path] = ParquetReader(path, storage_ip)
        # reader_cache[path] = ArrowParquetReader(path, storage_ip)
    reader_cache_lock.release()

    return reader_cache[path]

class ParquetReader(fastparquet.ParquetFile):
    def __init__(self, path: str, storage_ip=None):
        self.infile = WebHdfsFile(path, storage_ip=storage_ip)
        super(ParquetReader, self).__init__(self.infile)
        self.dtypes = self.dtypes.values()
        self.num_row_groups = len(self.row_groups)
        self.infile.close()

    def read_rg(self, index: int, columns: list):
        rg = self.row_groups[index]
        #  Allocate memory
        df, assign = self.pre_allocate(rg.num_rows, columns, {}, None)

        fin = self.infile.clone()

        for col in rg.columns:
            name = col.meta_data.path_in_schema[0]
            if name in columns:
                off = min((col.meta_data.dictionary_page_offset or col.meta_data.data_page_offset,
                           col.meta_data.data_page_offset))
                fin.prefetch(off, col.meta_data.total_compressed_size + 65536)
                core.read_col(col, self.schema, fin, assign=assign[name])
                print(f'RG {index} Column {name} total_uncompressed_size {col.meta_data.total_uncompressed_size} total_compressed_size {col.meta_data.total_compressed_size}')

        fin.close()
        del fin
        return df

class ArrowParquetReader(pyarrow.parquet.ParquetFile):
    def __init__(self, path: str, storage_ip=None):
        self.url = urllib.parse.urlparse(path)
        if self.url.scheme == 'webhdfs':
            self.infile = WebHdfsFile(path, storage_ip=storage_ip)
        else:
            self.infile = open(self.url.path, 'rb')

        super(ArrowParquetReader, self).__init__(self.infile)
        self.columns = self.metadata.schema.names
        self.dtypes = [numpy.dtype(d.to_pandas_dtype()) for d in self.schema_arrow.types]
        # self.num_row_groups = self.metadata.num_row_groups
        self.infile.close()

    def read_rg(self, index: int, columns: list):
        if self.url.scheme == 'webhdfs':
            fin = self.infile.clone()
        else:
            fin = open(self.url.path, 'rb')

        pfin = pyarrow.parquet.ParquetFile(fin, metadata=self.metadata)
        df = pfin.read_row_group(index, columns=columns)
        fin.close()
        return df

def get_storage_ip():
    result = subprocess.run('docker network inspect qflock-net'.split(' '), stdout=subprocess.PIPE)
    d = json.loads(result.stdout)

    for c in d[0]['Containers'].values():
        print(c['Name'], c['IPv4Address'].split('/')[0])
        if c['Name'] == 'qflock-storage':
            return c['IPv4Address'].split('/')[0]

def get_webhdfs_reader():
    fname = '/tpcds-parquet/store_sales.parquet/part-00000-40849ae6-37a1-4e22-848d-d4cfea002b2b-c000.snappy.parquet'

    # storage_ip = get_storage_ip()
    # return get_reader(f'webhdfs://{storage_ip}:9870/{fname}?user.name=peter', storage_ip)
    storage_ip = '127.0.0.1'
    return get_reader(f'webhdfs://{storage_ip}:9860/{fname}?user.name=peter', storage_ip)

def get_file_reader():
    fname = '/mnt/usb-SanDisk_Ultra_128GB/caerus-dikeHDFS/data/tpch-test-parquet-1g/lineitem.parquet/' \
            'part-00000-badcef81-d816-44c1-b936-db91dae4c15f-c000.snappy.parquet'

    return get_reader(f'file://dikehdfs:9870/{fname}?user.name=peter')

def get_network_stats(docker_name):
    # docker stats --no-stream --format "{{.NetIO}}" qflock-storage
    cmd = 'docker stats --no-stream --format "{{.NetIO}}" ' + docker_name
    result = subprocess.run(cmd.split(' '), stdout=subprocess.PIPE)
    stats = result.stdout.decode("utf-8")
    stats = stats.lower().replace('\n', '').replace('"', '').replace(' ', '').replace('b', '').split('/')
    print(stats)
    sizes = {'k': 1 << 10, 'm': 1 << 20, 'g': 1 << 30}

    for i in range(0, len(stats)):
        factor = 1
        if stats[i][-1] in sizes.keys():
            factor = sizes[stats[i][-1]]
            stats[i] = stats[i][:-1]

        stats[i] = float(stats[i]) * factor

    return stats

if __name__ == '__main__':
    # reader = get_file_reader()
    reader = get_webhdfs_reader();

    dtypes = [t.name for t in reader.dtypes]

    total_compressed = 0

    for c in range(0, len(reader.columns)):
        c_name = reader.columns[c]
        c_type = dtypes[c]
        c_total_uncompressed = reader.info['rows'] * 8  # We need to check c_type here
        c_total_compressed = 0
        for i in range(0, reader.num_row_groups):
            c_total_compressed += reader.row_groups[i].columns[c].meta_data.total_compressed_size

        total_compressed += c_total_compressed
        c_ratio = c_total_uncompressed / c_total_compressed
        print(f'{c_name:<24} {c_type:<10} {c_total_uncompressed:<10} {c_total_compressed:<10} {c_ratio:<10.1f}')

    delta = (reader.infile.size - total_compressed) * 100 / reader.infile.size
    print(f'Total compressed {total_compressed} File size {reader.infile.size} Delta {delta:.3f}%')

    # Read one column with Docker network stats
    stats_before = get_network_stats('qflock-storage')
    for i in range(0, reader.num_row_groups):
        reader.read_rg(i, ['ss_sold_date_sk'])
    stats_after = get_network_stats('qflock-storage')
    transfer_size = int(stats_after[1] - stats_before[1])
    print(f'Transfer size {transfer_size}')


#!/usr/bin/python

# This is not the more object oriented approach as the other with classes and such

from gi.repository import Gtk,WebKit

def on_quit(widget):
    Gtk.main_quit()

def refreshclick(widget):
    #Open this url
    webview.open('https://google.com/')

def runscript(*args):
    spinner.hide()
    #Run this script on page load
    #webview.execute_script("$('body').hide()")
    return

def end(*args):
    window.show_all()

window = Gtk.Window(type=Gtk.WindowType.TOPLEVEL)
window.set_position(Gtk.WindowPosition.CENTER)
window.set_default_size(1024, 768)
window.connect("destroy", on_quit)

toolbar = Gtk.Toolbar(toolbar_style=Gtk.ToolbarStyle.BOTH)
context = toolbar.get_style_context()
context.add_class(Gtk.STYLE_CLASS_PRIMARY_TOOLBAR)

refresh = Gtk.ToolButton()
refresh.set_label('Refresh')
refresh.connect('clicked',refreshclick)
toolbar.add(refresh)

webview = WebKit.WebView()
webview.props.settings.props.enable_default_context_menu = False
webview.connect('document_load_finished',runscript)
webview.connect('load_started',end)
spinner = Gtk.Spinner(margin_left=5, margin_right=5, margin_top=5, margin_bottom=5)
spinner.start()
fixed = Gtk.Fixed()
fixed.add(spinner)

fixed.set_halign(Gtk.Align.START)
fixed.set_valign(Gtk.Align.END)

overlay = Gtk.Overlay()
overlay.add(webview)
overlay.add_overlay(fixed)

vbox = Gtk.VBox()
vbox.pack_start(toolbar,False,False,0)
vbox.pack_start(overlay,True,True,0)

window.add(vbox)
window.show_all()
refreshclick(webview)

Gtk.main()


import itertools
import sys
sys.stdin = open('input.txt', 'r')

board = [list(input()) for _ in range(5)]
adj = [(-1, 0), (0, 1), (1, 0), (0, -1)]

students = []
for i in range(5):
    for j in range(5):
        students.append((i, j))

answer = 0

for group in itertools.combinations(students, 7):
    ds = 0
    for i, j in group:
        if board[i][j] == 'S':
            ds += 1
    if ds >= 4:
        visit = {}
        for pos in group:
            visit[pos] = False

        queue = [group[0]]
        visit[group[0]] = True

        while queue:
            pos1 = queue.pop(0)
            for di, dj in adj:
                pos2 = (pos1[0]+di, pos1[1]+dj)
                if pos2 in group and visit[pos2] is False:
                    visit[pos2] = True
                    queue.append(pos2)

        if False not in visit.values():
            answer += 1

print(answer)


from pyspark.sql import SparkSession
from pyspark.sql.types import *
import pyspark.sql.functions as F
# import logging
# logging.basicConfig(filename="C:\\Users\\aakan\\PycharmProjects\\logs\\pysparkloffile.log", filemode="w")
# log = logging.getLogger()
# log.setLevel(logging.WARN)

#creating function to create a spark session object
def create_sparkseesion():
    spark = SparkSession.builder\
                  .master("local[1]")\
                  .appName("assignment on pyspark")\
                  .getOrCreate()
    return spark

#creating a dataframe for product table
def product_dataframe(spark):
    product_schema=StructType([StructField(name='ProductName',dataType=StringType()),
                          StructField(name='Issue Date',dataType=StringType()),
                          StructField(name='Price',dataType=IntegerType()),
                          StructField(name='Brand',dataType=StringType()),
                          StructField(name='Country',dataType=StringType()),
                          StructField(name='Product_Number',dataType=IntegerType())
                          ])
    product_df = spark\
            .read\
            .format('csv')\
            .options(header=True)\
            .schema(product_schema)\
            .load("../../resource/product.csv")

    return product_df

#creating a function to convert  the issue date to timestamp format.
def issue_date_timestamp(product_df):
    df_timestamp_string = product_df.withColumn("Issue_Date_timsetampformat",F.from_unixtime(F.col("Issue Date")/1000,"yyyy-MM-dd'T'HH:mm:ssZZZZ"))
    return df_timestamp_string

#creating a function to convert  the issue date to date type.
def issue_date_datestamp(df_timestamp_string):
    df_date = df_timestamp_string.withColumn("Date",F.date_format(F.col("Issue_Date_timsetampformat"),"yyyy-MM-dd"))
    return df_date

#creating a function to remove starting space in Brand Column.
def remove_space(product_df):
    df_remove_space = product_df.withColumn("Brand_new",F.trim(F.col("Brand")))
    return df_remove_space

#creating a function to replace null values with empty spaces Country Column.
def replace_null(product_df):
    df_replace_null = product_df.withColumn("Country_new",F.when(F.col("Country") == "null", '').otherwise(F.col("Country")))
    return df_replace_null

#creating a function to create a dataframe for transaction table
def transaction_dataframe(spark):
    transaction_schema=StructType([StructField(name='SourceId',dataType=IntegerType()),
                          StructField(name='TransactionNumber',dataType=IntegerType()),
                          StructField(name='Language',dataType=StringType()),
                          StructField(name='ModelNumber',dataType=IntegerType()),
                          StructField(name='StartTime',dataType=StringType()),
                          StructField(name='ProductNumber',dataType=IntegerType())
                          ])
    transaction_df = spark\
            .read\
            .format('csv')\
            .options(header=True)\
            .schema(transaction_schema)\
            .load("../../resource/Product_transaction.csv")
    return transaction_df

# #creating a function to covert column names from camel case to snake case.

#creating a function to add a new column as start_time_ms and convert the StartTime column values to milliseconds.
def millisecond(transaction_df):
    df_millisecond = transaction_df.withColumn("StartTime_timestamp",F.to_timestamp(F.col("StartTime")))\
                                   .withColumn("start_time_ms",F.unix_timestamp(F.col("StartTime_timestamp")))\
                                   .drop("StartTime_timestamp")
    return df_millisecond

#creating a function to join both the dataframe
def joindataframe(df_replace_null,df_millisecond):
    joined_df = df_replace_null.join(df_millisecond,df_replace_null.Product_Number == df_millisecond.ProductNumber,"left")\
                                .drop("Product_Number")
    return joined_df

#creating a function to filter the country with language as EN
def filter_country(joined_df):
    df_filter = joined_df.filter(F.col("Language") == "EN")
    return df_filter
#


#!/usr/bin/env python
__author__ = "Simon Swolfs en Robbe de Budt"
__email__ = "simon.swolfs@student.kdg.be"
__version__ = "1.0"

# LIBRARIES
import time
import gpiozero
from gpiozero import Button
from gpiozero import LED

# VARIABLES
systemwork = True
howworks = {
    Button(26): LED(6),
    Button(16): LED(5),
    Button(25): LED(23),
    Button(17): LED(24)
}

while systemwork:
    for key, value in howworks.items():
        if key.value == 1:
            print('The button was pressed!')
            if value.value == 1:
                value.off()
            else:
                value.on()

        time.sleep(0.50)


from django.shortcuts import render
from .models import Room

# Create your views here.

home = [
    {'id':1, 'name':'Home One'},
    {'id':2, 'name':'Home Two'},
    {'id':3, 'name':'Home Three'},
]

roomData_niu = [
    {'id':1, 'name':'Frontend Developer'},
    {'id':2, 'name':'Backend Developer'},
    {'id':3, 'name':'Full Stack Developer'},
]

def home(request):
    roomData = Room.objects.all()
    print('---------')
    print(roomData)
    print('---------')
    context = {'roomData':roomData}
    return render(request, 'appPractice/home.html', context)

roomDetails = [
    {'id': 1, 'name':'Room One'},
    {'id': 2, 'name':'Room Two'},
    {'id': 3, 'name':'Room Three'},
]
def room(request, pk):
    roomName = None
    # context = {'roomData':roomData}
    # print('---------------')
    # print(pk)
    # print(int(pk))
    # print('---------------')
    # for i in roomData:
    #     print('---------------')
    #     print(i)
    #     print('---------------')
    #     if i['id'] == int(pk):
    #         roomName = i['name']
    roomName = Room.objects.get(id=pk)

    context = {'roomName':roomName}
    return render(request, 'appPractice/room.html', context)


import numpy as np
import tensorflow as tf
from gensim.models import Word2Vec

tf.reset_default_graph()
sess = tf.InteractiveSession()

class Seq2seq():
	def __init__(self, rnn_size, num_layers, batch_size, vocab_size, mode,
				 max_encoder_steps, max_decoder_steps, embedding_size, beam_size=0):
		self.rnn_size = rnn_size
		self.num_layers = num_layers
		self.batch_size = batch_size
		self.vocab_size = vocab_size
		self.mode = mode
		self.max_encoder_steps = max_encoder_steps
		self.max_decoder_steps = max_decoder_steps
		self.beam_size = beam_size
		self.embedding_size = embedding_size

		self.build_model()

		self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=50)

	def build_model(self):
		# Model input & output

		# shape = (batch_size, max_seq_len)
		self.encoder_inputs  = tf.placeholder(shape=(None, None), dtype=tf.int32, name='encoder_inputs')
		self.encoder_inputs_length = tf.placeholder(shape=[None], dtype=tf.int32, name='encoder_inputs_length')
		# shape = (batch_size, max_seq_len)
		self.decoder_inputs  = tf.placeholder(shape=(None, None), dtype=tf.int32, name='decoder_inputs')
		self.decoder_targets = tf.placeholder(shape=(None, None), dtype=tf.int32, name='decoder_targets')
		self.decoder_targets_length = tf.placeholder(shape=[None], dtype=tf.int32, name='decoder_targets_length')

		self.keep_prob_placeholder = tf.placeholder(dtype=tf.float32)
		self.sampling_probability = tf.placeholder(shape=[], dtype=tf.float32)

		# embeddings (gensin word2vec)
		w2v_model = Word2Vec.load("word2vec.model")
		w2v_embedding = np.zeros([len(w2v_model.wv.index2word), self.embedding_size])
		for idx, word in enumerate(w2v_model.wv.index2word):
			w2v_embedding[idx, :] = w2v_model[word]
		embeddings = tf.Variable(tf.convert_to_tensor(w2v_embedding, np.float32), trainable=False, name='embedding')

		encoder_inputs_embedded = tf.nn.embedding_lookup(embeddings, self.encoder_inputs)
		encoder_inputs_embedded = tf.transpose(encoder_inputs_embedded, [1,0,2])

		decoder_inputs_embedded = tf.nn.embedding_lookup(embeddings, self.decoder_inputs)
		decoder_inputs_embedded = tf.transpose(decoder_inputs_embedded, [1,0,2])

		# Encoder

		encoder_cell = self._create_cell()
		encoder_outputs, encoder_final_state = tf.nn.dynamic_rnn(encoder_cell, encoder_inputs_embedded, dtype = tf.float32, time_major = True)
#####???
		# encoder_outputs, encoder_final_state = tf.nn.bidirectional_dynamic_rnn(
		# 												cell_fw=self._create_cell(),
		# 												cell_bw=self._create_cell(),
		# 												inputs=encoder_inputs_embedded,
		# 												sequence_length=self.encoder_inputs_length,
		# 												dtype=tf.float32,
		# 												time_major=True )
#####???

		encoder_outputs = tf.transpose(encoder_outputs, [1,0,2]) # shape(batch_size,max_step,rnn_size)
		if self.mode == 'test':
			encoder_outputs = tf.contrib.seq2seq.tile_batch(encoder_outputs, multiplier=self.beam_size)

		# Decoder

		# Decoder cell with attention
		attention_mechanism = tf.contrib.seq2seq.BahdanauAttention(
													num_units=self.rnn_size,
													memory=encoder_outputs,
													normalize=True)

		decoder_cell = tf.contrib.seq2seq.AttentionWrapper(
												cell=self._create_cell(),
												attention_mechanism=attention_mechanism,
												attention_layer_size=self.rnn_size,
												name='Attention_Wrapper')

		output_layer = tf.layers.Dense(self.vocab_size, kernel_initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1))

		max_target_sequence_length = tf.reduce_max(self.decoder_targets_length, name='max_target_len')
		mask = tf.sequence_mask(self.decoder_targets_length, max_target_sequence_length, dtype=tf.float32, name='masks')

		# Define helper based on 'train' or 'infer' mode
		if self.mode == 'train':
			decoder_initial_state = decoder_cell.zero_state(batch_size=self.batch_size, dtype=tf.float32).clone(cell_state=encoder_final_state)
#####???
			training_helper = tf.contrib.seq2seq.ScheduledEmbeddingTrainingHelper(
											inputs=decoder_inputs_embedded,
											sequence_length=self.decoder_targets_length,
											embedding=embeddings,
											sampling_probability=self.sampling_probability,
											time_major=True)
#####???
			# training_helper = tf.contrib.seq2seq.TrainingHelper(
			# 								inputs=decoder_inputs_embedded,
			# 								sequence_length=self.decoder_targets_length,
			# 								# memory_sequence_length=?????????,
			# 								time_major=True, name='training_helper')

			# Decoder
			training_decoder = tf.contrib.seq2seq.BasicDecoder(
											cell=decoder_cell, helper=training_helper,
											initial_state=decoder_initial_state,
											output_layer=output_layer)

			decoder_outputs, _, _ = tf.contrib.seq2seq.dynamic_decode(
											decoder=training_decoder,
											impute_finished=True,
											maximum_iterations=max_target_sequence_length)

			# Calculate loss with sequence_loss
			decoder_logits_train = tf.identity(decoder_outputs.rnn_output)
			self.decoder_predict_train = tf.argmax(decoder_logits_train, axis=-1, name='decoder_pred_train')
			# decoder_logits_train = tf.transpose(decoder_logits_train, [1,0,2])

			self.loss = tf.contrib.seq2seq.sequence_loss(
									logits=decoder_logits_train,
									targets=self.decoder_targets,
									weights=mask)
			optimizer = tf.train.AdamOptimizer()

			# Clip gradient if gradient is too large
			gradients = optimizer.compute_gradients(self.loss)
			capped_gradients = [ (tf.clip_by_value(grad, -0.3, 0.3), var) for grad, var in gradients if grad is not None ]
			self.train_op = optimizer.apply_gradients(capped_gradients)

		else:
			# 如果使用beam_search，则需要将encoder的输出进行tile_batch，其实就是复制beam_size份。
			self.tiled_encoder_final_state = tf.contrib.seq2seq.tile_batch(encoder_final_state, multiplier=self.beam_size)

			batch_size = self.batch_size * self.beam_size

			decoder_initial_state = decoder_cell.zero_state(batch_size=batch_size, dtype=tf.float32).clone(cell_state=self.tiled_encoder_final_state)

			inference_decoder = tf.contrib.seq2seq.BeamSearchDecoder(
											cell=decoder_cell,
											embedding=embeddings,
											start_tokens=tf.fill([self.batch_size, ], 1),	# <BOS> = 1
											end_token=2, 									# <EOS> = 2
											initial_state=decoder_initial_state,
											beam_width=self.beam_size,
											output_layer=output_layer)

			inference_decoder_outputs, _, _ = tf.contrib.seq2seq.dynamic_decode(
															decoder=inference_decoder,
															maximum_iterations=self.max_decoder_steps)

			self.decoder_predict_decode = inference_decoder_outputs.predicted_ids
			self.decoder_predict_logits = inference_decoder_outputs.beam_search_decoder_output

	def _create_cell(self):
		def single_cell():
			cell = tf.contrib.rnn.LSTMCell(self.rnn_size)
			cell = tf.contrib.rnn.DropoutWrapper(cell, output_keep_prob=self.keep_prob_placeholder)
			return cell
		cell = tf.contrib.rnn.MultiRNNCell([ single_cell() for _ in range(self.num_layers) ])
		return cell

	def train(self, sess, encoder_inputs, encoder_inputs_length, decoder_inputs,
			  decoder_targets, decoder_targets_length, sampling_probability):
		feed_dict = { self.encoder_inputs : encoder_inputs,
					  self.encoder_inputs_length : encoder_inputs_length,
					  self.decoder_inputs : decoder_inputs,
					  self.decoder_targets : decoder_targets,
					  self.decoder_targets_length : decoder_targets_length,
					  self.sampling_probability : sampling_probability,
					  self.keep_prob_placeholder : 0.8 }

		_, loss = sess.run([self.train_op, self.loss], feed_dict=feed_dict)
		return loss

	def infer(self, sess, encoder_inputs):
		feed_dict = { self.encoder_inputs : encoder_inputs,
					  self.keep_prob_placeholder : 1.0 }

		predict, logits = sess.run([self.decoder_predict_decode, self.decoder_predict_logits], feed_dict=feed_dict)
		return predict, logits


# Project Imports
import torch
from torch import tensor
from torch.utils.data import DataLoader

from dataloaders.dataset_base import DatasetDepth, DatasetSemantic, DatasetRGB, \
    PathsHandlerDepthDense, PathsHandlerSemantic
from dataloaders.scritps.prepare_synthia import available_splits
from misc import transforms as tf_prep

# Packages
import matplotlib.pyplot as plt
import matplotlib.colors
import numpy as np
import random
import glob
import os
import PIL.Image as pil
import cv2
from torchvision import transforms
import re

from utils.constans import IGNORE_VALUE_DEPTH, IGNORE_INDEX_SEMANTIC

def decompose_rgb_path(path_file):
    """
    Decomposes the path to instance and splits the path of the file into meaningful components
    :param path_file: path to the file.
    :return:
        base_path: base path to the dataset
        information_type: e.g. 'DEPTH', 'RGB', 'MASK'
        class_name: e.g 'person', 'rider'
        instance_id: id of the instance
        file_format: extension of the file e.g. '.png'
    """
    path_components = {}
    par_dir = path_file

    # Go through the rgb path and decompose its parent directory names step for step and store in path_components
    for i in range(3):
        path_components[i] = os.path.basename(par_dir)
        par_dir = os.path.dirname(par_dir)

    # Take the path components out as needed

    file_format = re.search('(\.[A-Za-z0-9]*)', path_components[0]).group(0)
    instance_id = re.search('([A-Za-z0-9]*)\.', path_components[0]).group(0)
    instance_id = instance_id[:-1]  # remove '.' from id

    # example of return: [basepath, 'DENSE', 'person', '60', '.png']
    return [par_dir, path_components[2], path_components[1], instance_id, file_format]

class _PathsSynthiaInstances(PathsHandlerSemantic, PathsHandlerDepthDense):
    """
    Class for handling the paths to the images and labels.
    """

    def __init__(self, mode, _, cfg):
        """
        :param mode: mode of usage: 'train', 'val', 'test'
        :param cfg: configuration
        """

        super(_PathsSynthiaInstances, self).__init__(mode, None, cfg)

    def get_rgb_image_paths(self, mode, class_name='*', frame_ids='*', file_format=".png", *args, **kwargs):
        """
        Gets all the RGB image paths.
        :param mode: mode of usage: 'train', 'val', 'test'
        :param class_name: name of the class of the instance, e.g 'person'
        :param frame_ids: ids of frames to be selected, * = all images
        :param file_format: format of the files
        :param split: split name of the data
        """
        return sorted(
            glob.glob(
                self.get_rgb_image_path(self.path_base, class_name, frame_ids, file_format)
            )
        )

    def get_semantic_label_paths(self, mode, class_name='*', frame_ids='*', file_format=".png", *args, **kwargs):
        """
        Gets all the semantic label paths.
        :param mode: mode of usage: 'train', 'val', 'test'
        :param class_name: name of the class of the instance, e.g 'person'
        :param frame_ids: ids of frames to be selected, * = all images
        :param file_format: format of the files
        :param split: split name of the data
        """
        return sorted(
            glob.glob(
                self.get_semantic_label_path(self.path_base, class_name, frame_ids, file_format)
            )
        )

    def get_gt_depth_image_paths(self, mode, class_name='*', frame_ids='*', file_format=".png", *args, **kwargs):
        """
        Gets all the depth label paths.
        :param mode: mode of usage: 'train', 'val', 'test'
        :param class_name: name of the class of the instance, e.g 'person'
        :param frame_ids: ids of frames to be selected, * = all images
        :param file_format: format of the files
        :param split: split name of the data
        """
        return sorted(
            glob.glob(
                self.get_depth_label_path(self.path_base, class_name, frame_ids, file_format)
            )
        )

    @staticmethod
    def get_rgb_image_path(path_base, class_name, frame_ids, file_format):
        """
            Builds path to an rgb image.
            Default path: path_base/RGB/*.png
            Specific path: path_base/RGB/0000001.png
            :param path_base: path to directory containing images and labels
            :param frame_ids: 7-digit string id or '*'
            :param file_format: extensions of files '.png'
            :return: single path string
        """
        path = os.path.join(path_base, 'RGB', class_name, frame_ids + file_format)
        if os.path.exists(path) or not frame_ids.isnumeric():
            # Apply path exists check only if exact id is provided, no regex like '*'
            return path
        else:
            return None

    @staticmethod
    def get_semantic_label_path(path_base, class_name, frame_ids, file_format):
        """
            Builds path to an semantic label image.
            Default path: path_base/GT/LABEL/*.png
            Specific path: path_base/GT/LABEL/0000001.png
            :param path_base: path to directory containing images and labels
            :param frame_ids: 7-digit string id or '*'
            :param file_format: extensions of files '.png'
            :return: single path string
        """
        path = os.path.join(path_base, 'MASK', class_name, frame_ids + file_format)
        if os.path.exists(path) or not frame_ids.isnumeric():
            # Apply path exists check only if exact id is provided, no regex like '*'
            return path
        else:
            return None

    @staticmethod
    def get_depth_label_path(path_base, class_name, frame_ids, file_format):
        """
            Builds path to an rgb image.
            Default path: path_base/RGB/*.png
            Specific path: path_base/RGB/0000001.png
            :param path_base: path to directory containing images and labels
            :param frame_ids: 7-digit string id or '*'
            :param file_format: extensions of files '.png'
            :return: single path string
        """
        path = os.path.join(path_base, 'DEPTH', class_name, frame_ids + file_format)
        if os.path.exists(path) or not frame_ids.isnumeric():
            # Apply path exists check only if exact id is provided, no regex like '*'
            return path
        else:
            return None

class SynthiaInstancesDataset(DatasetRGB, DatasetSemantic, DatasetDepth):
    """
        RGB Images annotated with Depth and Semantic classes. This dataset does not provide any image sequences.
    """

    def __init__(self, mode, split, cfg):
        """
        :param mode: mode of usage, synthia only supports 'train'
        :param split: split name of the data, currently only None is supported
        :param cfg: configuration
        """
        # The dataset does not contain sequences therefore the offsets refer only to the selected RGB image(offset == 0)
        assert len(cfg.dataset.rgb_frame_offsets) == 1
        assert cfg.dataset.rgb_frame_offsets[0] == 0
        assert split is None, 'Synthia_Rand_Cityscapes currently does not support the spilt of the data!'
        assert mode == 'train', 'Synthia_Rand_Cityscapes does not support any other mode than train!'

        self.paths = _PathsSynthiaInstances(mode, None, cfg)

        super(SynthiaInstancesDataset, self).__init__(pathsObj=self.paths, cfg=cfg)

        # Stuff related to class encoding -------------------------------------------------------------

        # classes considered void have no equivalent valid label in cityscapes
        # See: https://www.cityscapes-dataset.com/dataset-overview/ for void and valid classes
        if self.cfg.dataset.num_classes != 16:
            raise ValueError(f'Synthia not implemented for {self.cfg.dataset.num_classes} classes. \n'
                             f'Only compatible with Cityscapes with 16 classes: ids from 0 to 15!')

        self.ignore_index = IGNORE_INDEX_SEMANTIC
        self.ignore_value = IGNORE_VALUE_DEPTH

        # Stuff related to augmentations --------------------------------------------------------------

        self.do_normalization = self.cfg.dataset.img_norm

        # todo: validate these values
        self.mean = torch.tensor([[[0.314747602, 0.277402550, 0.248091921]]]).transpose(0, 2)
        self.var = tensor([[[1.0, 1.0, 1.0]]]).transpose(0, 2)

        self.resize_factor = self.cfg.dataset.resize_factor

    def __len__(self):
        """
        Returns the number of rgb images for the mode in the split.
        :return:
        """
        return len(self.paths.paths_rgb)

    def __getitem__(self, index):
        """
            Collects the rgb images of the sequence and the label of the index image and returns them as a dict.
            Images can be accessed by dict_variable[("rgb", offset)] with e.g. offset = 0 for the index image.
            Labels can be accessed by dict_variable["gt"].
            :param index: Index of an image with offset 0 in the sequence.
            :return: dict of images and label
        """
        rgb_path_file = self.paths.paths_rgb[index]
        path_components = decompose_rgb_path(rgb_path_file)  # [basepath, 'DENSE', 'person', '60', '.png']
        base_path = path_components[0]
        class_name = path_components[2]
        instance_id = path_components[3]
        file_format = path_components[4]

        rgb_img = self.get_rgb(rgb_path_file)

        # Get all required training elements
        semantic_path_file = os.path.join(base_path, 'MASK', class_name, instance_id + file_format)
        gt_semantic = self.get_semantic(semantic_path_file)

        depth_path_file = os.path.join(base_path, 'DEPTH', class_name, instance_id + file_format)
        gt_depth_dense = self.get_depth(depth_path_file)

        if self.mode == 'train':
            rgb_img, gt_semantic, gt_depth_dense = self.transform_train(rgb_img, gt_semantic, gt_depth_dense)
        elif self.mode == 'val' or self.mode == 'test':
            raise ValueError("Synthia only available for training")
        else:
            assert False, "The mode {} is not defined!".format(self.mode)

        # Create the map to be outputted
        data = {}

        data["semantic"] = gt_semantic

        data["depth_dense"] = gt_depth_dense

        data['rgb'] = rgb_img

        return data

    def transform_train(self, rgb_img, gt_semantic, gt_depth_dense):
        """
            Transforms the rgb images and the semantic ground truth for training.
            :param gt_depth_dense: depth ground truth of image with offset = 0
            :param rgb_img: rgb image.
            :param gt_semantic: semantic ground truth of the image with offset = 0
            :return: dict of transformed rgb images and transformed label
        """
        if self.cfg.dataset.debug:
            do_flip = random.random() > 0.5
            do_aug = False
        else:
            do_flip = random.random() > 0.5
            do_aug = random.random() > 0.5

        # Get the transformation objects
        tf_rgb_train = self.tf_rgb_train(tgt_size=self.feed_img_size,
                                         do_flip=do_flip,
                                         do_normalization=self.do_normalization,
                                         mean=self.mean,
                                         var=self.var,
                                         resize_factor=self.resize_factor)

        tf_semantic_train = self.tf_semantic_train(tgt_size=self.feed_img_size,
                                                   do_flip=do_flip,
                                                   resize_factor=self.resize_factor)

        tf_depth_dense_train = self.tf_depth_train(tgt_size=self.feed_img_size,
                                                   do_flip=do_flip,
                                                   resize_factor=self.resize_factor)

        # Apply transformations

        rgb_img = tf_rgb_train(rgb_img)

        gt_semantic = tf_semantic_train(gt_semantic) if gt_semantic is not None else None

        gt_depth_dense = tf_depth_dense_train(gt_depth_dense) if gt_depth_dense is not None else None

        return rgb_img, gt_semantic, gt_depth_dense

    def get_rgb(self, path_file, *args):
        """
        Loads the PIL Image form the path_file.

        :param path_file: path to the image
        :return: PIL RGB Images.
        """
        if path_file is None:
            return None
        else:
            assert os.path.exists(path_file), "The file {} does not exist!".format(path_file)
            img = pil.open(path_file).convert('RGB')
            return img

    def get_semantic(self, path_file):
        """
            Loads the Labels with opencv (loading with pillow destroys encoding).
            :param path_file: path to the label
            :return: numpy ndarray of semantic ground truth.
        """
        if path_file is None:
            return None
        else:
            assert os.path.exists(path_file), "The file {} does not exist!".format(path_file)

            label = cv2.imread(path_file, cv2.IMREAD_UNCHANGED)
            return label

    def get_depth(self, path_file):
        """
            Loads the depth as PIL Image form the path_file.
            :param path_file: path to the label
            :return: PIL RGB Images.
        """
        if path_file is None:
            return None
        else:
            assert os.path.exists(path_file), "The file {} does not exist!".format(path_file)

            label = cv2.imread(path_file, cv2.IMREAD_UNCHANGED).astype(np.int32)
            return label

    @staticmethod
    def tf_rgb_train(tgt_size, do_flip, do_normalization, mean, var, resize_factor):
        """
        Transformations of the rgb image during training.
        :param resize_factor: factor by which to down sample the resolution, e.g 4 --> 1/4 resolution
        :param tgt_size: target size of the images after resize operation
        :param do_flip: True if the image should be horizontally flipped else False
        :param do_normalization: true if image should be normalized
        :param mean: mean of R,G,B of rgb images
        :param var: variance in R,G,B of rgb images

        :return: Transformation composition
        """
        return transforms.Compose(
            [
                tf_prep.PILResizeByFactor(pil.BICUBIC, resize_factor),
                tf_prep.PILHorizontalFlip(do_flip),
                tf_prep.PrepareForNet(do_normalization, mean, var)
            ]
        )

    @staticmethod
    def tf_semantic_train(tgt_size, do_flip, resize_factor):
        """
            Transformations of the label during training.
            :param tgt_size: target size of the labels after resize operation
            :param do_flip: True if the image should be horizontally flipped else False
            :param resize_factor: factor by which to down sample the resolution, e.g 4 --> 1/4 resolution

            :return: Transformation composition
        """
        return transforms.Compose(
            [
                tf_prep.CV2ResizeByFactor(interpolation=cv2.INTER_NEAREST, factor=resize_factor),
                tf_prep.CV2HorizontalFlip(do_flip=do_flip),
                tf_prep.ToInt64Array(),
            ]
        )

    @staticmethod
    def tf_depth_train(tgt_size, do_flip, resize_factor):
        """
            Transformations of the depth label during training.
            :param tgt_size: target size of the labels after resize operation
            :param do_flip: True if the image should be horizontally flipped else False
            :param resize_factor: factor by which to down sample the resolution, e.g 4 --> 1/4 resolution

            :return: Transformation composition
        """
        return transforms.Compose(
            [
                tf_prep.CV2ResizeByFactor(interpolation=cv2.INTER_NEAREST, factor=resize_factor),
                tf_prep.CV2HorizontalFlip(do_flip=do_flip),
                tf_prep.TransformToDepthSynthia(),
                tf_prep.PrepareForNet(do_normaliazion=False)
            ]
        )

    # ---------------------------functions below are useless but need to be defined--------------------------------------
    def transform_val(self, *args, **kwargs):
        pass

    @staticmethod
    def tf_rgb_val(*args, **kwargs):
        pass

    @staticmethod
    def tf_semantic_val(*args, **kwargs):
        pass

    @staticmethod
    def tf_depth_val(*args, **kwargs):
        pass

if __name__ == '__main__':
    from cfg.config_dataset import get_cfg_dataset_defaults
    path_16 = r'.\cfg\yaml_files\dataloaders\synthia_instances.yaml'
    cfg_16 = get_cfg_dataset_defaults()
    cfg_16.merge_from_file(path_16)
    cfg_16.freeze()

    dataset = SynthiaInstancesDataset('train', None, cfg_16)
    loader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0, pin_memory=True, drop_last=True)

    for idx, data in enumerate(loader):
        from utils.plotting_like_cityscapes_utils import visu_depth_prediction as vdp
        fig, (ax1, ax2, ax3) = plt.subplots(1, 3)
        data['rgb'][0] = data['rgb'][0] + torch.tensor([[[0.314747602, 0.277402550, 0.248091921]]]).transpose(0, 2)
        ax1.imshow(data['rgb'][0].numpy().transpose(1, 2, 0))
        print(data['rgb'][0].shape)
        ax2.imshow(data['semantic'][0])
        print(data['semantic'][0].shape)
        print(f'Min depth: {torch.min(data["depth_dense"])}')
        ax3.imshow(vdp(1 / data['depth_dense'][0]))
        print(data['depth_dense'][0].shape)
        plt.show()