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

from __future__ import absolute_import, division, print_function

import os
import tensorflow as tf
import tensorflow.contrib.slim as slim
import numpy as np

from libs.networks import resnet
from libs.networks import mobilenet_v2
from libs.box_utils import encode_and_decode
from libs.box_utils import boxes_utils
from libs.box_utils import anchor_utils
from libs.configs import cfgs
from libs.losses import losses
from libs.losses import tfapi_loss
from libs.box_utils import show_box_in_tensor
from libs.detection_oprations.proposal_opr import postprocess_rpn_proposals
from libs.detection_oprations.anchor_target_layer_without_boxweight import anchor_target_layer
from libs.detection_oprations.proposal_target_layer import proposal_target_layer
from tensorflow.python import pywrap_tensorflow

class DetectionNetwork(object):

    def __init__(self, base_network_name, is_training):

        self.base_network_name = base_network_name
        self.is_training = is_training
        self.num_anchors_per_location = len(cfgs.ANCHOR_SCALES) * len(cfgs.ANCHOR_RATIOS)

    def build_base_network(self, input_img_batch):

        if self.base_network_name.startswith('resnet_v1'):
            return resnet.resnet_base(input_img_batch, scope_name=self.base_network_name, is_training=self.is_training)

        elif self.base_network_name.startswith('MobilenetV2'):
            return mobilenet_v2.mobilenetv2_base(input_img_batch, is_training=self.is_training)

        else:
            raise ValueError('Sry, we only support resnet or mobilenet_v2')

    def postprocess_fastrcnn(self, rois, bbox_ppred, scores, img_shape, scope):
        '''

        :param rois:[-1, 4]
        :param bbox_ppred: [-1, (cfgs.Class_num+1) * 4]
        :param scores: [-1, cfgs.Class_num + 1]
        :return:
        '''

        with tf.name_scope('postprocess_fastrcnn_{}'.format(scope)):
            rois = tf.stop_gradient(rois)
            scores = tf.stop_gradient(scores)
            bbox_ppred = tf.reshape(bbox_ppred, [-1, cfgs.CLASS_NUM + 1, 4])
            bbox_ppred = tf.stop_gradient(bbox_ppred)

            bbox_pred_list = tf.unstack(bbox_ppred, axis=1)
            score_list = tf.unstack(scores, axis=1)

            allclasses_boxes = []
            allclasses_scores = []
            categories = []
            for i in range(1, cfgs.CLASS_NUM+1):

                # 1. decode boxes in each class
                tmp_encoded_box = bbox_pred_list[i]
                tmp_score = score_list[i]
                tmp_decoded_boxes = encode_and_decode.decode_boxes(encoded_boxes=tmp_encoded_box,
                                                                   reference_boxes=rois,
                                                                   scale_factors=cfgs.ROI_SCALE_FACTORS)
                # tmp_decoded_boxes = encode_and_decode.decode_boxes(boxes=rois,
                #                                                    deltas=tmp_encoded_box,
                #                                                    scale_factor=cfgs.ROI_SCALE_FACTORS)

                # 2. clip to img boundaries
                tmp_decoded_boxes = boxes_utils.clip_boxes_to_img_boundaries(decode_boxes=tmp_decoded_boxes,
                                                                             img_shape=img_shape)

                # 3. NMS
                keep = tf.image.non_max_suppression(
                    boxes=tmp_decoded_boxes,
                    scores=tmp_score,
                    max_output_size=cfgs.FAST_RCNN_NMS_MAX_BOXES_PER_CLASS,
                    iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD)

                perclass_boxes = tf.gather(tmp_decoded_boxes, keep)
                perclass_scores = tf.gather(tmp_score, keep)

                allclasses_boxes.append(perclass_boxes)
                allclasses_scores.append(perclass_scores)
                categories.append(tf.ones_like(perclass_scores) * i)

            final_boxes = tf.concat(allclasses_boxes, axis=0)
            final_scores = tf.concat(allclasses_scores, axis=0)
            final_category = tf.concat(categories, axis=0)

            if self.is_training:
                '''
                in training. We should show the detecitons in the tensorboard. So we add this.
                '''
                kept_indices = tf.reshape(tf.where(tf.greater_equal(final_scores, cfgs.SHOW_SCORE_THRSHOLD)), [-1])

                final_boxes = tf.gather(final_boxes, kept_indices)
                final_scores = tf.gather(final_scores, kept_indices)
                final_category = tf.gather(final_category, kept_indices)

        return final_boxes, final_scores, final_category

    def postprocess_cascade(self, rois, bbox_ppred, scope):
        '''

        :param rois:[-1, 4]
        :param bbox_ppred: bbox_ppred: [-1, 4]
        :param scores: [-1, 1]
        :return:

        '''

        with tf.name_scope('postprocess_cascade_{}'.format(scope)):
            rois = tf.stop_gradient(rois)
            bbox_ppred = tf.stop_gradient(bbox_ppred)

            # decode boxes
            decoded_boxes = encode_and_decode.decode_boxes(encoded_boxes=bbox_ppred,
                                                           reference_boxes=rois,
                                                           scale_factors=cfgs.ROI_SCALE_FACTORS)

        return decoded_boxes

    def roi_pooling(self, feature_maps, rois, img_shape):
        '''
        Here use roi warping as roi_pooling

        :param featuremaps_dict: feature map to crop
        :param rois: shape is [-1, 4]. [x1, y1, x2, y2]
        :return:
        '''

        with tf.variable_scope('ROI_Warping'):
            img_h, img_w = tf.cast(img_shape[1], tf.float32), tf.cast(img_shape[2], tf.float32)
            N = tf.shape(rois)[0]
            x1, y1, x2, y2 = tf.unstack(rois, axis=1)

            normalized_x1 = x1 / img_w
            normalized_x2 = x2 / img_w
            normalized_y1 = y1 / img_h
            normalized_y2 = y2 / img_h

            normalized_rois = tf.transpose(
                tf.stack([normalized_y1, normalized_x1, normalized_y2, normalized_x2]), name='get_normalized_rois')

            normalized_rois = tf.stop_gradient(normalized_rois)

            cropped_roi_features = tf.image.crop_and_resize(feature_maps, normalized_rois,
                                                            box_ind=tf.zeros(shape=[N, ],
                                                                             dtype=tf.int32),
                                                            crop_size=[cfgs.ROI_SIZE, cfgs.ROI_SIZE],
                                                            name='CROP_AND_RESIZE'
                                                            )
            roi_features = slim.max_pool2d(cropped_roi_features,
                                           [cfgs.ROI_POOL_KERNEL_SIZE, cfgs.ROI_POOL_KERNEL_SIZE],
                                           stride=cfgs.ROI_POOL_KERNEL_SIZE)

        return roi_features

    def build_fastrcnn(self, feature_to_cropped, rois, img_shape, scope):

        with tf.variable_scope('Fast-RCNN_{}'.format(scope)):
            # 5. ROI Pooling
            with tf.variable_scope('rois_pooling'):
                pooled_features = self.roi_pooling(feature_maps=feature_to_cropped, rois=rois, img_shape=img_shape)

            # 6. inferecne rois in Fast-RCNN to obtain fc_flatten features
            if self.base_network_name.startswith('resnet'):
                fc_flatten = resnet.restnet_head(input=pooled_features,
                                                 is_training=self.is_training,
                                                 scope_name=self.base_network_name,
                                                 stage=scope)
            elif self.base_network_name.startswith('Mobile'):
                fc_flatten = mobilenet_v2.mobilenetv2_head(inputs=pooled_features,
                                                           is_training=self.is_training,
                                                           stage=scope)
            else:
                raise NotImplementedError('only support resnet and mobilenet')

            # 7. cls and reg in Fast-RCNN
            # tf.variance_scaling_initializer()
            # tf.VarianceScaling()
            with slim.arg_scope([slim.fully_connected], weights_regularizer=slim.l2_regularizer(cfgs.WEIGHT_DECAY)):

                cls_score = slim.fully_connected(fc_flatten,
                                                 num_outputs=cfgs.CLASS_NUM+1,
                                                 weights_initializer=slim.variance_scaling_initializer(factor=1.0,
                                                                                mode='FAN_AVG',
                                                                                uniform=True),
                                                 activation_fn=None, trainable=self.is_training,
                                                 scope='cls_fc')

                bbox_pred = slim.fully_connected(fc_flatten,
                                                 num_outputs=(cfgs.CLASS_NUM+1)*4,
                                                 weights_initializer=slim.variance_scaling_initializer(factor=1.0,
                                                                                mode='FAN_AVG',
                                                                                uniform=True),
                                                activation_fn=None, trainable=self.is_training,
                                                 scope='reg_fc')
                if not scope == 'stage3':
                    # for convient. It also produce (cls_num +1) bboxes
                    cls_score = tf.reshape(cls_score, [-1, cfgs.CLASS_NUM+1])
                    bbox_pred = tf.reshape(bbox_pred, [-1, 4*(cfgs.CLASS_NUM+1)])
                    bbox_pred_ins = tf.reshape(bbox_pred, [-1, cfgs.CLASS_NUM+1, 4])

                    # only keep a box which score is the bigest
                    keep_abox = tf.argmax(cls_score, axis=1)
                    keep_inds = tf.reshape(tf.transpose(tf.stack([tf.cumsum(tf.ones_like(keep_abox))-1, keep_abox])), [-1, 2])
                    bbox_pred_fliter = tf.reshape(tf.gather_nd(bbox_pred_ins, keep_inds), [-1, 4])

                    return bbox_pred_fliter, bbox_pred, cls_score
                else:
                    cls_score = tf.reshape(cls_score, [-1, cfgs.CLASS_NUM + 1])
                    bbox_pred = tf.reshape(bbox_pred, [-1, 4 * (cfgs.CLASS_NUM + 1)])
                    return bbox_pred, cls_score

    def add_anchor_img_smry(self, img, anchors, labels):

        positive_anchor_indices = tf.reshape(tf.where(tf.greater_equal(labels, 1)), [-1])
        negative_anchor_indices = tf.reshape(tf.where(tf.equal(labels, 0)), [-1])

        positive_anchor = tf.gather(anchors, positive_anchor_indices)
        negative_anchor = tf.gather(anchors, negative_anchor_indices)

        pos_in_img = show_box_in_tensor.only_draw_boxes(img_batch=img,
                                                        boxes=positive_anchor)
        neg_in_img = show_box_in_tensor.only_draw_boxes(img_batch=img,
                                                        boxes=negative_anchor)

        tf.summary.image('positive_anchor', pos_in_img)
        tf.summary.image('negative_anchors', neg_in_img)

    def add_roi_batch_img_smry(self, img, rois, labels, scope):
        positive_roi_indices = tf.reshape(tf.where(tf.greater_equal(labels, 1)), [-1])

        negative_roi_indices = tf.reshape(tf.where(tf.equal(labels, 0)), [-1])

        pos_roi = tf.gather(rois, positive_roi_indices)
        neg_roi = tf.gather(rois, negative_roi_indices)

        pos_in_img = show_box_in_tensor.only_draw_boxes(img_batch=img,
                                                               boxes=pos_roi)
        neg_in_img = show_box_in_tensor.only_draw_boxes(img_batch=img,
                                                               boxes=neg_roi)
        tf.summary.image('pos_rois_{}'.format(scope), pos_in_img)
        tf.summary.image('neg_rois_{}'.format(scope), neg_in_img)

    def build_loss(self, rpn_box_pred, rpn_bbox_targets, rpn_cls_score, rpn_labels,
                   bbox_pred, bbox_targets, stage2_bbox_pred, stage2_bbox_targets,  stage1_bbox_pred, stage1_bbox_targets,
                   cls_score, labels, stage2_cls_score, stage2_labels, stage1_cls_score, stage1_labels):
        '''

        :param rpn_box_pred: [-1, 4]
        :param rpn_bbox_targets: [-1, 4]
        :param rpn_cls_score: [-1]
        :param rpn_labels: [-1]
        :param bbox_pred: [-1, 4*(cls_num+1)]
        :param bbox_targets: [-1, 4*(cls_num+1)]
        :param cls_score: [-1, cls_num+1]
        :param labels: [-1]
        :return:
        '''
        with tf.variable_scope('build_loss') as sc:
            with tf.variable_scope('rpn_loss'):

                rpn_bbox_loss = losses.smooth_l1_loss_rpn(bbox_pred=rpn_box_pred,
                                                          bbox_targets=rpn_bbox_targets,
                                                          label=rpn_labels,
                                                          sigma=cfgs.RPN_SIGMA)
                # rpn_bbox_loss = tfapi_loss.smooth_l1_loss_rpn(bbox_pred=rpn_box_pred,
                #                                               bbox_targets=rpn_bbox_targets,
                #                                               label=rpn_labels,
                #                                               sigma=cfgs.RPN_SIGMA)
                # rpn_cls_loss:
                # rpn_cls_score = tf.reshape(rpn_cls_score, [-1, 2])
                # rpn_labels = tf.reshape(rpn_labels, [-1])
                # ensure rpn_labels shape is [-1]
                rpn_select = tf.reshape(tf.where(tf.not_equal(rpn_labels, -1)), [-1])
                rpn_cls_score = tf.reshape(tf.gather(rpn_cls_score, rpn_select), [-1, 2])
                rpn_labels = tf.reshape(tf.gather(rpn_labels, rpn_select), [-1])

                rpn_cls_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=rpn_cls_score,
                                                                                labels=rpn_labels))

                rpn_cls_loss = rpn_cls_loss * cfgs.RPN_CLASSIFICATION_LOSS_WEIGHT
                rpn_bbox_loss = rpn_bbox_loss * cfgs.RPN_LOCATION_LOSS_WEIGHT

            with tf.variable_scope('FastRCNN_loss'):
                if not cfgs.FAST_RCNN_MINIBATCH_SIZE == -1:
                    bbox_loss = losses.smooth_l1_loss_rcnn(bbox_pred=bbox_pred,
                                                           bbox_targets=bbox_targets,
                                                           label=labels,
                                                           num_classes=cfgs.CLASS_NUM + 1,
                                                           sigma=cfgs.FASTRCNN_SIGMA)
                    stage2_bbox_loss = losses.smooth_l1_loss_rcnn(bbox_pred=stage2_bbox_pred,
                                                           bbox_targets=stage2_bbox_targets,
                                                           label=stage2_labels,
                                                           num_classes=cfgs.CLASS_NUM + 1,
                                                           sigma=cfgs.FASTRCNN_SIGMA)
                    stage1_bbox_loss = losses.smooth_l1_loss_rcnn(bbox_pred=stage1_bbox_pred,
                                                           bbox_targets=stage1_bbox_targets,
                                                           label=stage1_labels,
                                                           num_classes=cfgs.CLASS_NUM + 1,
                                                           sigma=cfgs.FASTRCNN_SIGMA)
                    cls_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
                        logits=cls_score, labels=labels))  # beacause already sample before
                    stage2_cls_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
                        logits=stage2_cls_score, labels=stage2_labels))  # beacause already sample before
                    stage1_cls_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
                        logits=stage1_cls_score, labels=stage1_labels))  # beacause already sample before

                    bbox_loss = bbox_loss + stage2_bbox_loss + stage1_bbox_loss
                    cls_loss = cls_loss + stage2_cls_loss + stage1_cls_loss
                else:
                    '''
                    applying OHEM here
                    '''
                    print(20 * "@@")
                    print("@@" + 10 * " " + "TRAIN WITH OHEM ...")
                    print(20 * "@@")
                    cls_loss, bbox_loss = losses.sum_ohem_loss(cls_score=cls_score,
                                                               label=labels,
                                                               bbox_targets=bbox_targets,
                                                               bbox_pred=bbox_pred,
                                                               num_ohem_samples=256,
                                                               num_classes=cfgs.CLASS_NUM + 1)
                cls_loss = cls_loss * cfgs.FAST_RCNN_CLASSIFICATION_LOSS_WEIGHT
                bbox_loss = bbox_loss * cfgs.FAST_RCNN_LOCATION_LOSS_WEIGHT
            loss_dict = {
                'rpn_cls_loss': rpn_cls_loss,
                'rpn_loc_loss': rpn_bbox_loss,
                'fastrcnn_cls_loss': cls_loss,
                'fastrcnn_loc_loss': bbox_loss
            }
        return loss_dict

    def build_whole_detection_network(self, input_img_batch, gtboxes_batch):

        if self.is_training:
            # ensure shape is [M, 5]
            gtboxes_batch = tf.reshape(gtboxes_batch, [-1, 5])
            gtboxes_batch = tf.cast(gtboxes_batch, tf.float32)

        img_shape = tf.shape(input_img_batch)

        # 1. build base network
        feature_to_cropped = self.build_base_network(input_img_batch)

        # 2. build rpn
        with tf.variable_scope('build_rpn',
                               regularizer=slim.l2_regularizer(cfgs.WEIGHT_DECAY)):

            rpn_conv3x3 = slim.conv2d(
                feature_to_cropped, 512, [3, 3],
                trainable=self.is_training, weights_initializer=cfgs.INITIALIZER,
                activation_fn=tf.nn.relu,
                scope='rpn_conv/3x3')
            rpn_cls_score = slim.conv2d(rpn_conv3x3, self.num_anchors_per_location*2, [1, 1],stride=1,
                                        trainable=self.is_training, weights_initializer=cfgs.INITIALIZER,
                                        activation_fn=None,
                                        scope='rpn_cls_score')
            rpn_box_pred = slim.conv2d(rpn_conv3x3, self.num_anchors_per_location*4, [1, 1], stride=1,
                                       trainable=self.is_training, weights_initializer=cfgs.BBOX_INITIALIZER,
                                       activation_fn=None,
                                       scope='rpn_bbox_pred')
            rpn_box_pred = tf.reshape(rpn_box_pred, [-1, 4])
            rpn_cls_score = tf.reshape(rpn_cls_score, [-1, 2])
            rpn_cls_prob = slim.softmax(rpn_cls_score, scope='rpn_cls_prob')

        # 3. generate_anchors
        featuremap_height, featuremap_width = tf.shape(feature_to_cropped)[1], tf.shape(feature_to_cropped)[2]
        featuremap_height = tf.cast(featuremap_height, tf.float32)
        featuremap_width = tf.cast(featuremap_width, tf.float32)

        anchors = anchor_utils.make_anchors(base_anchor_size=cfgs.BASE_ANCHOR_SIZE_LIST[0],
                                            anchor_scales=cfgs.ANCHOR_SCALES, anchor_ratios=cfgs.ANCHOR_RATIOS,
                                            featuremap_height=featuremap_height,
                                            featuremap_width=featuremap_width,
                                            stride=cfgs.ANCHOR_STRIDE,
                                            name="make_anchors_forRPN")

        # with tf.variable_scope('make_anchors'):
        #     anchors = anchor_utils.make_anchors(height=featuremap_height,
        #                                         width=featuremap_width,
        #                                         feat_stride=cfgs.ANCHOR_STRIDE[0],
        #                                         anchor_scales=cfgs.ANCHOR_SCALES,
        #                                         anchor_ratios=cfgs.ANCHOR_RATIOS, base_size=16
        #                                         )

        # 4. postprocess rpn proposals. such as: decode, clip, NMS
        with tf.variable_scope('postprocess_RPN'):
            # rpn_cls_prob = tf.reshape(rpn_cls_score, [-1, 2])
            # rpn_cls_prob = slim.softmax(rpn_cls_prob, scope='rpn_cls_prob')
            # rpn_box_pred = tf.reshape(rpn_box_pred, [-1, 4])
            rois, roi_scores = postprocess_rpn_proposals(rpn_bbox_pred=rpn_box_pred,
                                                         rpn_cls_prob=rpn_cls_prob,
                                                         img_shape=img_shape,
                                                         anchors=anchors,
                                                         is_training=self.is_training)
            # rois shape [-1, 4]
            # +++++++++++++++++++++++++++++++++++++add img smry+++++++++++++++++++++++++++++++++++++++++++++++++++++++

            if self.is_training:
                rois_in_img = show_box_in_tensor.draw_boxes_with_scores(img_batch=input_img_batch,
                                                                        boxes=rois,
                                                                        scores=roi_scores)
                tf.summary.image('all_rpn_rois', rois_in_img)

                score_gre_05 = tf.reshape(tf.where(tf.greater_equal(roi_scores, 0.5)), [-1])
                score_gre_05_rois = tf.gather(rois, score_gre_05)
                score_gre_05_score = tf.gather(roi_scores, score_gre_05)
                score_gre_05_in_img = show_box_in_tensor.draw_boxes_with_scores(img_batch=input_img_batch,
                                                                                boxes=score_gre_05_rois,
                                                                                scores=score_gre_05_score)
                tf.summary.image('score_greater_05_rois', score_gre_05_in_img)
            # ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

        if self.is_training:
            with tf.variable_scope('sample_anchors_minibatch'):
                rpn_labels, rpn_bbox_targets = \
                    tf.py_func(
                        anchor_target_layer,
                        [gtboxes_batch, img_shape, anchors],
                        [tf.float32, tf.float32])
                rpn_bbox_targets = tf.reshape(rpn_bbox_targets, [-1, 4])
                rpn_labels = tf.to_int32(rpn_labels, name="to_int32")
                rpn_labels = tf.reshape(rpn_labels, [-1])
                self.add_anchor_img_smry(input_img_batch, anchors, rpn_labels)

            # --------------------------------------add smry----------------------------------------------------------------

            rpn_cls_category = tf.argmax(rpn_cls_prob, axis=1)
            kept_rpppn = tf.reshape(tf.where(tf.not_equal(rpn_labels, -1)), [-1])
            rpn_cls_category = tf.gather(rpn_cls_category, kept_rpppn)
            acc = tf.reduce_mean(tf.to_float(tf.equal(rpn_cls_category, tf.to_int64(tf.gather(rpn_labels, kept_rpppn)))))
            tf.summary.scalar('ACC/rpn_accuracy', acc)

            with tf.control_dependencies([rpn_labels]):
                with tf.variable_scope('sample_RCNN_minibatch_stage1'):
                    stage1_rois, stage1_labels, stage1_bbox_targets = \
                    tf.py_func(proposal_target_layer,
                               [rois, gtboxes_batch, cfgs.FAST_RCNN_IOU_POSITIVE_THRESHOLD[0]],
                               [tf.float32, tf.float32, tf.float32])
                    stage1_rois = tf.reshape(stage1_rois, [-1, 4])
                    stage1_labels = tf.to_int32(stage1_labels)
                    stage1_labels = tf.reshape(stage1_labels, [-1])
                    stage1_bbox_targets = tf.reshape(stage1_bbox_targets, [-1, 4*(cfgs.CLASS_NUM+1)])
                    self.add_roi_batch_img_smry(input_img_batch, stage1_rois, stage1_labels, 'stage1')
        else:
            stage1_rois = rois

        # -------------------------------------------------------------------------------------------------------------#
        #                                            Fast-RCNN-before1                                                 #
        # -------------------------------------------------------------------------------------------------------------#

        # 5. build Fast-RCNN-before1
        # rois = tf.Print(rois, [tf.shape(rois)], 'rois shape', summarize=10)
        stage1_bbox_pred_fliter, stage1_bbox_pred, stage1_cls_score = self.build_fastrcnn(feature_to_cropped=feature_to_cropped,
                                                                                rois=stage1_rois,
                                                                                img_shape=img_shape,
                                                                                scope='stage1')
        # bbox_pred shape: [-1, 4*(cls_num+1)].
        # cls_score shape： [-1, cls_num+1]

        stage1_cls_prob = slim.softmax(stage1_cls_score, 'stage1_cls_prob')
        stage1_cls_category = tf.argmax(stage1_cls_prob, axis=1)
        # ----------------------------------------------add smry-------------------------------------------------------
        if self.is_training:
            stage1_fast_acc = tf.reduce_mean(tf.to_float(tf.equal(stage1_cls_category, tf.to_int64(stage1_labels))))
            tf.summary.scalar('ACC/stage1_fast_acc', stage1_fast_acc)

        #  postprocess_fastrcnn_before1
        stage1_bbox = self.postprocess_cascade(rois=stage1_rois,
                                               bbox_ppred=stage1_bbox_pred_fliter,
                                               scope='stage1')

        if self.is_training:
            with tf.control_dependencies([stage1_bbox]):
                with tf.variable_scope('sample_RCNN_minibatch_stage2'):
                    stage2_rois, stage2_labels, stage2_bbox_targets = \
                    tf.py_func(proposal_target_layer,
                               [stage1_bbox, gtboxes_batch, cfgs.FAST_RCNN_IOU_POSITIVE_THRESHOLD[1]],
                               [tf.float32, tf.float32, tf.float32])
                    stage2_rois = tf.reshape(stage2_rois, [-1, 4])
                    stage2_labels = tf.to_int32(stage2_labels)
                    stage2_labels = tf.reshape(stage2_labels, [-1])
                    stage2_bbox_targets = tf.reshape(stage2_bbox_targets, [-1, 4*(cfgs.CLASS_NUM+1)])
                    self.add_roi_batch_img_smry(input_img_batch, stage2_rois, stage2_labels, 'stage2')
        else:
            stage2_rois = stage1_bbox

        # -------------------------------------------------------------------------------------------------------------#
        #                                            Fast-RCNN-before2                                                 #
        # -------------------------------------------------------------------------------------------------------------#

        # 6. build Fast-RCNN-before2
        # rois = tf.Print(rois, [tf.shape(rois)], 'rois shape', summarize=10)
        # stage2_rois = tf.stop_gradient(stage2_rois)
        stage2_bbox_pred_fliter, stage2_bbox_pred, stage2_cls_score = self.build_fastrcnn(feature_to_cropped=feature_to_cropped,
                                                                                rois=stage2_rois,
                                                                                img_shape=img_shape,
                                                                                scope='stage2')
        # bbox_pred shape: [-1, 4*(cls_num+1)].
        # cls_score shape： [-1, cls_num+1]

        stage2_cls_prob = slim.softmax(stage2_cls_score, 'stage2_cls_prob')
        stage2_cls_category = tf.argmax(stage2_cls_prob, axis=1)
        # ----------------------------------------------add smry-------------------------------------------------------
        if self.is_training:
            stage2_fast_acc = tf.reduce_mean(tf.to_float(tf.equal(stage2_cls_category, tf.to_int64(stage2_labels))))
            tf.summary.scalar('ACC/stage2_fast_acc', stage2_fast_acc)

        #  postprocess_fastrcnn_before2
        stage2_bbox = self.postprocess_cascade(rois=stage2_rois,
                                               bbox_ppred=stage2_bbox_pred_fliter,
                                               scope='stage2')

        if self.is_training:
            with tf.control_dependencies([stage2_bbox]):
                with tf.variable_scope('sample_RCNN_minibatch_stage3'):
                    stage3_rois, stage3_labels, stage3_bbox_targets = \
                    tf.py_func(proposal_target_layer,
                               [stage2_rois, gtboxes_batch, cfgs.FAST_RCNN_IOU_POSITIVE_THRESHOLD[2]],
                               [tf.float32, tf.float32, tf.float32])
                    stage3_rois = tf.reshape(stage3_rois, [-1, 4])
                    stage3_labels = tf.to_int32(stage3_labels)
                    stage3_labels = tf.reshape(stage3_labels, [-1])
                    stage3_bbox_targets = tf.reshape(stage3_bbox_targets, [-1, 4*(cfgs.CLASS_NUM+1)])
                    self.add_roi_batch_img_smry(input_img_batch, stage3_rois, stage3_labels, 'stage3')
        else:
            stage3_rois = stage2_bbox

        # -------------------------------------------------------------------------------------------------------------#
        #                                            Fast-RCNN                                                         #
        # -------------------------------------------------------------------------------------------------------------#

        # 7. build Fast-RCNN
        # rois = tf.Print(rois, [tf.shape(rois)], 'rois shape', summarize=10)
        # stage3_rois = tf.stop_gradient(stage3_rois)
        stage3_bbox_pred, stage3_cls_score = self.build_fastrcnn(feature_to_cropped=feature_to_cropped,
                                                                 rois=stage3_rois,
                                                                 img_shape=img_shape,
                                                                 scope='stage3')
        # bbox_pred shape: [-1, 4*(cls_num+1)].
        # cls_score shape： [-1, cls_num+1]

        stage3_cls_prob = slim.softmax(stage3_cls_score, 'stage3_cls_prob')
        stage3_cls_category = tf.argmax(stage3_cls_prob, axis=1)

        # ----------------------------------------------add smry-------------------------------------------------------
        if self.is_training:
            fast_acc = tf.reduce_mean(tf.to_float(tf.equal(stage3_cls_category, tf.to_int64(stage3_labels))))
            tf.summary.scalar('ACC/fast_acc', fast_acc)

        #  postprocess_fastrcnn
        if not self.is_training:
            with slim.arg_scope([slim.conv2d, slim.conv2d_in_plane,
                                 slim.conv2d_transpose, slim.separable_conv2d, slim.fully_connected],
                                reuse=True):
                _, _, final_scores_stage2 = self.build_fastrcnn(feature_to_cropped=feature_to_cropped,
                                                                rois=stage3_rois,
                                                                img_shape=img_shape,
                                                                scope='stage2')
                final_scores_stage2= slim.softmax(final_scores_stage2, 'final_scores_stage2')

                _, _, final_scores_stage1 = self.build_fastrcnn(feature_to_cropped=feature_to_cropped,
                                                                rois=stage3_rois,
                                                                img_shape=img_shape,
                                                                scope='stage1')
                final_scores_stage1 = slim.softmax(final_scores_stage1, 'final_scores_stage1')
                # choose which stage to export
                cls_prob = tf.add(final_scores_stage2, final_scores_stage1)
                cls_prob = tf.add(cls_prob, stage3_cls_prob)/3
                return self.postprocess_fastrcnn(rois=stage3_rois, bbox_ppred=stage3_bbox_pred, scores=cls_prob, img_shape=img_shape, scope='stage3')
        else:
            '''
            when trian. We need build Loss
            '''
            loss_dict = self.build_loss(rpn_box_pred=rpn_box_pred,
                                        rpn_bbox_targets=rpn_bbox_targets,
                                        rpn_cls_score=rpn_cls_score,
                                        rpn_labels=rpn_labels,
                                        bbox_pred=stage3_bbox_pred,
                                        bbox_targets=stage3_bbox_targets,
                                        stage2_bbox_pred=stage2_bbox_pred,
                                        stage2_bbox_targets=stage2_bbox_targets,
                                        stage1_bbox_pred=stage1_bbox_pred,
                                        stage1_bbox_targets=stage1_bbox_targets,
                                        cls_score=stage3_cls_score,
                                        labels=stage3_labels,
                                        stage2_cls_score=stage2_cls_score,
                                        stage2_labels = stage2_labels,
                                        stage1_cls_score=stage1_cls_score,
                                        stage1_labels=stage1_labels)
            final_bbox, final_scores, final_category = self.postprocess_fastrcnn(rois=stage3_rois,
                                                                                bbox_ppred=stage3_bbox_pred,
                                                                                scores=stage3_cls_prob,
                                                                                img_shape=img_shape,
                                                                                scope='stage3')
            return final_bbox, final_scores, final_category, loss_dict

    def get_restorer(self):
        checkpoint_path = tf.train.latest_checkpoint(os.path.join(cfgs.TRAINED_CKPT, cfgs.VERSION))

        if checkpoint_path != None:
            if cfgs.RESTORE_FROM_RPN:
                print('___restore from rpn___')
                model_variables = slim.get_model_variables()
                restore_variables = [var for var in model_variables if not var.name.startswith('FastRCNN_Head')] + \
                                    [slim.get_or_create_global_step()]
                for var in restore_variables:
                    print(var.name)
                restorer = tf.train.Saver(restore_variables)
            else:
                restorer = tf.train.Saver()
                print("model restore from :", checkpoint_path)

        else:
            checkpoint_path = cfgs.PRETRAINED_CKPT
            print("model restore from pretrained mode, path is :", checkpoint_path)

            model_variables = slim.get_model_variables()
            # for var in model_variables:
            #     print(var.name)
            # print(20*"__++__++__")

            def name_in_ckpt_rpn(var):
                return var.op.name

            def name_in_ckpt_fastrcnn_head(var):
                '''
                Fast-RCNN/resnet_v1_50/block4 -->resnet_v1_50/block4
                Fast-RCNN/MobilenetV2/** -- > MobilenetV2 **
                :param var:
                :return:
                '''
                return '/'.join(var.op.name.split('/')[1:])

            nameInCkpt_Var_dict = {}
            for var in model_variables:
                if var.name.startswith('Fast-RCNN_stage1/'+self.base_network_name):  # +'/block4'
                    var_name_in_ckpt = name_in_ckpt_fastrcnn_head(var)
                    nameInCkpt_Var_dict[var_name_in_ckpt] = var
                else:
                    if var.name.startswith(self.base_network_name):
                        var_name_in_ckpt = name_in_ckpt_rpn(var)
                        nameInCkpt_Var_dict[var_name_in_ckpt] = var
                    else:
                        continue
            restore_variables = nameInCkpt_Var_dict
            for key, item in restore_variables.items():
                print("var_in_graph: ", item.name)
                print("var_in_ckpt: ", key)
                print(20*"___")
            restorer = tf.train.Saver(restore_variables)
            print(20 * "****")
            print("restore from pretrained_weighs in IMAGE_NET")
        return restorer, checkpoint_path

    def get_gradients(self, optimizer, loss):
        '''

        :param optimizer:
        :param loss:
        :return:

        return vars and grads that not be fixed
        '''

        # if cfgs.FIXED_BLOCKS > 0:
        #     trainable_vars = tf.trainable_variables()
        #     # trained_vars = slim.get_trainable_variables()
        #     start_names = [cfgs.NET_NAME + '/block%d'%i for i in range(1, cfgs.FIXED_BLOCKS+1)] + \
        #                   [cfgs.NET_NAME + '/conv1']
        #     start_names = tuple(start_names)
        #     trained_var_list = []
        #     for var in trainable_vars:
        #         if not var.name.startswith(start_names):
        #             trained_var_list.append(var)
        #     # slim.learning.train()
        #     grads = optimizer.compute_gradients(loss, var_list=trained_var_list)
        #     return grads
        # else:
        #     return optimizer.compute_gradients(loss)
        return optimizer.compute_gradients(loss)

    def enlarge_gradients_for_bias(self, gradients):

        final_gradients = []
        with tf.variable_scope("Gradient_Mult") as scope:
            for grad, var in gradients:
                scale = 1.0
                if cfgs.MUTILPY_BIAS_GRADIENT and './biases' in var.name:
                    scale = scale * cfgs.MUTILPY_BIAS_GRADIENT
                if not np.allclose(scale, 1.0):
                    grad = tf.multiply(grad, scale)
                final_gradients.append((grad, var))
        return final_gradients


