"""You Only Look Once Object Detection v3"""
# pylint: disable=arguments-differ
from __future__ import absolute_import
from __future__ import division

import os
import warnings
import numpy as np
import math
import mxnet as mx
from mxnet import gluon
from mxnet import autograd
from mxnet.gluon import nn
from mxnet.gluon.nn import BatchNorm
from models.definitions.darknet.darknet import get_darknet
from models.definitions.yolo.yolo_target import YOLOV3TargetMerger

from models.definitions.layers import TimeDistributed, Conv, Corr, _upsample, _temp_pad, _conv2d, _conv21d
from gluoncv.loss import YOLOV3Loss

__all__ = ['YOLOV3Temporal',
           'get_yolov3'
           ]

class YOLOOutputV3(gluon.HybridBlock):
    """YOLO output layer V3.
    Parameters
    ----------
    index : int
        Index of the yolo output layer, to avoid naming conflicts only.
    num_class : int
        Number of foreground objects.
    anchors : iterable
        The anchor setting. Reference: https://arxiv.org/pdf/1804.02767.pdf.
    stride : int
        Stride of feature map.
    alloc_size : tuple of int, default is (128, 128)
        For advanced users. Define `alloc_size` to generate large enough anchor
        maps, which will later saved in parameters. During inference, we support arbitrary
        input image by cropping corresponding area of the anchor map. This allow us
        to export to symbol so we can run it in c++, Scalar, etc.
    """
    def __init__(self, index, num_class, anchors, stride,
                 alloc_size=(128, 128), agnostic=False, **kwargs):
        super(YOLOOutputV3, self).__init__(**kwargs)
        anchors = np.array(anchors).astype('float32')
        self._classes = num_class
        self._num_pred = 1 + 4 + num_class  # 1 objness + 4 box + num_class
        self._num_anchors = anchors.size // 2
        self._stride = stride
        self._agnostic = agnostic
        with self.name_scope():
            all_pred = self._num_pred * self._num_anchors
            self.prediction = nn.Conv2D(all_pred, kernel_size=1, padding=0, strides=1)
            # anchors will be multiplied to predictions
            anchors = anchors.reshape(1, 1, -1, 2)
            self.anchors = self.params.get_constant('anchor_%d'%(index), anchors)
            # offsets will be added to predictions
            grid_x = np.arange(alloc_size[1])
            grid_y = np.arange(alloc_size[0])
            grid_x, grid_y = np.meshgrid(grid_x, grid_y)
            # stack to (n, n, 2)
            offsets = np.concatenate((grid_x[:, :, np.newaxis], grid_y[:, :, np.newaxis]), axis=-1)
            # expand dims to (1, 1, n, n, 2) so it's easier for broadcasting
            offsets = np.expand_dims(np.expand_dims(offsets, axis=0), axis=0)
            self.offsets = self.params.get_constant('offset_%d'%(index), offsets)

    def reset_class(self, classes, reuse_weights=None):
        """Reset class prediction.
        Parameters
        ----------
        classes : type
            Description of parameter `classes`.
        reuse_weights : dict
            A {new_integer : old_integer} mapping dict that allows the new predictor to reuse the
            previously trained weights specified by the integer index.
        Returns
        -------
        type
            Description of returned object.
        """
        self._clear_cached_op()
        # keep old records
        old_classes = self._classes
        old_pred = self.prediction
        old_num_pred = self._num_pred
        ctx = list(old_pred.params.values())[0].list_ctx()
        self._classes = len(classes)
        self._num_pred = 1 + 4 + len(classes)
        all_pred = self._num_pred * self._num_anchors
        # to avoid deferred init, number of in_channels must be defined
        in_channels = list(old_pred.params.values())[0].shape[1]

        self.prediction = nn.Conv2D(all_pred, kernel_size=1, padding=0, strides=1, in_channels=in_channels,
                                    prefix=old_pred.prefix)
        self.prediction.initialize(ctx=ctx)
        if reuse_weights:
            new_pred = self.prediction
            assert isinstance(reuse_weights, dict)
            for old_params, new_params in zip(old_pred.params.values(), new_pred.params.values()):
                old_data = old_params.data()
                new_data = new_params.data()
                for k, v in reuse_weights.items():
                    if k >= self._classes or v >= old_classes:
                        warnings.warn("reuse mapping {}/{} -> {}/{} out of range".format(
                            k, self._classes, v, old_classes))
                        continue
                    for i in range(self._num_anchors):
                        off_new = i * self._num_pred
                        off_old = i * old_num_pred
                        # copy along the first dimension
                        new_data[1 + 4 + k + off_new] = old_data[1 + 4 + v + off_old]
                        # copy non-class weights as well
                        new_data[off_new : 1 + 4 + off_new] = old_data[off_old : 1 + 4 + off_old]
                # set data to new conv layers
                new_params.set_data(new_data)

    def hybrid_forward(self, F, x, anchors, offsets):
        """Hybrid Forward of YOLOV3Output layer.
        Parameters
        ----------
        F : mxnet.nd or mxnet.sym
            `F` is mxnet.sym if hybridized or mxnet.nd if not.
        x : mxnet.nd.NDArray
            Input feature map.
        anchors : mxnet.nd.NDArray
            Anchors loaded from self, no need to supply.
        offsets : mxnet.nd.NDArray
            Offsets loaded from self, no need to supply.
        Returns
        -------
        (tuple of) mxnet.nd.NDArray
            During training, return (bbox, raw_box_centers, raw_box_scales, objness,
            class_pred, anchors, offsets).
            During inference, return detections.
        """
        # prediction flat to (batch, pred per pixel, height * width)
        pred = self.prediction(x)
        pred = pred.reshape((0, self._num_anchors * self._num_pred, -1))
        # transpose to (batch, height * width, num_anchor, num_pred)
        pred = pred.transpose(axes=(0, 2, 1)).reshape((0, -1, self._num_anchors, self._num_pred))
        # components
        raw_box_centers = pred.slice_axis(axis=-1, begin=0, end=2)
        raw_box_scales = pred.slice_axis(axis=-1, begin=2, end=4)
        objness = pred.slice_axis(axis=-1, begin=4, end=5)
        class_pred = pred.slice_axis(axis=-1, begin=5, end=None)

        # valid offsets, (1, 1, height, width, 2)
        offsets = F.slice_like(offsets, x * 0, axes=(2, 3))
        # reshape to (1, height*width, 1, 2)
        offsets = offsets.reshape((1, -1, 1, 2))

        box_centers = F.broadcast_add(F.sigmoid(raw_box_centers), offsets) * self._stride
        box_scales = F.broadcast_mul(F.exp(raw_box_scales), anchors)
        confidence = F.sigmoid(objness)
        class_score = F.broadcast_mul(F.sigmoid(class_pred), confidence)
        wh = box_scales / 2.0
        bbox = F.concat(box_centers - wh, box_centers + wh, dim=-1)

        if autograd.is_training():
            # during training, we don't need to convert whole bunch of info to detection results
            return (bbox.reshape((0, -1, 4)), raw_box_centers, raw_box_scales,
                    objness, class_pred, anchors, offsets)

        if self._agnostic:
            idsa = F.broadcast_add(confidence * 0, F.arange(0, 1).reshape((0, 1, 1, 1)))
            agnostic_detections = F.concat(idsa, confidence, bbox, dim=-1)
            agnostic_detections = F.reshape(agnostic_detections, (0, -1, 6))
            return agnostic_detections  # nms might merge some boxes as they now have same class

        # prediction per class
        bboxes = F.tile(bbox, reps=(self._classes, 1, 1, 1, 1))
        scores = F.transpose(class_score, axes=(3, 0, 1, 2)).expand_dims(axis=-1)
        ids = F.broadcast_add(scores * 0, F.arange(0, self._classes).reshape((0, 1, 1, 1, 1)))
        detections = F.concat(ids, scores, bboxes, dim=-1)
        # reshape to (B, xx, 6)
        detections = F.reshape(detections.transpose(axes=(1, 0, 2, 3, 4)), (0, -1, 6))

        return detections

class YOLODetectionBlockV3(gluon.HybridBlock):
    """YOLO V3 Detection Block which does the following:
    - add a few conv layers
    - return the output
    - have a branch that do yolo detection.
    Parameters
    ----------
    channel : int
        Number of channels for 1x1 conv. 3x3 Conv will have 2*channel.
    norm_layer : object
        Normalization layer used (default: :class:`mxnet.gluon.nn.BatchNorm`)
        Can be :class:`mxnet.gluon.nn.BatchNorm` or :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`.
    norm_kwargs : dict
        Additional `norm_layer` arguments, for example `num_devices=4`
        for :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`.
    """
    def __init__(self, channel, conv_type='2', norm_layer=BatchNorm, norm_kwargs=None, **kwargs):
        super(YOLODetectionBlockV3, self).__init__(**kwargs)

        self._conv_type = conv_type

        assert channel % 2 == 0, "channel {} cannot be divided by 2".format(channel)
        with self.name_scope():
            self.body = nn.HybridSequential(prefix='')
            for _ in range(2):
                # 1x1 reduce
                if conv_type in ['3', '21']:  # keep the expand as a normal 1x1x1 3d conv
                    self.body.add(Conv('3', channel, 1, 0, 1, norm_layer=norm_layer, norm_kwargs=norm_kwargs))
                else:
                    self.body.add(_conv2d(channel, 1, 0, 1, norm_layer=norm_layer, norm_kwargs=norm_kwargs))
                    # self.body.add(Conv('2', channel, 1, 0, 1, norm_layer=norm_layer, norm_kwargs=norm_kwargs))

                # 3x3 expand
                if conv_type in ['3', '21']:
                    self.body.add(Conv(conv_type, channel * 2, 3, 1, 1, norm_layer=norm_layer, norm_kwargs=norm_kwargs))
                else:
                    self.body.add(_conv2d(channel * 2, 3, 1, 1, norm_layer=norm_layer, norm_kwargs=norm_kwargs))

            if conv_type in ['3', '21']:
                self.body.add(Conv('3', channel, 1, 0, 1, norm_layer=norm_layer, norm_kwargs=norm_kwargs))
            else:
                self.body.add(_conv2d(channel, 1, 0, 1, norm_layer=norm_layer, norm_kwargs=norm_kwargs))
                # self.body.add(Conv('2', channel, 1, 0, 1, norm_layer=norm_layer, norm_kwargs=norm_kwargs))

            if conv_type in ['3', '21']:
                self.tip = Conv(conv_type, channel * 2, 3, 1, 1, norm_layer=norm_layer, norm_kwargs=norm_kwargs)
            else:
                self.tip = _conv2d(channel * 2, 3, 1, 1, norm_layer=norm_layer, norm_kwargs=norm_kwargs)

    def hybrid_forward(self, F, x):
        if self._conv_type in ['3', '21']:
            x = F.swapaxes(x, 1, 2)
        route = self.body(x)
        tip = self.tip(route)
        if self._conv_type in ['3', '21']:
            route = F.swapaxes(route, 1, 2)
            tip = F.swapaxes(tip, 1, 2)
        return route, tip

