import numpy as np
import torch
import torch.nn as nn
from mmcv.cnn import ConvModule, bias_init_with_prob, normal_init
from mmcv.ops import DeformConv2d

from mmdet.core import (PointGenerator, bbox_overlaps, build_assigner,
                        build_sampler, images_to_levels, multi_apply,
                        multiclass_nms, unmap)
from ..builder import HEADS, build_loss
from .anchor_free_head import AnchorFreeHead

@HEADS.register_module()
class RepPointsHead(AnchorFreeHead):
    """RepPoint head.

    Args:
        point_feat_channels (int): Number of channels of points features.
        gradient_mul (float): The multiplier to gradients from
            points refinement and recognition.
        point_strides (Iterable): points strides.
        point_base_scale (int): bbox scale for assigning labels.
        loss_cls (dict): Config of classification loss.
        loss_bbox_init (dict): Config of initial points loss.
        loss_bbox_refine (dict): Config of points loss in refinement.
        use_grid_points (bool): If we use bounding box representation, the
        reppoints is represented as grid points on the bounding box.
        center_init (bool): Whether to use center point assignment.
        transform_method (str): The methods to transform RepPoints to bbox.
    """  # noqa: W605

    def __init__(self,
                 num_classes,
                 in_channels,
                 point_feat_channels=256,
                 num_points=9,
                 gradient_mul=0.1,
                 point_strides=[8, 16, 32, 64, 128],
                 point_base_scale=4,
                 loss_cls=dict(
                     type='FocalLoss',
                     use_sigmoid=True,
                     gamma=2.0,
                     alpha=0.25,
                     loss_weight=1.0),
                 use_vfl=False,
                 loss_cls_vfl=dict(
                     type='VarifocalLoss',
                     use_sigmoid=True,
                     alpha=0.75,
                     gamma=2.0,
                     iou_weighted=True,
                     loss_weight=1.0),
                 loss_bbox_init=dict(
                     type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=0.5),
                 loss_bbox_refine=dict(
                     type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=1.0),
                 use_grid_points=False,
                 center_init=True,
                 transform_method='moment',
                 moment_mul=0.01,
                 **kwargs):
        self.num_points = num_points
        self.point_feat_channels = point_feat_channels
        self.use_grid_points = use_grid_points
        self.center_init = center_init

        # we use deform conv to extract points features
        self.dcn_kernel = int(np.sqrt(num_points))
        self.dcn_pad = int((self.dcn_kernel - 1) / 2)
        assert self.dcn_kernel * self.dcn_kernel == num_points, \
            'The points number should be a square number.'
        assert self.dcn_kernel % 2 == 1, \
            'The points number should be an odd square number.'
        dcn_base = np.arange(-self.dcn_pad,
                             self.dcn_pad + 1).astype(np.float64)
        dcn_base_y = np.repeat(dcn_base, self.dcn_kernel)
        dcn_base_x = np.tile(dcn_base, self.dcn_kernel)
        dcn_base_offset = np.stack([dcn_base_y, dcn_base_x], axis=1).reshape(
            (-1))
        self.dcn_base_offset = torch.tensor(dcn_base_offset).view(1, -1, 1, 1)

        super().__init__(num_classes, in_channels, loss_cls=loss_cls, **kwargs)

        self.gradient_mul = gradient_mul
        self.point_base_scale = point_base_scale
        self.point_strides = point_strides
        self.point_generators = [PointGenerator() for _ in self.point_strides]

        self.sampling = loss_cls['type'] not in ['FocalLoss']
        if self.train_cfg:
            self.init_assigner = build_assigner(self.train_cfg.init.assigner)
            self.refine_assigner = build_assigner(
                self.train_cfg.refine.assigner)
            # use PseudoSampler when sampling is False
            if self.sampling and hasattr(self.train_cfg, 'sampler'):
                sampler_cfg = self.train_cfg.sampler
            else:
                sampler_cfg = dict(type='PseudoSampler')
            self.sampler = build_sampler(sampler_cfg, context=self)
        self.transform_method = transform_method
        if self.transform_method == 'moment':
            self.moment_transfer = nn.Parameter(
                data=torch.zeros(2), requires_grad=True)
            self.moment_mul = moment_mul

        self.use_sigmoid_cls = loss_cls.get('use_sigmoid', False)
        if self.use_sigmoid_cls:
            self.cls_out_channels = self.num_classes
        else:
            self.cls_out_channels = self.num_classes + 1
        self.use_vfl = use_vfl
        if self.use_vfl:
            self.loss_cls = build_loss(loss_cls_vfl)
        else:
            self.loss_cls = build_loss(loss_cls)
        self.loss_bbox_init = build_loss(loss_bbox_init)
        self.loss_bbox_refine = build_loss(loss_bbox_refine)

    def _init_layers(self):
        """Initialize layers of the head."""
        self.relu = nn.ReLU(inplace=True)
        self.cls_convs = nn.ModuleList()
        self.reg_convs = nn.ModuleList()
        for i in range(self.stacked_convs):
            chn = self.in_channels if i == 0 else self.feat_channels
            self.cls_convs.append(
                ConvModule(
                    chn,
                    self.feat_channels,
                    3,
                    stride=1,
                    padding=1,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg))
            self.reg_convs.append(
                ConvModule(
                    chn,
                    self.feat_channels,
                    3,
                    stride=1,
                    padding=1,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg))
        pts_out_dim = 4 if self.use_grid_points else 2 * self.num_points
        self.reppoints_cls_conv = DeformConv2d(self.feat_channels,
                                               self.point_feat_channels,
                                               self.dcn_kernel, 1,
                                               self.dcn_pad)
        self.reppoints_cls_out = nn.Conv2d(self.point_feat_channels,
                                           self.cls_out_channels, 1, 1, 0)
        self.reppoints_pts_init_conv = nn.Conv2d(self.feat_channels,
                                                 self.point_feat_channels, 3,
                                                 1, 1)
        self.reppoints_pts_init_out = nn.Conv2d(self.point_feat_channels,
                                                pts_out_dim, 1, 1, 0)
        self.reppoints_pts_refine_conv = DeformConv2d(self.feat_channels,
                                                      self.point_feat_channels,
                                                      self.dcn_kernel, 1,
                                                      self.dcn_pad)
        self.reppoints_pts_refine_out = nn.Conv2d(self.point_feat_channels,
                                                  pts_out_dim, 1, 1, 0)

    def init_weights(self):
        """Initialize weights of the head."""
        for m in self.cls_convs:
            normal_init(m.conv, std=0.01)
        for m in self.reg_convs:
            normal_init(m.conv, std=0.01)
        bias_cls = bias_init_with_prob(0.01)
        normal_init(self.reppoints_cls_conv, std=0.01)
        normal_init(self.reppoints_cls_out, std=0.01, bias=bias_cls)
        normal_init(self.reppoints_pts_init_conv, std=0.01)
        normal_init(self.reppoints_pts_init_out, std=0.01)
        normal_init(self.reppoints_pts_refine_conv, std=0.01)
        normal_init(self.reppoints_pts_refine_out, std=0.01)

    def points2bbox(self, pts, y_first=True):
        """Converting the points set into bounding box.

        :param pts: the input points sets (fields), each points
            set (fields) is represented as 2n scalar.
        :param y_first: if y_fisrt=True, the point set is represented as
            [y1, x1, y2, x2 ... yn, xn], otherwise the point set is
            represented as [x1, y1, x2, y2 ... xn, yn].
        :return: each points set is converting to a bbox [x1, y1, x2, y2].
        """
        pts_reshape = pts.view(pts.shape[0], -1, 2, *pts.shape[2:])
        pts_y = pts_reshape[:, :, 0, ...] if y_first else pts_reshape[:, :, 1,
                                                                      ...]
        pts_x = pts_reshape[:, :, 1, ...] if y_first else pts_reshape[:, :, 0,
                                                                      ...]
        if self.transform_method == 'minmax':
            bbox_left = pts_x.min(dim=1, keepdim=True)[0]
            bbox_right = pts_x.max(dim=1, keepdim=True)[0]
            bbox_up = pts_y.min(dim=1, keepdim=True)[0]
            bbox_bottom = pts_y.max(dim=1, keepdim=True)[0]
            bbox = torch.cat([bbox_left, bbox_up, bbox_right, bbox_bottom],
                             dim=1)
        elif self.transform_method == 'partial_minmax':
            pts_y = pts_y[:, :4, ...]
            pts_x = pts_x[:, :4, ...]
            bbox_left = pts_x.min(dim=1, keepdim=True)[0]
            bbox_right = pts_x.max(dim=1, keepdim=True)[0]
            bbox_up = pts_y.min(dim=1, keepdim=True)[0]
            bbox_bottom = pts_y.max(dim=1, keepdim=True)[0]
            bbox = torch.cat([bbox_left, bbox_up, bbox_right, bbox_bottom],
                             dim=1)
        elif self.transform_method == 'moment':
            pts_y_mean = pts_y.mean(dim=1, keepdim=True)
            pts_x_mean = pts_x.mean(dim=1, keepdim=True)
            pts_y_std = torch.std(pts_y - pts_y_mean, dim=1, keepdim=True)
            pts_x_std = torch.std(pts_x - pts_x_mean, dim=1, keepdim=True)
            moment_transfer = (self.moment_transfer * self.moment_mul) + (
                self.moment_transfer.detach() * (1 - self.moment_mul))
            moment_width_transfer = moment_transfer[0]
            moment_height_transfer = moment_transfer[1]
            half_width = pts_x_std * torch.exp(moment_width_transfer)
            half_height = pts_y_std * torch.exp(moment_height_transfer)
            bbox = torch.cat([
                pts_x_mean - half_width, pts_y_mean - half_height,
                pts_x_mean + half_width, pts_y_mean + half_height
            ],
                             dim=1)
        else:
            raise NotImplementedError
        return bbox

    def gen_grid_from_reg(self, reg, previous_boxes):
        """Base on the previous bboxes and regression values, we compute the
        regressed bboxes and generate the grids on the bboxes.

        :param reg: the regression value to previous bboxes.
        :param previous_boxes: previous bboxes.
        :return: generate grids on the regressed bboxes.
        """
        b, _, h, w = reg.shape
        bxy = (previous_boxes[:, :2, ...] + previous_boxes[:, 2:, ...]) / 2.
        bwh = (previous_boxes[:, 2:, ...] -
               previous_boxes[:, :2, ...]).clamp(min=1e-6)
        grid_topleft = bxy + bwh * reg[:, :2, ...] - 0.5 * bwh * torch.exp(
            reg[:, 2:, ...])
        grid_wh = bwh * torch.exp(reg[:, 2:, ...])
        grid_left = grid_topleft[:, [0], ...]
        grid_top = grid_topleft[:, [1], ...]
        grid_width = grid_wh[:, [0], ...]
        grid_height = grid_wh[:, [1], ...]
        intervel = torch.linspace(0., 1., self.dcn_kernel).view(
            1, self.dcn_kernel, 1, 1).type_as(reg)
        grid_x = grid_left + grid_width * intervel
        grid_x = grid_x.unsqueeze(1).repeat(1, self.dcn_kernel, 1, 1, 1)
        grid_x = grid_x.view(b, -1, h, w)
        grid_y = grid_top + grid_height * intervel
        grid_y = grid_y.unsqueeze(2).repeat(1, 1, self.dcn_kernel, 1, 1)
        grid_y = grid_y.view(b, -1, h, w)
        grid_yx = torch.stack([grid_y, grid_x], dim=2)
        grid_yx = grid_yx.view(b, -1, h, w)
        regressed_bbox = torch.cat([
            grid_left, grid_top, grid_left + grid_width, grid_top + grid_height
        ], 1)
        return grid_yx, regressed_bbox

    def forward(self, feats):
        return multi_apply(self.forward_single, feats)

    def forward_single(self, x):
        """Forward feature map of a single FPN level."""
        dcn_base_offset = self.dcn_base_offset.type_as(x)
        # If we use center_init, the initial reppoints is from center points.
        # If we use bounding bbox representation, the initial reppoints is
        #   from regular grid placed on a pre-defined bbox.
        if self.use_grid_points or not self.center_init:
            scale = self.point_base_scale / 2
            points_init = dcn_base_offset / dcn_base_offset.max() * scale
            bbox_init = x.new_tensor([-scale, -scale, scale,
                                      scale]).view(1, 4, 1, 1)
        else:
            points_init = 0
        cls_feat = x
        pts_feat = x
        for cls_conv in self.cls_convs:
            cls_feat = cls_conv(cls_feat)
        for reg_conv in self.reg_convs:
            pts_feat = reg_conv(pts_feat)
        # initialize reppoints
        pts_out_init = self.reppoints_pts_init_out(
            self.relu(self.reppoints_pts_init_conv(pts_feat)))
        if self.use_grid_points:
            pts_out_init, bbox_out_init = self.gen_grid_from_reg(
                pts_out_init, bbox_init.detach())
        else:
            pts_out_init = pts_out_init + points_init
        # refine and classify reppoints
        pts_out_init_grad_mul = (1 - self.gradient_mul) * pts_out_init.detach(
        ) + self.gradient_mul * pts_out_init
        dcn_offset = pts_out_init_grad_mul - dcn_base_offset
        cls_out = self.reppoints_cls_out(
            self.relu(self.reppoints_cls_conv(cls_feat, dcn_offset)))
        pts_out_refine = self.reppoints_pts_refine_out(
            self.relu(self.reppoints_pts_refine_conv(pts_feat, dcn_offset)))
        if self.use_grid_points:
            pts_out_refine, bbox_out_refine = self.gen_grid_from_reg(
                pts_out_refine, bbox_out_init.detach())
        else:
            pts_out_refine = pts_out_refine + pts_out_init.detach()
        return cls_out, pts_out_init, pts_out_refine

    def get_points(self, featmap_sizes, img_metas, device):
        """Get points according to feature map sizes.

        Args:
            featmap_sizes (list[tuple]): Multi-level feature map sizes.
            img_metas (list[dict]): Image meta info.

        Returns:
            tuple: points of each image, valid flags of each image
        """
        num_imgs = len(img_metas)
        num_levels = len(featmap_sizes)

        # since feature map sizes of all images are the same, we only compute
        # points center for one time
        multi_level_points = []
        for i in range(num_levels):
            points = self.point_generators[i].grid_points(
                featmap_sizes[i], self.point_strides[i], device)
            multi_level_points.append(points)
        points_list = [[point.clone() for point in multi_level_points]
                       for _ in range(num_imgs)]

        # for each image, we compute valid flags of multi level grids
        valid_flag_list = []
        for img_id, img_meta in enumerate(img_metas):
            multi_level_flags = []
            for i in range(num_levels):
                point_stride = self.point_strides[i]
                feat_h, feat_w = featmap_sizes[i]
                h, w = img_meta['pad_shape'][:2]
                valid_feat_h = min(int(np.ceil(h / point_stride)), feat_h)
                valid_feat_w = min(int(np.ceil(w / point_stride)), feat_w)
                flags = self.point_generators[i].valid_flags(
                    (feat_h, feat_w), (valid_feat_h, valid_feat_w), device)
                multi_level_flags.append(flags)
            valid_flag_list.append(multi_level_flags)

        return points_list, valid_flag_list

    def centers_to_bboxes(self, point_list):
        """Get bboxes according to center points.

        Only used in :class:`MaxIoUAssigner`.
        """
        bbox_list = []
        for i_img, point in enumerate(point_list):
            bbox = []
            for i_lvl in range(len(self.point_strides)):
                scale = self.point_base_scale * self.point_strides[i_lvl] * 0.5
                bbox_shift = torch.Tensor([-scale, -scale, scale,
                                           scale]).view(1, 4).type_as(point[0])
                bbox_center = torch.cat(
                    [point[i_lvl][:, :2], point[i_lvl][:, :2]], dim=1)
                bbox.append(bbox_center + bbox_shift)
            bbox_list.append(bbox)
        return bbox_list

    def offset_to_pts(self, center_list, pred_list):
        """Change from point offset to point coordinate."""
        pts_list = []
        for i_lvl in range(len(self.point_strides)):
            pts_lvl = []
            for i_img in range(len(center_list)):
                pts_center = center_list[i_img][i_lvl][:, :2].repeat(
                    1, self.num_points)
                pts_shift = pred_list[i_lvl][i_img]
                yx_pts_shift = pts_shift.permute(1, 2, 0).view(
                    -1, 2 * self.num_points)
                y_pts_shift = yx_pts_shift[..., 0::2]
                x_pts_shift = yx_pts_shift[..., 1::2]
                xy_pts_shift = torch.stack([x_pts_shift, y_pts_shift], -1)
                xy_pts_shift = xy_pts_shift.view(*yx_pts_shift.shape[:-1], -1)
                pts = xy_pts_shift * self.point_strides[i_lvl] + pts_center
                pts_lvl.append(pts)
            pts_lvl = torch.stack(pts_lvl, 0)
            pts_list.append(pts_lvl)
        return pts_list

    def _point_target_single(self,
                             flat_proposals,
                             valid_flags,
                             gt_bboxes,
                             gt_bboxes_ignore,
                             gt_labels,
                             label_channels=1,
                             stage='init',
                             unmap_outputs=True):
        inside_flags = valid_flags
        if not inside_flags.any():
            return (None, ) * 7
        # assign gt and sample proposals
        proposals = flat_proposals[inside_flags, :]

        if stage == 'init':
            assigner = self.init_assigner
            pos_weight = self.train_cfg.init.pos_weight
        else:
            assigner = self.refine_assigner
            pos_weight = self.train_cfg.refine.pos_weight
        assign_result = assigner.assign(proposals, gt_bboxes, gt_bboxes_ignore,
                                        None if self.sampling else gt_labels)
        sampling_result = self.sampler.sample(assign_result, proposals,
                                              gt_bboxes)

        num_valid_proposals = proposals.shape[0]
        bbox_gt = proposals.new_zeros([num_valid_proposals, 4])
        pos_proposals = torch.zeros_like(proposals)
        proposals_weights = proposals.new_zeros([num_valid_proposals, 4])
        labels = proposals.new_full((num_valid_proposals, ),
                                    self.num_classes,
                                    dtype=torch.long)
        label_weights = proposals.new_zeros(
            num_valid_proposals, dtype=torch.float)

        pos_inds = sampling_result.pos_inds
        neg_inds = sampling_result.neg_inds
        if len(pos_inds) > 0:
            pos_gt_bboxes = sampling_result.pos_gt_bboxes
            bbox_gt[pos_inds, :] = pos_gt_bboxes
            pos_proposals[pos_inds, :] = proposals[pos_inds, :]
            proposals_weights[pos_inds, :] = 1.0
            if gt_labels is None:
                # Only rpn gives gt_labels as None
                # Foreground is the first class
                labels[pos_inds] = 0
            else:
                labels[pos_inds] = gt_labels[
                    sampling_result.pos_assigned_gt_inds]
            if pos_weight <= 0:
                label_weights[pos_inds] = 1.0
            else:
                label_weights[pos_inds] = pos_weight
        if len(neg_inds) > 0:
            label_weights[neg_inds] = 1.0

        # map up to original set of proposals
        if unmap_outputs:
            num_total_proposals = flat_proposals.size(0)
            labels = unmap(labels, num_total_proposals, inside_flags)
            label_weights = unmap(label_weights, num_total_proposals,
                                  inside_flags)
            bbox_gt = unmap(bbox_gt, num_total_proposals, inside_flags)
            pos_proposals = unmap(pos_proposals, num_total_proposals,
                                  inside_flags)
            proposals_weights = unmap(proposals_weights, num_total_proposals,
                                      inside_flags)

        return (labels, label_weights, bbox_gt, pos_proposals,
                proposals_weights, pos_inds, neg_inds)

    def get_targets(self,
                    proposals_list,
                    valid_flag_list,
                    gt_bboxes_list,
                    img_metas,
                    gt_bboxes_ignore_list=None,
                    gt_labels_list=None,
                    stage='init',
                    label_channels=1,
                    unmap_outputs=True):
        """Compute corresponding GT box and classification targets for
        proposals.

        Args:
            proposals_list (list[list]): Multi level points/bboxes of each
                image.
            valid_flag_list (list[list]): Multi level valid flags of each
                image.
            gt_bboxes_list (list[Tensor]): Ground truth bboxes of each image.
            img_metas (list[dict]): Meta info of each image.
            gt_bboxes_ignore_list (list[Tensor]): Ground truth bboxes to be
                ignored.
            gt_bboxes_list (list[Tensor]): Ground truth labels of each box.
            stage (str): `init` or `refine`. Generate target for init stage or
                refine stage
            label_channels (int): Channel of label.
            unmap_outputs (bool): Whether to map outputs back to the original
                set of anchors.

        Returns:
            tuple:
                - labels_list (list[Tensor]): Labels of each level.
                - label_weights_list (list[Tensor]): Label weights of each level.  # noqa: E501
                - bbox_gt_list (list[Tensor]): Ground truth bbox of each level.
                - proposal_list (list[Tensor]): Proposals(points/bboxes) of each level.  # noqa: E501
                - proposal_weights_list (list[Tensor]): Proposal weights of each level.  # noqa: E501
                - num_total_pos (int): Number of positive samples in all images.  # noqa: E501
                - num_total_neg (int): Number of negative samples in all images.  # noqa: E501
        """
        assert stage in ['init', 'refine']
        num_imgs = len(img_metas)
        assert len(proposals_list) == len(valid_flag_list) == num_imgs

        # points number of multi levels
        num_level_proposals = [points.size(0) for points in proposals_list[0]]

        # concat all level points and flags to a single tensor
        for i in range(num_imgs):
            assert len(proposals_list[i]) == len(valid_flag_list[i])
            proposals_list[i] = torch.cat(proposals_list[i])
            valid_flag_list[i] = torch.cat(valid_flag_list[i])

        # compute targets for each image
        if gt_bboxes_ignore_list is None:
            gt_bboxes_ignore_list = [None for _ in range(num_imgs)]
        if gt_labels_list is None:
            gt_labels_list = [None for _ in range(num_imgs)]
        (all_labels, all_label_weights, all_bbox_gt, all_proposals,
         all_proposal_weights, pos_inds_list, neg_inds_list) = multi_apply(
             self._point_target_single,
             proposals_list,
             valid_flag_list,
             gt_bboxes_list,
             gt_bboxes_ignore_list,
             gt_labels_list,
             stage=stage,
             label_channels=label_channels,
             unmap_outputs=unmap_outputs)
        # no valid points
        if any([labels is None for labels in all_labels]):
            return None
        # sampled points of all images
        num_total_pos = sum([max(inds.numel(), 1) for inds in pos_inds_list])
        num_total_neg = sum([max(inds.numel(), 1) for inds in neg_inds_list])
        labels_list = images_to_levels(all_labels, num_level_proposals)
        label_weights_list = images_to_levels(all_label_weights,
                                              num_level_proposals)
        bbox_gt_list = images_to_levels(all_bbox_gt, num_level_proposals)
        proposals_list = images_to_levels(all_proposals, num_level_proposals)
        proposal_weights_list = images_to_levels(all_proposal_weights,
                                                 num_level_proposals)
        return (labels_list, label_weights_list, bbox_gt_list, proposals_list,
                proposal_weights_list, num_total_pos, num_total_neg)

    def loss_single(self, cls_score, pts_pred_init, pts_pred_refine, labels,
                    label_weights, bbox_gt_init, bbox_weights_init,
                    bbox_gt_refine, bbox_weights_refine, stride,
                    num_total_samples_init, num_total_samples_refine):
        # classification loss
        labels = labels.reshape(-1)
        label_weights = label_weights.reshape(-1)
        cls_score = cls_score.permute(0, 2, 3,
                                      1).reshape(-1, self.cls_out_channels)
        cls_score = cls_score.contiguous()
        # points loss
        bbox_gt_init = bbox_gt_init.reshape(-1, 4)
        bbox_weights_init = bbox_weights_init.reshape(-1, 4)
        bbox_pred_init = self.points2bbox(
            pts_pred_init.reshape(-1, 2 * self.num_points), y_first=False)
        bbox_gt_refine = bbox_gt_refine.reshape(-1, 4)
        bbox_weights_refine = bbox_weights_refine.reshape(-1, 4)
        bbox_pred_refine = self.points2bbox(
            pts_pred_refine.reshape(-1, 2 * self.num_points), y_first=False)
        normalize_term = self.point_base_scale * stride
        loss_pts_init = self.loss_bbox_init(
            bbox_pred_init / normalize_term,
            bbox_gt_init / normalize_term,
            bbox_weights_init,
            avg_factor=num_total_samples_init)
        loss_pts_refine = self.loss_bbox_refine(
            bbox_pred_refine / normalize_term,
            bbox_gt_refine / normalize_term,
            bbox_weights_refine,
            avg_factor=num_total_samples_refine)
        if self.use_vfl:
            pos_inds = ((labels >= 0)
                        & (labels < self.num_classes)).nonzero().reshape(-1)
            pos_labels = labels[pos_inds]
            ious = bbox_overlaps(
                bbox_pred_refine.detach(),
                bbox_gt_refine.detach(),
                is_aligned=True)
            pos_ious = ious[pos_inds]
            cls_iou_targets = torch.zeros_like(cls_score)
            cls_iou_targets[pos_inds, pos_labels] = pos_ious
            loss_cls = self.loss_cls(
                cls_score,
                cls_iou_targets,
                label_weights.unsqueeze(1),
                avg_factor=num_total_samples_refine)
        else:
            loss_cls = self.loss_cls(
                cls_score,
                labels,
                label_weights,
                avg_factor=num_total_samples_refine)

        return loss_cls, loss_pts_init, loss_pts_refine

    def loss(self,
             cls_scores,
             pts_preds_init,
             pts_preds_refine,
             gt_bboxes,
             gt_labels,
             img_metas,
             gt_bboxes_ignore=None):
        featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
        assert len(featmap_sizes) == len(self.point_generators)
        device = cls_scores[0].device
        label_channels = self.cls_out_channels if self.use_sigmoid_cls else 1

        # target for initial stage
        center_list, valid_flag_list = self.get_points(featmap_sizes,
                                                       img_metas, device)
        pts_coordinate_preds_init = self.offset_to_pts(center_list,
                                                       pts_preds_init)
        if self.train_cfg.init.assigner['type'] == 'PointAssigner':
            # Assign target for center list
            candidate_list = center_list
        else:
            # transform center list to bbox list and
            #   assign target for bbox list
            bbox_list = self.centers_to_bboxes(center_list)
            candidate_list = bbox_list
        cls_reg_targets_init = self.get_targets(
            candidate_list,
            valid_flag_list,
            gt_bboxes,
            img_metas,
            gt_bboxes_ignore_list=gt_bboxes_ignore,
            gt_labels_list=gt_labels,
            stage='init',
            label_channels=label_channels)
        (*_, bbox_gt_list_init, candidate_list_init, bbox_weights_list_init,
         num_total_pos_init, num_total_neg_init) = cls_reg_targets_init
        num_total_samples_init = (
            num_total_pos_init +
            num_total_neg_init if self.sampling else num_total_pos_init)

        # target for refinement stage
        center_list, valid_flag_list = self.get_points(featmap_sizes,
                                                       img_metas, device)
        pts_coordinate_preds_refine = self.offset_to_pts(
            center_list, pts_preds_refine)
        bbox_list = []
        for i_img, center in enumerate(center_list):
            bbox = []
            for i_lvl in range(len(pts_preds_refine)):
                bbox_preds_init = self.points2bbox(
                    pts_preds_init[i_lvl].detach())
                bbox_shift = bbox_preds_init * self.point_strides[i_lvl]
                bbox_center = torch.cat(
                    [center[i_lvl][:, :2], center[i_lvl][:, :2]], dim=1)
                bbox.append(bbox_center +
                            bbox_shift[i_img].permute(1, 2, 0).reshape(-1, 4))
            bbox_list.append(bbox)
        cls_reg_targets_refine = self.get_targets(
            bbox_list,
            valid_flag_list,
            gt_bboxes,
            img_metas,
            gt_bboxes_ignore_list=gt_bboxes_ignore,
            gt_labels_list=gt_labels,
            stage='refine',
            label_channels=label_channels)
        (labels_list, label_weights_list, bbox_gt_list_refine,
         candidate_list_refine, bbox_weights_list_refine, num_total_pos_refine,
         num_total_neg_refine) = cls_reg_targets_refine
        num_total_samples_refine = (
            num_total_pos_refine +
            num_total_neg_refine if self.sampling else num_total_pos_refine)

        # compute loss
        losses_cls, losses_pts_init, losses_pts_refine = multi_apply(
            self.loss_single,
            cls_scores,
            pts_coordinate_preds_init,
            pts_coordinate_preds_refine,
            labels_list,
            label_weights_list,
            bbox_gt_list_init,
            bbox_weights_list_init,
            bbox_gt_list_refine,
            bbox_weights_list_refine,
            self.point_strides,
            num_total_samples_init=num_total_samples_init,
            num_total_samples_refine=num_total_samples_refine)
        loss_dict_all = {
            'loss_cls': losses_cls,
            'loss_pts_init': losses_pts_init,
            'loss_pts_refine': losses_pts_refine
        }
        return loss_dict_all

    def get_bboxes(self,
                   cls_scores,
                   pts_preds_init,
                   pts_preds_refine,
                   img_metas,
                   cfg=None,
                   rescale=False,
                   with_nms=True):
        assert len(cls_scores) == len(pts_preds_refine)
        device = cls_scores[0].device
        bbox_preds_refine = [
            self.points2bbox(pts_pred_refine)
            for pts_pred_refine in pts_preds_refine
        ]
        num_levels = len(cls_scores)
        mlvl_points = [
            self.point_generators[i].grid_points(cls_scores[i].size()[-2:],
                                                 self.point_strides[i], device)
            for i in range(num_levels)
        ]
        result_list = []
        for img_id in range(len(img_metas)):
            cls_score_list = [
                cls_scores[i][img_id].detach() for i in range(num_levels)
            ]
            bbox_pred_list = [
                bbox_preds_refine[i][img_id].detach()
                for i in range(num_levels)
            ]
            img_shape = img_metas[img_id]['img_shape']
            scale_factor = img_metas[img_id]['scale_factor']
            proposals = self._get_bboxes_single(cls_score_list, bbox_pred_list,
                                                mlvl_points, img_shape,
                                                scale_factor, cfg, rescale,
                                                with_nms)
            result_list.append(proposals)
        return result_list

    def _get_bboxes_single(self,
                           cls_scores,
                           bbox_preds,
                           mlvl_points,
                           img_shape,
                           scale_factor,
                           cfg,
                           rescale=False,
                           with_nms=True):
        cfg = self.test_cfg if cfg is None else cfg
        assert len(cls_scores) == len(bbox_preds) == len(mlvl_points)
        mlvl_bboxes = []
        mlvl_scores = []
        for i_lvl, (cls_score, bbox_pred, points) in enumerate(
                zip(cls_scores, bbox_preds, mlvl_points)):
            assert cls_score.size()[-2:] == bbox_pred.size()[-2:]
            cls_score = cls_score.permute(1, 2,
                                          0).reshape(-1, self.cls_out_channels)
            if self.use_sigmoid_cls:
                scores = cls_score.sigmoid()
            else:
                scores = cls_score.softmax(-1)
            bbox_pred = bbox_pred.permute(1, 2, 0).reshape(-1, 4)
            nms_pre = cfg.get('nms_pre', -1)
            if nms_pre > 0 and scores.shape[0] > nms_pre:
                if self.use_sigmoid_cls:
                    max_scores, _ = scores.max(dim=1)
                else:
                    # remind that we set FG labels to [0, num_class-1]
                    # since mmdet v2.0
                    # BG cat_id: num_class
                    max_scores, _ = scores[:, :-1].max(dim=1)
                _, topk_inds = max_scores.topk(nms_pre)
                points = points[topk_inds, :]
                bbox_pred = bbox_pred[topk_inds, :]
                scores = scores[topk_inds, :]
            bbox_pos_center = torch.cat([points[:, :2], points[:, :2]], dim=1)
            bboxes = bbox_pred * self.point_strides[i_lvl] + bbox_pos_center
            x1 = bboxes[:, 0].clamp(min=0, max=img_shape[1])
            y1 = bboxes[:, 1].clamp(min=0, max=img_shape[0])
            x2 = bboxes[:, 2].clamp(min=0, max=img_shape[1])
            y2 = bboxes[:, 3].clamp(min=0, max=img_shape[0])
            bboxes = torch.stack([x1, y1, x2, y2], dim=-1)
            mlvl_bboxes.append(bboxes)
            mlvl_scores.append(scores)
        mlvl_bboxes = torch.cat(mlvl_bboxes)
        if rescale:
            mlvl_bboxes /= mlvl_bboxes.new_tensor(scale_factor)
        mlvl_scores = torch.cat(mlvl_scores)
        if self.use_sigmoid_cls:
            # Add a dummy background class to the backend when using sigmoid
            # remind that we set FG labels to [0, num_class-1] since mmdet v2.0
            # BG cat_id: num_class
            padding = mlvl_scores.new_zeros(mlvl_scores.shape[0], 1)
            mlvl_scores = torch.cat([mlvl_scores, padding], dim=1)
        if with_nms:
            det_bboxes, det_labels = multiclass_nms(mlvl_bboxes, mlvl_scores,
                                                    cfg.score_thr, cfg.nms,
                                                    cfg.max_per_img)
            return det_bboxes, det_labels
        else:
            return mlvl_bboxes, mlvl_scores