# import the necessary packages
import CONST
import numpy as np
import cv2
from keras.models import load_model
from keras.preprocessing.image import img_to_array
import os
import time
from keras import applications
from gtts import gTTS
from pygame import mixer
import shutil

#load previously trained model
model = applications.VGG16(include_top=False, weights='imagenet', input_shape=(CONST.IMAGE_SIZE,CONST.IMAGE_SIZE,3))
top_model = load_model(os.path.join(CONST.SAVE_DIR,CONST.BOTTLENECK_MODEL))

#load labels
labels = os.listdir('./dataset')

#store the sound of each label
if os.path.exists("./sounds") == True:
    shutil.rmtree("./sounds")

os.makedirs("./sounds")
(gTTS(text=labels[0], lang='en')).save("./sounds/0.mp3")
(gTTS(text=labels[1], lang='en')).save("./sounds/1.mp3")
(gTTS(text=labels[2], lang='en')).save("./sounds/2.mp3")
mixer.init()

# Turn on the webcam
cap = cv2.VideoCapture(0)
time.sleep(2)

print ('Press q to exit')
y_pred_old = '-1'
while True:
    # Capture frame-by-frame
    ret, frame = cap.read()

    #preprocessing frame to predict its label
    frame2 = cv2.resize(frame, (CONST.IMAGE_SIZE, CONST.IMAGE_SIZE))
    frame2 = img_to_array(frame2)
    frame2 = np.array(frame2, dtype="float32") / 255
    # generating a prdiction of the frame
    #y_pred = model.predict_classes(frame2[None,:,:,:])

    y_pred = top_model.predict_classes(model.predict(frame2[None,:,:,:]))

    if(y_pred[0] != y_pred_old):
        mixer.music.load("./sounds/"+str(y_pred[0])+'.mp3')
        mixer.music.play()

    y_pred_old = y_pred[0]

    cv2.putText(frame, labels[y_pred[0]] , (10, 30), CONST.FONT, 1, (0, 255, 0), 2, cv2.LINE_AA)
    cv2.putText(frame, "Press q to exit", (10, 450), CONST.FONT, 1, (0, 255, 0), 2, cv2.LINE_AA)

    # Display the resulting frame
    cv2.imshow('frame',frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# release the capture
cap.release()
cv2.destroyAllWindows()
#remove sounds folder


import time
import numpy as np
import tensorflow as tf
import tensorlayer as tl
import utils
from models_resnet import generator_txt2img_resnet, discriminator_txt2img_resnet

t_dim = 256
n_epoch = 500 # "Epoch to train [25]"
z_dim = 100 # "Num of noise value]"
lr = 0.0002 # "Learning rate of for adam [0.0002]")
beta1 = 0.5 # "Momentum term of adam [0.5]")
batch_size = 100 # "The number of batch images [64]")
output_size = 64 # "The size of the output images to produce [64]")
sample_size = 64 # "The number of sample images [64]")
c_dim = 3 # "Number of image channels. [3]")
save_every_epoch = 5 # "The interval of saving checkpoints.")

checkpoint_dir = "Output/checkpoint" # "Directory name to save the checkpoints [checkpoint]")
sample_dir = "Output/output" # "Directory name to save the image samples [samples]")
log_dir = "Output/Log_loss"

tl.files.exists_or_mkdir(checkpoint_dir) # save model
tl.files.exists_or_mkdir(sample_dir) # save generated image
tl.files.exists_or_mkdir(log_dir) # save generated image

# get vectorized images and caption
# skipvectors_sample.npy created in generate_thought_vectors.py
array_caption_vector = np.load('Data/vectors_files/skipvectors_sample.npy')
#imgvectors_sample.npy created in process_images.py
array_image_vector = np.load('Data/vectors_files/imgvectors_sample.npy')

tensor_captions = tf.convert_to_tensor(array_caption_vector)

tensor_images = tf.convert_to_tensor(array_image_vector)

def train():
    reduced_text_embedding = utils.lrelu(utils.linear(tensor_captions, 256))
    #G = get_generator([None, z_dim+reduced_text_embedding.shape[1]])
    #D = get_discriminator([None, output_size, output_size, c_dim], input_rnn_embed = [None, reduced_text_embedding.shape[1]])

    G = generator_txt2img_resnet([None, z_dim+reduced_text_embedding.shape[1]])
    #tl.files.load_hdf5_to_weights(checkpoint_dir+'/G.h5', G)
    D = discriminator_txt2img_resnet([None, output_size, output_size, c_dim],t_txt=[None, reduced_text_embedding.shape[1]])
    #tl.files.load_hdf5_to_weights(checkpoint_dir+'/D.h5', D)

    G.train()
    D.train()

    d_optimizer = tf.optimizers.Adam(lr, beta_1=beta1)
    g_optimizer = tf.optimizers.Adam(lr, beta_1=beta1)

    n_step_epoch = int(len(tensor_images) // batch_size)
    dataset_img = tf.data.Dataset.from_tensor_slices(tensor_images)
    batched_ds_img = dataset_img.batch(batch_size, drop_remainder=True)

    dataset_txt = tf.data.Dataset.from_tensor_slices(reduced_text_embedding)
    batched_ds_txt = dataset_txt.batch(batch_size, drop_remainder=True)
    list_batch_txt = [batch for batch in batched_ds_txt]

    text_file = open(log_dir+"/log.txt", "a")
    text_file.write('epoch, d_loss, g_loss\n')
    text_file.close()

    for epoch in range(n_epoch):
        for step, batch_images in enumerate(batched_ds_img):
            batch_txt = list_batch_txt[step]
            step_time = time.time()
            with tf.GradientTape(persistent=True) as tape:
                z = np.random.normal(loc=0.0, scale=1.0, size=[batch_size, z_dim]).astype(np.float32)
                z = tf.convert_to_tensor(z)
                z_text_concat = tf.concat([z, batch_txt], 1)
                net_fake_image = G(z_text_concat)

                cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=False)
                disc_fake_image_logits = D([net_fake_image, batch_txt])
                g_loss = cross_entropy(tf.ones_like(disc_fake_image_logits),disc_fake_image_logits)

                disc_real_image_logits = D([batch_images, batch_txt])
                d_loss1 = cross_entropy(tf.ones_like(disc_real_image_logits),disc_real_image_logits)

                idexs = utils.get_random_int(min=0, max=len(batch_txt) - 1, number=batch_size)
                wrong_caption = [batch_txt[i]for i in idexs]
                wrong_caption = tf.stack(wrong_caption)
                disc_mismatch_logits = D([batch_images, wrong_caption])
                d_loss2 = cross_entropy(tf.zeros_like(disc_mismatch_logits), disc_mismatch_logits)
                d_loss3 = cross_entropy(tf.zeros_like(disc_fake_image_logits), disc_fake_image_logits)
                d_loss = d_loss1 + (d_loss2 + d_loss3) * 0.5

            grad = tape.gradient(g_loss, G.trainable_weights)
            g_optimizer.apply_gradients(zip(grad, G.trainable_weights))
            if(step%2 == 0):
                grad = tape.gradient(d_loss, D.trainable_weights)
                d_optimizer.apply_gradients(zip(grad, D.trainable_weights))
            del tape

            print("Epoch: [{}/{}] [{}/{}] took: {:.3f}, d_loss: {:.5f}, g_loss: {:.5f}".format(epoch, \
                                                                                n_epoch, step,
                                                                                n_step_epoch,
                                                                                time.time() - step_time,
                                                                                d_loss, g_loss))
            text_file = open(log_dir+"/log.txt", "a")
            text_file.write('{}, {:.5f}, {:.5f} \n'.format(epoch, d_loss, g_loss))
            text_file.close()

        if np.mod(epoch, save_every_epoch) == 0:
            #G.save_weights('{}/G_weights.npz'.format(checkpoint_dir), format='npz')
            #D.save_weights('{}/D_weights.npz'.format(checkpoint_dir), format='npz')
            G.save('{}/G.h5'.format(checkpoint_dir), save_weights=True)
            G.eval()
            D.save('{}/D.h5'.format(checkpoint_dir), save_weights=True)
            D.eval()

            z = np.random.normal(loc=0.0, scale=1.0, size=[4, z_dim]).astype(np.float32)
            z = tf.convert_to_tensor(z)

            test_descr = np.reshape(batch_txt[0:4], (4,256))
            z_text_concat = tf.concat([z, test_descr], 1)

            result = G(z_text_concat)

            G.train()
            D.train()

            img = result.numpy().squeeze().astype(np.uint8)

            #Sauvegarde plusieurs images
            tl.visualize.save_images(img, [2, 2], '{}/train_{:02d}.png'.format(sample_dir, epoch))

            #Sauvegarde une image
            #tl.visualize.save_image(img, '{}/train_{:02d}.png'.format(sample_dir, epoch))

if __name__ == '__main__':
    train()


from nltk.tokenize import word_tokenize
from nltk.tag import pos_tag
import networkx as nx
import numpy as np
from collections import defaultdict, Counter
import re
from textblob import TextBlob
import operator
import requests
import json
from sklearn.cluster import KMeans
from bs4 import UnicodeDammit

#### MODEL parameters

parameters = {
"SIGMA_VSN": 15,
"SIGMA_R": 5,
"GAP": 4,
"SIGMA_SIM": 0.3,
"MAX_SENTENCE_LENGTH": 15,
"NNEIGH": 7,
"OUT_CNN": 10,
"SENTIMENT_DIFF": 0.2
}

####

def get_graph(reviews):
    tokenized_reviews = []
    for review in reviews:
        review = (review.encode('ascii', 'ignore')).decode("utf-8")
        tokenized_reviews.append(" ".join([word.lower()+"/"+tag for word,tag in pos_tag(word_tokenize(review))]))

    filename = "tagged_reviews.tmp"

    with open(filename, "w") as f:
        for review in tokenized_reviews:
            f.write(review+"\n")
    f.close()
    #
    with open(filename,"r") as f:
        bigrams = []; nodes_PRI = defaultdict(list)
        for doc_id, line in enumerate(f):
            try:
                words = (line.strip()).split()
            except:
                continue
            #
            words = [word.replace("'s/VBZ","is/VBZ").replace("n't","not") for word in words]
            bigrams_in_document = list(zip(words, words[1:]))
            bigrams.extend(bigrams_in_document)
            # Positional Reference Information (PRI) for each node
            for j, word in enumerate(words):
                nodes_PRI[word].append([doc_id,j])
    #
    edges = Counter(bigrams)
    #
    with open("review_edges.tmp", "w") as f:
        for bigram in edges:
            f.write(" ".join([bigram[0], bigram[1]]) +
                    " " + str(edges[bigram]))
            f.write("\n")
    #
    graph = nx.read_edgelist("review_edges.tmp", create_using=nx.DiGraph(),
                         data=(("count",int),))
    return graph, nodes_PRI

def is_valid_start(node,PRI,SIGMA_VSN):
    '''
        Check if a node is a Valid Starting Node (VSN)

        A node is VSN is the average PID <= SIGMA_VSN and

        its tag is one in the list above

        Also, if a node is common start word, we don't check the PID

    '''

    # check if a node is a common start word
    common_start_word = ["the","a","an","this","these","that","those",\
    					"when","if","for"]
    #
    if(node[:-2] in common_start_word):
    	return True
    # if not look at the average PDI
    PID = np.array(PRI)[:,1]
    valid_start_tags = ["JJ", "RB", "PRP$", "VBG", "NN", "DT"]
    if((np.mean(PID) <= SIGMA_VSN) and (node[-2:] in valid_start_tags)):
        return True
    else:
        return False

def is_valid_end(node):
    '''
        Check if a node is a Valid End Node (VEN)

        A node is a VEN is a period, a comma, or a coordinating conjunction

    '''

    if("/." in node or "/," in node or "/CC" in node):
        return True
    else:
        return False

def is_collapsible(graph,node,OUT_CNN):
    '''
        Check if a node is Collapsible

        A node is collapsible is it's a verb or if it has more than OUT_CNN outwards connections

    '''
    if(is_valid_end(node) is False):
	    if(len(graph[node])>OUT_CNN):
	    	return True

    if re.match(".*(/VB[A-Z]|/IN)", node) is not None:
        return True
    else:
        return False

def is_valid_path(sentence):
    '''
	    Determine if the sentence is a valid candidate.
    '''
    sentence = sentence.split()
    sent = " ".join(sentence)
    last = sentence[-1]
    w, t = last.split("/")
    if t in set(["TO", "VBZ", "IN", "CC", "WDT", "PRP", "DT", ","]):
        return False
    if re.match(".*(/DT)*.*(/VB)+.*(/PRP)+.*(/VB)+.*", sent): # This is it is
        return False
    if re.match(".*(very/RB very/RB).*", sent): # Very very very
        return False
    if re.match(".*(/JJ)*.*(/NN)+.*(/VB)+.*(/JJ)+.*", sent):
        return True
    elif re.match(".*(/RB)*.*(/JJ)+.*(/NN)+.*", sent) and not re.match(".*(/DT).*", sent):
        return True
    elif re.match(".*(/PRP|/DT)+.*(/VB)+.*(/RB|/JJ)+.*(/NN)+.*", sent):
        return True
    elif re.match(".*(/JJ)+.*(/TO)+.*(/VB).*", sent):
        return True
    elif re.match(".*(/RB)+.*(/IN)+.*(/NN)+.*", sent):
        return True
    else:
        return False

def PathScore(redundancy,length):
    '''
        Path Score
    '''
    return np.log2(length) * redundancy

def intersection(PRI, PRI_neigh, GAP):
    '''
        Compute the intersection between the PRIs of two nodes

        It returns the PRIs of the neighboring node that have the same SID

        and a different in PIDs less or equal than GAP

    '''

    PRI_intersect = []
    for the_pri in PRI:
        sid, pid = the_pri
        for the_pri_neigh in PRI_neigh:
            sid_n, pid_n = the_pri_neigh
            if sid == sid_n and abs(pid_n - pid) <= GAP:
                PRI_intersect.append(the_pri_neigh)
                break
    #
    return PRI_intersect

def rreplace(s, old, new, occurrence):
	li = s.rsplit(old, occurrence)
	return new.join(li)

# better way to stitch more sentences to an anchor
# 1) join the sentences ordered by score
# 2) sentences are connected based on their sentiment

def stitch(c_anchor, cc, SENTIMENT_DIFF):
	joined_sentence = " ".join(c_anchor)+" "
	#
	dummy = np.array(sorted([(key,score) for key, score in list(cc.items())],key=operator.itemgetter(1),reverse=True))
	all_ccs, scores = dummy[:,0], dummy[:,1].astype("float")
	#
	max_length = len(all_ccs)
	if(max_length==1):
		joined_sentence += all_ccs[0]
		return joined_sentence

	# otherwise just stitch all the pieces with commas and a `and/DT` if more than three pieces

	stitched_sentences = " , ".join(all_ccs)
	joined_sentence += rreplace(stitched_sentences," , "," and/DT ",1)
	#
	return joined_sentence

def jaccard(sent1, sent2):
    '''
        Jaccard similarity score between two strings
    '''
    sent1 = set(sent1.split())
    sent2 = set(sent2.split())
    inter = sent1 & sent2
    return len(inter) / (len(sent1) + len(sent2) - len(inter))

#jaccard = np.vectorize(jaccard)

def remove_duplicates(tmp,SIGMA_SIM):
    '''

        tmp is a dictionary whose keys are sentences and values their scores.

        This functions prunes the keys who have high Jaccard similarity and only retains those with the highest score.

    '''
    sentences = np.array(list(tmp.keys()))

    if(len(sentences)==1):
        return tmp

    scores = np.array(list(tmp.values()))

    sentences_to_examine = list(sentences)
    sentences_to_keep = []
    i=0
    while(len(sentences_to_examine)>0):
        similar_sentences = np.array([jaccard(sentences_to_examine[i],dummy) >= SIGMA_SIM for dummy in sentences])
        already_kept = np.array([ss not in sentences_to_keep for ss in sentences])
        similar_sentences = similar_sentences & already_kept
        order = np.argsort(scores[similar_sentences])[::-1]
        sentences_to_keep.append(sentences[similar_sentences][order][0])
        # removing the other sentence
        for sentence_to_remove in np.unique(sentences[similar_sentences][order]):
            sentences_to_examine = list([x for x in sentences_to_examine if x!=sentence_to_remove])

    tmp_pruned = {key:tmp[key] for key in sentences_to_keep}
    return tmp_pruned

def untag(sentence):
    '''
        Removing the tags from a sentence and formatting it into a nice sentence
    '''

    untagged_sentence = " ".join([word.split("/")[0] for word in sentence.split()])
    untagged_sentence = untagged_sentence.capitalize()
    untagged_sentence = untagged_sentence.replace(" , ",", ")
    untagged_sentence = untagged_sentence.replace(" n't"," not")
    untagged_sentence = untagged_sentence.replace(" 's","'s")
    return untagged_sentence

def traverse(graph,nodes_PRI,cList,node,score,PRI,sentence,parent=[],parameters=parameters):
    '''
        Traversing the graph and building the candidate sentences

    '''
    # We only allow sentences that are at most MAX_SENTENCE_LENGTH long
    if(len(sentence)>parameters["MAX_SENTENCE_LENGTH"]):
        return
    #
    redundancy = len(PRI)
    #
    # traverse the graph only if a node meets the rudandancy criterion
    if(redundancy >= parameters["SIGMA_R"] or is_valid_end(node)):
        if(is_valid_end(node)):
            # remove the /CC or the punctation at the end of the sentence if present
            del sentence[-1]
            # check if the sentence is valid
            if(is_valid_path(" ".join(parent + sentence))):
                final_score = score/len(parent + sentence)
                cList[" ".join(sentence)] = score
            return
        #
        # if not keep traversing the graph
        neighbors = graph[node]
        for neighbor in neighbors:
            if (neighbor in (sentence + parent) and neighbor[-2:] not in ["DT"]):
            	continue
            new_sentence = sentence[:]
            new_sentence.append(neighbor)
            PRI_new = intersection(PRI,nodes_PRI[neighbor],parameters["GAP"])
            redundancy = len(PRI_new)
            new_score = score + PathScore(redundancy,len(new_sentence))
            if(is_collapsible(graph,neighbor,parameters["OUT_CNN"]) and not parent):
                c_anchor = new_sentence
                tmp = defaultdict(int)
                best_neighbors = np.array(sorted([(nn,len(nodes_PRI[nn])) for nn in graph[neighbor]],key=operator.itemgetter(1),reverse=True))[:,0][:parameters["NNEIGH"]]
                for neighbor_an in best_neighbors:#graph[neighbor]:
                    traverse(graph,nodes_PRI,tmp,neighbor_an,0,PRI_new,[neighbor_an], new_sentence,parameters=parameters)
                    if("" in tmp): del tmp[""]
                    if(tmp):
                        tmp = remove_duplicates(tmp,parameters["SIGMA_SIM"])
                        CCPath_score = np.mean(list(tmp.values()))
                        final_score = new_score + CCPath_score
                        if(np.isnan(final_score)):
                            print(("NAN=",new_score, CCPath_score,tmp))
                        stitched_sent = stitch(c_anchor, tmp, parameters["SENTIMENT_DIFF"])
                        if(is_valid_path(stitched_sent)):
                            cList[stitched_sent] = final_score
            else:
                traverse(graph,nodes_PRI,cList,neighbor,new_score,PRI_new,new_sentence,parent,parameters=parameters)

def summarizer(graph,nodes_PRI,parameters=parameters):
	# summarization algorithm
	candidates = defaultdict(int)
	for node, PRI in list(nodes_PRI.items()):
	    if(is_valid_start(node,PRI,parameters["SIGMA_VSN"])):
	        PRI = nodes_PRI[node]
	        score = 0
	        cList = defaultdict(int)
	        # sentence formed so far
	        sentence = [node]
	        # recursively traverse the graph and create candidate sentences
	        traverse(graph,nodes_PRI,cList,node,score,PRI,sentence,[],parameters=parameters)
	        candidates.update(cList)
	return candidates

### CHECK IF THE REVIEWS MAKE SENSE
### Using the Microsoft Oxford Language Model API

# ---> Please use your key <---

OXFORD_LANGUAGE_MODEL_ENDPOINT = "https://od-api.oxforddictionaries.com/api/v1"
SUBSCRIPTION_KEY_NAME = ""
DEFAULT_SUBSCRIPTION_KEY = ""

def conditional_probability(query):
    '''

    Calculate the Conditional Probability that a particular word will follow a given sequence of words

    Examples:

    query = {"words": "this is", "word": "sparta"}

    '''
    data = json.dumps({"queries" : [query]})
    headers = {SUBSCRIPTION_KEY_NAME: DEFAULT_SUBSCRIPTION_KEY,
                "Content-Type": "application/json",
                "Content-Length": str(len(data))}
    #
    params = {
             "model" : "body",
             "order" : 1
            }
    #
    r = requests.post(OXFORD_LANGUAGE_MODEL_ENDPOINT + "calculateConditionalProbability", headers=headers,  params=params, data=data)
    #
    if r.status_code == 200:
        return 10.**(r.json()["results"][0]["probability"])
    else:
        return 10.**-30.

def readability_score(input_string, n_grams_order = 3):
    '''

    Readability score from Eq. 4 in Ganesan et al. 2012

    S_red (w_1...w_n) = 1/K * PROD_q^n P(w_k,w_{k-q+1}...w_{k-1})

    '''

    input_string = input_string.lower()
    words = input_string.split()

    k = list(range(n_grams_order-1,len(words)))

    probability = 1
    for kk in k:
        partial_probability = conditional_probability({"word": words[kk],"words": words[kk-n_grams_order+1] + " " + words[kk-n_grams_order+2]})
        probability *= partial_probability

    return np.log2(probability) / len(k)

def joint_probability(query):
    '''

    Calculate the Conditional Probability that a particular word will follow a given sequence of words

    Examples:

    query = "this is sparta"

    '''
    data = json.dumps({"queries" : [query]})
    headers = {SUBSCRIPTION_KEY_NAME: DEFAULT_SUBSCRIPTION_KEY,
                "Content-Type": "application/json",
                "Content-Length": str(len(data))}
    #
    params = {
             "model" : "body",
             "order" : 1
            }
    #
    r = requests.post(OXFORD_LANGUAGE_MODEL_ENDPOINT + "calculateJointProbability", headers=headers,  params=params, data=data)
    #
    if r.status_code == 200:
        return r.json()["results"][0]["probability"]
    else:
        return -99


# Standard Library
import builtins
import json
import copy
import numpy as np

# Import from third library
from up.utils.general.log_helper import default_logger as logger
from up.utils.general.registry_factory import EVALUATOR_REGISTRY
from up.data.metrics.base_evaluator import Evaluator
from up.utils.general.petrel_helper import PetrelHelper
from up.data.metrics import Metric

# fix pycocotools py2-style bug
builtins.unicode = str

__all__ = ['KittiEvaluator']

@EVALUATOR_REGISTRY.register('kitti')
class KittiEvaluator(Evaluator):
    def __init__(self, gt_file, recall_thresh_list=[0.3, 0.5, 0.7]):
        """
        Arguments:
            gt_file (str): directory or json file of annotations
            iou_types (str): list of iou types of [keypoints, bbox, segm]
        """
        super(KittiEvaluator, self).__init__()
        self.gt_file = gt_file
        self.recall_thresh_list = recall_thresh_list

    def load_dts(self, res_file, res):
        out = []
        if res is not None:
            out = res
        else:
            logger.info(f'loading res from {res_file}')
            out = []
            with PetrelHelper.open(res_file) as f:
                for line in f:
                    out.append(json.loads(line))
        out = [x for res_gpus in out for res_bs in res_gpus for x in res_bs]
        for idx in range(len(out)):
            for k, v in out[idx].items():
                if isinstance(v, list):
                    out[idx][k] = np.array(v)
        return out

    def get_metric(self, ret):
        metric = {
            'gt_num': 0,
        }
        for cur_thresh in self.recall_thresh_list:
            metric['recall_roi_%s' % str(cur_thresh)] = 0
            metric['recall_rcnn_%s' % str(cur_thresh)] = 0

        for i in range(len(ret)):
            recall_dict = ret[i]['recall_dict']
            for cur_thresh in self.recall_thresh_list:
                metric['recall_roi_%s' % str(cur_thresh)] += recall_dict.get('roi_%s' % str(cur_thresh), 0)
                metric['recall_rcnn_%s' % str(cur_thresh)] += recall_dict.get('rcnn_%s' % str(cur_thresh), 0)
            metric['gt_num'] += recall_dict.get('gt', 0)

            disp_dict = {}
            min_thresh = self.recall_thresh_list[0]
            disp_dict['recall_%s' % str(min_thresh)] = \
                '(%d, %d) / %d' % (metric['recall_roi_%s' % str(min_thresh)],
                                   metric['recall_rcnn_%s' % str(min_thresh)], metric['gt_num'])
        gt_num_cnt = metric['gt_num']
        recall_dict = {}
        for cur_thresh in self.recall_thresh_list:
            cur_roi_recall = metric['recall_roi_%s' % str(cur_thresh)] / max(gt_num_cnt, 1)
            cur_rcnn_recall = metric['recall_rcnn_%s' % str(cur_thresh)] / max(gt_num_cnt, 1)
            logger.info('recall_roi_%s: %f' % (cur_thresh, cur_roi_recall))
            logger.info('recall_rcnn_%s: %f' % (cur_thresh, cur_rcnn_recall))
            recall_dict['recall/roi_%s' % str(cur_thresh)] = cur_roi_recall
            recall_dict['recall/rcnn_%s' % str(cur_thresh)] = cur_rcnn_recall

        return recall_dict

    def eval(self, res_file, class_names, kitti_infos, res, **kwargs):
        if 'annos' not in kitti_infos[0].keys():
            return None, Metric
        det_annos = self.load_dts(res_file, res)
        from .kitti_object_eval_python import eval as kitti_eval
        eval_det_annos = copy.deepcopy(det_annos)
        eval_det_annos = sorted(eval_det_annos, key=lambda e: e['frame_id'])
        eval_gt_annos = [copy.deepcopy(info['annos']) for info in sorted(kitti_infos,
                                                                         key=lambda e:e['point_cloud']['lidar_idx'])]
        recall_dict = self.get_metric(eval_det_annos)
        result, recall_dict = kitti_eval.get_official_eval_result(eval_gt_annos, eval_det_annos, class_names)
        ave_recall = []
        for k, v in recall_dict.items():
            if "3d/moderate" in k:
                ave_recall.append(v)
        recall_dict["AVE_3d/moderate"] = np.array(ave_recall).mean()
        recall_dict = Metric(recall_dict)
        recall_dict.set_cmp_key("AVE_3d/moderate")
        logger.info(json.dumps(result, indent=2))
        return recall_dict

    @staticmethod
    def add_subparser(name, subparsers):
        subparser = subparsers.add_parser(
            name,
            help='subcommand for kitty evaluation')
        subparser.add_argument(
            '--anno_dir',
            required=True,
            help='directory holding kitty annotations')

        subparser.add_argument(
            '--res_file',
            required=True,
            help='file with each line of a result in json string format')
        return subparser


import os
import shutil

import numpy as np
import tensorflow as tf
import tensorflow_addons as tfa
import sys

def main():
    checkpoint_filename = sys.argv[1]
    out_path = sys.argv[2]
    tf.compat.v1.disable_eager_execution()
    for gpu in tf.config.experimental.list_physical_devices('GPU'):
        tf.config.experimental.set_memory_growth(gpu, True)

    with tf.compat.v1.Session() as sess:
        with tf.compat.v1.gfile.GFile(checkpoint_filename, 'rb') as file_handle:
            graph_def = tf.compat.v1.GraphDef()
            graph_def.ParseFromString(file_handle.read())

        tf.import_graph_def(graph_def, name='')
        graph = tf.compat.v1.get_default_graph()
        image_t = graph.get_tensor_by_name('images:0')
        features_t = graph.get_tensor_by_name('features:0')
        signature_def = tf.compat.v1.saved_model.signature_def_utils.predict_signature_def(
            inputs=dict(image=image_t), outputs=dict(features=features_t))
        os.mkdir(out_path)
        builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(out_path)
        builder.add_meta_graph_and_variables(
            sess, ['serve'], signature_def_map=dict(serving_default=signature_def))
        builder.save()
    tf.compat.v1.reset_default_graph()
    tf.compat.v1.enable_eager_execution()
    crop_model = tf.saved_model.load(out_path)
    shutil.rmtree(out_path)
    wrapped_model = ReIDModel(crop_model)
    tf.saved_model.save(wrapped_model, out_path)

class ReIDModel(tf.Module):
    def __init__(self, crop_model):
        super().__init__()
        self.crop_model = crop_model
        self.predict_crop = self.crop_model.signatures['serving_default']
        self.crop_size = (64, 128)

    @tf.function(input_signature=[
        tf.TensorSpec(shape=(None, None, 3), dtype=tf.uint8),
        tf.TensorSpec(shape=(3, 3), dtype=tf.float32),
        tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
        tf.TensorSpec(shape=(), dtype=tf.int32)])
    def predict_single_image(self, image, intrinsic_matrix, boxes, internal_batch_size=64):
        if tf.size(boxes) == 0:
            return tf.zeros(shape=(0, 128), dtype=tf.float32)
        ragged_boxes = tf.RaggedTensor.from_tensor(boxes[np.newaxis])
        return self.predict_multi_image(
            image[np.newaxis], intrinsic_matrix[np.newaxis], ragged_boxes, internal_batch_size)[0]

    @tf.function(input_signature=[
        tf.TensorSpec(shape=(None, None, None, 3), dtype=tf.uint8),
        tf.TensorSpec(shape=(None, 3, 3), dtype=tf.float32),
        tf.RaggedTensorSpec(shape=(None, None, 4), ragged_rank=1, dtype=tf.float32),
        tf.TensorSpec(shape=(), dtype=tf.int32)])
    def predict_multi_image(self, image, intrinsic_matrix, boxes, internal_batch_size=64):
        """Obtain ReID features for multiple bounding boxes specified for an image."""
        n_images = tf.shape(image)[0]
        if tf.size(boxes) == 0:
            # Special case for zero boxes provided
            result_flat = tf.zeros(shape=(0, 128), dtype=tf.float32)
            return tf.RaggedTensor.from_row_lengths(result_flat, tf.zeros(n_images, tf.int64))

        boxes_flat = boxes.flat_values
        n_box_per_image = boxes.row_lengths()
        image_id_per_box = boxes.value_rowids()
        n_total_boxes = tf.shape(boxes_flat)[0]

        if tf.shape(intrinsic_matrix)[0] == 1:
            intrinsic_matrix = tf.repeat(intrinsic_matrix, n_images, axis=0)

        Ks = tf.repeat(intrinsic_matrix, n_box_per_image, axis=0)

        if internal_batch_size == 0:
            # No batching
            results_flat = self.predict_single_batch(image, Ks, boxes_flat, image_id_per_box)
            return tf.RaggedTensor.from_row_lengths(results_flat, n_box_per_image)

        n_batches = tf.cast(tf.math.ceil(n_total_boxes / internal_batch_size), tf.int32)
        result_batches = tf.TensorArray(
            tf.float32, size=n_batches, element_shape=(None, 128), infer_shape=False)

        for i in tf.range(n_batches):
            box_batch = boxes_flat[i * internal_batch_size:(i + 1) * internal_batch_size]
            image_ids = image_id_per_box[i * internal_batch_size:(i + 1) * internal_batch_size]
            K_batch = Ks[i * internal_batch_size:(i + 1) * internal_batch_size]
            features = self.predict_single_batch(image, K_batch, box_batch, image_ids)
            result_batches = result_batches.write(i, features)

        results_flat = result_batches.concat()
        return tf.RaggedTensor.from_row_lengths(results_flat, n_box_per_image)

    def predict_single_batch(self, images, K, boxes, image_ids):
        n_box = tf.shape(boxes)[0]
        center_points = boxes[:, :2] + boxes[:, 2:4] / 2
        box_center_camspace = transf(center_points - K[:, :2, 2], tf.linalg.inv(K[:, :2, :2]))
        box_center_camspace = tf.concat(
            [box_center_camspace, tf.ones_like(box_center_camspace[:, :1])], axis=1)

        new_z = box_center_camspace / tf.linalg.norm(box_center_camspace, axis=-1, keepdims=True)
        new_x = tf.stack([new_z[:, 2], tf.zeros_like(new_z[:, 2]), -new_z[:, 0]], axis=1)
        new_y = tf.linalg.cross(new_z, new_x)
        new_R = tf.stack([new_x, new_y, new_z], axis=1)
        box_scales = tf.reduce_min(np.array(self.crop_size) / boxes[:, 2:4], axis=-1)
        new_K_mid = (tf.reshape(box_scales, [1, -1, 1, 1]) *
                     tf.reshape(K[:, :2, :2], [1, -1, 2, 2]))

        center = tf.concat([tf.fill((1, n_box, 1, 1), tf.cast(side, tf.float32))
                            for side in self.crop_size], axis=2) / 2
        intrinsic_matrix = tf.concat([
            tf.concat([new_K_mid, center], axis=3),
            tf.concat([tf.zeros((1, n_box, 1, 2), tf.float32),
                       tf.ones((1, n_box, 1, 1), tf.float32)], axis=3)], axis=2)
        new_proj_matrix = intrinsic_matrix @ new_R
        homography = K @ tf.linalg.inv(new_proj_matrix)
        homography = tf.reshape(homography, [1, n_box, 3, 3])
        homography = tf.reshape(homography, [-1, 9])
        homography = homography[:, :8] / homography[:, 8:]

        crops = perspective_transform(
            images, homography, (self.crop_size[1], self.crop_size[0]), 'BILINEAR', image_ids)
        crops = tf.reshape(crops, [1, n_box * self.crop_size[1], self.crop_size[0], 3])
        crops = tf.reshape(crops, [-1, self.crop_size[1], self.crop_size[0], 3])
        return self.predict_crop(image=crops)['features']

def transf(points, matrices):
    return tf.einsum('...k,...jk->...j', points, matrices)

def perspective_transform(images, homographies, output_shape, interpolation, image_ids):
    n_crops = tf.cast(tf.shape(homographies)[0], tf.int32)
    result_crops = tf.TensorArray(
        images.dtype, size=n_crops, element_shape=(*output_shape, 3), infer_shape=False)
    for i in tf.range(n_crops):
        tf.autograph.experimental.set_loop_options(parallel_iterations=1000)
        crop = tfa.image.transform(
            images[image_ids[i]], homographies[i], interpolation=interpolation,
            output_shape=output_shape)
        result_crops = result_crops.write(i, crop)
    return result_crops.stack()

def linspace_noend(start, stop, num):
    # Like np.linspace(endpoint=False)

    start = tf.convert_to_tensor(start)
    stop = tf.convert_to_tensor(stop, dtype=start.dtype)

    if num > 1:
        step = (stop - start) / tf.cast(num, start.dtype)
        new_stop = stop - step
        return tf.linspace(start, new_stop, num)
    else:
        return tf.linspace(start, stop, num)

if __name__ == '__main__':
    main()


import os
os.environ["LD_LIBRARY_PATH"] = ":/home/ztan/.mujoco/mujoco200/bin"
os.environ.get("LD_LIBRARY_PATH", "")

from gym.envs.mujoco import HalfCheetahEnv
from gym.envs.box2d import CarRacing

import rlkit.torch.pytorch_util as ptu
from rlkit.data_management.env_replay_buffer import EnvReplayBuffer
from rlkit.envs.wrappers import NormalizedBoxEnv
from rlkit.launchers.launcher_util import setup_logger
from rlkit.samplers.data_collector import MdpPathCollector
from rlkit.torch.sac.policies import (
    TanhGaussianPolicy,
    MakeDeterministic,
    TanhCNNGaussianPolicy,
    GaussianCNNPolicy,
)
from rlkit.torch.sac.sac import SACTrainer
from rlkit.torch.networks import ConcatMlp, PretrainedCNN, CNN
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm

import torch
import torchvision.models as models

from absl import app, flags
from typing import Sequence
import sys
from absl import app
from dm_control import viewer
from dm_robotics.moma import action_spaces

import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)