class YOLOV3Temporal(gluon.HybridBlock):
    """YOLO V3 detection network.
    Reference: https://arxiv.org/pdf/1804.02767.pdf.
    Parameters
    ----------
    stages : mxnet.gluon.HybridBlock
        Staged feature extraction blocks.
        For example, 3 stages and 3 YOLO output layers are used original paper.
    channels : iterable
        Number of conv channels for each appended stage.
        `len(channels)` should match `len(stages)`.
    num_class : int
        Number of foreground objects.
    anchors : iterable
        The anchor setting. `len(anchors)` should match `len(stages)`.
    strides : iterable
        Strides of feature map. `len(strides)` should match `len(stages)`.
    alloc_size : tuple of int, default is (128, 128)
        For advanced users. Define `alloc_size` to generate large enough anchor
        maps, which will later saved in parameters. During inference, we support arbitrary
        input image by cropping corresponding area of the anchor map. This allow us
        to export to symbol so we can run it in c++, Scalar, etc.
    nms_thresh : float, default is 0.45.
        Non-maximum suppression threshold. You can specify < 0 or > 1 to disable NMS.
    nms_topk : int, default is 400
        Apply NMS to top k detection results, use -1 to disable so that every Detection
         result is used in NMS.
    post_nms : int, default is 100
        Only return top `post_nms` detection results, the rest is discarded. The number is
        based on COCO dataset which has maximum 100 objects per image. You can adjust this
        number if expecting more objects. You can use -1 to return all detections.
    pos_iou_thresh : float, default is 1.0
        IOU threshold for true anchors that match real objects.
        'pos_iou_thresh < 1' is not implemented.
    ignore_iou_thresh : float
        Anchors that has IOU in `range(ignore_iou_thresh, pos_iou_thresh)` don't get
        penalized of objectness score.
    norm_layer : object
        Normalization layer used (default: :class:`mxnet.gluon.nn.BatchNorm`)
        Can be :class:`mxnet.gluon.nn.BatchNorm` or :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`.
    norm_kwargs : dict
        Additional `norm_layer` arguments, for example `num_devices=4`
        for :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`.
    """
    def __init__(self, stages, channels, anchors, strides, classes, alloc_size=(128, 128),
                 nms_thresh=0.45, nms_topk=400, post_nms=100, pos_iou_thresh=1.0,
                 ignore_iou_thresh=0.7, norm_layer=BatchNorm, norm_kwargs=None, agnostic=False, t_out=True, t=1, conv=2, corr_d=0,
                 **kwargs):
        super(YOLOV3Temporal, self).__init__(**kwargs)
        self._classes = classes
        self.nms_thresh = nms_thresh
        self.nms_topk = nms_topk
        self.post_nms = post_nms
        self.t_out = t_out
        self.t = t
        self.first_gap = 2*int(math.floor(self.t/2)/2)
        self.second_gap = 2*int(math.ceil(self.t/2)/2)

        self.conv = conv
        self.corr_d = corr_d
        self._pos_iou_thresh = pos_iou_thresh
        self._ignore_iou_thresh = ignore_iou_thresh
        if pos_iou_thresh >= 1:
            self._target_generator = YOLOV3TargetMerger(len(classes), ignore_iou_thresh)
        else:
            raise NotImplementedError(
                "pos_iou_thresh({}) < 1.0 is not implemented!".format(pos_iou_thresh))
        self._loss = YOLOV3Loss()
        with self.name_scope():
            self.stages = nn.HybridSequential()
            self.transitions = nn.HybridSequential()
            self.yolo_blocks = nn.HybridSequential()
            self.yolo_outputs = nn.HybridSequential()
            if t_out and corr_d:
                self.corr = Corr(corr_d, t=5, kernal_size=3, stride=1, keep='none', comp_mid=True)
                self.convs1 = nn.HybridSequential()
                self.convs1.add(_conv2d(channel=128, kernel=3, stride=1, padding=1,
                                        norm_layer=norm_layer, norm_kwargs=norm_kwargs))
                self.convs2 = nn.HybridSequential()
                self.convs2.add(_conv2d(channel=128, kernel=3, stride=2, padding=1,
                                        norm_layer=norm_layer, norm_kwargs=norm_kwargs))
                self.convs3 = nn.HybridSequential()
                self.convs3.add(_conv2d(channel=128, kernel=3, stride=2, padding=1,
                                        norm_layer=norm_layer, norm_kwargs=norm_kwargs))
            if not t_out:
                self.convs1 = nn.HybridSequential()
                self.convs2 = nn.HybridSequential()
                # for _ in range(1):
                self.convs1.add(_conv21d(channel=512, t=3, d=3, m=256, padding=[1, 0], stride=[(1, 2, 2), 1],
                                         norm_layer=norm_layer, norm_kwargs=norm_kwargs))
                self.convs2.add(_conv21d(channel=1024, t=3, d=3, m=512, padding=[1, 0], stride=[(1, 2, 2), 1],
                                         norm_layer=norm_layer, norm_kwargs=norm_kwargs))
            # note that anchors and strides should be used in reverse order
            for i, stage, channel, anchor, stride in zip(range(len(stages)), stages, channels, anchors[::-1], strides[::-1]):
                self.stages.add(stage)

                block = YOLODetectionBlockV3(channel, conv_type=str(conv), norm_layer=norm_layer, norm_kwargs=norm_kwargs)
                self.yolo_blocks.add(block)

                output = YOLOOutputV3(i, len(classes), anchor, stride, alloc_size=alloc_size, agnostic=agnostic)
                self.yolo_outputs.add(output)

                if i > 0:
                    self.transitions.add(_conv2d(channel, 1, 0, 1, norm_layer=norm_layer, norm_kwargs=norm_kwargs))

    @property
    def num_class(self):
        """Number of (non-background) categories.
        Returns
        -------
        int
            Number of (non-background) categories.
        """
        return self._num_class

    @property
    def classes(self):
        """Return names of (non-background) categories.
        Returns
        -------
        iterable of str
            Names of (non-background) categories.
        """
        return self._classes

    def hybrid_forward(self, F, x, *args):
        """YOLOV3 network hybrid forward.
        Parameters
        ----------
        F : mxnet.nd or mxnet.sym
            `F` is mxnet.sym if hybridized or mxnet.nd if not.
        x : mxnet.nd.NDArray
            Input data.
        *args : optional, mxnet.nd.NDArray
            During training, extra inputs are required:
            (gt_boxes, obj_t, centers_t, scales_t, weights_t, clas_t)
            These are generated by YOLOV3PrefetchTargetGenerator in dataloader transform function.
        Returns
        -------
        (tuple of) mxnet.nd.NDArray
            During inference, return detections in shape (B, N, 6)
            with format (cid, score, xmin, ymin, xmax, ymax)
            During training, return losses only: (obj_loss, center_loss, scale_loss, cls_loss).
        """
        all_box_centers = []
        all_box_scales = []
        all_objectness = []
        all_class_pred = []
        all_anchors = []
        all_offsets = []
        all_feat_maps = []
        all_detections = []
        routes = []
        if self.t == 1:
            for stage in self.stages:
                x = stage(x)
                routes.append(x)
        else:
            assert self.t == 5, 'Currently only support t=5 but will increase to more later'

            if self.t_out:
                if self.corr_d:
                    x = TimeDistributed(self.stages[0])(x)

                    # get middle feature for further Darknet processing
                    mid = F.squeeze(x.slice_axis(axis=1, begin=int(self.t / 2), end=int(self.t / 2)+1), axis=1)

                    x = self.corr(x)  # perform correlations across all timesteps

                    x = TimeDistributed(self.convs1)(x)  # do first conv
                    mid_rep = F.repeat(F.expand_dims(mid, axis=1), axis=1, repeats=self.t)  # repeat the mid feats t times
                    routes.append(F.concat(mid_rep, x, dim=2))  # concat and pass to YOLO (a,d)

                    mid = self.stages[1](mid)  # pass middle frame through more darknet
                    mid_rep = F.repeat(F.expand_dims(mid, axis=1), axis=1, repeats=self.t)  # repeat the mid feats t times
                    x = TimeDistributed(self.convs2)(x)  # downscale x with another conv
                    routes.append(F.concat(mid_rep, x, dim=2))  # concat and pass to YOLO (b,e)

                    mid = self.stages[2](mid)  # pass middle frame through last bit of darknet
                    mid_rep = F.repeat(F.expand_dims(mid, axis=1), axis=1, repeats=self.t)  # repeat the mid feats t times
                    x = TimeDistributed(self.convs3)(x)  # downscale x with another conv
                    x = F.concat(mid_rep, x, dim=2)
                    routes.append(x)  # concat and pass to YOLO (c,f)
                else:
                    x = TimeDistributed(self.stages[0])(x)
                    routes.append(x)
                    x = TimeDistributed(self.stages[1])(x)
                    # x = TimeDistributed(self.stages[1])(x.slice_axis(axis=1, begin=1, end=4))  # old code when did heir
                    routes.append(x)
                    x = TimeDistributed(self.stages[2])(x)
                    # x = TimeDistributed(self.stages[2])(x.slice_axis(axis=1, begin=1, end=2))  # old code when did heir
                    routes.append(x)
            else:
                x = TimeDistributed(self.stages[0])(x)
                routes.append(x.slice_axis(axis=1, begin=2, end=3).squeeze(axis=1))
                cx = F.swapaxes(self.convs1(F.swapaxes(x, 1, 2)), 1, 2)
                x = TimeDistributed(self.stages[1])(x.slice_axis(axis=1, begin=1, end=4))
                x = x + cx
                routes.append(x.slice_axis(axis=1, begin=1, end=2).squeeze(axis=1))
                cx = F.swapaxes(self.convs2(F.swapaxes(x, 1, 2)), 1, 2)
                x = TimeDistributed(self.stages[2])(x.slice_axis(axis=1, begin=1, end=2))
                x = x + cx
                x = x.squeeze(axis=1)
                routes.append(x)

        # the YOLO output layers are used in reverse order, i.e., from very deep layers to shallow
        for i, block, output in zip(range(len(routes)), self.yolo_blocks, self.yolo_outputs):
            if self.t > 1 and self.conv == 2 and self.t_out:
                x, tip = TimeDistributed(block)(x)
            else:
                x, tip = block(x)

            if self.t > 1 and self.t_out:
                if autograd.is_training():
                    dets, box_centers, box_scales, objness, class_pred, anchors, offsets = TimeDistributed(output, style='for')(tip)
                    all_box_centers.append(box_centers.reshape((0, 0, -3, -1)))
                    all_box_scales.append(box_scales.reshape((0, 0, -3, -1)))
                    all_objectness.append(objness.reshape((0, 0, -3, -1)))
                    all_class_pred.append(class_pred.reshape((0, 0, -3, -1)))
                    all_anchors.append(anchors)
                    all_offsets.append(offsets)
                    # here we use fake featmap to reduce memory consumption, only shape[2, 3] is used
                    fake_featmap = F.zeros_like(tip.slice_axis(
                        axis=0, begin=0, end=1).slice_axis(axis=2, begin=0, end=1))
                    all_feat_maps.append(fake_featmap)
                else:
                    dets = TimeDistributed(output)(tip)
            else:
                if autograd.is_training():
                    dets, box_centers, box_scales, objness, class_pred, anchors, offsets = output(tip)
                    all_box_centers.append(box_centers.reshape((0, -3, -1)))
                    all_box_scales.append(box_scales.reshape((0, -3, -1)))
                    all_objectness.append(objness.reshape((0, -3, -1)))
                    all_class_pred.append(class_pred.reshape((0, -3, -1)))
                    all_anchors.append(anchors)
                    all_offsets.append(offsets)
                    # here we use fake featmap to reduce memory consumption, only shape[2, 3] is used
                    fake_featmap = F.zeros_like(tip.slice_axis(
                        axis=0, begin=0, end=1).slice_axis(axis=1, begin=0, end=1))
                    all_feat_maps.append(fake_featmap)
                else:
                    dets = output(tip)

            all_detections.append(dets)

            if i >= len(routes) - 1:
                break

            # add transition layers
            if self.t > 1 and self.t_out:
                x = TimeDistributed(self.transitions[i])(x)
            else:
                x = self.transitions[i](x)

            # upsample feature map reverse to shallow layers
            upsample = _upsample(x, stride=2)

            route_now = routes[::-1][i + 1]

            x = F.concat(F.slice_like(upsample, route_now * 0, axes=(-2, -1)), route_now, dim=-3)  # concat to darknet

        if autograd.is_training():
            # during training, the network behaves differently since we don't need detection results
            if autograd.is_recording():
                box_preds = F.concat(*all_detections, dim=-2)

                if self.t == 1 or not self.t_out:  # the original if no temporal
                    all_targets = self._target_generator(box_preds, *args)
                    all_preds = [F.concat(*p, dim=1) for p in [
                        all_objectness, all_box_centers, all_box_scales, all_class_pred]]
                    return self._loss(*(all_preds + all_targets))

                losses = [[], [], [], []]
                for t in range(self.t):

                    all_preds = [F.slice_axis(F.concat(*p, dim=-2), axis=1, begin=t, end=t + 1).squeeze(axis=1)
                                 for p in [all_objectness, all_box_centers, all_box_scales, all_class_pred]]
                    box_preds_t = F.slice_axis(box_preds, axis=1, begin=t, end=t+1).squeeze(axis=1)
                    argst = [F.slice_axis(a, axis=1, begin=t, end=t+1).squeeze(axis=1) for a in args]
                    all_targets = self._target_generator(box_preds_t, *argst)
                    ls = self._loss(*(all_preds + all_targets))

                    for i, l in enumerate(ls):
                        losses[i].append(l)

                return [F.mean(F.concat(*l, dim=0)) for l in losses]

            if self.t > 1 and self.t_out:
                all_anchors = [F.slice_axis(a, axis=1, begin=1, end=2).squeeze(axis=1) for a in all_anchors]
                all_offsets = [F.slice_axis(a, axis=1, begin=1, end=2).squeeze(axis=1) for a in all_offsets]
                all_feat_maps = [F.slice_axis(a, axis=1, begin=1, end=2).squeeze(axis=1) for a in all_feat_maps]

            # orig 2d:
            return (F.concat(*all_detections, dim=-2), all_anchors, all_offsets, all_feat_maps,
                    F.concat(*all_box_centers, dim=-2), F.concat(*all_box_scales, dim=-2),
                    F.concat(*all_objectness, dim=-2), F.concat(*all_class_pred, dim=-2))

        # concat all detection results from different stages
        result = F.concat(*all_detections, dim=-2)

        if self.nms_thresh > 0 and self.nms_thresh < 1:  # todo check this works for the extra dim
            result = F.contrib.box_nms(
                result, overlap_thresh=self.nms_thresh, valid_thresh=0.01,
                topk=self.nms_topk, id_index=0, score_index=1, coord_start=2, force_suppress=False)

            if self.post_nms > 0:
                result = result.slice_axis(axis=-2, begin=0, end=self.post_nms)

        ids = result.slice_axis(axis=-1, begin=0, end=1)
        scores = result.slice_axis(axis=-1, begin=1, end=2)
        bboxes = result.slice_axis(axis=-1, begin=2, end=None)
        return ids, scores, bboxes

    def set_nms(self, nms_thresh=0.45, nms_topk=400, post_nms=100):
        """Set non-maximum suppression parameters.
        Parameters
        ----------
        nms_thresh : float, default is 0.45.
            Non-maximum suppression threshold. You can specify < 0 or > 1 to disable NMS.
        nms_topk : int, default is 400
            Apply NMS to top k detection results, use -1 to disable so that every Detection
             result is used in NMS.
        post_nms : int, default is 100
            Only return top `post_nms` detection results, the rest is discarded. The number is
            based on COCO dataset which has maximum 100 objects per image. You can adjust this
            number if expecting more objects. You can use -1 to return all detections.
        Returns
        -------
        None
        """
        self._clear_cached_op()
        self.nms_thresh = nms_thresh
        self.nms_topk = nms_topk
        self.post_nms = post_nms

    def reset_class(self, classes, reuse_weights=None):
        """Reset class categories and class predictors.
        Parameters
        ----------
        classes : iterable of str
            The new categories. ['apple', 'orange'] for example.
        reuse_weights : dict
            A {new_integer : old_integer} or mapping dict or {new_name : old_name} mapping dict,
            or a list of [name0, name1,...] if class names don't change.
            This allows the new predictor to reuse the
            previously trained weights specified.

        Example
        -------
        >>> net = gluoncv.model_zoo.get_model('yolo3_darknet53_voc', pretrained=True)
        >>> # use direct name to name mapping to reuse weights
        >>> net.reset_class(classes=['person'], reuse_weights={'person':'person'})
        >>> # or use interger mapping, person is the 14th category in VOC
        >>> net.reset_class(classes=['person'], reuse_weights={0:14})
        >>> # you can even mix them
        >>> net.reset_class(classes=['person'], reuse_weights={'person':14})
        >>> # or use a list of string if class name don't change
        >>> net.reset_class(classes=['person'], reuse_weights=['person'])

        """
        self._clear_cached_op()
        old_classes = self._classes
        self._classes = classes
        if self._pos_iou_thresh >= 1:
            self._target_generator = YOLOV3TargetMerger(len(classes), self._ignore_iou_thresh)
        if isinstance(reuse_weights, (dict, list)):
            if isinstance(reuse_weights, dict):
                # trying to replace str with indices
                new_keys = []
                new_vals = []
                for k, v in reuse_weights.items():
                    if isinstance(v, str):
                        try:
                            new_vals.append(old_classes.index(v))  # raise ValueError if not found
                        except ValueError:
                            raise ValueError(
                                "{} not found in old class names {}".format(v, old_classes))
                    else:
                        if v < 0 or v >= len(old_classes):
                            raise ValueError(
                                "Index {} out of bounds for old class names".format(v))
                        new_vals.append(v)
                    if isinstance(k, str):
                        try:
                            new_keys.append(self.classes.index(k))  # raise ValueError if not found
                        except ValueError:
                            raise ValueError(
                                "{} not found in new class names {}".format(k, self.classes))
                    else:
                        if k < 0 or k >= len(self.classes):
                            raise ValueError(
                                "Index {} out of bounds for new class names".format(k))
                        new_keys.append(k)
                reuse_weights = dict(zip(new_keys, new_vals))
            else:
                new_map = {}
                for x in reuse_weights:
                    try:
                        new_idx = self._classes.index(x)
                        old_idx = old_classes.index(x)
                        new_map[new_idx] = old_idx
                    except ValueError:
                        warnings.warn("{} not found in old: {} or new class names: {}".format(
                            x, old_classes, self._classes))
                reuse_weights = new_map

        for outputs in self.yolo_outputs:
            outputs.reset_class(classes, reuse_weights=reuse_weights)