#! /usr/bin/env python

import collections
import io

import math
import numpy as np

OneKelvinInKeV = 0.00000008617328149741  # 1 K = 8.617328149741e-8 keV
T9inKeV = 1e9 * OneKelvinInKeV  # 1 GK = 86.17328149741 keV

class ReaclibLibrary:
    def __init__(self, s=None):
        self.chapters = {}
        if s is not None:
            self.parse(s)

    def parse(self, s):
        lines = io.StringIO(s).readlines()
        chunks = {}
        thischunk = ''
        chapter = 0
        while lines:
            line = lines.pop(0)
            if (line[:1]).isdigit():
                if thischunk.strip() != '':
                    chunks.setdefault(chapter, []).append(thischunk)
                chapter = int(line[:2])
                thischunk = ''
            thischunk += line
        if thischunk.strip() != '':
            chunks.setdefault(chapter, []).append(thischunk)
        for i in chunks:
            # print( ''.join(chunks[i]) )
            w = ReaclibChapter(''.join(chunks[i]))
            self.chapters[i] = w

    def __str__(self):
        result = ''
        for i in self.chapters:
            result += str(self.chapters[i])
        return result

    def find(self, targ=None, proj=None, reaction=None, reference=None, chapter=None):
        for chapter in self.chapters:
            pass
        raise NotImplementedError("FIXME: write me")

    def add_inverse_reactions_detailed_balance(self):
        raise NotImplementedError("FIXME: write me")