import dmc2gym
from gym import core, spaces
from dm_control import suite
from dm_env import specs
import numpy as np
from rgb_stacking import environment

ptu.set_gpu_mode(True)

def main(_argv):
    variant = dict(
        algorithm="SAC",
        version="normal",
        layer_size=256,
        replay_buffer_size=int(2e4),
        algorithm_kwargs=dict(
            num_epochs=int(1e4),
            # num_eval_steps_per_epoch=2000,
            # num_trains_per_train_loop=2000,
            # num_expl_steps_per_train_loop=2000,
            # min_num_steps_before_training=2000,
            num_eval_steps_per_epoch=4,
            num_trains_per_train_loop=4,
            num_expl_steps_per_train_loop=4,
            min_num_steps_before_training=4,
            max_path_length=400,
            batch_size=12,
        ),
        trainer_kwargs=dict(
            discount=0.99,
            soft_target_tau=5e-3,
            target_update_period=1,
            policy_lr=3e-4,
            qf_lr=3e-4,
            reward_scale=1,
            use_automatic_entropy_tuning=True,
        ),
    )

    expl_env = NormalizedBoxEnv(dmc2gym.make(domain_name="rgb_stacking", task_name='rgb_test_triplet3'))
    eval_env = NormalizedBoxEnv(dmc2gym.make(domain_name="rgb_stacking", task_name='rgb_test_triplet3'))
    obs_dim = expl_env.observation_space.low.size
    action_dim = eval_env.action_space.low.size
    M = variant["layer_size"]

    # "size" will return the desired product of dimensions
    print(expl_env.observation_space.shape)

    input_width, input_height, input_channels = expl_env.observation_space.shape  # channel last!!!
    qf1 = PretrainedCNN(
        input_width,
        input_height,
        input_channels,
        output_size=1,
        hidden_sizes=[128, 64],  # this is the hidden sizes of FC layers after the CNN
        added_fc_input_size=action_dim,  # layer used to merge image output and action input
        batch_norm_fc=False,
        init_w=1e-4,
        # hidden_init=nn.init.xavier_uniform_,
        # hidden_activation=nn.ReLU(),
        # output_activation=identity,
        output_conv_channels=False,
        model_architecture=models.efficientnet_b0,
        model_pretrained=True,
        model_freeze=False,
    )
    qf2 = PretrainedCNN(
        input_width,
        input_height,
        input_channels,
        output_size=1,
        hidden_sizes=[128, 64],  # this is the hidden sizes of FC layers after the CNN
        added_fc_input_size=action_dim,  # layer used to merge image output and action input
        batch_norm_fc=False,
        init_w=1e-4,
        # hidden_init=nn.init.xavier_uniform_,
        # hidden_activation=nn.ReLU(),
        # output_activation=identity,
        output_conv_channels=False,
        model_architecture=models.efficientnet_b0,
        model_pretrained=True,
        model_freeze=False,
    )
    target_qf1 = PretrainedCNN(
        input_width,
        input_height,
        input_channels,
        output_size=1,
        hidden_sizes=[128, 64],  # this is the hidden sizes of FC layers after the CNN
        added_fc_input_size=action_dim,  # layer used to merge image output and action input
        batch_norm_fc=False,
        init_w=1e-4,
        # hidden_init=nn.init.xavier_uniform_,
        # hidden_activation=nn.ReLU(),
        # output_activation=identity,
        output_conv_channels=False,
        model_architecture=models.efficientnet_b0,
        model_pretrained=True,
        model_freeze=False,
    )
    target_qf2 = PretrainedCNN(
        input_width,
        input_height,
        input_channels,
        output_size=1,
        hidden_sizes=[128, 64],  # this is the hidden sizes of FC layers after the CNN
        added_fc_input_size=action_dim,  # layer used to merge image output and action input
        batch_norm_fc=False,
        init_w=1e-4,
        # hidden_init=nn.init.xavier_uniform_,
        # hidden_activation=nn.ReLU(),
        # output_activation=identity,
        output_conv_channels=False,
        model_architecture=models.efficientnet_b0,
        model_pretrained=True,
        model_freeze=False,
    )
    policy = GaussianCNNPolicy(
        hidden_sizes=[128, 64],  # hidden size of FC after CNN; it uses "return_last_activations" to skip the last FC
        obs_dim=obs_dim,
        action_dim=action_dim,
        std=None,
        init_w=1e-3,
        min_log_std=-20,
        max_log_std=2,
        std_architecture="shared",
        **{
            "input_width": input_width,
            "input_height": input_height,
            "input_channels": input_channels,
            "kernel_sizes": [5, 5, 5],
            "n_channels": [32, 64, 128],
            "strides": [1] * 3,
            "paddings": ["same"] * 3,
        },
    )

    # now: also use pretrainedCNN

    # self.conv_output_flat_size: 1280 is the CNN output (effnet for example!)
    eval_policy = MakeDeterministic(policy)
    eval_path_collector = MdpPathCollector(eval_env, eval_policy,)
    expl_path_collector = MdpPathCollector(expl_env, policy,)
    replay_buffer = EnvReplayBuffer(variant["replay_buffer_size"], expl_env,)

    trainer = SACTrainer(
            env=eval_env,
            policy=policy,
            qf1=qf1,
            qf2=qf2,
            target_qf1=target_qf1,
            target_qf2=target_qf2,
            **variant["trainer_kwargs"]
        )
    algorithm = TorchBatchRLAlgorithm(
        trainer=trainer,
        exploration_env=expl_env,
        evaluation_env=eval_env,
        exploration_data_collector=expl_path_collector,
        evaluation_data_collector=eval_path_collector,
        replay_buffer=replay_buffer,
        **variant["algorithm_kwargs"]
    )

    setup_logger("rgb_stacking", variant=variant)
    algorithm.to(ptu.device)
    algorithm.train()

    expl_env.close()
    eval_env.close()