# class YOLOV3Temporal(gluon.HybridBlock):
#     """YOLO V3 detection network.
#     Reference: https://arxiv.org/pdf/1804.02767.pdf.
#     Parameters
#     ----------
#     stages : mxnet.gluon.HybridBlock
#         Staged feature extraction blocks.
#         For example, 3 stages and 3 YOLO output layers are used original paper.
#     channels : iterable
#         Number of conv channels for each appended stage.
#         `len(channels)` should match `len(stages)`.
#     num_class : int
#         Number of foreground objects.
#     anchors : iterable
#         The anchor setting. `len(anchors)` should match `len(stages)`.
#     strides : iterable
#         Strides of feature map. `len(strides)` should match `len(stages)`.
#     alloc_size : tuple of int, default is (128, 128)
#         For advanced users. Define `alloc_size` to generate large enough anchor
#         maps, which will later saved in parameters. During inference, we support arbitrary
#         input image by cropping corresponding area of the anchor map. This allow us
#         to export to symbol so we can run it in c++, Scalar, etc.
#     nms_thresh : float, default is 0.45.
#         Non-maximum suppression threshold. You can specify < 0 or > 1 to disable NMS.
#     nms_topk : int, default is 400
#         Apply NMS to top k detection results, use -1 to disable so that every Detection
#          result is used in NMS.
#     post_nms : int, default is 100
#         Only return top `post_nms` detection results, the rest is discarded. The number is
#         based on COCO dataset which has maximum 100 objects per image. You can adjust this
#         number if expecting more objects. You can use -1 to return all detections.
#     pos_iou_thresh : float, default is 1.0
#         IOU threshold for true anchors that match real objects.
#         'pos_iou_thresh < 1' is not implemented.
#     ignore_iou_thresh : float
#         Anchors that has IOU in `range(ignore_iou_thresh, pos_iou_thresh)` don't get
#         penalized of objectness score.
#     norm_layer : object
#         Normalization layer used (default: :class:`mxnet.gluon.nn.BatchNorm`)
#         Can be :class:`mxnet.gluon.nn.BatchNorm` or :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`.
#     norm_kwargs : dict
#         Additional `norm_layer` arguments, for example `num_devices=4`
#         for :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`.
#     """
#     def __init__(self, d_model, channels, anchors, strides, classes, alloc_size=(128, 128),
#                  nms_thresh=0.45, nms_topk=400, post_nms=100, pos_iou_thresh=1.0,
#                  ignore_iou_thresh=0.7, norm_layer=BatchNorm, norm_kwargs=None,
#                  k=None, k_join_type=None, k_join_pos=None, block_conv_type='2',
#                  corr_pos=None, corr_d=None, agnostic=False, **kwargs):
#         super(YOLOV3Temporal, self).__init__(**kwargs)
#         self.d_model = d_model
#         self._classes = classes
#         self.nms_thresh = nms_thresh
#         self.nms_topk = nms_topk
#         self.post_nms = post_nms
#         self._pos_iou_thresh = pos_iou_thresh
#         self._ignore_iou_thresh = ignore_iou_thresh
#         self._k = k
#         self._k_join_type = k_join_type
#         self._k_join_pos = k_join_pos
#         self._block_conv_type = block_conv_type
#         self._corr_pos = corr_pos
#         self._corr_d = corr_d
#         if block_conv_type in ['3', '21']:
#             assert k > 1, "k must be greater than 1 to use 3D or 2+1D convolutions"
#             assert k_join_pos == 'late', "only 'late' pooling can be used when using 3D or 2+1D convolutions"
#             assert k_join_type is not None, "please specify a k_join_type: max, mean, or cat"
#         assert k_join_type in [None, 'max', 'mean', 'cat']
#         assert k_join_pos in [None, 'early', 'late']
#         assert corr_pos in [None, 'early', 'late']
#         if pos_iou_thresh >= 1:
#             self._target_generator = YOLOV3TargetMerger(len(classes), ignore_iou_thresh)
#         else:
#             raise NotImplementedError(
#                 "pos_iou_thresh({}) < 1.0 is not implemented!".format(pos_iou_thresh))
#         self._loss = YOLOV3Loss()
#         with self.name_scope():
#
#             if k > 1 and k_join_type in ['max', 'mean']:
#                 self.pool = TemporalPooling(k=k, type=k_join_type)
#
#             if k > 1 and corr_pos is not None:
#                 self.corr = Corr(corr_d, k, kernal_size=1, stride=1, keep='all')
#
#             self.stages = nn.HybridSequential()
#             self.transitions = nn.HybridSequential()
#             self.yolo_blocks = nn.HybridSequential()
#             self.yolo_tips = (None, None, None)
#             self.yolo_outputs = nn.HybridSequential()
#
#             # note that anchors and strides should be used in reverse order
#             for i, channel, anchor, stride in zip(range(3), channels, anchors[::-1], strides[::-1]):
#
#                 # with late rnn_pos we split block and tip into sep
#                 block = YOLODetectionBlockV3(channel, block_conv_type, norm_layer=norm_layer, norm_kwargs=norm_kwargs)
#
#                 if self._k > 1 and block_conv_type == '2' and \
#                         (self._k_join_pos == 'late' or self._corr_pos == 'late'):
#                     self.yolo_blocks.add(TimeDistributed(block))
#                 else:
#                     self.yolo_blocks.add(block)
#
#                 output = YOLOOutputV3(i, len(classes), anchor, stride, alloc_size=alloc_size, agnostic=agnostic)
#                 self.yolo_outputs.add(output)
#
#                 if i > 0:
#                     if self._k > 1 and \
#                             (self._k_join_pos == 'late' or self._corr_pos == 'late'):
#                         self.transitions.add(TimeDistributed(_conv2d(channel, 1, 0, 1,
#                                                              norm_layer=norm_layer, norm_kwargs=norm_kwargs)))
#                     else:
#                         self.transitions.add(_conv2d(channel, 1, 0, 1,
#                                                      norm_layer=norm_layer, norm_kwargs=norm_kwargs))
#
#     @property
#     def num_class(self):
#         """Number of (non-background) categories.
#         Returns
#         -------
#         int
#             Number of (non-background) categories.
#         """
#         return self._num_class
#
#     @property
#     def classes(self):
#         """Return names of (non-background) categories.
#         Returns
#         -------
#         iterable of str
#             Names of (non-background) categories.
#         """
#         return self._classes
#
#     def hybrid_forward(self, F, x, *args):
#         """YOLOV3 network hybrid forward.
#         Parameters
#         ----------
#         F : mxnet.nd or mxnet.sym
#             `F` is mxnet.sym if hybridized or mxnet.nd if not.
#         x : mxnet.nd.NDArray
#             Input data.
#         *args : optional, mxnet.nd.NDArray
#             During training, extra inputs are required:
#             (gt_boxes, obj_t, centers_t, scales_t, weights_t, clas_t)
#             These are generated by YOLOV3PrefetchTargetGenerator in dataloader transform function.
#         Returns
#         -------
#         (tuple of) mxnet.nd.NDArray
#             During inference, return detections in shape (B, N, 6)
#             with format (cid, score, xmin, ymin, xmax, ymax)
#             During training, return losses only: (obj_loss, center_loss, scale_loss, cls_loss).
#         """
#         all_box_centers = []
#         all_box_scales = []
#         all_objectness = []
#         all_class_pred = []
#         all_anchors = []
#         all_offsets = []
#         all_feat_maps = []
#         all_detections = []
#
#         routes = []
#
#         mid = int(self._k/2)
#         if self._k > 1:
#             for r in TimeDistributed(self.d_model)(x):
#                 if self._k_join_pos == 'early':
#                     if self._k_join_type == 'cat':
#                         r = F.reshape(r, (0, -3, -2))  # B,K,C,H,W -> B,K*C,H,W
#                     elif self._k_join_type in ['max', 'mean']:
#                         r = self.pool(r)
#                 elif self._corr_pos == 'early':
#                     r = self.corr(r)
#                 routes.append(r)
#         else:
#             routes = self.d_model(x)
#
#         x = routes[-1]
#
#         # the YOLO output layers are used in reverse order, i.e., from very deep layers to shallow
#         for i, block, tips, output in zip(range(len(routes)), self.yolo_blocks, self.yolo_tips, self.yolo_outputs):
#
#             x, tip = block(x)
#
#             if self._k > 1 and self._k_join_pos == 'late':
#                 if self._k_join_type == 'cat':
#                     tip = F.reshape(tip, (0, -3, -2))  # B,K,C,H,W -> B,K*C,H,W
#                 elif self._k_join_type in ['max', 'mean']:
#                     tip = self.pool(tip)
#             elif self._k > 1 and self._corr_pos == 'late':
#                 tip = self.corr(tip)
#
#             if autograd.is_training():
#                 dets, box_centers, box_scales, objness, class_pred, anchors, offsets = output(tip)
#                 all_box_centers.append(box_centers.reshape((0, -3, -1)))
#                 all_box_scales.append(box_scales.reshape((0, -3, -1)))
#                 all_objectness.append(objness.reshape((0, -3, -1)))
#                 all_class_pred.append(class_pred.reshape((0, -3, -1)))
#                 all_anchors.append(anchors)
#                 all_offsets.append(offsets)
#                 # here we use fake featmap to reduce memory consumption, only shape[2, 3] is used
#                 fake_featmap = F.zeros_like(tip.slice_axis(axis=0, begin=0, end=1).slice_axis(axis=1, begin=0, end=1))
#                 all_feat_maps.append(fake_featmap)
#             else:
#                 dets = output(tip)
#
#             all_detections.append(dets)
#
#             if i >= len(routes) - 1: # last output layer scale
#                 break
#
#             # add transition layers
#             x = self.transitions[i](x)
#
#             # upsample feature map reverse to shallow layers
#             upsample = _upsample(x, stride=2)
#             route_now = routes[::-1][i + 1]
#
#             # concat
#             if self._k > 1 and (self._k_join_pos == 'late' or self._corr_pos == 'late'):
#                 x = F.concat(F.slice_like(upsample, route_now * 0, axes=(3, 4)), route_now, dim=2)
#             else:
#                 x = F.concat(F.slice_like(upsample, route_now * 0, axes=(2, 3)), route_now, dim=1)
#
#         if autograd.is_training():
#             # during training, the network behaves differently since we don't need detection results
#             if autograd.is_recording():
#                 # generate losses and return them directly
#                 box_preds = F.concat(*all_detections, dim=1)
#                 all_preds = [F.concat(*p, dim=1) for p in [
#                     all_objectness, all_box_centers, all_box_scales, all_class_pred]]
#                 all_targets = self._target_generator(box_preds, *args)
#                 return self._loss(*(all_preds + all_targets))
#
#             # return raw predictions, this is only used in DataLoader transform function.
#             return (F.concat(*all_detections, dim=1), all_anchors, all_offsets, all_feat_maps,
#                     F.concat(*all_box_centers, dim=1), F.concat(*all_box_scales, dim=1),
#                     F.concat(*all_objectness, dim=1), F.concat(*all_class_pred, dim=1))
#
#         # concat all detection results from different stages
#         result = F.concat(*all_detections, dim=1)
#         # apply nms per class
#         if self.nms_thresh > 0 and self.nms_thresh < 1:
#             result = F.contrib.box_nms(
#                 result, overlap_thresh=self.nms_thresh, valid_thresh=0.01,
#                 topk=self.nms_topk, id_index=0, score_index=1, coord_start=2, force_suppress=False)
#             if self.post_nms > 0:
#                 result = result.slice_axis(axis=1, begin=0, end=self.post_nms)
#         ids = result.slice_axis(axis=-1, begin=0, end=1)
#         scores = result.slice_axis(axis=-1, begin=1, end=2)
#         bboxes = result.slice_axis(axis=-1, begin=2, end=None)
#         return ids, scores, bboxes
#
#     def set_nms(self, nms_thresh=0.45, nms_topk=400, post_nms=100):
#         """Set non-maximum suppression parameters.
#         Parameters
#         ----------
#         nms_thresh : float, default is 0.45.
#             Non-maximum suppression threshold. You can specify < 0 or > 1 to disable NMS.
#         nms_topk : int, default is 400
#             Apply NMS to top k detection results, use -1 to disable so that every Detection
#              result is used in NMS.
#         post_nms : int, default is 100
#             Only return top `post_nms` detection results, the rest is discarded. The number is
#             based on COCO dataset which has maximum 100 objects per image. You can adjust this
#             number if expecting more objects. You can use -1 to return all detections.
#         Returns
#         -------
#         None
#         """
#         self._clear_cached_op()
#         self.nms_thresh = nms_thresh
#         self.nms_topk = nms_topk
#         self.post_nms = post_nms
#
#     def reset_class(self, classes, reuse_weights=None):
#         """Reset class categories and class predictors.
#         Parameters
#         ----------
#         classes : iterable of str
#             The new categories. ['apple', 'orange'] for example.
#         reuse_weights : dict
#             A {new_integer : old_integer} or mapping dict or {new_name : old_name} mapping dict,
#             or a list of [name0, name1,...] if class names don't change.
#             This allows the new predictor to reuse the
#             previously trained weights specified.
#
#         Example
#         -------
#         >>> net = gluoncv.model_zoo.get_model('yolo3_darknet53_voc', pretrained=True)
#         >>> # use direct name to name mapping to reuse weights
#         >>> net.reset_class(classes=['person'], reuse_weights={'person':'person'})
#         >>> # or use interger mapping, person is the 14th category in VOC
#         >>> net.reset_class(classes=['person'], reuse_weights={0:14})
#         >>> # you can even mix them
#         >>> net.reset_class(classes=['person'], reuse_weights={'person':14})
#         >>> # or use a list of string if class name don't change
#         >>> net.reset_class(classes=['person'], reuse_weights=['person'])
#
#         """
#         self._clear_cached_op()
#         old_classes = self._classes
#         self._classes = classes
#         if self._pos_iou_thresh >= 1:
#             self._target_generator = YOLOV3TargetMerger(len(classes), self._ignore_iou_thresh)
#         if isinstance(reuse_weights, (dict, list)):
#             if isinstance(reuse_weights, dict):
#                 # trying to replace str with indices
#                 new_keys = []
#                 new_vals = []
#                 for k, v in reuse_weights.items():
#                     if isinstance(v, str):
#                         try:
#                             new_vals.append(old_classes.index(v))  # raise ValueError if not found
#                         except ValueError:
#                             raise ValueError(
#                                 "{} not found in old class names {}".format(v, old_classes))
#                     else:
#                         if v < 0 or v >= len(old_classes):
#                             raise ValueError(
#                                 "Index {} out of bounds for old class names".format(v))
#                         new_vals.append(v)
#                     if isinstance(k, str):
#                         try:
#                             new_keys.append(self.classes.index(k))  # raise ValueError if not found
#                         except ValueError:
#                             raise ValueError(
#                                 "{} not found in new class names {}".format(k, self.classes))
#                     else:
#                         if k < 0 or k >= len(self.classes):
#                             raise ValueError(
#                                 "Index {} out of bounds for new class names".format(k))
#                         new_keys.append(k)
#                 reuse_weights = dict(zip(new_keys, new_vals))
#             else:
#                 new_map = {}
#                 for x in reuse_weights:
#                     try:
#                         new_idx = self._classes.index(x)
#                         old_idx = old_classes.index(x)
#                         new_map[new_idx] = old_idx
#                     except ValueError:
#                         warnings.warn("{} not found in old: {} or new class names: {}".format(
#                             x, old_classes, self._classes))
#                 reuse_weights = new_map
#
#         for outputs in self.yolo_outputs:
#             outputs.reset_class(classes, reuse_weights=reuse_weights)