class ReaclibChapter:
    def __init__(self, s=None):
        self.reacts = []
        self.chapter = None
        if s is not None:
            self.parse(s)

    def parse(self, s):
        split = s.splitlines()
        nchunks = len(split) / 4
        for ichunk in range(0, len(split), 4):
            chunk = split[ichunk:ichunk + 4]
            header = chunk.pop(0)
            set1 = chunk.pop(0)
            set2 = chunk.pop(0)
            set3 = chunk.pop(0)
            r = ReaclibReact(set1 + set2 + set3)
            self.reacts.append(r)
            self.chapter = header

    def __str__(self):
        mymap2 = ''
        for r in self.reacts:
            mymap2 += '\n'.join(list(map(str, [self.chapter, r]))) + '\n'
        return mymap2

class ReaclibReact:
    def __init__(self, s=None):
        self.original = s
        self.set_label = None
        self.rate = None
        self.reverserate = None
        self.Qval = None
        self.e = collections.OrderedDict(((i, None) for i in range(1, 6)))
        if s is not None:
            self.parse(s)

    def add_particle(self, p, slot):
        pstring = p.rjust(5)
        self.e[slot] = pstring

    def add_coefficient(self, c, slot):
        if slot == 0:
            self.a_zero = ['', c]
        elif slot == 1:
            self.a_one = ['', c]
        elif slot == 2:
            self.a_two = ['', c]
        elif slot == 3:
            self.a_three = ['', c]
        elif slot == 4:
            self.a_four = ['', c]
        elif slot == 5:
            self.a_five = ['', c]
        elif slot == 6:
            self.a_six = ['', c]
        else:
            raise ValueError("Coefficient slots are 0-6, you entered %s" % str(slot))

    def parse(self, s):
        self.e[1] = s[5:10]
        self.e[2] = s[10:15]
        self.e[3] = s[15:20]
        self.e[4] = s[20:25]
        self.e[5] = s[25:30]
        self.e[6] = s[30:35]
        self.set_label = s[43:47]
        self.rate = s[47]
        self.reverserate = s[48]
        self.Qval = [str(s[52:64]), float(s[52:64])]
        #        print( '/'+s[74:87]+'/' )
        #        print( '/'+s[87:100]+'/' )
        self.a_zero = [str(s[74:87]), float(s[74:87])]
        self.a_one = [str(s[87:100]), float(s[87:100])]
        self.a_two = [str(s[100:113]), float(s[100:113])]
        self.a_three = [str(s[113:126]), float(s[113:126])]
        self.a_four = [str(s[148:161]), float(s[148:161])]
        self.a_five = [str(s[161:174]), float(s[161:174])]
        self.a_six = [str(s[174:187]), float(s[174:187])]

    def __str__(self):
        line1 = "     " + self.e[1] + self.e[2] + self.e[3] + self.e[4] + self.e[5] + self.e[
            6] + "        " + self.set_label + self.rate + self.reverserate + "   " + self.Qval[0] + "          "
        line2 = self.a_zero[0] + self.a_one[0] + self.a_two[0] + self.a_three[0] + "                      "
        line3 = self.a_four[0] + self.a_five[0] + self.a_six[0] + "                                   "
        mymap = list(map(str, [line1, line2, line3]))
        return "\n".join(mymap)

    def evaluate(self, x, unit='T9'):
        if unit == 'T9':
            T = x
        elif unit == 'keV':
            T = x / T9inKeV
        else:
            raise NotImplementedError("unit must be 'keV' or 'T9', got %s" % unit)
        logrrate = (
                self.a_zero[1] +
                (self.a_one[1] * pow(T, -1.0)) +
                (self.a_two[1] * pow(T, -0.33333333333)) +
                (self.a_three[1] * pow(T, 0.3333333333)) +
                (self.a_four[1] * pow(T, 1)) +
                (self.a_five[1] * pow(T, 1.66666666666)) +
                (self.a_six[1] * math.log(T)))
        return math.exp(logrrate)