if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass


from tensorflow import keras
from tensorflow.keras.layers import Dense, Input, LayerNormalization
from vit_keras.layers.transformer import TransformerEncoderLayer, TransformerInputConv2DLayer
from tensorflow.io import gfile
import numpy as np

class VitBuilder():
    def __init__(self, image_size=(384, 384, 3),
                 patch_size=16,
                 num_heads=12,
                 num_layers=12,
                 num_classes=1000,
                 load_pretrain=True,
                 batch_size=1):
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.num_classes = num_classes
        self.load_pretrain = load_pretrain
        self.batch_size = batch_size

    def build(self):
        input_layer = Input(shape=self.image_size, batch_size=self.batch_size)
        x = TransformerInputConv2DLayer(self.image_size, self.patch_size)(input_layer)
        for i in range(self.num_layers):
            x = TransformerEncoderLayer(f'Transformer/encoderblock_{i}', self.image_size, self.patch_size, self.num_heads)(x)
        x = LayerNormalization(name='encoder_norm')(x)
        output = Dense(self.num_classes, name='head')(x[:, 0])

        vit_model = keras.Model(input_layer, output)

        if self.load_pretrain:
            vit_model = self.load_pretrain_weights(vit_model)

        return vit_model

    def load_pretrain_weights(self, vit):
        with gfile.GFile('ViT-B_16_imagenet2012.npz', 'rb') as f:
            ckpt_dict = np.load(f, allow_pickle=False)

        for layer in vit.get_layer('Transformer/posembed_input').trainable_weights:
            if 'cls:0' == layer.name:
                layer.assign(ckpt_dict['cls'])
            if 'position_embedding:0' == layer.name:
                layer.assign(ckpt_dict['Transformer/posembed_input/pos_embedding'])
            if 'Transformer/posembed_input/embedding/kernel:0' == layer.name:
                layer.assign(ckpt_dict['embedding/kernel'].reshape(layer.shape))
            if 'Transformer/posembed_input/embedding/bias:0' == layer.name:
                layer.assign(ckpt_dict['embedding/bias'])

        for i in range(self.num_layers):
            for layer in vit.get_layer(f'Transformer/encoderblock_{i}').trainable_weights:
                for k in [0, 2]:
                    if layer.name.endswith(f'/LayerNorm_{k}/gamma:0'):
                        layer.assign(ckpt_dict[f'Transformer/encoderblock_{i}/LayerNorm_{k}/scale'])
                    if layer.name.endswith(f'/LayerNorm_{k}/beta:0'):
                        layer.assign(ckpt_dict[f'Transformer/encoderblock_{i}/LayerNorm_{k}/bias'])

                for qkv in ['query', 'key', 'value', 'out']:
                    if layer.name.endswith(f'/{qkv}_kernel:0'):
                        layer.assign(
                            ckpt_dict[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/{qkv}/kernel'])
                    if layer.name.endswith(f'/{qkv}_bias:0'):
                        layer.assign(
                            ckpt_dict[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/{qkv}/bias'])

                for k in range(2):
                    if layer.name.endswith(f'/Dense_{k}/kernel:0'):
                        layer.assign(ckpt_dict[f'Transformer/encoderblock_{i}/MlpBlock_3/Dense_{k}/kernel'])
                    if layer.name.endswith(f'/Dense_{k}/bias:0'):
                        layer.assign(ckpt_dict[f'Transformer/encoderblock_{i}/MlpBlock_3/Dense_{k}/bias'])

        for layer in vit.get_layer('encoder_norm').trainable_weights:
            if 'encoder_norm/gamma:0' == layer.name:
                layer.assign(ckpt_dict['Transformer/encoder_norm/scale'])
            if 'encoder_norm/beta:0' == layer.name:
                layer.assign(ckpt_dict['Transformer/encoder_norm/bias'])

        for layer in vit.get_layer('head').trainable_weights:
            if 'head/kernel:0' == layer.name:
                layer.assign(ckpt_dict['head/kernel'])
            if 'head/bias:0' == layer.name:
                layer.assign(ckpt_dict['head/bias'])

        return vit