def get_yolov3(stages, filters, anchors, strides, classes, **kwargs):
    """Get YOLOV3 models.
    Parameters
    ----------
    stages : iterable of str or `HybridBlock`
        List of network internal output names, in order to specify which layers are
        used for predicting bbox values.
        If `name` is `None`, `features` must be a `HybridBlock` which generate multiple
        outputs for prediction.
    filters : iterable of float or None
        List of convolution layer channels which is going to be appended to the base
        network feature extractor. If `name` is `None`, this is ignored.
    sizes : iterable fo float
        Sizes of anchor boxes, this should be a list of floats, in incremental order.
        The length of `sizes` must be len(layers) + 1. For example, a two stage SSD
        model can have ``sizes = [30, 60, 90]``, and it converts to `[30, 60]` and
        `[60, 90]` for the two stages, respectively. For more details, please refer
        to original paper.
    ratios : iterable of list
        Aspect ratios of anchors in each output layer. Its length must be equals
        to the number of SSD output layers.
    steps : list of int
        Step size of anchor boxes in each output layer.
    classes : iterable of str
        Names of categories.
    dataset : str
        Name of dataset. This is used to identify model name because models trained on
        different datasets are going to be very different.
    pretrained : bool or str
        Boolean value controls whether to load the default pretrained weights for model.
        String value represents the hashtag for a certain version of pretrained weights.
    pretrained_base : bool or str, optional, default is True
        Load pretrained base network, the extra layers are randomized. Note that
        if pretrained is `True`, this has no effect.
    ctx : mxnet.Context
        Context such as mx.cpu(), mx.gpu(0).
    root : str
        Model weights storing path.
    norm_layer : object
        Normalization layer used (default: :class:`mxnet.gluon.nn.BatchNorm`)
        Can be :class:`mxnet.gluon.nn.BatchNorm` or :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`.
    norm_kwargs : dict
        Additional `norm_layer` arguments, for example `num_devices=4`
        for :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`.
    Returns
    -------
    HybridBlock
        A YOLOV3 detection network.
    """
    net = YOLOV3Temporal(stages, filters, anchors, strides, classes=classes, **kwargs)
    return net