def gnds_crossSection_to_Reaclib_ARR_coefficients(xs, useCovariance=False, verbose=False, minTemp=0.1, maxTemp=10.0,
                                                  numTemps=15):
    from brownies.BNL.inter.metrics import computeAstrophysicalReactionRate
    import pqu.PQU as pqu

    # Set up temperature grid
    Temp_inT9 = np.logspace(start=math.log(minTemp), stop=math.log(maxTemp),
                            num=numTemps)  # [.1, .2, .3, .4, .5, 1, 1.5, 2.0, 5.0, 10.0]
    Temp_inkeV = [T9inKeV * (T) for T in Temp_inT9]

    # Set up array of reaction rates
    b_matrix = []
    for T in Temp_inkeV:
        ARR = computeAstrophysicalReactionRate(xs, pqu.PQU(T, 'keV'), useCovariance=useCovariance)
        b_matrix.append(np.log(ARR.getValue()))
    b = np.array(b_matrix)

    # Set up matrix of powers of temperatures so we can fit the coefficients
    a_matrix = []
    for T in Temp_inT9:
        a_matrix_row = [1, pow(T, -1.0), pow(T, -1.0 / 3.0), pow(T, 1.0 / 3.0), pow(T, 1.0), pow(T, 5.0 / 3.0),
                        np.log(T)]
        a_matrix.append(a_matrix_row)
    a = np.array(a_matrix)

    # a*x = b, solve for vector x
    if verbose: print('b:', b)
    if verbose: print('a:', a)
    x, residual, rank, s = np.linalg.lstsq(a, b, rcond=None)
    if verbose: print('a*x', np.dot(a, x))
    if verbose: print(residual)
    return x

def gnds_reactionSuite_to_ReaclibLibrary(rs, maxEThreshold=0.5e6, skipCrossSectionSums=True, verbose=False):
    import fudge
    projectile = str(rs.projectile)
    target = str(rs.target)
    resultChapters = []
    for r in rs:
        # Get rid of all reactions that are not plain reactions or sums of reactions
        # (e.g. no fission components of production stuff)
        if not isinstance(r, (fudge.reactions.reaction.reaction, fudge.sums.crossSectionSum)): continue
        if isinstance(r, fudge.sums.crossSectionSum) and skipCrossSectionSums: continue
        if not hasattr(r, 'outputChannel'): continue

        # Compute outgoing particle names
        prods = [str(p) for p in r.outputChannel]
        reactionName = str(r)

        # Get Ethreshold
        if hasattr(r, "getThreshold"):
            EThreshold = r.getThreshold(unit='eV')
        else:
            EThreshold = 0.0

        # Skip over some reactions for which an astrophysical reaction rate is useless
        if projectile in prods and target in prods: continue  # skip elastic
        if EThreshold > maxEThreshold: continue  # skip high threshold reactions
        isInelastic = False
        for prod in prods:
            if '(' in prod:
                isInelastic = True  # skip anything that's a discrete level excitation, we're not ready for isomers
                break
        if isInelastic: continue

        print(20 * '-', str(r), 20 * '-')

        # Get Q
        if hasattr(r, 'getQ'):
            Q = r.getQ(unit='eV')
        else:
            Q = 0.0

        # Get coefficients in ARR parameterization
        a = gnds_crossSection_to_Reaclib_ARR_coefficients(r.crossSection, verbose=False)

        # Figure out relevant products and compute chapter
        if 'photon' in prods: prods.remove('photon')
        if len(prods) == 1:
            chapter = 4
        elif len(prods) == 2:
            chapter = 5
        elif len(prods) == 3:
            chapter = 6
        else:
            continue

        thisReaction = ReaclibReact()
        thisReaction.set_label = 'endf'
        thisReaction.rate = None
        thisReaction.reverserate = ''
        thisReaction.Qval = ['', Q]
        thisReaction.add_particle(target, 1)
        thisReaction.add_particle(projectile, 2)
        for ip, product in enumerate(prods): thisReaction.add_particle(product, ip + 3)
        for ic, c in enumerate(a): thisReaction.add_coefficient(c, ic)

        print(prods, reactionName, a, chapter, Q, EThreshold)
        #       print( thisReaction )
        thisChapter = ReaclibChapter()
        thisChapter.chapter = chapter
        thisChapter.reacts.append(thisReaction)
        resultChapters.append(thisChapter)
    return resultChapters
    raise NotImplementedError("FIXME: write me")

def endfFile_to_ReaclibLibrary(filename, verbose=False):
    from brownies.legacy.converting import endfFileToGNDS
    rs = endfFileToGNDS.endfFileToGNDS(filename,
                                       toStdOut=int(verbose) * 10,
                                       skipBadData=True,
                                       reconstructResonances=True,
                                       continuumSpectraFix=True,
                                       doCovariances=False,
                                       verboseWarnings=verbose,
                                       printBadNK14=False,
                                       ignoreBadDate=True,
                                       acceptBadMF10FissionZAP=True)['reactionSuite']
    return gnds_reactionSuite_to_ReaclibLibrary(rs, verbose=verbose)

def plot_rates(d, temperatureGrid=[.1, .2, .3, .4, .5, 1, 1.5, 2, 5, 10], title="Astrophysical Reaction Rate"):
    import matplotlib.pyplot as plt
    ''' d={ "ENDF/B-VIII.0":reactlibreac, 'KaDoNiS':...}'''
    for k in d:
        plt.plot(temperatureGrid, [d[k].evaluate(T) for T in temperatureGrid], label=k)
    plt.xlabel('$Temperature (GK)$')
    plt.ylabel('$ Astrophysical Reaction Rate(1e6)(cm^{3} s^{-1} mol^{-1})$')
    plt.title(title)
    plt.legend(loc='upper center', shadow=True, fontsize='x-large')
    plt.show()