if __name__ == '__main__':

    t = 5
    freeze_base = True
    filters = [512, 256, 128]
    anchors = [
        [10, 13, 16, 30, 33, 23],
        [30, 61, 62, 45, 59, 119],
        [116, 90, 156, 198, 373, 326]]
    strides = [8, 16, 32]

    darknet = get_darknet(pretrained=True, norm_layer=BatchNorm)
    if freeze_base:
        for param in darknet.collect_params().values():
            param.grad_req = 'null'
    stages = [darknet.features[:15], darknet.features[15:24], darknet.features[24:]]

    net = get_yolov3(stages, filters, anchors, strides, classes=list(range(30)), t=t)
    net.initialize()
    net.summary(mx.nd.random_normal(shape=(1, t, 3, 384, 384)))  # b,t,c,w,h

    with autograd.record():
        net(mx.nd.random_normal(shape=(1, t, 3, 384, 384)))


from keras.models import Model
from keras.utils import multi_gpu_model
from keras.callbacks import ModelCheckpoint, EarlyStopping
from argparse import ArgumentParser
from datagen import TripletsGenerator
from models import create_skip
import numpy as np
import tensorflow as tf
import keras
import datagen
import utils
import sys

# Set random seeds for reproducibility
import random
random.seed(333)
np.random.seed(333)
tf.set_random_seed(333)

# Default values
EMBEDDING_DIM   = 600
ENCODER_SIZE    = 2400
DECODER_SIZE    = 2400
CELL            = 'gru'
VOCAB_SIZE      = 20000
MAX_LEN         = 40
WINDOW          = 1

EPOCHS      = 100
BATCH_SIZE  = 128
STEPS_EPOCH = 1e3

if __name__ == "__main__":
    parser = ArgumentParser("Train sent2vec model")
    parser.add_argument('-g', '--gpu' ,type=int, default=0, help="GPU device to be used")
    parser.add_argument('-c','--corpus', required=True, help="Corpus file for the training")
    parser.add_argument('--dev', required=True, help="Development set file")
    parser.add_argument('-t','--tokenizer', default='tknzr', required=False, help="File name to save the tokenizer")
    parser.add_argument('-m','--model', required=True, help="Model name")
    parser.add_argument('-s','--size', type=int, default=ENCODER_SIZE, help="Size of encoder and decoder")
    parser.add_argument('--cell', type=str, default=CELL, choices=['gru', 'lstm'], help="Cell type of the recurrent netowrk: GRU or LSTM")
    parser.add_argument('-v','--vocab-size', type=int, default=VOCAB_SIZE, help="Size of the vocabulary")
    parser.add_argument('--embedding-dim', type=int, default=EMBEDDING_DIM, help="Emmbedding vector dimensions")
    parser.add_argument('-b','--batch_size', type=int, default=BATCH_SIZE, help="Batch size")
    parser.add_argument('-e','--epochs', type=int, default=EPOCHS, help="Number of epochs")
    parser.add_argument('--max-len', type=int, default=MAX_LEN, help="Max sequence length")
    parser.add_argument('-sp','--steps', type=int, default=STEPS_EPOCH, help='Number of steps (batches) per epoch')
    parser.add_argument('-w','--window', type=int, default=WINDOW, help='Window of context. Number of sentences to use on backward and forward')
    parser.add_argument('--no-filters', dest='filters', action='store_false')
    args = parser.parse_args(sys.argv[1:])
    print(args.__dict__)

    tripgen = TripletsGenerator(args.corpus,
                            batch_size=args.batch_size,
                            max_len=args.max_len,
                            vocab_size=args.vocab_size,
                            window=args.window,
                            filters=args.filters)
    devgen = TripletsGenerator(args.dev,
                            tokenizer=tripgen.tknzr,
                            batch_size=args.batch_size,
                            window=args.window,
                            max_len=MAX_LEN)
    utils.oov_pct(tripgen.data, tripgen.tknzr.word_index[datagen.OOV])

    if 'bpe' in args.corpus or 'BPE' in args.corpus:
        preprocess = '-bpe'
    else:
        preprocess = ''
    if args.filters:
        preprocess += '-filters'
    else:
        preprocess += '-nofilters'

    suffix = '-c' + utils.num_to_str(len(tripgen.data)) \
                            + '-' + args.cell + str(args.size) \
                            + '-w' + str(args.window) \
                            + '-e' + str(args.embedding_dim) \
                            + '-v' + utils.num_to_str(args.vocab_size) \
                            + preprocess
    modelname = 'models/' + args.model \
                            + suffix \
                            + '.best.h5'
    tokenname = 'models/' + args.tokenizer \
                            + suffix \
                            +'.pkl'
    datagen.save_tokenizer(tripgen.tknzr,tokenname)
    print('Vocab size:',args.vocab_size)
    print('Num samples:',len(tripgen.data))

with tf.device('/gpu:'+str(args.gpu)):
    model = create_skip(embedding_dim=args.embedding_dim,
                        encoder_size=args.size,
                        decoder_size=args.size,
                        cell=args.cell,
                        vocab_size=args.vocab_size,
                        max_len=args.max_len,
                        window=args.window,
                        )
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
    model.summary()

    checkpoint = ModelCheckpoint(modelname,
                                monitor='val_loss',
                                save_best_only=True,
                                mode='min',
                                period=1)
    earlystop = EarlyStopping(monitor='val_loss',
                            patience=10,
                            mode='min')
    history = model.fit_generator(generator=tripgen,
                use_multiprocessing=True,
                workers=8,
                shuffle=True,
                epochs=args.epochs,
                steps_per_epoch=args.steps,
                callbacks=[checkpoint,earlystop],
                validation_data=devgen,
                verbose=1,
                )
    utils.plot(history,args.model,modelname)


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