if __name__ == "__main__":
    import unittest

    class Test_ReaclibReact(unittest.TestCase):
        def setUp(self):
            self.s = "         n    p                            wc12w     7.82300e-01          \n" \
                     "-6.781610e+00 0.000000e+00 0.000000e+00 0.000000e+00                      \n" \
                     " 0.000000e+00 0.000000e+00 0.000000e+00                                   \n"

        # self.s=open('reaclibv1.txt').read()
        def test_output(self):
            r = ReaclibReact(self.s)
            self.assertEqual(str(r), self.s)

        def test_evaluate(self):
            r = ReaclibReact(self.s)
            self.assertAlmostEqual(r.evaluate(1.0), 4.0)

    class Test_ReaclibChapter(unittest.TestCase):
        def setUp(self):
            self.s = \
                """1
                         n    p                            wc12w     7.82300e-01
                -6.781610e+00 0.000000e+00 0.000000e+00 0.000000e+00
                 0.000000e+00 0.000000e+00 0.000000e+00
                1
                         t  he3                            wc12w     1.86000e-02
                -2.014560e+01 0.000000e+00 0.000000e+00 0.000000e+00
                 0.000000e+00 0.000000e+00 0.000000e+00
                1
                       he3    t                              ecw    -1.90000e-02
                -3.246200e+01-2.133800e-01-8.215810e-01 1.112410e+01
                -5.773380e-01 2.904710e-02-2.627050e-01
                """

        def test_output(self):
            r = ReaclibChapter(self.s)
            self.assertEqual(str(r), self.s)

    class Test_ReaclibLibrary(unittest.TestCase):
        def setUp(self):
            self.s = ""

        def test_output(self):
            r = ReaclibLibrary(self.s)
            self.assertEqual(str(r), self.s)

    unittest.main()


#!/usr/bin/env python
# coding: utf-8

# this code is a modification of:
# notes:
# todo : I canceled the randomize weights for the last layer + freezed the weights for all of the layers (some weights were trained anyway).
#todo : mayb -- save to fule during evaluate function the outputs
# **Outline of Steps**
#     + Initialization
#         + Download COCO detection data from http://cocodataset.org/#download
#             + http://images.cocodataset.org/zips/train2014.zip <= train images
#             + http://images.cocodataset.org/zips/val2014.zip <= validation images
#             + http://images.cocodataset.org/annotations/annotations_trainval2014.zip <= train and validation annotations
#         + Run this script to convert annotations in COCO format to VOC format
#             + https://gist.github.com/chicham/6ed3842d0d2014987186#file-coco2pascal-py
#         + Download pre-trained weights from https://pjreddie.com/darknet/yolo/
#             + https://pjreddie.com/media/files/yolo.weights
#         + Specify the directory of train annotations (train_annot_folder) and train images (train_image_folder)
#         + Specify the directory of validation annotations (valid_annot_folder) and validation images (valid_image_folder)
#         + Specity the path of pre-trained weights by setting variable *wt_path*
#     + Construct equivalent network in Keras
#         + Network arch from https://github.com/pjreddie/darknet/blob/master/cfg/yolo-voc.cfg
#     + Load the pretrained weights
#     + Perform training
#     + Perform detection on an image with newly trained weights
#     + Perform detection on an video with newly trained weights

# # Initialization

# In[51]:
#from IPython import get_ipython
from keras.models import Sequential, Model
from keras.layers import Reshape, Activation, Conv2D, Input, MaxPooling2D, BatchNormalization, Flatten, Dense, Lambda, \
    UpSampling2D, TimeDistributed, LSTM
from keras.layers.advanced_activations import LeakyReLU
from keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard
from keras.optimizers import SGD, Adam, RMSprop
from keras.layers.merge import concatenate
import matplotlib.pyplot as plt
import keras.backend as K
import tensorflow as tf
import imgaug as ia
from tqdm import tqdm
from imgaug import augmenters as iaa
import numpy as np
import pickle
import os, cv2
from preprocessing import parse_annotation, BatchGenerator, LSTMBatchGenerator
from utils import WeightReader, decode_netout, draw_boxes

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ""

# get_ipython().run_line_magic('matplotlib', 'inline')

# In[52]:
SUP_NUM_IMAGES   = 3
UNSUP_NUM_IMAGES = 3
EVAL_NUM_IMAGES  = 3

LABELS = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
          'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
          'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
          'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
          'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
          'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
          'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
          'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
          'hair drier', 'toothbrush']

IMAGE_H, IMAGE_W = 416, 416
GRID_H, GRID_W = 13, 13
BOX = 5
CLASS = len(LABELS)
CLASS_WEIGHTS = np.ones(CLASS, dtype='float32')
OBJ_THRESHOLD = 0.3  # 0.5
NMS_THRESHOLD = 0.3  # 0.45
ANCHORS = [0.57273, 0.677385, 1.87446, 2.06253, 3.33843, 5.47434, 7.88282, 3.52778, 9.77052, 9.16828]

NO_OBJECT_SCALE = 1.0
OBJECT_SCALE = 5.0
COORD_SCALE = 1.0
CLASS_SCALE = 1.0

BATCH_SIZE = 16
WARM_UP_BATCHES = 0
TRUE_BOX_BUFFER = 50

MAX_BOX_PER_IMAGE = 10

# In[53]:

wt_path = 'yolov2.weights'
train_image_folder = './data/images/train2014/'
train_annot_folder = './data/train_converted/'
valid_image_folder = './data/images/val2014/'
valid_annot_folder = './data/val_converted/'

# # Construct the network

# the function to implement the orgnization layer (thanks to github.com/allanzelener/YAD2K)
def space_to_depth_x2(x):
    return tf.space_to_depth(x, block_size=2)

import frontend
""" creates a new dir names coco_x with the results, weights, and all the relevant files"""
# TB_COUNT = len([d for d in os.listdir(os.path.expanduser('./results_lstm/')) if 'coco_' in d]) + 1
# PATH = os.path.expanduser('./results_lstm/') + 'coco_' + '_' + str(TB_COUNT)
# os.makedirs(PATH)
PATH = './lstm/'
print("=================== Directory " , PATH ,  " Created ")
# PATH = "./results/coco__25"

class ToharGenerator2(BatchGenerator):
    def __getitem__(self, item):
        # t= [x_batch,b_batch],y_batch
        #    [input,goutndtruth],desired network output]
        t = super().__getitem__(item)
        x_batch = t[0][0] #the input
        GT = t[0][1]
        y_batch = t[1]

        new_x_batch = predict(model,x_batch) #instead of input img vector we want the YOLO's output vector
        t[0][0]= new_x_batch
        return [new_x_batch, GT], y_batch

input_image = Input(shape=(IMAGE_H, IMAGE_W, 3))
true_boxes = Input(shape=(1, 1, 1, TRUE_BOX_BUFFER, 4))

# Layer 1
x = Conv2D(32, (3, 3), strides=(1, 1), padding='same', name='conv_1', use_bias=False)(input_image)
x = BatchNormalization(name='norm_1')(x)
x = LeakyReLU(alpha=0.1)(x)
encoded = MaxPooling2D(pool_size=(2, 2))(x)

# Layer 2
x = Conv2D(64, (3, 3), strides=(1, 1), padding='same', name='conv_2', use_bias=False, trainable=False)(encoded)
x = BatchNormalization(name='norm_2', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)
x = MaxPooling2D(pool_size=(2, 2))(x)

# Layer 3
x = Conv2D(128, (3, 3), strides=(1, 1), padding='same', name='conv_3', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_3', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 4
x = Conv2D(64, (1, 1), strides=(1, 1), padding='same', name='conv_4', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_4', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 5
x = Conv2D(128, (3, 3), strides=(1, 1), padding='same', name='conv_5', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_5', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)
x = MaxPooling2D(pool_size=(2, 2))(x)

# Layer 6
x = Conv2D(256, (3, 3), strides=(1, 1), padding='same', name='conv_6', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_6', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 7
x = Conv2D(128, (1, 1), strides=(1, 1), padding='same', name='conv_7', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_7', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 8
x = Conv2D(256, (3, 3), strides=(1, 1), padding='same', name='conv_8', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_8', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)
x = MaxPooling2D(pool_size=(2, 2))(x)

# Layer 9
x = Conv2D(512, (3, 3), strides=(1, 1), padding='same', name='conv_9', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_9', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 10
x = Conv2D(256, (1, 1), strides=(1, 1), padding='same', name='conv_10', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_10', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 11
x = Conv2D(512, (3, 3), strides=(1, 1), padding='same', name='conv_11', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_11', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 12
x = Conv2D(256, (1, 1), strides=(1, 1), padding='same', name='conv_12', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_12', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 13
x = Conv2D(512, (3, 3), strides=(1, 1), padding='same', name='conv_13', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_13', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

skip_connection = x

x = MaxPooling2D(pool_size=(2, 2))(x)

# Layer 14
x = Conv2D(1024, (3, 3), strides=(1, 1), padding='same', name='conv_14', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_14', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 15
x = Conv2D(512, (1, 1), strides=(1, 1), padding='same', name='conv_15', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_15', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 16
x = Conv2D(1024, (3, 3), strides=(1, 1), padding='same', name='conv_16', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_16', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 17
x = Conv2D(512, (1, 1), strides=(1, 1), padding='same', name='conv_17', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_17', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 18
x = Conv2D(1024, (3, 3), strides=(1, 1), padding='same', name='conv_18', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_18', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 19
x = Conv2D(1024, (3, 3), strides=(1, 1), padding='same', name='conv_19', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_19', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 20
x = Conv2D(1024, (3, 3), strides=(1, 1), padding='same', name='conv_20', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_20', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 21
skip_connection = Conv2D(64, (1, 1), strides=(1, 1), padding='same', name='conv_21', use_bias=False, trainable=False)(
    skip_connection)
skip_connection = BatchNormalization(name='norm_21', trainable=False)(skip_connection)
skip_connection = LeakyReLU(alpha=0.1)(skip_connection)
skip_connection = Lambda(space_to_depth_x2)(skip_connection)

x = concatenate([skip_connection, x])

# Layer 22
x = Conv2D(1024, (3, 3), strides=(1, 1), padding='same', name='conv_22', use_bias=False, trainable=False)(x)
x = BatchNormalization(name='norm_22', trainable=False)(x)
x = LeakyReLU(alpha=0.1)(x)

# Layer 23
x = Conv2D(BOX * (4 + 1 + CLASS), (1, 1), strides=(1, 1), padding='same', name='conv_23')(x)
output = Reshape((GRID_H, GRID_W, BOX, 4 + 1 + CLASS))(x)

# small hack to allow true_boxes to be registered when Keras build the model
# for more information: https://github.com/fchollet/keras/issues/2790
output = Lambda(lambda args: args[0])([output, true_boxes])

model = Model([input_image, true_boxes], output)

# model.summary()
print("output=====")
print(output.shape)
'''build lstm model: '''
lstm_input = Input(shape=(GRID_H, GRID_W, BOX, 4 + 1 + CLASS))

input_dim = GRID_H * GRID_W * BOX * (4 + 1 + CLASS)
# input_dim=(GRID_H,GRID_W, BOX, 4 + 1 + CLASS, 1, 1, 1, TRUE_BOX_BUFFER, 4)
print(input_dim)

timesteps = EVAL_NUM_IMAGES

# lstm.add(units= Dense(input_shape=(GRID_H, GRID_W, BOX, 4 + 1 + CLASS)))
# l=Lambda(lambda x: K.batch_flatten(x))(lstm_input)
# l=LSTM(input_dim, batch_input_shape= (None, timesteps, input_dim), activation='sigmoid',recurrent_activation='hard_sigmoid',return_sequences=True)(l)
# # l = (Dense(output_dim=input_dim, activation="relu"))(lstm)
# #
# # # l = LSTM(input_dim)(l)
# # # # hidden_layer = Dense(output_dim=input_shape, activation="relu")(x)
# # # # outputs = Dense(output_dim=input_shape, activation="softmax")(hidden_layer)
# #
# loutput = Reshape((GRID_H, GRID_W, BOX, 4 + 1 + CLASS))(l)
# #
# # # small hack to allow true_boxes to be registered when Keras build the model
# # # for more information: https://github.com/fchollet/keras/issues/2790
# out = Lambda(lambda args: args[0])([loutput, true_boxes])
#
#
#
# lstm = Model([lstm_input, true_boxes], out)
# lstm.summary()

input_dim = GRID_H * GRID_W * BOX * (4 + 1 + CLASS)

#take 5 frames every time
frames = Input(shape=(5, IMAGE_H, IMAGE_W, 3))
x = TimeDistributed(model)(frames)
x = TimeDistributed(Flatten())(x)
#now- timestamsp=5
x = LSTM(input_dim, name='lstm')(x)
out = Dense(input_dim, name='out')(x)
lstm = Model(inputs=frames, outputs=out)

exit()
# # Load pretrained weights

# **Load the weights originally provided by YOLO**
print("**Load the weights originally provided by YOLO**")
weight_reader = WeightReader(wt_path)

weight_reader.reset()  # don't worry! it doesn't delete the weights.
nb_conv = 23

for i in range(1, nb_conv + 1):
    conv_layer = model.get_layer('conv_' + str(i))

    if i < nb_conv:
        norm_layer = model.get_layer('norm_' + str(i))

        size = np.prod(norm_layer.get_weights()[0].shape)

        beta = weight_reader.read_bytes(size)
        gamma = weight_reader.read_bytes(size)
        mean = weight_reader.read_bytes(size)
        var = weight_reader.read_bytes(size)

        weights = norm_layer.set_weights([gamma, beta, mean, var])

    if len(conv_layer.get_weights()) > 1:
        bias = weight_reader.read_bytes(np.prod(conv_layer.get_weights()[1].shape))
        kernel = weight_reader.read_bytes(np.prod(conv_layer.get_weights()[0].shape))
        kernel = kernel.reshape(list(reversed(conv_layer.get_weights()[0].shape)))
        kernel = kernel.transpose([2, 3, 1, 0])
        conv_layer.set_weights([kernel, bias])
    else:
        kernel = weight_reader.read_bytes(np.prod(conv_layer.get_weights()[0].shape))
        kernel = kernel.reshape(list(reversed(conv_layer.get_weights()[0].shape)))
        kernel = kernel.transpose([2, 3, 1, 0])
        conv_layer.set_weights([kernel])

# model_t = model  #model that trained but not pre-trained
# model_un = model #model without training at all

# **Randomize weights of the last layer**

# In[ ]:

# print("========randomize last layer")
# layer   = model.layers[-4] # the last convolutional layer
# weights = layer.get_weights()
#
# new_kernel = np.random.normal(size=weights[0].shape)/(GRID_H*GRID_W)
# new_bias   = np.random.normal(size=weights[1].shape)/(GRID_H*GRID_W)
#
# layer.set_weights([new_kernel, new_bias])

# # Perform training

# **Loss function**

# $$\begin{multline}
# \lambda_\textbf{coord}
# \sum_{i = 0}^{S^2}
#     \sum_{j = 0}^{B}
#      L_{ij}^{\text{obj}}
#             \left[
#             \left(
#                 x_i - \hat{x}_i
#             \right)^2 +
#             \left(
#                 y_i - \hat{y}_i
#             \right)^2
#             \right]
# \\
# + \lambda_\textbf{coord}
# \sum_{i = 0}^{S^2}
#     \sum_{j = 0}^{B}
#          L_{ij}^{\text{obj}}
#          \left[
#         \left(
#             \sqrt{w_i} - \sqrt{\hat{w}_i}
#         \right)^2 +
#         \left(
#             \sqrt{h_i} - \sqrt{\hat{h}_i}
#         \right)^2
#         \right]
# \\
# + \sum_{i = 0}^{S^2}
#     \sum_{j = 0}^{B}
#         L_{ij}^{\text{obj}}
#         \left(
#             C_i - \hat{C}_i
#         \right)^2
# \\
# + \lambda_\textrm{noobj}
# \sum_{i = 0}^{S^2}
#     \sum_{j = 0}^{B}
#     L_{ij}^{\text{noobj}}
#         \left(
#             C_i - \hat{C}_i
#         \right)^2
# \\
# + \sum_{i = 0}^{S^2}
# L_i^{\text{obj}}
#     \sum_{c \in \textrm{classes}}
#         \left(
#             p_i(c) - \hat{p}_i(c)
#         \right)^2
# \end{multline}$$

# In[ ]:

import backend

def predict(model, image, i, img_name, path=""):
    """
    input_size = IMAGE_H

    image_h, image_w, _ = image.shape
    feature_extractor = backend.FullYoloFeature()
    image = cv2.resize(image, (input_size, input_size))
    image =feature_extractor.normalize(image)

    input_image = image[:,:,::-1]
    input_image = np.expand_dims(input_image, 0)
    dummy_array = np.zeros((1,1,1,1, MAX_BOX_PER_IMAGE,4))

    netout = model.predict([input_image, dummy_array])[0]
    boxes  = decode_netout(netout, ANCHORS, len(LABELS))
    """
    dummy_array = np.zeros((1, 1, 1, 1, TRUE_BOX_BUFFER, 4))
    # print("dummy array:", dummy_array)
    plt.figure(figsize=(10, 10))

    input_image = cv2.resize(image, (416, 416))
    input_image = input_image / 255.
    input_image = input_image[:, :, ::-1]
    input_image = np.expand_dims(input_image, 0)

    netout = model.predict([input_image, dummy_array])

    boxes = decode_netout(netout[0],
                          obj_threshold=OBJ_THRESHOLD,
                          nms_threshold=NMS_THRESHOLD,
                          anchors=ANCHORS,
                          nb_class=CLASS)
    image = draw_boxes(image, boxes, labels=LABELS)

    plt.imshow(image[:, :, ::-1])
    path = str(path)
    if i <= 100:
        # Create target directory & all intermediate directories if don't exists
        if not os.path.exists(path):
            os.makedirs(path)
            print("Directory ", path, " Created ")
        else:
            pass
            # print("Directory ", path, " already exists")
        #os.makedirs(path) # create the directory on given path, also if any intermediate-level directory don’t exists then it will create that too.
        plt.savefig(path+ "/" + img_name)

    return boxes

from utils import decode_netout, compute_overlap, compute_ap

from os.path import normpath, basename

def evaluate(model, generator,
             iou_threshold=0.3,
             score_threshold=0.3,
             max_detections=100,
             save_path=None):
    """ Evaluate a given dataset using a given model.
    code originally from https://github.com/fizyr/keras-retinanet

    # Arguments
        generator       : The generator that represents the dataset to evaluate.
        model           : The model to evaluate.
        iou_threshold   : The threshold used to consider when a detection is positive or negative.
        score_threshold : The score confidence threshold to use for detections.
        max_detections  : The maximum number of detections to use per image.
        save_path       : The path to save images with visualized detections to.
    # Returns
        A dict mapping class names to mAP scores.
    """
    # gather all detections and annotations
    all_detections = [[None for i in range(generator.num_classes())] for j in range(generator.size())]
    all_annotations = [[None for i in range(generator.num_classes())] for j in range(generator.size())]

    for i in range(generator.size()):
        raw_image = generator.load_image(i)

        path = generator.images[i]['filename']
        img_name = basename(normpath(path))

        raw_height, raw_width, raw_channels = raw_image.shape

        # make the boxes and the labels
        pred_boxes = predict(model, raw_image, i, img_name, path=save_path)

        score = np.array([box.score for box in pred_boxes])
        pred_labels = np.array([box.label for box in pred_boxes])

        if len(pred_boxes) > 0:
            pred_boxes = np.array([[box.xmin * raw_width, box.ymin * raw_height, box.xmax * raw_width,
                                    box.ymax * raw_height, box.score] for box in pred_boxes])
        else:
            pred_boxes = np.array([[]])

            # sort the boxes and the labels according to scores
        score_sort = np.argsort(-score)
        pred_labels = pred_labels[score_sort]
        pred_boxes = pred_boxes[score_sort]

        # copy detections to all_detections
        for label in range(generator.num_classes()):
            all_detections[i][label] = pred_boxes[pred_labels == label, :]

        annotations = generator.load_annotation(i)

        # copy detections to all_annotations
        for label in range(generator.num_classes()):
            all_annotations[i][label] = annotations[annotations[:, 4] == label, :4].copy()

    # compute mAP by comparing all detections and all annotations
    average_precisions = {}

    for label in range(generator.num_classes()):
        false_positives = np.zeros((0,))
        true_positives = np.zeros((0,))
        scores = np.zeros((0,))
        num_annotations = 0.0

        for i in range(generator.size()):
            detections = all_detections[i][label]
            annotations = all_annotations[i][label]
            num_annotations += annotations.shape[0]
            detected_annotations = []

            for d in detections:
                scores = np.append(scores, d[4])

                if annotations.shape[0] == 0:
                    false_positives = np.append(false_positives, 1)
                    true_positives = np.append(true_positives, 0)
                    continue

                overlaps = compute_overlap(np.expand_dims(d, axis=0), annotations)
                assigned_annotation = np.argmax(overlaps, axis=1)
                max_overlap = overlaps[0, assigned_annotation]

                if max_overlap >= iou_threshold and assigned_annotation not in detected_annotations:
                    false_positives = np.append(false_positives, 0)
                    true_positives = np.append(true_positives, 1)
                    detected_annotations.append(assigned_annotation)
                else:
                    false_positives = np.append(false_positives, 1)
                    true_positives = np.append(true_positives, 0)

        # no annotations -> AP for this class is 0 (is this correct?)
        if num_annotations == 0:
            average_precisions[label] = 0
            continue

        # sort by score
        indices = np.argsort(-scores)
        false_positives = false_positives[indices]
        true_positives = true_positives[indices]

        # compute false positives and true positives
        false_positives = np.cumsum(false_positives)
        true_positives = np.cumsum(true_positives)

        # compute recall and precision
        recall = true_positives / num_annotations
        precision = true_positives / np.maximum(true_positives + false_positives, np.finfo(np.float64).eps)

        # compute average precision
        average_precision = compute_ap(recall, precision)
        average_precisions[label] = average_precision

    import pickle
    f = open(save_path+"/mAP.pkl", "wb")
    pickle.dump(average_precisions, f)
    f.close()

    return average_precisions

def custom_loss(y_true, y_pred):
    mask_shape = tf.shape(y_true)[:4]

    cell_x = tf.to_float(tf.reshape(tf.tile(tf.range(GRID_W), [GRID_H]), (1, GRID_H, GRID_W, 1, 1)))
    cell_y = tf.transpose(cell_x, (0, 2, 1, 3, 4))

    cell_grid = tf.tile(tf.concat([cell_x, cell_y], -1), [BATCH_SIZE, 1, 1, 5, 1])

    coord_mask = tf.zeros(mask_shape)
    conf_mask = tf.zeros(mask_shape)
    class_mask = tf.zeros(mask_shape)

    seen = tf.Variable(0.)
    total_recall = tf.Variable(0.)

    """
    Adjust prediction
    """
    ### adjust x and y
    pred_box_xy = tf.sigmoid(y_pred[..., :2]) + cell_grid

    ### adjust w and h
    pred_box_wh = tf.exp(y_pred[..., 2:4]) * np.reshape(ANCHORS, [1, 1, 1, BOX, 2])

    ### adjust confidence
    pred_box_conf = tf.sigmoid(y_pred[..., 4])

    ### adjust class probabilities
    pred_box_class = y_pred[..., 5:]

    """
    Adjust ground truth
    """
    ### adjust x and y
    true_box_xy = y_true[..., 0:2]  # relative position to the containing cell

    ### adjust w and h
    true_box_wh = y_true[..., 2:4]  # number of cells accross, horizontally and vertically

    ### adjust confidence
    true_wh_half = true_box_wh / 2.
    true_mins = true_box_xy - true_wh_half
    true_maxes = true_box_xy + true_wh_half

    pred_wh_half = pred_box_wh / 2.
    pred_mins = pred_box_xy - pred_wh_half
    pred_maxes = pred_box_xy + pred_wh_half

    intersect_mins = tf.maximum(pred_mins, true_mins)
    intersect_maxes = tf.minimum(pred_maxes, true_maxes)
    intersect_wh = tf.maximum(intersect_maxes - intersect_mins, 0.)
    intersect_areas = intersect_wh[..., 0] * intersect_wh[..., 1]

    true_areas = true_box_wh[..., 0] * true_box_wh[..., 1]
    pred_areas = pred_box_wh[..., 0] * pred_box_wh[..., 1]

    union_areas = pred_areas + true_areas - intersect_areas
    iou_scores = tf.truediv(intersect_areas, union_areas)

    true_box_conf = iou_scores * y_true[..., 4]

    ### adjust class probabilities
    true_box_class = tf.argmax(y_true[..., 5:], -1)

    """
    Determine the masks
    """
    ### coordinate mask: simply the position of the ground truth boxes (the predictors)
    coord_mask = tf.expand_dims(y_true[..., 4], axis=-1) * COORD_SCALE

    ### confidence mask: penelize predictors + penalize boxes with low IOU
    # penalize the confidence of the boxes, which have IOU with some ground truth box < 0.6
    true_xy = true_boxes[..., 0:2]
    true_wh = true_boxes[..., 2:4]

    true_wh_half = true_wh / 2.
    true_mins = true_xy - true_wh_half
    true_maxes = true_xy + true_wh_half

    pred_xy = tf.expand_dims(pred_box_xy, 4)
    pred_wh = tf.expand_dims(pred_box_wh, 4)

    pred_wh_half = pred_wh / 2.
    pred_mins = pred_xy - pred_wh_half
    pred_maxes = pred_xy + pred_wh_half

    intersect_mins = tf.maximum(pred_mins, true_mins)
    intersect_maxes = tf.minimum(pred_maxes, true_maxes)
    intersect_wh = tf.maximum(intersect_maxes - intersect_mins, 0.)
    intersect_areas = intersect_wh[..., 0] * intersect_wh[..., 1]

    true_areas = true_wh[..., 0] * true_wh[..., 1]
    pred_areas = pred_wh[..., 0] * pred_wh[..., 1]

    union_areas = pred_areas + true_areas - intersect_areas
    iou_scores = tf.truediv(intersect_areas, union_areas)

    best_ious = tf.reduce_max(iou_scores, axis=4)
    conf_mask = conf_mask + tf.to_float(best_ious < 0.6) * (1 - y_true[..., 4]) * NO_OBJECT_SCALE

    # penalize the confidence of the boxes, which are reponsible for corresponding ground truth box
    conf_mask = conf_mask + y_true[..., 4] * OBJECT_SCALE

    ### class mask: simply the position of the ground truth boxes (the predictors)
    class_mask = y_true[..., 4] * tf.gather(CLASS_WEIGHTS, true_box_class) * CLASS_SCALE

    """
    Warm-up training
    """
    no_boxes_mask = tf.to_float(coord_mask < COORD_SCALE / 2.)
    seen = tf.assign_add(seen, 1.)

    true_box_xy, true_box_wh, coord_mask = tf.cond(tf.less(seen, WARM_UP_BATCHES),
                                                   lambda: [true_box_xy + (0.5 + cell_grid) * no_boxes_mask,
                                                            true_box_wh + tf.ones_like(true_box_wh) * np.reshape(
                                                                ANCHORS, [1, 1, 1, BOX, 2]) * no_boxes_mask,
                                                            tf.ones_like(coord_mask)],
                                                   lambda: [true_box_xy,
                                                            true_box_wh,
                                                            coord_mask])

    """
    Finalize the loss
    """
    nb_coord_box = tf.reduce_sum(tf.to_float(coord_mask > 0.0))
    nb_conf_box = tf.reduce_sum(tf.to_float(conf_mask > 0.0))
    nb_class_box = tf.reduce_sum(tf.to_float(class_mask > 0.0))

    loss_xy = tf.reduce_sum(tf.square(true_box_xy - pred_box_xy) * coord_mask) / (nb_coord_box + 1e-6) / 2.
    loss_wh = tf.reduce_sum(tf.square(true_box_wh - pred_box_wh) * coord_mask) / (nb_coord_box + 1e-6) / 2.
    loss_conf = tf.reduce_sum(tf.square(true_box_conf - pred_box_conf) * conf_mask) / (nb_conf_box + 1e-6) / 2.
    loss_class = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=true_box_class, logits=pred_box_class)
    loss_class = tf.reduce_sum(loss_class * class_mask) / (nb_class_box + 1e-6)

    loss = loss_xy + loss_wh + loss_conf + loss_class

    nb_true_box = tf.reduce_sum(y_true[..., 4])
    nb_pred_box = tf.reduce_sum(tf.to_float(true_box_conf > 0.5) * tf.to_float(pred_box_conf > 0.3))

    """
    Debugging code
    """
    current_recall = nb_pred_box / (nb_true_box + 1e-6)
    total_recall = tf.assign_add(total_recall, current_recall)

    loss = tf.Print(loss, [tf.zeros((1))], message='Dummy Line \t', summarize=1000)
    loss = tf.Print(loss, [loss_xy], message='Loss XY \t', summarize=1000)
    loss = tf.Print(loss, [loss_wh], message='Loss WH \t', summarize=1000)
    loss = tf.Print(loss, [loss_conf], message='Loss Conf \t', summarize=1000)
    loss = tf.Print(loss, [loss_class], message='Loss Class \t', summarize=1000)
    loss = tf.Print(loss, [loss], message='Total Loss \t', summarize=1000)
    loss = tf.Print(loss, [current_recall], message='Current Recall \t', summarize=1000)
    loss = tf.Print(loss, [total_recall / seen], message='Average Recall \t', summarize=1000)

    return loss

# **Parse the annotations to construct train generator and validation generator**

generator_config = {
    'IMAGE_H': IMAGE_H,
    'IMAGE_W': IMAGE_W,
    'GRID_H': GRID_H,
    'GRID_W': GRID_W,
    'BOX': BOX,
    'LABELS': LABELS,
    'CLASS': len(LABELS),
    'ANCHORS': ANCHORS,
    'BATCH_SIZE': BATCH_SIZE,
    'TRUE_BOX_BUFFER': 50,
}

def normalize(image):
    return image / 255.

# train_imgs, seen_train_labels = parse_annotation(train_annot_folder, train_image_folder, labels=LABELS)
# ## write parsed annotations to pickle for fast retrieval next time
# with open('train_imgs', 'wb') as fp:
#    pickle.dump(train_imgs, fp)

# ## read saved pickle of parsed annotations
# with open('train_imgs', 'rb') as fp:
#     train_imgs = pickle.load(fp)
#
# from random import shuffle
# shuffle(train_imgs)
#
# with open('train_imgs_shuffled', 'wb') as fp:
#    pickle.dump(train_imgs, fp)

with open('train_imgs_shuffled', 'rb') as fp:
    train_imgs = pickle.load(fp)

# valid_imgs, seen_valid_labels = parse_annotation(valid_annot_folder, valid_image_folder, labels=LABELS)
# ## write parsed annotations to pickle for fast retrieval next time
# with open('valid_imgs', 'wb') as fp:
#    pickle.dump(valid_imgs, fp)
## read saved pickle of parsed annotations
with open('valid_imgs', 'rb') as fp:
    valid_imgs = pickle.load(fp)

sup_train_imgs = train_imgs[:SUP_NUM_IMAGES]

# split the training set (supervised date) into train and validation 80%, 20% respectively:
train = sup_train_imgs[:int(SUP_NUM_IMAGES*0.8)]
val = sup_train_imgs[-int(SUP_NUM_IMAGES*0.2):] #takes the last 20% images from the training

train_batch = BatchGenerator(train, generator_config, norm=normalize)

eval_imgs = valid_imgs[:EVAL_NUM_IMAGES] #we use the valid_imgs as our evaluation set (testing). while we use 20% of the training for validation.

valid_batch = BatchGenerator(val, generator_config, norm=normalize, jitter=False)

"""we evaluate the model on the validation set"""
tohar_eval_batch = BatchGenerator(eval_imgs, generator_config, norm=normalize, jitter=False,
                                  shuffle=False)

# **Setup a few callbacks and start the training**

early_stop = EarlyStopping(monitor='val_loss',
                           min_delta=0.001,
                           patience=3,
                           mode='min',
                           verbose=1)

checkpoint = ModelCheckpoint(PATH+'/LSTM_weights_coco.h5',
                             monitor='val_loss',
                             verbose=1,
                             save_best_only=True,
                             mode='min',
                             period=1)
org_checkpoint = ModelCheckpoint(PATH+'/original_weights_coco.h5',
                                monitor='val_loss',
                                verbose=1,
                                save_best_only=True,
                                mode='min',
                                period=1)

# In[ ]:

tb_counter = len([log for log in os.listdir(os.path.expanduser('./lstm/')) if 'coco_' in log]) + 1
tensorboard = TensorBoard(log_dir=os.path.expanduser('./lstm/') + 'coco_' + '_' + str(tb_counter),
                          histogram_freq=0,
                          write_graph=True,
                          write_images=False)

optimizer = Adam(lr=0.5e-4, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)
# optimizer = SGD(lr=1e-4, decay=0.0005, momentum=0.9)
# optimizer = RMSprop(lr=1e-4, rho=0.9, epsilon=1e-08, decay=0.0)

model.compile(loss=custom_loss, optimizer=optimizer)
# model_t.compile(loss=custom_loss, optimizer=optimizer)

from keras.callbacks import TensorBoard

"""evaluating on original YOLO (no training at all)"""
model.load_weights("yolo.h5")
# YOLO = evaluate(model, tohar_eval_batch, save_path=PATH+"/YOLO")
# print("YOLO:\n", YOLO)
# print(np.average(list(YOLO.values())))

'''creating a modified batch to the lstm:'''
# [x_batch, GT], y_batch
# [x_batch, GT], \
lstm_batch = LSTMBatchGenerator(eval_imgs, generator_config, model, norm=None, jitter=False, shuffle=False)

print(len(lstm_batch))

exit()

"""X_train2 should be YOLO's output vectors
y_train2 should be the ground truth in the exact same format of YOLO's output
"""

# autoencoder.fit_generator(generator=train_batch_lstm,  #(input, input)
#                           steps_per_epoch=len(train_batch_lstm),
#                           epochs=100,
#                           verbose=1,
#                           # validation_data=tohar_valid_batch,
#                           # validation_steps=len(tohar_valid_batch),
#                           callbacks=[early_stop, ae_checkpoint, tensorboard],
#                           max_queue_size=3)
# print("===================== Done training AE")
# print("===================== Save weights to AE_weights_coco.h5")
# autoencoder.save_weights(PATH+"/AE_weights_coco.h5")  # save the autoencoder's weights in this file
# print("===================== Load weights from AE_weights_coco.h5")
# model.load_weights(PATH+'/AE_weights_coco.h5',
#                    by_name=True)  # copy the AE's weights to the "YOLO model" weights, only to layers with the same name as the AE

## end ae
##uncomment for training:

# Perform detection on image

# print("===================== load YOLO model's weights to weights_coco.h5")

# evaluate:

# train_batch_lstm = ToharGenerator2(train, generator_config, norm=normalize)

""" Add lstm on top of the trained YOLO model. the lstm should have many to many sturcture. each latm cell predict 1 output . help:"""
# https://stackoverflow.com/questions/49535488/lstm-on-top-of-a-pre-trained-cnn
# https://github.com/keras-team/keras/issues/5527
''' Freeze previous layers '''
for layer in model.layers:
    layer.trainable = False

from keras.applications.vgg16 import VGG16
from keras.models import Model
from keras.layers import Dense, Input
from keras.layers.pooling import GlobalAveragePooling2D
from keras.layers.recurrent import LSTM
from keras.layers.wrappers import TimeDistributed
from keras.optimizers import Nadam

frames = len(tohar_eval_batch)
print(frames)

units = GRID_H * GRID_W * BOX * (4 + 1 + CLASS)
print("==========",units)
length=5 #todo:batch size

#todo: input dim is problematic.
# input_images = Input(shape=( None, frames ,IMAGE_H, IMAGE_W, 3))
#https://riptutorial.com/keras/example/29812/vgg-16-cnn-and-lstm-for-video-classification

# frames, rows, columns, channels = 10, IMAGE_H, IMAGE_W, 3
# video = Input(shape=(frames,
#                      rows,
#                      columns,
#                      channels))
#
# # cnn_base = VGG16(input_shape=(rows, columns, channels),
# #                  weights="imagenet",
# #                  include_top=False)
# # cnn_out = GlobalAveragePooling2D()(cnn_base.output)
# # cnn = Model(input=cnn_base.input, output=cnn_out)
#
# model.trainable = False
#
# encoded_frames = TimeDistributed(model)(video)
# encoded_sequence = LSTM(256)(encoded_frames)
# hidden_layer = Dense(output_dim=1024, activation="relu")(encoded_sequence)
# outputs = Dense(output_dim=units, activation="softmax")(hidden_layer)
# lstm = Model([video], outputs)

#
# # x = Reshape((len(train_batch)*10 ,IMAGE_H, IMAGE_W, 3))(input_images)
# x = TimeDistributed(model)(x)
# x = TimeDistributed(Flatten())(x)
# x = LSTM(units, name='lstm')(x) # This has the effect of each LSTM unit returning a sequence of 1 output, one for each time step in the input data
# # x = Dense( n_output,name='lstm_out')(x)
# # x = Conv2D(BOX * (4 + 1 + CLASS), (1, 1), strides=(1, 1), padding='same', name='lstm_conv')(x)
# out = Reshape((GRID_H, GRID_W, BOX, 4 + 1 + CLASS))(x)

print("======== lstm:")
lstm.summary()

lstm.compile(loss=custom_loss, optimizer=optimizer)

exit()
lstm.fit_generator(generator=train_batch,  # train_batch #(input, ground_truth)
                    steps_per_epoch=len(train_batch),
                    epochs=3,
                    verbose=1,
                    validation_data=valid_batch,
                    validation_steps=len(valid_batch),
                    callbacks=[early_stop, checkpoint, tensorboard],
                    max_queue_size=3)

"""evaluating on LSTM YOLO """
LSTM = evaluate(model, tohar_eval_batch, save_path=PATH+"/LSTM")
print("LSTM:\n",LSTM)
print(np.average(list(LSTM.values())))

# """evaluating on original YOLO (no training at all)"""
# model.load_weights("yolo.h5")
# YOLO = evaluate(model, tohar_eval_batch, save_path=PATH+"/YOLO")
# print("YOLO:\n", YOLO)
# print(np.average(list(YOLO.values())))
#
#
# """evaluating on original YOLO (no training at all) """
# model_t.load_weights(PATH+"/T_weights_coco.h5")
# NO_AE = evaluate(model_t, tohar_eval_batch, save_path=PATH+"/NO_AE")
# print("NO_AE:\n", NO_AE)
# print(np.average(list(NO_AE.values())))

params={"SUP_NUM_IMAGES:": SUP_NUM_IMAGES,
"UNSUP_NUM_IMAGES:":UNSUP_NUM_IMAGES,
"EVAL_NUM_IMAGES:":EVAL_NUM_IMAGES}

f = open(PATH + "/mAP.txt", "w")
f.write("LSTM:\n")
f.write(str(LSTM)+"\n")
f.write("NO_AE:\n")
# f.write(str(NO_AE)+"\n")
f.write("YOLO:\n")
# f.write(str(YOLO)+"\n")
f.write("AVG:"+"\n")
f.write(str(np.average(list(LSTM.values())))+"\n")
# f.write(str(np.average(list(NO_AE.values())))+"\n")
# f.write(str(np.average(list(YOLO.values())))+"\n")
f.write("LOG:"+"\n")
f.write(str(params) )

f.close()

# image = cv2.imread('images/giraffe.jpg')
# dummy_array = np.zeros((1,1,1,1,TRUE_BOX_BUFFER,4))
# plt.figure(figsize=(10,10))
#
# input_image = cv2.resize(image, (416, 416))
# input_image = input_image / 255.
# input_image = input_image[:,:,::-1]
# input_image = np.expand_dims(input_image, 0)
#
# netout = model.predict([input_image, dummy_array])
#
# boxes = decode_netout(netout[0],
#                       obj_threshold=OBJ_THRESHOLD,
#                       nms_threshold=NMS_THRESHOLD,
#                       anchors=ANCHORS,
#                       nb_class=CLASS)
#
# image = draw_boxes(image, boxes, labels=LABELS)
#
# plt.imshow(image[:,:,::-1]); #plt.show()
# i=0
# plt.savefig("./predictions/figure"+str(i))
print('\a')
print('\a')

exit()

# # Perform detection on video

# In[ ]:

model.load_weights("weights_coco.h5")

dummy_array = np.zeros((1, 1, 1, 1, TRUE_BOX_BUFFER, 4))

# In[ ]:

video_inp = '../basic-yolo-keras/images/phnom_penh.mp4'
video_out = '../basic-yolo-keras/images/phnom_penh_bbox.mp4'

video_reader = cv2.VideoCapture(video_inp)

nb_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
frame_h = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT))
frame_w = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH))

video_writer = cv2.VideoWriter(video_out,
                               cv2.VideoWriter_fourcc(*'XVID'),
                               50.0,
                               (frame_w, frame_h))

for i in tqdm(range(nb_frames)):
    ret, image = video_reader.read()

    input_image = cv2.resize(image, (416, 416))
    input_image = input_image / 255.
    input_image = input_image[:, :, ::-1]
    input_image = np.expand_dims(input_image, 0)

    netout = model.predict([input_image, dummy_array])

    boxes = decode_netout(netout[0],
                          obj_threshold=0.3,
                          nms_threshold=NMS_THRESHOLD,
                          anchors=ANCHORS,
                          nb_class=CLASS)
    image = draw_boxes(image, boxes, labels=LABELS)

    video_writer.write(np.uint8(image))

video_reader.release()
video_writer.release()


import urllib3, json
import os
import pandas as pd
import numpy as np
import math
from pathlib import Path

def get_annotation_file(dst=None, version='1.3'):
    if version == '1.3':
        # Get ActivityNet Annotation file for v1.3 and save it in a json file
        url = "http://ec2-52-25-205-214.us-west-2.compute.amazonaws.com/files/activity_net.v1-3.min.json"
        response = urllib3.PoolManager().request('GET', url).data
        data = json.loads(response)
        # Prepare destination path
        if dst == None:
            dst = os.path.dirname(os.path.realpath(__file__))
        dst_path = os.path.join(dst, 'activity_net.v1-3.min.json')
        with open(dst_path, 'w') as outfile:
            json.dump(data, outfile)
    else:
        raise KeyError(f"Version {version} not configured in function get_activitynet_annotation_file.")

def rename_df_col(col):
    if isinstance(col, tuple):
        col = '_'.join(str(c) for c in col)
    return col

def annotation_stats(src):
    ###Create dataframe with annotation info
    annotation_path = Path(src)
    with annotation_path.open('r') as f:
        data = json.load(f)
    # Define columns
    cols = ['video_id', 'url', 'subset', 'resolution', 'fps', 'duration', 'label', 'segment_low', 'segment_high', 'frame_low', 'frame_high', 'frame_qty']
    # Create empty DF
    df_ann = pd.DataFrame(columns=cols)
    # Loop over annotation file
    i = 0       # Counter
    total = len(data['database'].keys())
    for video_id in data['database'].keys():
        i+=1                                        # Increment counter
        print(f"Loop {i}/{total}: {video_id}")      # print info
        # Prepare base vec
        try:
            base_vec = {'video_id': video_id,
                        'url': data['database'][video_id]['url'],
                        'subset': data['database'][video_id]['subset'],
                        'resolution': data['database'][video_id]['resolution'],
                        'fps': data['database'][video_id]['fps'],
                        'duration': data['database'][video_id]['duration']}
            # Loop over annotations of the video_id
            for ann in data['database'][video_id]['annotations']:
                # Add info to the final vec, that will be append to the dataframe
                final_vec = base_vec
                final_vec['label'] = ann['label']
                final_vec['segment_low'] = ann['segment'][0]
                final_vec['segment_high'] = ann['segment'][1]
                final_vec['frame_low'] = math.floor(ann['segment'][0] * base_vec['fps']) + 1
                final_vec['frame_high'] = math.floor(ann['segment'][1] * base_vec['fps']) + 1
                final_vec['frame_qty'] = (final_vec['frame_high']-final_vec['frame_low']) if (final_vec['frame_high']-final_vec['frame_low']) > 8 else 0
                # Append to the dataframe
                df_ann = df_ann.append(final_vec, ignore_index=True)
        except KeyError as e:
            print(f"KeyError in Video: {video_id} Key: {e}")
    # Generate annotation stats
    print("Generating annotation stats")
    df_res = df_ann.groupby(['subset', 'video_id', 'label']).agg({'fps':'mean', 'duration':'mean', 'frame_qty': 'sum'}).groupby(['subset', 'label']).agg({'fps':'mean', 'duration':'mean', 'frame_qty': ['sum', 'count']})
    df_res[('frame_qty', 'mean')] = df_res[('frame_qty',   'sum')]/df_res[('frame_qty', 'count')]
    df_res.columns = map(rename_df_col, df_res.columns)
    df_res.reset_index(inplace=True)
    # Save in excel file
    print("Saving stats in excel file")
    dst = os.path.dirname(os.path.realpath(__file__))
    dst_path = os.path.join(dst, 'annotation_stats.xlsx')
    df_res.to_excel(dst_path)
    # Plot distribuitions
    print("Training set's video count distribuition over classes")
    df_res[df_res['subset']=='training']['frame_qty_count'].plot.kde()
    print("Validation set's video count distribuition over classes")
    df_res[df_res['subset']=='validation']['frame_qty_count'].plot.kde()
    print("Training set's frame count distribuition over classes")
    df_res[df_res['subset']=='training']['frame_qty_mean'].plot.kde()
    print("Validation set's frame count distribuition over classes")
    df_res[df_res['subset']=='validation']['frame_qty_mean'].plot.kde()