"""
import tensorflow as tf
import numpy as np
import tflearn

from ddpg.ddpg import build_summaries, ActorNetwork, CriticNetwork, OrnsteinUhlenbeckActionNoise, ReplayBuffer, \
    getReward

from src.BallTracker import ballTracker
from src.Pepper import Pepper
from src.Pepper.Pepper import readAngle
from src.Settings import *

# ===========================
#   Agent Training
# ===========================

def train(sess, session, thread, args, actor, critic, actor_noise, update_model, saver):
    # Set up summary Ops
    summary_ops, summary_vars = build_summaries()

    if update_model == False:
        sess.run(tf.global_variables_initializer())
        # Initialize target network weights
        actor.update_target_network()
        critic.update_target_network()

    writer = tf.summary.FileWriter(args['summary_dir'], sess.graph)

    # Initialize replay memory
    replay_buffer = ReplayBuffer(int(args['buffer_size']), int(args['random_seed']))

    # Needed to enable BatchNorm.
    # This hurts the performance on Pendulum but could be useful
    # in other environments.
    tflearn.is_training(True)

    for i in range(int(args['max_episodes'])):

        # s = env.reset()

        ep_reward = 0
        ep_ave_max_q = 0
        for j in range(int(args['max_episode_len'])):
            service = session.service("ALMotion")
            params = dict()
            # Hole Anfangszustand
            delta1 = thread.delta[0]
            winkel1 = readAngle(session)
            s = [winkel1, delta1]

            # Hole action
            a = actor.predict(np.reshape(s, (1, 2))) + (1. / (1. + i))
            # ITERATE THORUGH SAMPLED DATA AND ADD TO REPLAY BUFFER

            a = actor.predict(np.reshape(s, (1, actor.s_dim))) + actor_noise()

            # a[0] = ((a[0] - (1 - action_bound)) / (action_bound - (1 - action_bound))) * (
            #            OBERE_GRENZE - UNTERE_GRENZE) + UNTERE_GRENZE
            # a[0] = a[0]
            rewardTMP = 0
            if a[0] < UNTERE_GRENZE:
                # print("Winkel zu klein :" + str(a[0]))
                a[0] = UNTERE_GRENZE
                rewardTMP = -1000

            if a[0] > OBERE_GRENZE:
                # print("Winkel zu gross :" + str(a[0]))
                a[0] = OBERE_GRENZE
                rewardTMP = -1000

            # Fuehre Action aus
            params[args['motor']] = [a[0], TIME_TO_MOVE]
            Pepper.move(params, service)

            # Hole Folgezustand
            delta2 = thread.delta[0]
            winkel2 = readAngle(session)
            s2 = [winkel2, delta2]

            # Hole Reward
            r = getReward(delta2) + rewardTMP
            terminal = False
            print(str(a[0]) + "\t" + str(
                r))

            replay_buffer.add(np.reshape(s, (actor.s_dim,)), np.reshape(a, (actor.a_dim,)), r,
                              terminal, np.reshape(s2, (actor.s_dim,)))

            # export actor Model somewhere

            # Keep adding experience to the memory until
            # there are at least minibatch size samples
            if replay_buffer.size() > int(args['minibatch_size']):
                s_batch, a_batch, r_batch, t_batch, s2_batch = \
                    replay_buffer.sample_batch(int(args['minibatch_size']))

                # Calculate targets
                target_q = critic.predict_target(
                    s2_batch, actor.predict_target(s2_batch))
                y_i = []
                for k in range(int(args['minibatch_size'])):
                    if t_batch[k]:
                        y_i.append(r_batch[k])
                    else:
                        y_i.append(r_batch[k] + critic.gamma * target_q[k])

                # Update the critic given the targets
                predicted_q_value, _ = critic.train(
                    s_batch, a_batch, np.reshape(y_i, (int(args['minibatch_size']), 1)))

                ep_ave_max_q += np.amax(predicted_q_value)

                # Update the actor policy using the sampled gradient
                a_outs = actor.predict(s_batch)
                grads = critic.action_gradients(s_batch, a_outs)
                actor.train(s_batch, grads[0])

                # Update target networks
                actor.update_target_network()
                critic.update_target_network()

            s = s2
            ep_reward += r

            if terminal:
                summary_str = sess.run(summary_ops, feed_dict={
                    summary_vars[0]: ep_reward,
                    summary_vars[1]: ep_ave_max_q / float(j)
                })

                writer.add_summary(summary_str, i)
                writer.flush()

                print('| Reward: {:d} | Episode: {:d} | Qmax: {:.4f}'.format(int(ep_reward), \
                                                                             i, (ep_ave_max_q / float(j))))
                break

def main():
    with tf.Session() as sess:
        print("Starte BallTrackerThread")
        global delta

        thread1 = ballTracker.BallTrackerThread()
        thread1.start()

        print("Main running...")
        session = Pepper.init(ip, port)
        Pepper.roboInit(session)

        # Ensure action bound is symmetric
        # assert (env.action_space.high == -env.action_space.low)

        actor = ActorNetwork(sess, state_dim, action_dim, action_bound,
                             float(args['actor_lr']), float(args['tau']),
                             int(args['minibatch_size']))

        critic = CriticNetwork(sess, state_dim, action_dim,
                               float(args['critic_lr']), float(args['tau']),
                               float(args['gamma']),
                               actor.get_num_trainable_vars())

        actor_noise = OrnsteinUhlenbeckActionNoise(mu=np.zeros(action_dim))

        if args['mode'] == 'INIT':
            saver = tf.train.Saver()
            train(sess, args, actor, critic, actor_noise, False, saver)
        elif args['mode'] == 'TRAIN':
            saver = tf.train.Saver()
            saver.restore(sess, args['model'] + "/model")
            train(sess, args, actor, critic, actor_noise, True, saver)
        elif args['mode'] == 'TEST':
            saver = tf.train.Saver()
            saver.restore(sess, args['model'] + "/model")
            testDDPG(sess, args, actor, critic, actor_noise)
        else:
            print('No mode defined!')
        #train(sess, session, thread1, args, actor, critic, actor_noise)
        print('Terminated')

if __name__ == '__main__':
    main()


from base64 import b64encode
import dbm
import hashlib
import logging
import os
import sys
import tempfile

import numpy as np
import faiss
import tensorflow as tf

logger = logging.getLogger(__name__)

class ImageRetriever(object):
    """
    Image retrieval manager for PixelVAE embeddings.

    Example::

        encoder = tf.keras.models.load_model('../data/pixelvae/encoder')
        image_db = ImageRetriever(
            dataset_path='../data/crello-image/train*',
            cache_dir='../tmp/retriever/train_images',
        )
        image_db.build(encoder)

        query = np.zeros((1, image_db.dim), np.float32)
        image_url = image_db.search(query)
    """
    def __init__(self, dataset_path, cache_dir=None):
        self._dataset_path = dataset_path

        if cache_dir is None:
            self._temp_dir = tempfile.TemporaryDirectory()
            self._cache_dir = self._temp_dir.name
        else:
            self._cache_dir = cache_dir
            if not os.path.exists(self._cache_dir):
                logger.info(f'Creating cache dir at {self._cache_dir}')
                os.makedirs(self._cache_dir)

        self._image_cache_path = os.path.join(self._cache_dir, 'image')
        self._feature_cache_path = os.path.join(self._cache_dir, 'feature')
        self._key_cache_path = os.path.join(self._cache_dir, 'keys')

        self._image_keys = None
        self._image_cache = None
        self._feature_index = None

    def build(self, encoder, batch_size=16):
        """Compute and build feature index."""

        image_cache_file = self._image_cache_path
        if sys.platform == 'linux':
            image_cache_file += '.dat'

        if all(
                os.path.exists(p) for p in (
                    image_cache_file,
                    self._feature_cache_path,
                    self._key_cache_path,
                )):
            logger.info(f'Loading cached retriever from {self._cache_dir}')
            with open(self._key_cache_path, 'r') as f:
                self._image_keys = [x.strip() for x in f.readlines()]
            self._feature_index = faiss.read_index(self._feature_cache_path)

        else:
            logger.info(f'No cache found at {self._cache_dir}')

            logger.info('Computing image features...')
            data_path = os.path.join(self._dataset_path)
            dataset = tf.data.Dataset.list_files(data_path)
            dataset = tf.data.TFRecordDataset(dataset)
            dataset = dataset.batch(batch_size)
            dataset = dataset.map(
                lambda x: _compute_feature(encoder, x),
                num_parallel_calls=tf.data.experimental.AUTOTUNE,
            )
            dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)

            features = []
            self._image_keys = []
            with dbm.open(self._image_cache_path, 'n') as db:
                for batch in dataset.as_numpy_iterator():
                    features.append(batch[0])

                    for x in batch[1]:
                        image_key = hashlib.md5(x).hexdigest()
                        db[image_key] = x
                        self._image_keys.append(image_key)

                features = np.concatenate(features, axis=0)

            logger.info('Building feature index')
            self._feature_index = faiss.IndexFlatL2(features.shape[1])
            self._feature_index.add(features)

            logger.info('Caching data')
            with open(self._key_cache_path, 'w') as f:
                f.writelines(x + '\n' for x in self._image_keys)
            faiss.write_index(self._feature_index, self._feature_cache_path)

        logger.info(f'Indexd {len(self._image_keys)} features of '
                    f'dim={self._feature_index.d}')

        self._image_cache = dbm.open(self._image_cache_path, 'r')

    @property
    def dim(self):
        return self._feature_index.d if self._feature_index else None

    def len(self):
        return len(self._image_keys) if self._image_keys else 0

    def search(self, query, k=1):
        assert self._feature_index is not None, 'Index is not built yet'

        if isinstance(query, list):
            query = np.array([query], dtype=np.float32)

        if query.ndim == 1:
            query = np.expand_dims(query, axis=0)
        assert query.dtype == np.float32
        assert query.shape[1] == self._feature_index.d

        def _make_data_url(i):
            key = self._image_keys[i]
            image_bytes = self._image_cache[key]
            encoded = b64encode(image_bytes).decode('ascii')
            return f'data:image/png;base64,{encoded}'

        _, index = self._feature_index.search(query, k)
        urls = [_make_data_url(i) for i in index[0].tolist()]
        if k == 1:
            return urls[0]
        return urls

@tf.function
def _compute_feature(encoder, serialized):
    image_bytes = tf.io.parse_example(
        serialized, {
            'image': tf.io.FixedLenFeature((), dtype=tf.string),
        })['image']
    images = tf.map_fn(
        lambda x: tf.io.decode_png(x, channels=4),
        image_bytes,
        fn_output_signature=tf.uint8,
    )
    features = encoder(images)
    return features, image_bytes


# -*- coding:utf-8 -*-
# author: gaochen
# date: 2018.06.04

import numpy as np
import tensorflow as tf
import os
from tensorlayer.layers import *

class gan_model:
    def __init__(self, norm_epsilon, norm_decay, anchors_path, classes_path, pre_train):
        """
        Introduction
        ------------
            初始化函数
        Parameters
        ----------
            norm_decay: 在预测时计算moving average时的衰减率
            norm_epsilon: 方差加上极小的数，防止除以0的情况
            anchors_path: yolo anchor 文件路径
            classes_path: 数据集类别对应文件
            pre_train: 是否使用预训练darknet53模型
        """
        self.norm_epsilon = norm_epsilon
        self.norm_decay = norm_decay
        self.anchors_path = anchors_path
        self.classes_path = classes_path
        self.pre_train = pre_train
        self.anchors = self._get_anchors()
        self.classes = self._get_class()

    def GAN_g1(self, t_image, is_train=False, reuse=False):
        w_init = tf.random_normal_initializer(stddev=0.02)
        b_init = None  # tf.constant_initializer(value=0.0)
        g_init = tf.random_normal_initializer(1., 0.02)

        with tf.variable_scope("GAN_g1", reuse=reuse) as vs:
            n = InputLayer(t_image, name='in')
            n = Conv2d(n, 64, (3, 3), (1, 1), act=None, padding='SAME', W_init=w_init, name='n64s1/c')
            temp = n

            for i in range(8):
                nn = Conv2d(n, 64, (3, 3), (1, 1), act=None, padding='SAME', W_init=w_init, b_init=b_init,
                            name='n64s1/c1/%s' % i)
                nn = BatchNormLayer(nn, act=tf.nn.relu, is_train=is_train, gamma_init=g_init, name='n64s1/b1/%s' % i)
                nn = Conv2d(nn, 64, (3, 3), (1, 1), act=None, padding='SAME', W_init=w_init, b_init=b_init,
                            name='n64s1/c2/%s' % i)
                nn = BatchNormLayer(nn, is_train=is_train, gamma_init=g_init, name='n64s1/b2/%s' % i)
                nn = ElementwiseLayer([n, nn], tf.add, name='b_residual_add/%s' % i)
                n = nn

            n = Conv2d(n, 64, (3, 3), (1, 1), act=None, padding='SAME', W_init=w_init, b_init=b_init, name='n64s1/c/m')
            n = DeConv2d(n, 256, (3, 3), (2, 2), act=None, padding='SAME', W_init=w_init, b_init=b_init,
                         name='n256s2/dc/m')
            n = DeConv2d(n, 256, (3, 3), (2, 2), act=None, padding='SAME', W_init=w_init, b_init=b_init,
                         name='n256s3/dc/m')
            n = Conv2d(n, 3, (1, 1), (1, 1), act=None, padding='SAME', W_init=w_init, name='n3s1/c')

            return (n)

    def GAN_g2(self, t_image, is_train=False, reuse=False):
        w_init = tf.random_normal_initializer(stddev=0.02)
        b_init = None  # tf.constant_initializer(value=0.0)
        g_init = tf.random_normal_initializer(1., 0.02)

        with tf.variable_scope("GAN_g2", reuse=reuse) as vs:
            # n = InputLayer(t_image, name='in')
            n = Conv2d(t_image, 64, (3, 3), (1, 1), act=None, padding='SAME', W_init=w_init, name='n64s1/c')
            temp = n

            for i in range(8):
                nn = Conv2d(n, 64, (3, 3), (1, 1), act=None, padding='SAME', W_init=w_init, b_init=b_init,
                            name='n64s1/c1/%s' % i)
                nn = BatchNormLayer(nn, act=tf.nn.relu, is_train=is_train, gamma_init=g_init, name='n64s1/b1/%s' % i)
                nn = Conv2d(nn, 64, (3, 3), (1, 1), act=None, padding='SAME', W_init=w_init, b_init=b_init,
                            name='n64s1/c2/%s' % i)
                nn = BatchNormLayer(nn, is_train=is_train, gamma_init=g_init, name='n64s1/b2/%s' % i)
                nn = ElementwiseLayer([n, nn], tf.add, name='b_residual_add/%s' % i)
                n = nn

            n = Conv2d(n, 64, (3, 3), (1, 1), act=None, padding='SAME', W_init=w_init, b_init=b_init, name='n64s1/c/m')
            n = Conv2d(n, 256, (3, 3), (1, 1), act=None, padding='SAME', W_init=w_init, b_init=b_init,
                       name='n256s1/c/m')
            n = Conv2d(n, 256, (3, 3), (1, 1), act=None, padding='SAME', W_init=w_init, b_init=b_init,
                       name='n256s1/c/m')
            n = Conv2d(n, 3, (1, 1), (1, 1), act=None, padding='SAME', W_init=w_init, name='n3s1/c')
            return (n)

    def GAN_g(self, t_image, is_train=False, reuse=False):
        with tf.variable_scope("GAN_g", reuse=reuse) as vs:
            n = self.GAN_g1(t_image, is_train, reuse)
            n = self.GAN_g2(n, is_train, reuse)
            return n

    def _get_class(self):
        """
        Introduction
        ------------
            获取类别名字
        Returns
        -------
            class_names: coco数据集类别对应的名字
        """
        classes_path = os.path.expanduser(self.classes_path)
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

    def _get_anchors(self):
        """
        Introduction
        ------------
            获取anchors
        """
        anchors_path = os.path.expanduser(self.anchors_path)
        with open(anchors_path) as f:
            anchors = f.readline()
        anchors = [float(x) for x in anchors.split(',')]
        return np.array(anchors).reshape(-1, 2)

    def _batch_normalization_layer(self, input_layer, name = None, training = True, norm_decay = 0.99, norm_epsilon = 1e-3):
        '''
        Introduction
        ------------
            对卷积层提取的feature map使用batch normalization
        Parameters
        ----------
            input_layer: 输入的四维tensor
            name: batchnorm层的名字
            trainging: 是否为训练过程
            norm_decay: 在预测时计算moving average时的衰减率
            norm_epsilon: 方差加上极小的数，防止除以0的情况
        Returns
        -------
            bn_layer: batch normalization处理之后的feature map
        '''
        bn_layer = tf.layers.batch_normalization(inputs = input_layer,
            momentum = norm_decay, epsilon = norm_epsilon, center = True,
            scale = True, training = training, name = name)
        return tf.nn.leaky_relu(bn_layer, alpha = 0.1)

    def _conv2d_layer(self, inputs, filters_num, kernel_size, name, use_bias = False, strides = 1):
        """
        Introduction
        ------------
            使用tf.layers.conv2d减少权重和偏置矩阵初始化过程，以及卷积后加上偏置项的操作
            经过卷积之后需要进行batch norm，最后使用leaky ReLU激活函数
            根据卷积时的步长，如果卷积的步长为2，则对图像进行降采样
            比如，输入图片的大小为416*416，卷积核大小为3，若stride为2时，（416 - 3 + 2）/ 2 + 1， 计算结果为208，相当于做了池化层处理
            因此需要对stride大于1的时候，先进行一个padding操作, 采用四周都padding一维代替'same'方式
        Parameters
        ----------
            inputs: 输入变量
            filters_num: 卷积核数量
            strides: 卷积步长
            name: 卷积层名字
            trainging: 是否为训练过程
            use_bias: 是否使用偏置项
            kernel_size: 卷积核大小
        Returns
        -------
            conv: 卷积之后的feature map
        """
        conv = tf.layers.conv2d(
            inputs = inputs, filters = filters_num,
            kernel_size = kernel_size, strides = [strides, strides], kernel_initializer = tf.glorot_uniform_initializer(),
            padding = ('SAME' if strides == 1 else 'VALID'), kernel_regularizer = tf.contrib.layers.l2_regularizer(scale = 5e-4), use_bias = use_bias, name = name)
        return conv

    def _Residual_block(self, inputs, filters_num, blocks_num, conv_index, training = True, norm_decay = 0.99, norm_epsilon = 1e-3):
        """
        Introduction
        ------------
            Darknet的残差block，类似resnet的两层卷积结构，分别采用1x1和3x3的卷积核，使用1x1是为了减少channel的维度
        Parameters
        ----------
            inputs: 输入变量
            filters_num: 卷积核数量
            trainging: 是否为训练过程
            blocks_num: block的数量
            conv_index: 为了方便加载预训练权重，统一命名序号
            weights_dict: 加载预训练模型的权重
            norm_decay: 在预测时计算moving average时的衰减率
            norm_epsilon: 方差加上极小的数，防止除以0的情况
        Returns
        -------
            inputs: 经过残差网络处理后的结果
        """
        # 在输入feature map的长宽维度进行padding
        inputs = tf.pad(inputs, paddings=[[0, 0], [1, 0], [1, 0], [0, 0]], mode='CONSTANT')
        layer = self._conv2d_layer(inputs, filters_num, kernel_size = 3, strides = 2, name = "conv2d_" + str(conv_index))
        layer = self._batch_normalization_layer(layer, name = "batch_normalization_" + str(conv_index), training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
        conv_index += 1
        for _ in range(blocks_num):
            shortcut = layer
            layer = self._conv2d_layer(layer, filters_num // 2, kernel_size = 1, strides = 1, name = "conv2d_" + str(conv_index))
            layer = self._batch_normalization_layer(layer, name = "batch_normalization_" + str(conv_index), training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
            conv_index += 1
            layer = self._conv2d_layer(layer, filters_num, kernel_size = 3, strides = 1, name = "conv2d_" + str(conv_index))
            layer = self._batch_normalization_layer(layer, name = "batch_normalization_" + str(conv_index), training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
            conv_index += 1
            layer += shortcut
        return layer, conv_index

    def _darknet53(self, inputs, conv_index, training = True, norm_decay = 0.99, norm_epsilon = 1e-3):
        """
        Introduction
        ------------
            构建yolo3使用的darknet53网络结构
        Parameters
        ----------
            inputs: 模型输入变量
            conv_index: 卷积层数序号，方便根据名字加载预训练权重
            weights_dict: 预训练权重
            training: 是否为训练
            norm_decay: 在预测时计算moving average时的衰减率
            norm_epsilon: 方差加上极小的数，防止除以0的情况
        Returns
        -------
            conv: 经过52层卷积计算之后的结果, 输入图片为416x416x3，则此时输出的结果shape为13x13x1024
            route1: 返回第26层卷积计算结果52x52x256, 供后续使用
            route2: 返回第43层卷积计算结果26x26x512, 供后续使用
            conv_index: 卷积层计数，方便在加载预训练模型时使用
        """
        with tf.variable_scope('darknet53'):
            conv = self._conv2d_layer(inputs, filters_num = 32, kernel_size = 3, strides = 1, name = "conv2d_" + str(conv_index))
            conv = self._batch_normalization_layer(conv, name = "batch_normalization_" + str(conv_index), training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
            conv_index += 1
            conv, conv_index = self._Residual_block(conv, conv_index = conv_index, filters_num = 64, blocks_num = 1, training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
            conv, conv_index = self._Residual_block(conv, conv_index = conv_index, filters_num = 128, blocks_num = 2, training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
            conv, conv_index = self._Residual_block(conv, conv_index = conv_index, filters_num = 256, blocks_num = 8, training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
            route1 = conv
            conv, conv_index = self._Residual_block(conv, conv_index = conv_index, filters_num = 512, blocks_num = 8, training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
            route2 = conv
            conv, conv_index = self._Residual_block(conv, conv_index = conv_index,  filters_num = 1024, blocks_num = 4, training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
        return  route1, route2, conv, conv_index

    def _d_block(self, inputs, filters_num, out_filters, conv_index, training = True, norm_decay = 0.99, norm_epsilon = 1e-3):
        """
        Introduction
        ------------
            yolo3在Darknet53提取的特征层基础上，又加了针对3种不同比例的feature map的block，这样来提高对小物体的检测率
        Parameters
        ----------
            inputs: 输入特征
            filters_num: 卷积核数量
            out_filters: 最后输出层的卷积核数量
            conv_index: 卷积层数序号，方便根据名字加载预训练权重
            training: 是否为训练
            norm_decay: 在预测时计算moving average时的衰减率
            norm_epsilon: 方差加上极小的数，防止除以0的情况
        Returns
        -------
            route: 返回最后一层卷积的前一层结果
            conv: 返回最后一层卷积的结果
            conv_index: conv层计数
        """
        conv = self._conv2d_layer(inputs, filters_num = filters_num, kernel_size = 1, strides = 1, name = "conv2d_" + str(conv_index))
        conv = self._batch_normalization_layer(conv, name = "batch_normalization_" + str(conv_index), training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
        conv_index += 1
        conv = self._conv2d_layer(conv, filters_num = filters_num * 2, kernel_size = 3, strides = 1, name = "conv2d_" + str(conv_index))
        conv = self._batch_normalization_layer(conv, name = "batch_normalization_" + str(conv_index), training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
        conv_index += 1
        conv = self._conv2d_layer(conv, filters_num = filters_num, kernel_size = 1, strides = 1, name = "conv2d_" + str(conv_index))
        conv = self._batch_normalization_layer(conv, name = "batch_normalization_" + str(conv_index), training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
        conv_index += 1
        conv = self._conv2d_layer(conv, filters_num = filters_num * 2, kernel_size = 3, strides = 1, name = "conv2d_" + str(conv_index))
        conv = self._batch_normalization_layer(conv, name = "batch_normalization_" + str(conv_index), training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
        conv_index += 1
        conv = self._conv2d_layer(conv, filters_num = filters_num, kernel_size = 1, strides = 1, name = "conv2d_" + str(conv_index))
        conv = self._batch_normalization_layer(conv, name = "batch_normalization_" + str(conv_index), training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
        conv_index += 1
        route = conv
        conv = self._conv2d_layer(conv, filters_num = filters_num * 2, kernel_size = 3, strides = 1, name = "conv2d_" + str(conv_index))
        conv = self._batch_normalization_layer(conv, name = "batch_normalization_" + str(conv_index), training = training, norm_decay = norm_decay, norm_epsilon = norm_epsilon)
        conv_index += 1
        conv = self._conv2d_layer(conv, filters_num = out_filters, kernel_size = 1, strides = 1, name = "conv2d_" + str(conv_index), use_bias = True)
        conv_index += 1
        return route, conv, conv_index

    def d_inference(self, inputs, num_anchors, num_classes, training = True):
        """
        Introduction
        ------------
            构建yolo模型结构
        Parameters
        ----------
            inputs: 模型的输入变量
            num_anchors: 每个grid cell负责检测的anchor数量
            num_classes: 类别数量
            training: 是否为训练模式
        """
        conv_index = 1
        conv2d_26, conv2d_43, conv, conv_index = self._darknet53(inputs, conv_index, training = training, norm_decay = self.norm_decay, norm_epsilon = self.norm_epsilon)
        with tf.variable_scope('yolo'):
            out1 = self._conv2d_layer(conv, filters_num=2, kernel_size=1, strides=1, name="out1")  #HR/SR
            conv2d_57, conv2d_59, conv_index = self._d_block(conv, 512, num_anchors * (num_classes + 5), conv_index = conv_index, training = training, norm_decay = self.norm_decay, norm_epsilon = self.norm_epsilon)
            conv2d_60 = self._conv2d_layer(conv2d_57, filters_num = 256, kernel_size = 1, strides = 1, name = "conv2d_" + str(conv_index))
            conv2d_60 = self._batch_normalization_layer(conv2d_60, name = "batch_normalization_" + str(conv_index),training = training, norm_decay = self.norm_decay, norm_epsilon = self.norm_epsilon)
            conv_index += 1
            unSample_0 = tf.image.resize_nearest_neighbor(conv2d_60, [2 * tf.shape(conv2d_60)[1], 2 * tf.shape(conv2d_60)[1]], name='upSample_0')
            route0 = tf.concat([unSample_0, conv2d_43], axis = -1, name = 'route_0')
            conv2d_65, conv2d_67, conv_index = self._d_block(route0, 256, num_anchors * (num_classes + 5), conv_index = conv_index, training = training, norm_decay = self.norm_decay, norm_epsilon = self.norm_epsilon)
            conv2d_68 = self._conv2d_layer(conv2d_65, filters_num = 128, kernel_size = 1, strides = 1, name = "conv2d_" + str(conv_index))
            conv2d_68 = self._batch_normalization_layer(conv2d_68, name = "batch_normalization_" + str(conv_index), training=training, norm_decay=self.norm_decay, norm_epsilon = self.norm_epsilon)
            conv_index += 1
            unSample_1 = tf.image.resize_nearest_neighbor(conv2d_68, [2 * tf.shape(conv2d_68)[1], 2 * tf.shape(conv2d_68)[1]], name='upSample_1')
            route1 = tf.concat([unSample_1, conv2d_26], axis = -1, name = 'route_1')
            _, conv2d_75, _ = self._d_block(route1, 128, num_anchors * (num_classes + 5), conv_index = conv_index, training = training, norm_decay = self.norm_decay, norm_epsilon = self.norm_epsilon)

        return [conv2d_59, conv2d_67, conv2d_75, out1]

    def d_head(self, feats, anchors, num_classes, input_shape, training = True):
        """
        Introduction
        ------------
            根据不同大小的feature map做多尺度的检测，三种feature map大小分别为13x13x1024, 26x26x512, 52x52x256
        Parameters
        ----------
            feats: 输入的特征feature map
            anchors: 针对不同大小的feature map的anchor
            num_classes: 类别的数量
            input_shape: 图像的输入大小，一般为416
            trainging: 是否训练，用来控制返回不同的值
        Returns
        -------
        """
        num_anchors = len(anchors)
        anchors_tensor = tf.reshape(tf.constant(anchors, dtype = tf.float32), [1, 1, 1, num_anchors, 2])
        grid_size = tf.shape(feats)[1:3]
        predictions = tf.reshape(feats, [-1, grid_size[0], grid_size[1], num_anchors, num_classes + 5])
        # 这里构建13*13*1*2的矩阵，对应每个格子加上对应的坐标
        grid_y = tf.tile(tf.reshape(tf.range(grid_size[0]), [-1, 1, 1, 1]), [1, grid_size[1], 1, 1])
        grid_x = tf.tile(tf.reshape(tf.range(grid_size[1]), [1, -1, 1, 1]), [grid_size[0], 1, 1, 1])
        grid = tf.concat([grid_x, grid_y], axis = -1)
        grid = tf.cast(grid, tf.float32)
        #将x,y坐标归一化为占416的比例
        box_xy = (tf.sigmoid(predictions[..., :2]) + grid) / tf.cast(grid_size[::-1], tf.float32)
        #将w,h也归一化为占416的比例
        box_wh = tf.exp(predictions[..., 2:4]) * anchors_tensor / input_shape[::-1]
        box_confidence = tf.sigmoid(predictions[..., 4:5])
        box_class_probs = tf.sigmoid(predictions[..., 5:])
        if training == True:
            return grid, predictions, box_xy, box_wh
        return box_xy, box_wh, box_confidence, box_class_probs

    def d_boxes_scores(self, feats, anchors, num_classes, input_shape, image_shape):
        """
        Introduction
        ------------
            该函数是将box的坐标修正，除去之前按照长宽比缩放填充的部分，最后将box的坐标还原成相对原始图片的
        Parameters
        ----------
            feats: 模型输出feature map
            anchors: 模型anchors
            num_classes: 数据集类别数
            input_shape: 训练输入图片大小
            image_shape: 原始图片的大小
        """
        input_shape = tf.cast(input_shape, tf.float32)
        image_shape = tf.cast(image_shape, tf.float32)
        box_xy, box_wh, box_confidence, box_class_probs = self.d_head(feats, anchors, num_classes, input_shape, training = False)
        box_yx = box_xy[..., ::-1]
        box_hw = box_wh[..., ::-1]
        new_shape = tf.round(image_shape * tf.reduce_min(input_shape / image_shape))
        offset = (input_shape - new_shape) / 2. / input_shape
        scale = input_shape / new_shape
        box_yx = (box_yx - offset) * scale
        box_hw = box_hw * scale

        box_min = box_yx - box_hw / 2.
        box_max = box_yx + box_hw / 2.
        boxes = tf.concat(
            [box_min[..., 0:1],
             box_min[..., 1:2],
             box_max[..., 0:1],
             box_max[..., 1:2]],
            axis = -1
        )
        boxes *= tf.concat([image_shape, image_shape], axis = -1)
        boxes = tf.reshape(boxes, [-1, 4])
        boxes_scores = box_confidence * box_class_probs
        boxes_scores = tf.reshape(boxes_scores, [-1, num_classes])
        return boxes, boxes_scores

    def box_iou(self, box1, box2):
        """
        Introduction
        ------------
            计算box tensor之间的iou
        Parameters
        ----------
            box1: shape=[grid_size, grid_size, anchors, xywh]
            box2: shape=[box_num, xywh]
        Returns
        -------
            iou:
        """
        box1 = tf.expand_dims(box1, -2)
        box1_xy = box1[..., :2]
        box1_wh = box1[..., 2:4]
        box1_mins = box1_xy - box1_wh / 2.
        box1_maxs = box1_xy + box1_wh / 2.

        box2 = tf.expand_dims(box2, 0)
        box2_xy = box2[..., :2]
        box2_wh = box2[..., 2:4]
        box2_mins = box2_xy - box2_wh / 2.
        box2_maxs = box2_xy + box2_wh / 2.

        intersect_mins = tf.maximum(box1_mins, box2_mins)
        intersect_maxs = tf.minimum(box1_maxs, box2_maxs)
        intersect_wh = tf.maximum(intersect_maxs - intersect_mins, 0.)
        intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
        box1_area = box1_wh[..., 0] * box1_wh[..., 1]
        box2_area = box2_wh[..., 0] * box2_wh[..., 1]
        iou = intersect_area / (box1_area + box2_area - intersect_area)
        return iou

    def yolo_loss(self, yolo_output, y_true, anchors, num_classes, tf_mask, ignore_thresh = .5):
        """
        Introduction
        ------------
            yolo模型的损失函数
        Parameters
        ----------
            yolo_output: yolo模型的输出
            y_true: 经过预处理的真实标签，shape为[batch, grid_size, grid_size, 5 + num_classes]
            anchors: yolo模型对应的anchors
            num_classes: 类别数量
            ignore_thresh: 小于该阈值的box我们认为没有物体
        Returns
        -------
            loss: 每个batch的平均损失值
            accuracy
        """
        loss = 0
        anchor_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
        input_shape = [416.0, 416.0]
        grid_shapes = [tf.cast(tf.shape(yolo_output[l])[1:3], tf.float32) for l in range(3)]
        for index in range(1):   #33333333333333~~~~
            # 只有负责预测ground truth box的grid对应的为1, 才计算相对应的loss
            # object_mask的shape为[batch_size, grid_size, grid_size, 3, 1]
            object_mask = y_true[index][..., 4:5]
            class_probs = y_true[index][..., 5:]
            grid, predictions, pred_xy, pred_wh = self.d_head(yolo_output[index], anchors[anchor_mask[index]], num_classes, input_shape, training = True)
            # pred_box的shape为[batch, box_num, 4]
            pred_box = tf.concat([pred_xy, pred_wh], axis = -1)
            raw_true_xy = y_true[index][..., :2] * grid_shapes[index][::-1] - grid
            object_mask_bool = tf.cast(object_mask, dtype = tf.bool)
            raw_true_wh = tf.log(tf.where(tf.equal(y_true[index][..., 2:4] / anchors[anchor_mask[index]] * input_shape[::-1], 0), tf.ones_like(y_true[index][..., 2:4]), y_true[index][..., 2:4] / anchors[anchor_mask[index]] * input_shape[::-1]))
            # 该系数是用来调整box坐标loss的系数
            box_loss_scale = 2 - y_true[index][..., 2:3] * y_true[index][..., 3:4]
            ignore_mask = tf.TensorArray(dtype = tf.float32, size=1, dynamic_size = True)
            def loop_body(internal_index, ignore_mask):
                # true_box的shape为[box_num, 4]
                true_box = tf.boolean_mask(y_true[index][internal_index, ..., 0:4], object_mask_bool[internal_index, ..., 0])
                iou = self.box_iou(pred_box[internal_index], true_box)
                # 计算每个true_box对应的预测的iou最大的box
                best_iou = tf.reduce_max(iou, axis = -1)
                ignore_mask = ignore_mask.write(internal_index, tf.cast(best_iou < ignore_thresh, tf.float32))
                return internal_index + 1, ignore_mask
            _, ignore_mask = tf.while_loop(lambda internal_index, ignore_mask : internal_index < tf.shape(yolo_output[0])[0], loop_body, [0, ignore_mask])
            ignore_mask = ignore_mask.stack()
            ignore_mask = tf.expand_dims(ignore_mask, axis = -1)
            # 计算四个部分的loss
            # xy_loss = object_mask * box_loss_scale * tf.nn.sigmoid_cross_entropy_with_logits(labels = raw_true_xy, logits = predictions[..., 0:2])
            # wh_loss = object_mask * box_loss_scale * 0.5 * tf.square(raw_true_wh - predictions[..., 2:4])
            # confidence_loss = object_mask * tf.nn.sigmoid_cross_entropy_with_logits(labels = object_mask, logits = predictions[..., 4:5]) + (1 - object_mask) * tf.nn.sigmoid_cross_entropy_with_logits(labels = object_mask, logits = predictions[..., 4:5]) * ignore_mask
            # class_loss = object_mask * tf.nn.sigmoid_cross_entropy_with_logits(labels =  class_probs, logits = predictions[..., 5:])

            xy_loss = object_mask * box_loss_scale * tf.square(raw_true_xy - predictions[..., 0:2])
            wh_loss = object_mask * box_loss_scale * 0.5 * tf.square(raw_true_wh - predictions[..., 2:4])
            confidence_loss = object_mask * tf.square(object_mask - predictions[..., 4:5]) + (1 - object_mask) * tf.square(object_mask - predictions[..., 4:5]) * ignore_mask
            class_loss = object_mask * tf.square(class_probs - predictions[..., 5:])

            xy_loss = tf.reduce_sum(xy_loss) / tf.cast(tf.shape(yolo_output[0])[0], tf.float32)
            wh_loss = tf.reduce_sum(wh_loss) / tf.cast(tf.shape(yolo_output[0])[0], tf.float32)
            confidence_loss = tf.reduce_sum(confidence_loss) / tf.cast(tf.shape(yolo_output[0])[0], tf.float32)
            class_loss = tf.reduce_sum(class_loss) / tf.cast(tf.shape(yolo_output[0])[0], tf.float32)

            loss += xy_loss + wh_loss + confidence_loss + class_loss
        if tf_mask == 1:
            d_loss1 = tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(yolo_output[3]), logits=yolo_output[3])
        if tf_mask == 0:
            d_loss1 = tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.zeros_like(yolo_output[3]), logits=yolo_output[3])

        d_loss = d_loss1 + loss

        return d_loss