"""
RPN:
data =
    {'data': [num_images, c, h, w],
     'im_info': [num_images, 4] (optional)}
label =
    {'gt_boxes': [num_boxes, 5] (optional),
     'label': [batch_size, 1] <- [batch_size, num_anchors, feat_height, feat_width],
     'bbox_target': [batch_size, num_anchors, feat_height, feat_width],
     'bbox_weight': [batch_size, num_anchors, feat_height, feat_width]}
"""

from __future__ import print_function
import sys
import logging
import datetime
import numpy as np
import numpy.random as npr

from ..logger import logger
from ..config import config
from .image import get_image, tensor_vstack, get_crop_image
from ..processing.generate_anchor import generate_anchors, anchors_plane
from ..processing.bbox_transform import bbox_overlaps, bbox_transform, landmark_transform

STAT = {0:0, 8:0, 16:0, 32:0}

def get_rpn_testbatch(roidb):
    """
    return a dict of testbatch
    :param roidb: ['image', 'flipped']
    :return: data, label, im_info
    """
    assert len(roidb) == 1, 'Single batch only'
    imgs, roidb = get_image(roidb)
    im_array = imgs[0]
    im_info = np.array([roidb[0]['im_info']], dtype=np.float32)

    data = {'data': im_array,
            'im_info': im_info}
    label = {}

    return data, label, im_info

def get_rpn_batch(roidb):
    """
    prototype for rpn batch: data, im_info, gt_boxes
    :param roidb: ['image', 'flipped'] + ['gt_boxes', 'boxes', 'gt_classes']
    :return: data, label
    """
    assert len(roidb) == 1, 'Single batch only'
    imgs, roidb = get_image(roidb)
    im_array = imgs[0]
    im_info = np.array([roidb[0]['im_info']], dtype=np.float32)

    # gt boxes: (x1, y1, x2, y2, cls)
    if roidb[0]['gt_classes'].size > 0:
        gt_inds = np.where(roidb[0]['gt_classes'] != 0)[0]
        gt_boxes = np.empty((roidb[0]['boxes'].shape[0], 5), dtype=np.float32)
        gt_boxes[:, 0:4] = roidb[0]['boxes'][gt_inds, :]
        gt_boxes[:, 4] = roidb[0]['gt_classes'][gt_inds]
    else:
        gt_boxes = np.empty((0, 5), dtype=np.float32)

    data = {'data': im_array,
            'im_info': im_info}
    label = {'gt_boxes': gt_boxes}

    return data, label

def get_crop_batch(roidb):
    """
    prototype for rpn batch: data, im_info, gt_boxes
    :param roidb: ['image', 'flipped'] + ['gt_boxes', 'boxes', 'gt_classes']
    :return: data, label
    """
    #assert len(roidb) == 1, 'Single batch only'
    data_list = []
    label_list = []
    imgs, roidb = get_crop_image(roidb)
    assert len(imgs)==len(roidb)
    for i in range(len(imgs)):
      im_array = imgs[i]
      im_info = np.array([roidb[i]['im_info']], dtype=np.float32)

      # gt boxes: (x1, y1, x2, y2, cls)
      if roidb[i]['gt_classes'].size > 0:
          gt_inds = np.where(roidb[i]['gt_classes'] != 0)[0]
          gt_boxes = np.empty((roidb[i]['boxes'].shape[0], 5), dtype=np.float32)
          gt_boxes[:, 0:4] = roidb[i]['boxes'][gt_inds, :]
          gt_boxes[:, 4] = roidb[i]['gt_classes'][gt_inds]
          if config.USE_BLUR:
            gt_blur = roidb[i]['blur']
          if config.FACE_LANDMARK:
            #gt_landmarks = np.empty((roidb[i]['landmarks'].shape[0], 11), dtype=np.float32)
            gt_landmarks = roidb[i]['landmarks'][gt_inds,:,:]
          if config.HEAD_BOX:
            gt_boxes_head = np.empty((roidb[i]['boxes_head'].shape[0], 5), dtype=np.float32)
            gt_boxes_head[:, 0:4] = roidb[i]['boxes_head'][gt_inds, :]
            gt_boxes_head[:, 4] = roidb[i]['gt_classes'][gt_inds]
      else:
          gt_boxes = np.empty((0, 5), dtype=np.float32)
          if config.USE_BLUR:
            gt_blur = np.empty((0,), dtype=np.float32)
          if config.FACE_LANDMARK:
            gt_landmarks = np.empty((0, 5, 3), dtype=np.float32)
          if config.HEAD_BOX:
            gt_boxes_head = np.empty((0, 5), dtype=np.float32)

      data = {'data': im_array,
              'im_info': im_info}
      label = {'gt_boxes': gt_boxes}
      if config.USE_BLUR:
        label['gt_blur'] = gt_blur
      if config.FACE_LANDMARK:
        label['gt_landmarks'] = gt_landmarks
      if config.HEAD_BOX:
        label['gt_boxes_head'] = gt_boxes_head
      data_list.append(data)
      label_list.append(label)

    return data_list, label_list

def assign_anchor_fpn(feat_shape, gt_label, im_info, landmark=False, prefix='face', select_stride=0):
    """
    assign ground truth boxes to anchor positions
    :param feat_shape: infer output shape
    :param gt_boxes: assign ground truth
    :param im_info: filter out anchors overlapped with edges
    :return: tuple
    labels: of shape (batch_size, 1) <- (batch_size, num_anchors, feat_height, feat_width)
    bbox_targets: of shape (batch_size, num_anchors * 4, feat_height, feat_width)
    bbox_weights: mark the assigned anchors
    """
    def _unmap(data, count, inds, fill=0):
        """" unmap a subset inds of data into original data of size count """
        if len(data.shape) == 1:
            ret = np.empty((count,), dtype=np.float32)
            ret.fill(fill)
            ret[inds] = data
        else:
            ret = np.empty((count,) + data.shape[1:], dtype=np.float32)
            ret.fill(fill)
            ret[inds, :] = data
        return ret

    global STAT
    DEBUG = False

    im_info = im_info[0]
    gt_boxes = gt_label['gt_boxes']
    # clean up boxes
    nonneg = np.where(gt_boxes[:, 4] != -1)[0]
    gt_boxes = gt_boxes[nonneg]
    if config.USE_BLUR:
      gt_blur = gt_label['gt_blur']
      gt_blur = gt_blur[nonneg]
    if landmark:
      gt_landmarks = gt_label['gt_landmarks']
      gt_landmarks = gt_landmarks[nonneg]
      assert gt_boxes.shape[0]==gt_landmarks.shape[0]
    #scales = np.array(scales, dtype=np.float32)
    feat_strides = config.RPN_FEAT_STRIDE
    bbox_pred_len = 4
    landmark_pred_len = 10
    if config.USE_BLUR:
      gt_boxes[:,4] = gt_blur
      bbox_pred_len = 5
    if config.USE_OCCLUSION:
      landmark_pred_len = 15

    anchors_list = []
    anchors_num_list = []
    inds_inside_list = []
    feat_infos = []
    A_list = []
    for i in range(len(feat_strides)):
        stride = feat_strides[i]
        sstride = str(stride)
        base_size = config.RPN_ANCHOR_CFG[sstride]['BASE_SIZE']
        allowed_border = config.RPN_ANCHOR_CFG[sstride]['ALLOWED_BORDER']
        ratios = config.RPN_ANCHOR_CFG[sstride]['RATIOS']
        scales = config.RPN_ANCHOR_CFG[sstride]['SCALES']
        base_anchors = generate_anchors(base_size=base_size, ratios=list(ratios), scales=np.array(scales, dtype=np.float32), stride = stride, dense_anchor = config.DENSE_ANCHOR)
        num_anchors = base_anchors.shape[0]
        feat_height, feat_width = feat_shape[i][-2:]
        feat_stride = feat_strides[i]
        feat_infos.append([feat_height, feat_width])

        A = num_anchors
        A_list.append(A)
        K = feat_height * feat_width

        all_anchors = anchors_plane(feat_height, feat_width, feat_stride, base_anchors)
        all_anchors = all_anchors.reshape((K * A, 4))
        #print('anchor0', stride, all_anchors[0])

        total_anchors = int(K * A)
        anchors_num_list.append(total_anchors)
        # only keep anchors inside the image
        inds_inside = np.where((all_anchors[:, 0] >= -allowed_border) &
                               (all_anchors[:, 1] >= -allowed_border) &
                               (all_anchors[:, 2] < im_info[1] + allowed_border) &
                               (all_anchors[:, 3] < im_info[0] + allowed_border))[0]
        if DEBUG:
            print('total_anchors', total_anchors)
            print('inds_inside', len(inds_inside))

        # keep only inside anchors
        anchors = all_anchors[inds_inside, :]
        #print('AA', anchors.shape, len(inds_inside))

        anchors_list.append(anchors)
        inds_inside_list.append(inds_inside)

    # Concat anchors from each level
    anchors = np.concatenate(anchors_list)
    for i in range(1, len(inds_inside_list)):
        inds_inside_list[i] = inds_inside_list[i] + sum(anchors_num_list[:i])
    inds_inside = np.concatenate(inds_inside_list)
    total_anchors = sum(anchors_num_list)
    #print('total_anchors', anchors.shape[0], len(inds_inside), file=sys.stderr)

    # label: 1 is positive, 0 is negative, -1 is dont care
    labels = np.empty((len(inds_inside),), dtype=np.float32)
    labels.fill(-1)
    #print('BB', anchors.shape, len(inds_inside))
    #print('gt_boxes', gt_boxes.shape, file=sys.stderr)

    if gt_boxes.size > 0:
        # overlap between the anchors and the gt boxes
        # overlaps (ex, gt)
        overlaps = bbox_overlaps(anchors.astype(np.float), gt_boxes.astype(np.float))
        argmax_overlaps = overlaps.argmax(axis=1)
        #print('AAA', argmax_overlaps.shape)
        max_overlaps = overlaps[np.arange(len(inds_inside)), argmax_overlaps]
        gt_argmax_overlaps = overlaps.argmax(axis=0)
        gt_max_overlaps = overlaps[gt_argmax_overlaps, np.arange(overlaps.shape[1])]
        gt_argmax_overlaps = np.where(overlaps == gt_max_overlaps)[0]

        if not config.TRAIN.RPN_CLOBBER_POSITIVES:
            # assign bg labels first so that positive labels can clobber them
            labels[max_overlaps < config.TRAIN.RPN_NEGATIVE_OVERLAP] = 0

        # fg label: for each gt, anchor with highest overlap
        if config.TRAIN.RPN_FORCE_POSITIVE:
          labels[gt_argmax_overlaps] = 1

        # fg label: above threshold IoU
        labels[max_overlaps >= config.TRAIN.RPN_POSITIVE_OVERLAP] = 1

        if config.TRAIN.RPN_CLOBBER_POSITIVES:
            # assign bg labels last so that negative labels can clobber positives
            labels[max_overlaps < config.TRAIN.RPN_NEGATIVE_OVERLAP] = 0
    else:
        labels[:] = 0
    fg_inds = np.where(labels == 1)[0]
    #print('fg count', len(fg_inds))

    # subsample positive labels if we have too many
    if config.TRAIN.RPN_ENABLE_OHEM==0:
      fg_inds = np.where(labels == 1)[0]
      num_fg = int(config.TRAIN.RPN_FG_FRACTION * config.TRAIN.RPN_BATCH_SIZE)
      if len(fg_inds) > num_fg:
          disable_inds = npr.choice(fg_inds, size=(len(fg_inds) - num_fg), replace=False)
          if DEBUG:
              disable_inds = fg_inds[:(len(fg_inds) - num_fg)]
          labels[disable_inds] = -1

      # subsample negative labels if we have too many
      num_bg = config.TRAIN.RPN_BATCH_SIZE - np.sum(labels == 1)
      bg_inds = np.where(labels == 0)[0]
      if len(bg_inds) > num_bg:
          disable_inds = npr.choice(bg_inds, size=(len(bg_inds) - num_bg), replace=False)
          if DEBUG:
              disable_inds = bg_inds[:(len(bg_inds) - num_bg)]
          labels[disable_inds] = -1

      #fg_inds = np.where(labels == 1)[0]
      #num_fg = len(fg_inds)
      #num_bg = num_fg*int(1.0/config.TRAIN.RPN_FG_FRACTION-1)

      #bg_inds = np.where(labels == 0)[0]
      #if len(bg_inds) > num_bg:
      #    disable_inds = npr.choice(bg_inds, size=(len(bg_inds) - num_bg), replace=False)
      #    if DEBUG:
      #        disable_inds = bg_inds[:(len(bg_inds) - num_bg)]
      #    labels[disable_inds] = -1
    else:
      fg_inds = np.where(labels == 1)[0]
      num_fg = len(fg_inds)
      bg_inds = np.where(labels == 0)[0]
      num_bg = len(bg_inds)

    #print('anchor stat', num_fg, num_bg)

    bbox_targets = np.zeros((len(inds_inside), bbox_pred_len), dtype=np.float32)
    if gt_boxes.size > 0:
        #print('GT', gt_boxes.shape, gt_boxes[argmax_overlaps, :4].shape)
        bbox_targets[:,:] = bbox_transform(anchors, gt_boxes[argmax_overlaps, :])
        #bbox_targets[:,4] = gt_blur

    bbox_weights = np.zeros((len(inds_inside), bbox_pred_len), dtype=np.float32)
    #bbox_weights[labels == 1, :] = np.array(config.TRAIN.RPN_BBOX_WEIGHTS)
    bbox_weights[labels == 1, 0:4] = 1.0
    if bbox_pred_len>4:
      bbox_weights[labels == 1, 4:bbox_pred_len] = 0.1

    if landmark:
      landmark_targets = np.zeros((len(inds_inside), landmark_pred_len), dtype=np.float32)
      #landmark_weights = np.zeros((len(inds_inside), 10), dtype=np.float32)
      landmark_weights = np.zeros((len(inds_inside), landmark_pred_len), dtype=np.float32)
      #landmark_weights[labels == 1, :] = np.array(config.TRAIN.RPN_LANDMARK_WEIGHTS)
      if landmark_pred_len==10:
        landmark_weights[labels == 1, :] = 1.0
      elif landmark_pred_len==15:
        v = [1.0, 1.0, 0.1] * 5
        assert len(v)==15
        landmark_weights[labels == 1, :] = np.array(v)
      else:
        assert False
      #TODO here
      if gt_landmarks.size > 0:
        #print('AAA',argmax_overlaps)
        a_landmarks = gt_landmarks[argmax_overlaps,:,:]
        landmark_targets[:] = landmark_transform(anchors, a_landmarks)
        invalid = np.where(a_landmarks[:,0,2]<0.0)[0]
        #assert len(invalid)==0
        #landmark_weights[invalid, :] = np.array(config.TRAIN.RPN_INVALID_LANDMARK_WEIGHTS)
        landmark_weights[invalid, :] = 0.0

    #if DEBUG:
    #    _sums = bbox_targets[labels == 1, :].sum(axis=0)
    #    _squared_sums = (bbox_targets[labels == 1, :] ** 2).sum(axis=0)
    #    _counts = np.sum(labels == 1)
    #    means = _sums / (_counts + 1e-14)
    #    stds = np.sqrt(_squared_sums / _counts - means ** 2)
    #    print 'means', means
    #    print 'stdevs', stds
    # map up to original set of anchors
    #print(labels.shape, total_anchors, inds_inside.shape, inds_inside[0], inds_inside[-1])
    labels = _unmap(labels, total_anchors, inds_inside, fill=-1)
    bbox_targets = _unmap(bbox_targets, total_anchors, inds_inside, fill=0)
    bbox_weights = _unmap(bbox_weights, total_anchors, inds_inside, fill=0)
    if landmark:
      landmark_targets = _unmap(landmark_targets, total_anchors, inds_inside, fill=0)
      landmark_weights = _unmap(landmark_weights, total_anchors, inds_inside, fill=0)
    #print('CC', anchors.shape, len(inds_inside))

    #if DEBUG:
    #    if gt_boxes.size > 0:
    #        print 'rpn: max max_overlaps', np.max(max_overlaps)
    #    print 'rpn: num_positives', np.sum(labels == 1)
    #    print 'rpn: num_negatives', np.sum(labels == 0)
    #    _fg_sum = np.sum(labels == 1)
    #    _bg_sum = np.sum(labels == 0)
    #    _count = 1
    #    print 'rpn: num_positive avg', _fg_sum / _count
    #    print 'rpn: num_negative avg', _bg_sum / _count

    # resahpe
    label_list = list()
    bbox_target_list = list()
    bbox_weight_list = list()
    if landmark:
      landmark_target_list = list()
      landmark_weight_list = list()
    anchors_num_range = [0] + anchors_num_list
    label = {}
    for i in range(len(feat_strides)):
        stride = feat_strides[i]
        feat_height, feat_width = feat_infos[i]
        A = A_list[i]
        _label = labels[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        if select_stride>0 and stride!=select_stride:
          #print('set', stride, select_stride)
          _label[:] = -1
        #print('_label', _label.shape, select_stride)
        #_fg_inds = np.where(_label == 1)[0]
        #n_fg = len(_fg_inds)
        #STAT[0]+=1
        #STAT[stride]+=n_fg
        #if STAT[0]%100==0:
        #  print('rpn_stat', STAT, file=sys.stderr)
        bbox_target = bbox_targets[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        bbox_weight = bbox_weights[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        if landmark:
          landmark_target = landmark_targets[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
          landmark_weight = landmark_weights[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]

        _label = _label.reshape((1, feat_height, feat_width, A)).transpose(0, 3, 1, 2)
        _label = _label.reshape((1, A * feat_height * feat_width))
        bbox_target = bbox_target.reshape((1, feat_height*feat_width, A * bbox_pred_len)).transpose(0, 2, 1)
        bbox_weight = bbox_weight.reshape((1, feat_height*feat_width, A * bbox_pred_len)).transpose((0, 2, 1))
        label['%s_label_stride%d'%(prefix, stride)] = _label
        label['%s_bbox_target_stride%d'%(prefix,stride)] = bbox_target
        label['%s_bbox_weight_stride%d'%(prefix,stride)] = bbox_weight
        if landmark:
          landmark_target = landmark_target.reshape((1, feat_height*feat_width, A * landmark_pred_len)).transpose(0, 2, 1)
          landmark_weight = landmark_weight.reshape((1, feat_height*feat_width, A * landmark_pred_len)).transpose((0, 2, 1))
          label['%s_landmark_target_stride%d'%(prefix,stride)] = landmark_target
          label['%s_landmark_weight_stride%d'%(prefix,stride)] = landmark_weight
        #print('in_rpn', stride,_label.shape, bbox_target.shape, bbox_weight.shape, file=sys.stderr)
        label_list.append(_label)
        #print('DD', _label.shape)
        bbox_target_list.append(bbox_target)
        bbox_weight_list.append(bbox_weight)
        if landmark:
          landmark_target_list.append(landmark_target)
          landmark_weight_list.append(landmark_weight)

    label_concat = np.concatenate(label_list, axis=1)
    bbox_target_concat = np.concatenate(bbox_target_list, axis=2)
    bbox_weight_concat = np.concatenate(bbox_weight_list, axis=2)
    #fg_inds = np.where(label_concat[0] == 1)[0]
    #print('fg_inds_in_rpn2', fg_inds, file=sys.stderr)

    label.update({'%s_label'%prefix: label_concat,
            '%s_bbox_target'%prefix: bbox_target_concat,
            '%s_bbox_weight'%prefix: bbox_weight_concat}
            )
    if landmark:
      landmark_target_concat = np.concatenate(landmark_target_list, axis=2)
      landmark_weight_concat = np.concatenate(landmark_weight_list, axis=2)
      label['%s_landmark_target'%prefix] = landmark_target_concat
      label['%s_landmark_weight'%prefix] = landmark_weight_concat
    return label

class AA:
  def __init__(self, feat_shape):
    self.feat_shape = feat_shape
    feat_strides = config.RPN_FEAT_STRIDE
    anchors_list = []
    anchors_num_list = []
    inds_inside_list = []
    feat_infos = []
    A_list = []
    DEBUG = False
    for i in range(len(feat_strides)):
        stride = feat_strides[i]
        sstride = str(stride)
        base_size = config.RPN_ANCHOR_CFG[sstride]['BASE_SIZE']
        allowed_border = config.RPN_ANCHOR_CFG[sstride]['ALLOWED_BORDER']
        ratios = config.RPN_ANCHOR_CFG[sstride]['RATIOS']
        scales = config.RPN_ANCHOR_CFG[sstride]['SCALES']
        base_anchors = generate_anchors(base_size=base_size, ratios=list(ratios), scales=np.array(scales, dtype=np.float32), stride = stride, dense_anchor = config.DENSE_ANCHOR)
        num_anchors = base_anchors.shape[0]
        feat_height, feat_width = feat_shape[i][-2:]
        feat_stride = feat_strides[i]
        feat_infos.append([feat_height, feat_width])

        A = num_anchors
        A_list.append(A)
        K = feat_height * feat_width

        all_anchors = anchors_plane(feat_height, feat_width, feat_stride, base_anchors)
        all_anchors = all_anchors.reshape((K * A, 4))
        #print('anchor0', stride, all_anchors[0])

        total_anchors = int(K * A)
        anchors_num_list.append(total_anchors)
        # only keep anchors inside the image
        inds_inside = np.where((all_anchors[:, 0] >= -allowed_border) &
                               (all_anchors[:, 1] >= -allowed_border) &
                               (all_anchors[:, 2] < config.SCALES[0][1] + allowed_border) &
                               (all_anchors[:, 3] < config.SCALES[0][1] + allowed_border))[0]
        if DEBUG:
            print('total_anchors', total_anchors)
            print('inds_inside', len(inds_inside))

        # keep only inside anchors
        anchors = all_anchors[inds_inside, :]
        #print('AA', anchors.shape, len(inds_inside))

        anchors_list.append(anchors)
        inds_inside_list.append(inds_inside)
    anchors = np.concatenate(anchors_list)
    for i in range(1, len(inds_inside_list)):
        inds_inside_list[i] = inds_inside_list[i] + sum(anchors_num_list[:i])
    inds_inside = np.concatenate(inds_inside_list)
    #self.anchors_list = anchors_list
    #self.inds_inside_list = inds_inside_list
    self.anchors = anchors
    self.inds_inside = inds_inside
    self.anchors_num_list = anchors_num_list
    self.feat_infos = feat_infos
    self.A_list = A_list
    self._times = [0.0, 0.0, 0.0, 0.0]

  @staticmethod
  def _unmap(data, count, inds, fill=0):
      """" unmap a subset inds of data into original data of size count """
      if len(data.shape) == 1:
          ret = np.empty((count,), dtype=np.float32)
          ret.fill(fill)
          ret[inds] = data
      else:
          ret = np.empty((count,) + data.shape[1:], dtype=np.float32)
          ret.fill(fill)
          ret[inds, :] = data
      return ret

  def assign_anchor_fpn(self, gt_label, im_info, landmark=False, prefix='face', select_stride=0):

    #ta = datetime.datetime.now()
    im_info = im_info[0]
    gt_boxes = gt_label['gt_boxes']
    # clean up boxes
    nonneg = np.where(gt_boxes[:, 4] != -1)[0]
    gt_boxes = gt_boxes[nonneg]
    if config.USE_BLUR:
      gt_blur = gt_label['gt_blur']
      gt_blur = gt_blur[nonneg]
    if landmark:
      gt_landmarks = gt_label['gt_landmarks']
      gt_landmarks = gt_landmarks[nonneg]
      assert gt_boxes.shape[0]==gt_landmarks.shape[0]
    #scales = np.array(scales, dtype=np.float32)
    feat_strides = config.RPN_FEAT_STRIDE
    bbox_pred_len = 4
    landmark_pred_len = 10
    if config.USE_BLUR:
      gt_boxes[:,4] = gt_blur
      bbox_pred_len = 5
    if config.USE_OCCLUSION:
      landmark_pred_len = 15

    #anchors_list = self.anchors_list
    #inds_inside_list = self.inds_inside_list
    anchors = self.anchors
    inds_inside = self.inds_inside
    anchors_num_list = self.anchors_num_list
    feat_infos = self.feat_infos
    A_list = self.A_list

    total_anchors = sum(anchors_num_list)
    #print('total_anchors', anchors.shape[0], len(inds_inside), file=sys.stderr)

    # label: 1 is positive, 0 is negative, -1 is dont care
    labels = np.empty((len(inds_inside),), dtype=np.float32)
    labels.fill(-1)
    #print('BB', anchors.shape, len(inds_inside))
    #print('gt_boxes', gt_boxes.shape, file=sys.stderr)
    #tb = datetime.datetime.now()
    #self._times[0] += (tb-ta).total_seconds()
    #ta = datetime.datetime.now()

    if gt_boxes.size > 0:
        # overlap between the anchors and the gt boxes
        # overlaps (ex, gt)
        overlaps = bbox_overlaps(anchors.astype(np.float), gt_boxes.astype(np.float))
        argmax_overlaps = overlaps.argmax(axis=1)
        #print('AAA', argmax_overlaps.shape)
        max_overlaps = overlaps[np.arange(len(inds_inside)), argmax_overlaps]
        gt_argmax_overlaps = overlaps.argmax(axis=0)
        gt_max_overlaps = overlaps[gt_argmax_overlaps, np.arange(overlaps.shape[1])]
        gt_argmax_overlaps = np.where(overlaps == gt_max_overlaps)[0]

        if not config.TRAIN.RPN_CLOBBER_POSITIVES:
            # assign bg labels first so that positive labels can clobber them
            labels[max_overlaps < config.TRAIN.RPN_NEGATIVE_OVERLAP] = 0

        # fg label: for each gt, anchor with highest overlap
        if config.TRAIN.RPN_FORCE_POSITIVE:
          labels[gt_argmax_overlaps] = 1

        # fg label: above threshold IoU
        labels[max_overlaps >= config.TRAIN.RPN_POSITIVE_OVERLAP] = 1

        if config.TRAIN.RPN_CLOBBER_POSITIVES:
            # assign bg labels last so that negative labels can clobber positives
            labels[max_overlaps < config.TRAIN.RPN_NEGATIVE_OVERLAP] = 0
    else:
        labels[:] = 0
    fg_inds = np.where(labels == 1)[0]
    #print('fg count', len(fg_inds))

    # subsample positive labels if we have too many
    if config.TRAIN.RPN_ENABLE_OHEM==0:
      fg_inds = np.where(labels == 1)[0]
      num_fg = int(config.TRAIN.RPN_FG_FRACTION * config.TRAIN.RPN_BATCH_SIZE)
      if len(fg_inds) > num_fg:
          disable_inds = npr.choice(fg_inds, size=(len(fg_inds) - num_fg), replace=False)
          if DEBUG:
              disable_inds = fg_inds[:(len(fg_inds) - num_fg)]
          labels[disable_inds] = -1

      # subsample negative labels if we have too many
      num_bg = config.TRAIN.RPN_BATCH_SIZE - np.sum(labels == 1)
      bg_inds = np.where(labels == 0)[0]
      if len(bg_inds) > num_bg:
          disable_inds = npr.choice(bg_inds, size=(len(bg_inds) - num_bg), replace=False)
          if DEBUG:
              disable_inds = bg_inds[:(len(bg_inds) - num_bg)]
          labels[disable_inds] = -1

      #fg_inds = np.where(labels == 1)[0]
      #num_fg = len(fg_inds)
      #num_bg = num_fg*int(1.0/config.TRAIN.RPN_FG_FRACTION-1)

      #bg_inds = np.where(labels == 0)[0]
      #if len(bg_inds) > num_bg:
      #    disable_inds = npr.choice(bg_inds, size=(len(bg_inds) - num_bg), replace=False)
      #    if DEBUG:
      #        disable_inds = bg_inds[:(len(bg_inds) - num_bg)]
      #    labels[disable_inds] = -1
    else:
      fg_inds = np.where(labels == 1)[0]
      num_fg = len(fg_inds)
      bg_inds = np.where(labels == 0)[0]
      num_bg = len(bg_inds)

    #print('anchor stat', num_fg, num_bg)

    bbox_targets = np.zeros((len(inds_inside), bbox_pred_len), dtype=np.float32)
    if gt_boxes.size > 0:
        #print('GT', gt_boxes.shape, gt_boxes[argmax_overlaps, :4].shape)
        bbox_targets[:,:] = bbox_transform(anchors, gt_boxes[argmax_overlaps, :])
        #bbox_targets[:,4] = gt_blur
    #tb = datetime.datetime.now()
    #self._times[1] += (tb-ta).total_seconds()
    #ta = datetime.datetime.now()

    bbox_weights = np.zeros((len(inds_inside), bbox_pred_len), dtype=np.float32)
    #bbox_weights[labels == 1, :] = np.array(config.TRAIN.RPN_BBOX_WEIGHTS)
    bbox_weights[labels == 1, 0:4] = 1.0
    if bbox_pred_len>4:
      bbox_weights[labels == 1, 4:bbox_pred_len] = 0.1

    if landmark:
      landmark_targets = np.zeros((len(inds_inside), landmark_pred_len), dtype=np.float32)
      #landmark_weights = np.zeros((len(inds_inside), 10), dtype=np.float32)
      landmark_weights = np.zeros((len(inds_inside), landmark_pred_len), dtype=np.float32)
      #landmark_weights[labels == 1, :] = np.array(config.TRAIN.RPN_LANDMARK_WEIGHTS)
      if landmark_pred_len==10:
        landmark_weights[labels == 1, :] = 1.0
      elif landmark_pred_len==15:
        v = [1.0, 1.0, 0.1] * 5
        assert len(v)==15
        landmark_weights[labels == 1, :] = np.array(v)
      else:
        assert False
      #TODO here
      if gt_landmarks.size > 0:
        #print('AAA',argmax_overlaps)
        a_landmarks = gt_landmarks[argmax_overlaps,:,:]
        landmark_targets[:] = landmark_transform(anchors, a_landmarks)
        # landmark 不存在就是0，存在为1
        invalid = np.where(a_landmarks[:,0,2]<=0.0)[0]
        #assert len(invalid)==0
        #landmark_weights[invalid, :] = np.array(config.TRAIN.RPN_INVALID_LANDMARK_WEIGHTS)
        landmark_weights[invalid, :] = 0.0
    #tb = datetime.datetime.now()
    #self._times[2] += (tb-ta).total_seconds()
    #ta = datetime.datetime.now()

    #if DEBUG:
    #    _sums = bbox_targets[labels == 1, :].sum(axis=0)
    #    _squared_sums = (bbox_targets[labels == 1, :] ** 2).sum(axis=0)
    #    _counts = np.sum(labels == 1)
    #    means = _sums / (_counts + 1e-14)
    #    stds = np.sqrt(_squared_sums / _counts - means ** 2)
    #    print 'means', means
    #    print 'stdevs', stds
    # map up to original set of anchors
    #print(labels.shape, total_anchors, inds_inside.shape, inds_inside[0], inds_inside[-1])
    labels = AA._unmap(labels, total_anchors, inds_inside, fill=-1)
    bbox_targets = AA._unmap(bbox_targets, total_anchors, inds_inside, fill=0)
    bbox_weights = AA._unmap(bbox_weights, total_anchors, inds_inside, fill=0)
    if landmark:
      landmark_targets = AA._unmap(landmark_targets, total_anchors, inds_inside, fill=0)
      landmark_weights = AA._unmap(landmark_weights, total_anchors, inds_inside, fill=0)
    #print('CC', anchors.shape, len(inds_inside))

    #if DEBUG:
    #    if gt_boxes.size > 0:
    #        print 'rpn: max max_overlaps', np.max(max_overlaps)
    #    print 'rpn: num_positives', np.sum(labels == 1)
    #    print 'rpn: num_negatives', np.sum(labels == 0)
    #    _fg_sum = np.sum(labels == 1)
    #    _bg_sum = np.sum(labels == 0)
    #    _count = 1
    #    print 'rpn: num_positive avg', _fg_sum / _count
    #    print 'rpn: num_negative avg', _bg_sum / _count

    # resahpe
    label_list = list()
    bbox_target_list = list()
    bbox_weight_list = list()
    if landmark:
      landmark_target_list = list()
      landmark_weight_list = list()
    anchors_num_range = [0] + anchors_num_list
    label = {}
    for i in range(len(feat_strides)):
        stride = feat_strides[i]
        feat_height, feat_width = feat_infos[i]
        A = A_list[i]
        _label = labels[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        if select_stride>0 and stride!=select_stride:
          #print('set', stride, select_stride)
          _label[:] = -1
        #print('_label', _label.shape, select_stride)
        #_fg_inds = np.where(_label == 1)[0]
        #n_fg = len(_fg_inds)
        #STAT[0]+=1
        #STAT[stride]+=n_fg
        #if STAT[0]%100==0:
        #  print('rpn_stat', STAT, file=sys.stderr)
        bbox_target = bbox_targets[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        bbox_weight = bbox_weights[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        if landmark:
          landmark_target = landmark_targets[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
          landmark_weight = landmark_weights[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]

        _label = _label.reshape((1, feat_height, feat_width, A)).transpose(0, 3, 1, 2)
        _label = _label.reshape((1, A * feat_height * feat_width))
        bbox_target = bbox_target.reshape((1, feat_height*feat_width, A * bbox_pred_len)).transpose(0, 2, 1)
        bbox_weight = bbox_weight.reshape((1, feat_height*feat_width, A * bbox_pred_len)).transpose((0, 2, 1))
        label['%s_label_stride%d'%(prefix, stride)] = _label
        label['%s_bbox_target_stride%d'%(prefix,stride)] = bbox_target
        label['%s_bbox_weight_stride%d'%(prefix,stride)] = bbox_weight
        if landmark:
          landmark_target = landmark_target.reshape((1, feat_height*feat_width, A * landmark_pred_len)).transpose(0, 2, 1)
          landmark_weight = landmark_weight.reshape((1, feat_height*feat_width, A * landmark_pred_len)).transpose((0, 2, 1))
          label['%s_landmark_target_stride%d'%(prefix,stride)] = landmark_target
          label['%s_landmark_weight_stride%d'%(prefix,stride)] = landmark_weight
        #print('in_rpn', stride,_label.shape, bbox_target.shape, bbox_weight.shape, file=sys.stderr)
        label_list.append(_label)
        #print('DD', _label.shape)
        bbox_target_list.append(bbox_target)
        bbox_weight_list.append(bbox_weight)
        if landmark:
          landmark_target_list.append(landmark_target)
          landmark_weight_list.append(landmark_weight)

    label_concat = np.concatenate(label_list, axis=1)
    bbox_target_concat = np.concatenate(bbox_target_list, axis=2)
    bbox_weight_concat = np.concatenate(bbox_weight_list, axis=2)
    #fg_inds = np.where(label_concat[0] == 1)[0]
    #print('fg_inds_in_rpn2', fg_inds, file=sys.stderr)

    label.update({'%s_label'%prefix: label_concat,
            '%s_bbox_target'%prefix: bbox_target_concat,
            '%s_bbox_weight'%prefix: bbox_weight_concat}
            )
    if landmark:
      landmark_target_concat = np.concatenate(landmark_target_list, axis=2)
      landmark_weight_concat = np.concatenate(landmark_weight_list, axis=2)
      label['%s_landmark_target'%prefix] = landmark_target_concat
      label['%s_landmark_weight'%prefix] = landmark_weight_concat
    #tb = datetime.datetime.now()
    #self._times[3] += (tb-ta).total_seconds()
    #ta = datetime.datetime.now()
    #print(self._times)
    return label


''' Visualization code for point clouds and 3D bounding boxes with mayavi.

Modified by Charles R. Qi
Date: September 2017

Ref: https://github.com/hengck23/didi-udacity-2017/blob/master/baseline-04/kitti_data/draw.py

Modified by Haotian Tang
Date: August 2020
'''

import argparse
import numpy as np
import mayavi.mlab as mlab
import os

import torch
from torchsparse.utils import sparse_quantize, sparse_collate
from torchsparse import SparseTensor
from model_zoo import minkunet, spvcnn, spvnas_specialized

def process_point_cloud(input_point_cloud, input_labels=None, voxel_size=0.05):
    input_point_cloud[:, 3] = input_point_cloud[:, 3]
    pc_ = np.round(input_point_cloud[:, :3] / voxel_size)
    pc_ -= pc_.min(0, keepdims=1)

    label_map = create_label_map()
    if input_labels is not None:
        labels_ = label_map[input_labels & 0xFFFF].astype(
            np.int64)  # semantic labels
    else:
        labels_ = np.zeros(pc_.shape[0], dtype=np.int64)

    feat_ = input_point_cloud

    if input_labels is not None:
        out_pc = input_point_cloud[labels_ != labels_.max(), :3]
        pc_ = pc_[labels_ != labels_.max()]
        feat_ = feat_[labels_ != labels_.max()]
        labels_ = labels_[labels_ != labels_.max()]
    else:
        out_pc = input_point_cloud
        pc_ = pc_

    inds, labels, inverse_map = sparse_quantize(pc_,
                                                feat_,
                                                labels_,
                                                return_index=True,
                                                return_invs=True)
    pc = np.zeros((inds.shape[0], 4))
    pc[:, :3] = pc_[inds]

    feat = feat_[inds]
    labels = labels_[inds]
    lidar = SparseTensor(
        torch.from_numpy(feat).float(),
        torch.from_numpy(pc).int()
    )
    return {
        'pc': out_pc,
        'lidar': lidar,
        'targets': labels,
        'targets_mapped': labels_,
        'inverse_map': inverse_map
    }

mlab.options.offscreen = True

def create_label_map(num_classes=19):
    name_label_mapping = {
        'unlabeled': 0, 'outlier': 1, 'car': 10, 'bicycle': 11,
        'bus': 13, 'motorcycle': 15, 'on-rails': 16, 'truck': 18,
        'other-vehicle': 20, 'person': 30, 'bicyclist': 31,
        'motorcyclist': 32, 'road': 40, 'parking': 44,
        'sidewalk': 48, 'other-ground': 49, 'building': 50,
        'fence': 51, 'other-structure': 52, 'lane-marking': 60,
        'vegetation': 70, 'trunk': 71, 'terrain': 72, 'pole': 80,
        'traffic-sign': 81, 'other-object': 99, 'moving-car': 252,
        'moving-bicyclist': 253, 'moving-person': 254, 'moving-motorcyclist': 255,
        'moving-on-rails': 256, 'moving-bus': 257, 'moving-truck': 258,
        'moving-other-vehicle': 259
    }

    for k in name_label_mapping:
        name_label_mapping[k] = name_label_mapping[k.replace('moving-', '')]
    train_label_name_mapping = {
        0: 'car', 1: 'bicycle', 2: 'motorcycle', 3: 'truck', 4:
        'other-vehicle', 5: 'person', 6: 'bicyclist', 7: 'motorcyclist',
        8: 'road', 9: 'parking', 10: 'sidewalk', 11: 'other-ground',
        12: 'building', 13: 'fence', 14: 'vegetation', 15: 'trunk',
        16: 'terrain', 17: 'pole', 18: 'traffic-sign'
    }

    label_map = np.zeros(260)+num_classes
    for i in range(num_classes):
        cls_name = train_label_name_mapping[i]
        label_map[name_label_mapping[cls_name]] = min(num_classes,i)
    return label_map.astype(np.int64)

cmap = np.array([
    [245, 150, 100, 255],#[0,0,0] for teaser1
    [245, 230, 100, 255],
    [150, 60, 30, 255],
    [180, 30, 80, 255],
    [255, 0, 0, 255],
    [30, 30, 255, 255],#[255,0,0] for teaser1
    [200, 40, 255, 255],
    [90, 30, 150, 255],
    [255, 0, 255, 255],
    [255, 150, 255, 255],
    [75, 0, 75, 255],
    [75, 0, 175, 255],
    [0, 200, 255, 255],
    [50, 120, 255, 255],
    [0, 175, 0, 255],
    [0, 60, 135, 255],
    [80, 240, 150, 255],
    [150, 240, 255, 255],
    [0, 0, 255, 255],
    #[255, 255, 255, 0]
])

#cmap[:, -1] = 255

def draw_lidar(pc, color=None, fig=None, bgcolor=(1,1,1), pts_scale=0.06, pts_mode='2dcircle', pts_color=None):
    ''' Draw lidar points
    Args:
        pc: numpy array (n,3) of XYZ
        color: numpy array (n) of intensity or whatever
        fig: mayavi figure handler, if None create new one otherwise will use it
    Returns:
        fig: created or used fig
    '''
    if fig is None: fig = mlab.figure(figure=None, bgcolor=bgcolor, fgcolor=None, engine=None, size=(800, 500))
    if color is None: color = pc[:,2]
    pts = mlab.points3d(pc[:,0], pc[:,1], pc[:,2], color, mode=pts_mode, scale_factor=pts_scale, figure=fig)
    pts.glyph.scale_mode = 'scale_by_vector'
    pts.glyph.color_mode = 'color_by_scalar' # Color by scalar
    pts.module_manager.scalar_lut_manager.lut.table = cmap
    pts.module_manager.scalar_lut_manager.lut.number_of_colors = cmap.shape[0]

    mlab.view(azimuth=180, elevation=70, focalpoint=[ 12.0909996 , -1.04700089, -2.03249991], distance=62, figure=fig)

    return fig

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--velodyne-dir', type=str, default='sample_data')
    parser.add_argument('--model', type=str, default='SemanticKITTI_val_SPVNAS@65GMACs')
    args = parser.parse_args()
    output_dir = os.path.join(args.velodyne_dir, 'outputs')
    os.makedirs(output_dir, exist_ok=True)

    if torch.cuda.is_available():
        device = 'cuda:0'
    else:
        device = 'cpu'

    if 'MinkUNet' in args.model:
        model = minkunet(args.model, pretrained=True)
    elif 'SPVCNN' in args.model:
        model = spvcnn(args.model, pretrained=True)
    elif 'SPVNAS' in args.model:
        model = spvnas_specialized(args.model, pretrained=True)
    else:
        raise NotImplementedError

    model = model.to(device)

    input_point_clouds = sorted(os.listdir(args.velodyne_dir))
    for point_cloud_name in input_point_clouds:
        if not point_cloud_name.endswith('.bin'):
            continue
        label_file_name = point_cloud_name.replace('.bin', '.label')
        vis_file_name = point_cloud_name.replace('.bin', '.png')
        gt_file_name = point_cloud_name.replace('.bin', '_GT.png')

        pc = np.fromfile('%s/%s'%(args.velodyne_dir, point_cloud_name), dtype=np.float32).reshape(-1,4)
        if os.path.exists(label_file_name):
            label = np.fromfile('%s/%s'%(args.velodyne_dir, label_file_name), dtype=np.int32)
        else:
            label = None
        feed_dict = process_point_cloud(pc, label)
        inputs = feed_dict['lidar'].to(device)
        outputs = model(inputs)
        predictions = outputs.argmax(1).cpu().numpy()
        predictions = predictions[feed_dict['inverse_map']]
        fig = draw_lidar(feed_dict['pc'], predictions.astype(np.int32))
        mlab.savefig('%s/%s'%(output_dir, vis_file_name))
        if label is not None:
            fig = draw_lidar(feed_dict['pc'], feed_dict['targets_mapped'])
            mlab.savefig('%s/%s'%(output_dir, gt_file_name))


#!/usr/bin/env python2
# -*- coding: utf-8 -*-
from __future__ import division
import numpy as np
import serial
import cv2
import crcmod.predefined
import os
from struct import unpack

import rospy
from cv_bridge import CvBridge
from sensor_msgs.msg import Image
from std_msgs.msg import Float64, Float64MultiArray, MultiArrayLayout, MultiArrayDimension
from teraranger.cfg import EvoThermalConfig
from dynamic_reconfigure.server import Server

class EvoThermal(object):
    def __init__(self):
        # ROS INIT
        rospy.init_node("evo_thermal")
        self.rgb_publisher = rospy.Publisher("teraranger_evo_thermal/rgb_image", Image, queue_size=1)
        self.raw_publisher = rospy.Publisher("teraranger_evo_thermal/raw_temp_array", Float64MultiArray, queue_size=1)
        self.ptat_publisher = rospy.Publisher("teraranger_evo_thermal/ptat", Float64, queue_size=1)

        self.window_size = rospy.get_param("~window_size", 5)
        self.portname = rospy.get_param("~portname", "/dev/ttyACM0")
        self.baudrate = rospy.get_param("~baudrate", 115200)
        self.evo_thermal_frame = "evo_thermal_frame"
        self.bridge = CvBridge()

        # Init variables
        self.cmap_number = 0
        self.minima_list = []
        self.maxima_list = []

        # Get colormap from files
        colormap_files = (
            "resources/colormap/dave.txt",
            "resources/colormap/ice.txt",
            "resources/colormap/ironbow.txt",
            "resources/colormap/high_contrast.txt",
            "resources/colormap/whot.txt",
        )
        self.cmap_list = self.load_colormap_files(colormap_files)
        # We also append opencv colormaps
        self.cmap_list.append(cv2.COLORMAP_HOT)
        self.cmap_list.append(cv2.COLORMAP_JET)
        self.selected_cmap = self.cmap_list[self.cmap_number]

        # Converter frame param
        self.scaling = 10.0
        self.celsius_offset = 273.15

        # Initialize reconfiguration server
        self.evo_thermal_cfg_server = Server(EvoThermalConfig, self.evo_thermal_callback)

        # Reconfigure parameters initialization
        self.thermal_image_flip_v = False
        self.thermal_image_flip_h = False
        self.thermal_image_interpolate = False
        self.thermal_image_autoscale = True
        self.manual_min_scaling = 20.0
        self.manual_max_scaling = 30.0

        self.port = serial.Serial(
            port=self.portname,
            baudrate=self.baudrate,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            bytesize=serial.EIGHTBITS,
            timeout=1.0
        )

        self.port.isOpen()
        self.crc32 = crcmod.predefined.mkPredefinedCrcFun('crc-32-mpeg')
        self.crc8 = crcmod.predefined.mkPredefinedCrcFun('crc-8')

    def get_frame(self):
        got_frame = False
        while not got_frame:
            # Polls for header
            temp = self.port.read(2)
            got_header = False
            while not got_header:
                if len(temp) < 2:  # Check for serial timeout
                    return None
                if ord(temp[0]) == 13 and ord(temp[1]) == 0:
                    got_header = True
                else:
                    temp = temp[1:]
                    temp += self.port.read(1)

            # Header received, now read rest of frame
            data = self.port.read(2068)
            if len(data) < 2068:  # Check for serial timeout
                return None

            # Calculate CRC for frame (except CRC value and header)
            calculatedCRC = self.crc32(data[:2064])
            data = unpack("H" * 1034, str(data))
            receivedCRC = (data[1032] & 0xFFFF) << 16
            receivedCRC |= data[1033] & 0xFFFF
            self.ptat = data[1024]
            data = data[:1024]
            data = np.reshape(data, (32, 32))
            # Compare calculated CRC to received CRC
            if calculatedCRC == receivedCRC:
                got_frame = True
            else:
                rospy.logwarn("Bad CRC")
                return None

        self.port.flushInput()
        return data

    def convert_frame(self, input_frame):
        array = np.copy(input_frame)
        # Data is sent in dK, need to convert to celsius
        array = (array / self.scaling) - self.celsius_offset

        # Mirror frame if necessary
        if self.thermal_image_flip_h:
            array = np.flip(array, 0)  # flip the image horizontally
        if self.thermal_image_flip_v:
            array = np.flip(array, 1)  # flip the image vertically

        return array

    def auto_scaling(self, input_frame):
        data = np.copy(input_frame)

        # Get min/maxfor averaging
        frame_min, frame_max = data.min(), data.max()
        self.minima_list.append(frame_min)
        self.maxima_list.append(frame_max)

        # Need at least 10 frames for better average
        if len(self.maxima_list) >= 10:
            avg_max = sum(self.maxima_list) / len(self.maxima_list)
            avg_min = sum(self.minima_list) / len(self.minima_list)
            # Delete oldest insertions
            self.maxima_list.pop(0)
            self.minima_list.pop(0)
        else:
            # Until list fills, use current frame min/max/ptat
            avg_max = frame_max
            avg_min = frame_min

        # Scale based on boolean toggled by event #
        if self.thermal_image_autoscale:
            # Auto scale based on avg min/max values for last 10 frames
            data[data <= avg_min] = avg_min
            data[data >= avg_max] = avg_max
            multiplier = 255 / (avg_max - avg_min)
            data = data - avg_min
            data = data * multiplier
        else:
            # Manual scale based on user setting
            data[data <= self.manual_min_scaling] = self.manual_min_scaling
            data[data >= self.manual_max_scaling] = self.manual_max_scaling
            multiplier = 255 / (self.manual_max_scaling - self.manual_min_scaling)
            data = data - self.manual_min_scaling
            data = data * multiplier
        frame = data.astype(np.uint8)

        # Check if interpolation is toggled on
        if self.thermal_image_interpolate:
            # If on, interpolate and blur
            frame = cv2.resize(frame, (128, 128), interpolation=cv2.INTER_LINEAR)
            #frame = cv2.GaussianBlur(frame, (9, 9), 9)

        # Resize the frame
        frame = cv2.resize(frame, (512, 512), interpolation=cv2.INTER_NEAREST)

        # Apply selected colormap and stamp FPS
        frame = cv2.applyColorMap(frame, self.selected_cmap)

        return frame

    def load_colormap_files(self, files):
        # Get colormap from files
        # Created using http://jdherman.github.io/colormap/
        colormap_list = []
        for rel_file_path in files:
            full_path = os.path.join(os.path.dirname(__file__), os.pardir, rel_file_path)
            with open(full_path, 'r') as f:
                r = []
                g = []
                b = []
                for i in range(256):
                    x, y, z = f.readline().split(',')
                    r.append(x)
                    g.append(y)
                    b.append(z.replace(";\n", ""))
            colormap = np.zeros((256, 1, 3), dtype=np.uint8)
            # We use BGR because that's default for openCV
            colormap[:, 0, 0] = b
            colormap[:, 0, 1] = g
            colormap[:, 0, 2] = r
            colormap_list.append(colormap)

        return colormap_list

    def evo_thermal_callback(self, config, level):
        rospy.logdebug("Evo Thermal parameters reconfigure request".format(**config))
        if level == -1:
            return config
        if level == 0:
            self.thermal_image_flip_h = config["thermal_image_flip_h"]
            self.thermal_image_flip_v = config["thermal_image_flip_v"]
            self.thermal_image_interpolate = config["thermal_image_interpolate"]
        elif level == 1:
            self.thermal_image_autoscale = config["thermal_image_autoscale"]
        elif level == 2:
            if config["manual_max_scaling"] <= config["manual_min_scaling"]:
                config["manual_max_scaling"] = config["manual_min_scaling"] + 1
            if config["manual_max_scaling"] <= config["manual_min_scaling"]:
                config["manual_min_scaling"] = config["manual_max_scaling"] - 1
            self.manual_min_scaling = config["manual_min_scaling"]
            self.manual_max_scaling = config["manual_max_scaling"]
        elif level == 3:
            self.reconfigure_color_map(config)
        else:
            rospy.loginfo("Invalid reconfiguration level")
        return config

    def reconfigure_color_map(self, config):
        if config["Map"] == EvoThermalConfig.EvoThermal_Dave:
            self.selected_cmap = self.cmap_list[0]
            rospy.loginfo("Changed colormap to Dave")
        elif config["Map"] == EvoThermalConfig.EvoThermal_Ice:
            self.selected_cmap = self.cmap_list[1]
            rospy.loginfo("Changed colormap to Ice")
        elif config["Map"] == EvoThermalConfig.EvoThermal_Ironbow:
            self.selected_cmap = self.cmap_list[2]
            rospy.loginfo("Changed colormap to Ironbow")
        elif config["Map"] == EvoThermalConfig.EvoThermal_High_contrast:
            self.selected_cmap = self.cmap_list[3]
            rospy.loginfo("Changed colormap to High Contrast")
        elif config["Map"] == EvoThermalConfig.EvoThermal_Whot:
            self.selected_cmap = self.cmap_list[4]
            rospy.loginfo("Changed colormap to White Hot")
        else:
            rospy.logerr("Unknown colormap index")

    def publish(self, rgb, temp_array):
        # Publish RGB image
        rgb_img_msg = self.bridge.cv2_to_imgmsg(rgb)
        rgb_img_msg.header.frame_id = self.evo_thermal_frame
        rgb_img_msg.header.stamp = rospy.Time.now()
        self.rgb_publisher.publish(rgb_img_msg)

        # Publish temperatures array
        raw_temp_array = Float64MultiArray()
        raw_temp_array.layout.dim.append(MultiArrayDimension())
        raw_temp_array.layout.dim[0].size = 32
        raw_temp_array.layout.dim[0].stride = 1
        raw_temp_array.layout.dim[0].label = "x-axis"
        raw_temp_array.layout.dim.append(MultiArrayDimension())
        raw_temp_array.layout.dim[1].size = 32
        raw_temp_array.layout.dim[1].stride = 1
        raw_temp_array.layout.dim[1].label = "y-axis"
        data = np.reshape(temp_array, (1024))
        np.round(data, 2)
        raw_temp_array.data = data
        self.raw_publisher.publish(raw_temp_array)

        # Publish internal sensor temperature
        self.ptat_publisher.publish(self.ptat/self.scaling - self.celsius_offset)

    def send_command(self, command):
        self.port.flushInput()
        self.port.write(command)

        # This loop discards buffered frames until a valid ACK frame is reached
        temp_ack = self.port.read(4)
        got_ack = False
        while not got_ack:
            if len(temp_ack) < 4:  # Check for serial timeout
                rospy.logwarn("Timeout when reading ACK")
                return False
            if ord(temp_ack[0]) == 20:
                if self.crc8(temp_ack[:3]) == ord(temp_ack[3]):
                    got_ack = True
                    break
            temp_ack = temp_ack[1:]
            temp_ack += self.port.read(1)

        # Check if ACK or NACK
        if ord(temp_ack[2]) == 0:
            return True
        else:
            rospy.logerr("Command refused by device")
            return False

    def start_sensor(self):
        rospy.loginfo("Starting sensor output...")
        res = self.send_command("\x00\x52\x02\x01\xDF")
        if res:
            rospy.loginfo("Sensor output started successfully")
            return True
        else:
            rospy.logerr("Failed to start sensor output")
            return False

    def stop_sensor(self):
        rospy.loginfo("Stopping sensor output...")
        res = self.send_command("\x00\x52\x02\x00\xD8")
        if res:
            rospy.loginfo("Sensor output stopped successfully")
            return True
        else:
            rospy.logerr("Failed to stop sensor output")
            return False

    def run(self):
        self.port.flushInput()
        if self.baudrate == 115200:  # Sending VCP start when connected via USB
            if not self.start_sensor():
                return

        while not rospy.is_shutdown():
            frame = self.get_frame()
            if frame is not None:
                temp_array = self.convert_frame(frame)
                thermal_image = self.auto_scaling(temp_array)

                self.publish(thermal_image, temp_array)
        else:
            if self.baudrate == 115200:  # Sending VCP stop when connected via USB
                self.stop_sensor()
            rospy.logwarn("Node shutting down")

if __name__ == "__main__":
    evo_thermal = EvoThermal()
    try:
        evo_thermal.run()
    except rospy.ROSInterruptException:
        exit()


######################################################################################
## Code is based on a facenet implementation: https://github.com/sainimohit23/FaceNet-Real-Time-face-recognition
## Code was amended for the use of this project

import os
import cv2
import numpy as np
import pandas as pd
from parameters import *
import pickle
from pathlib import Path

image_dirs = Path(TARGET_IMAGES_DIR)
images_trn = pd.DataFrame(columns=['label', 'path', 'name', 'img'], dtype=object)

labels =[]
for image_dir in image_dirs.glob('*'):
    label = image_dir.stem
    labels.append(label)
    for file in image_dir.glob('*'+trn_img+'*'):
        basename = os.path.basename(file)
        f_name, f_ext = os.path.splitext(basename)
        if f_ext.lower() != ".png" or f_name[0:4] !="pre1": continue
        images_trn.loc[len(images_trn)] = [label, file, f_name, ""]
    if len(images_trn) >= MAX_FILES: break

paths={}
fins=[]
for i, label in enumerate(labels):
    paths[label] = TARGET_IMAGES_DIR.replace("\\", "/")+"/"+label
    fins.append(label)

input_shape = (3, IMAGE_SIZE, IMAGE_SIZE)

images = {}

for key in paths.keys():
    li = []
    for i, file in enumerate(images_trn[images_trn.label==key].itertuples()):
        img1 = cv2.imread(str(file.path))
        img2 = img1[...,::-1]
        li.append(np.around(np.transpose(img2, (2,0,1))/255.0, decimals=12))
    images[key] = np.array(li)

def batch_generator(batch_size=16):
    y_val = np.zeros((batch_size, 2, 1))
    anchors = np.zeros((batch_size, input_shape[0], input_shape[1], input_shape[2]))
    positives = np.zeros((batch_size, input_shape[0], input_shape[1], input_shape[2]))
    negatives = np.zeros((batch_size, input_shape[0], input_shape[1], input_shape[2]))

    while True:
        for i in range(batch_size):
            positivefin = fins[np.random.randint(len(fins))]
            negativefin = fins[np.random.randint(len(fins))]
            while positivefin == negativefin:
                negativefin = fins[np.random.randint(len(fins))]

            positives[i] = images[positivefin][np.random.randint(len(images[positivefin]))]
            anchors[i] = images[positivefin][np.random.randint(len(images[positivefin]))]
            negatives[i] = images[negativefin][np.random.randint(len(images[negativefin]))]

        x_data = {'anchor': anchors,
                  'anchorPositive': positives,
                  'anchorNegative': negatives
                  }

        yield (x_data, [y_val, y_val, y_val])


import numpy as np
import numpy.random as npr
import scipy.io as sio
import os, sys
import _init_paths
from utils.timer import Timer
# Make sure that caffe is on the python path:
caffe_root = './caffe-fast-rcnn/'
os.chdir(caffe_root)
sys.path.insert(0, os.path.join(caffe_root, 'python'))
import caffe
import cv2
import time

NumTest = 20000
TH=0.7
if __name__ == '__main__':

    caffe.set_mode_cpu()
    inputSize = 128
    mean = np.array([104, 117, 123])
#    classify_net = caffe.Net("/Users/momo/Desktop/Rain/rawmodels/hand_gesture_cls.prototxt",
#                             "/Users/momo/Desktop/Rain/rawmodels/hand_gesture_cls.caffemodel", caffe.TEST)
    classify_net = caffe.Net("/Users/momo/Desktop/gesture/fromAli/cls/cls_bareness.prototxt",
                            "/Users/momo/Desktop/gesture/fromAli/cls/ba_lmdb_b512_iter_40000.caffemodel", caffe.TEST)

    fid = open("/Users/momo/wkspace/Data/gesture/tests/128obj4test_0920RS13x10_5/pos_128obj4test_0920RS13x10_5.txt","r")
    subdirlists = ['bg', 'heart', 'yearh', 'one', 'baoquan', 'five', 'bainian', 'zan', 'fingerheart', 'ok', 'call', 'rock', 'big_v']
    tp_dict = {}
    gt_dict = {}
    re_dict = {}
    for gname in subdirlists:
        tp_dict[gname] = 0
        gt_dict[gname] = 0
        re_dict[gname] = 0
    TP=0
    err = 0
    lines = fid.readlines()
    fid.close()
    cur_=0
    sum_=len(lines)
    regloss = np.array([])
    probs = np.array([])
    cls = 0
    totalTime = 0
    for line in lines:
        cur_+=1
        if not line or cur_ == NumTest:
            break;
        words = line.split()
        image_file_name = "/Users/momo/wkspace/Data/gesture/tests/" + words[0]
#        print words, cur_

        if cur_%500 == 0:
            print cur_,
            sys.stdout.flush()
        im = cv2.imread(image_file_name)
        h,w,ch = im.shape
        if h!=inputSize or w!=inputSize:
            im = cv2.resize(im,(int(inputSize),int(inputSize)))

        im = im.astype(np.int)
        im -= mean
        im = np.swapaxes(im, 0, 2)
        im = np.swapaxes(im, 1, 2)
        label    = int(words[1])
        gt_dict[subdirlists[label]] += 1

        classify_net.blobs['data'].reshape(1,3,inputSize,inputSize)
        classify_net.blobs['data'].data[...]=im

        startT = time.clock()
        out_ = classify_net.forward()
        endT = time.clock()
        totalTime += (endT-startT)
        prob = out_['loss'][0]
#        print prob
#        cls = np.where(np.max(prob))
#        print np.max(prob), np.where(prob==np.max(prob))[0][0]
        cls_prob = np.max(prob)
        cls = np.where(prob==np.max(prob))[0][0]
        re_dict[subdirlists[cls]] += 1
        if not cls == label:
            err += 1
#            print words, subdirlists[label], cls, subdirlists[cls], cls_prob
        else:
            tp_dict[subdirlists[cls]]+=1
    print err, sum_
    print 'tp:', tp_dict
    print 're', re_dict
    print 'gt', gt_dict
    reTotal = 0
    gtTotal = 0
    tpTotal = 0
    for gname in tp_dict:
        print "%12s"%(gname)," recall:%.2f"%( float(tp_dict[gname])/float(gt_dict[gname]+1) ), " precision:%.2f"%( float(tp_dict[gname])/float(re_dict[gname]+1) )
        reTotal += re_dict[gname]
        gtTotal += gt_dict[gname]
        tpTotal += tp_dict[gname]
    print "total recall:%.2f"%(float(tpTotal)/float(gtTotal)), "total precision:%.2f"%(float(tpTotal)/float(reTotal))

#    print "TP/total:", float(TP)/float(cur_),TP, cur_,
#    print "reg loss mean=", np.mean(regloss),"reg loss std=", np.std(regloss),"time:", totalTime*1000/cur_, "ms"
#    print "porb mean=", np.mean(probs),"prob std=", np.std(probs)


#coding:utf-8
import numpy as np
import os
import sys
import cv2
import csv
caffe_root = '/home/wangchao/ali-tianchi/'
sys.path.append(caffe_root+'python')

#导入进度条
from tqdm import tqdm

import caffe
#caffe.set_device(0)
#caffe.set_mode_gpu()

model_def = 'deploy.prototxt'
model_weights = 'resnet101_solver_adam_iter_90000.caffemodel'
net = caffe.Net(model_def, 'models/'+model_weights, caffe.TEST)
#net.blobs['data'].reshape(1,3,h,w) # (batch_size,c,h,w)
net.blobs['data'].reshape(1,3,224,224) # (batch_size,c,h,w)

transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
transformer.set_transpose('data', (2,0,1))  # (h,w,c)--->(c,h,w)
transformer.set_mean('data', np.array([104,117,123])) #注意是 BGR

transformer.set_raw_scale('data', 255)      # rescale from [0, 1] to [0, 255]
transformer.set_channel_swap('data', (2,1,0))# RGB--->BGR

'''
可能你对均值设置成BGR有疑问，不是到后面才把RGB转为BGR吗?
其实transformer.setXXX()这些只是设置属性，实际执行顺序是参考附录preprocess函数定义
'''

with open("test.txt","r") as f:
    image_names = f.readlines()

label_warp = {'正常': 0,
              '不导电': 1,
              '擦花': 2,
              '横条压凹': 3,
              '桔皮': 4,
              '漏底': 5,
              '碰伤': 6,
              '起坑': 7,
              '凸粉': 8,
              '涂层开裂': 9,
              '脏点': 10,
              '其他': 11,
              }
submit = ['norm','defect1','defect2','defect3','defect4','defect5','defect6','defect7','defect8','defect9','defect10','defect11']

cnt = 0
pbar = tqdm(total=440)

with open('submit.txt','w') as f:
  for image_name in image_names:
      pbar.update(1)
      cnt = cnt+1
      image_name = image_name.strip()
      image = caffe.io.load_image(image_name)# 用的是skimage库，见附录
      # 利用刚刚的设置进行图片预处理
      transformed_image = transformer.preprocess('data', image)
      net.blobs['data'].data[...] = transformed_image
      # 网络前传（测试无后传）
      output = net.forward()

      output_prob= output['prob'][0].argmax()  # 概率最大的label
      #image_name,output_prob)#
      #print (output_prob)
      f.write(image_name.split("/")[-1]+',' + submit[output_prob]+'\n')

os.rename('submit.txt', 'submit.csv')

pbar.close()

'''
标签对照表
正常-norm 不导电-defect1 擦花-defect2 横条压凹-defect3 桔皮-defect4 漏底-defect5 碰伤-defect6
起坑-defect7 凸粉-defect8 涂层开裂-defect9 脏点-defect10 其他-defect11
name_list = ['zhengchang','budaodian', 'cahua', 'hengtiaoyaao', 'jiepi', 'loudi', 'pengshang', 'qikeng', 'tufen', 'tucengkailie', 'zangdian', 'qita']
            ['norm','defect1','defect2','defect3','defect4','defect5','defect6','defect7','defect8','defect9','defect10','defect11',]

'''

'''
正常	norm
不导电	defect1
擦花	defect2
横条压凹	defect3
桔皮	defect4
漏底	defect5
碰伤	defect6
起坑	defect7
凸粉	defect8
涂层开裂	defect9
脏点	defect10
其他	defect11
'''


import os
import numpy as np
from moviepy.editor import *

fps = 12
time_length = 45
worker_id = int(os.environ['WORKER_ID'])
if os.environ.get('FPS') is not None:
    fps = int(os.environ['FPS'])
# max time
base_grid_size=(4, 8)

grid_sizes = [
    (4, 8),
    (4, 8),
    (4, 8),
    (5, 10),
]

downscale_factor = 2

video_width = 256
scaled_video_width = video_width // downscale_factor

target_resolutions = [(grid_size[0] * scaled_video_width,
                       grid_size[1] * scaled_video_width)
                      for grid_size in grid_sizes]

classes = ['frosh', 'smore', 'junior', 'senior']
my_class = classes[worker_id % 4]

data_dir = 'data'

time_id = str((worker_id % 4) * time_length) + '_' + str((worker_id % 4) * time_length + time_length)
result_video_dir = os.path.join(data_dir, my_class + '_videos')
all_video_names = [video_name for video_name in os.listdir(result_video_dir)
                   if '.mp4' in video_name]
my_video_names = []
for video_name in all_video_names:
    detected_fps = int(video_name.split('_')[4])
    detected_time_id = '_'.join(video_name.split('_')[2:4])
    if fps == detected_fps and time_id == detected_time_id:
        my_video_names.append(video_name)

"""
- Initialize big boy video
- Resize videos to (128, 128), then slot them in
- Pad them appropriately (so have horizontal offset)
- Get audio from official source
- Profit
"""

grid_size = grid_sizes[worker_id % 4]
target_resolution = target_resolutions[worker_id % 4]

movie_tensor = np.zeros((fps * time_length, target_resolution[0],
                  target_resolution[1], 3))

for i, video_name in enumerate(my_video_names):
    video_path = os.path.join(result_video_dir, video_name)
    clip = VideoFileClip(video_path)
    clip = clip.resize(width=clip.size[0] // downscale_factor)
    w, h = clip.size
    grid_y = i // grid_size[1]
    grid_x = i % grid_size[1]
    for j, frame in enumerate(clip.iter_frames()):
        if j == 0:
            print("frame shape:", frame.shape)
        if j >= fps * time_length:
            break
        y_start = grid_y * scaled_video_width
        x_start = grid_x * scaled_video_width
        movie_tensor[j, y_start:y_start + h,
                     x_start:x_start+w] = frame

numpy_backup_name = 'movie_tensor_' + str(worker_id % 4) + '.npy'

movie_name = 'still_alive_movie_' + str(worker_id % 4) + '.mp4'

# with open(os.path.join(data_dir, numpy_backup_name), 'wb+') as f:
#     np.save(f, movie_tensor)

movie = ImageSequenceClip(list(movie_tensor), fps=fps)
cropped_audio = AudioFileClip(os.path.join(
    data_dir, 'still_alive_good_cropped.ogg'
))

print('durations:', movie.duration, cropped_audio.duration)

cropped_audio = cropped_audio.subclip(t_start=(worker_id % 4) * time_length,
    t_end=min((worker_id % 4) * time_length + time_length, cropped_audio.duration))
movie = movie.set_audio(cropped_audio)
movie = movie.resize(width=base_grid_size[1] * scaled_video_width)
print(movie_name)
movie.write_videofile(os.path.join(data_dir, movie_name), audio_codec='aac')


import gym
from gym import error,spaces,utils
from gym.utils import seeding
import os
import pybullet as p
import pybullet_data
import math
import numpy as np
from time import sleep
from pprint import pprint
import random

p.connect(p.GUI)
p.resetSimulation()
p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 0)
p.setGravity(0, 0, -10)

p.setAdditionalSearchPath(pybullet_data.getDataPath())
planeUid = p.loadURDF("plane/plane.urdf")
tiago = p.loadURDF("tiago_dualhand/tiago_dual_modified.urdf", basePosition=[0, 0, 0.35], useFixedBase=False)
# tiago = p.loadURDF("PR2/pr2_no_torso_lift_tall.urdf", basePosition=[0, 0, 0.35], useFixedBase=False)
joint_num = p.getNumJoints(tiago)
prismatic_joints_indices = [i for i in range(joint_num) if p.getJointInfo(tiago, i)[2] == p.JOINT_PRISMATIC]
revolute_joints_indices = [i for i in range(joint_num) if p.getJointInfo(tiago, i)[2] == p.JOINT_REVOLUTE]
fixed_joints_indices = [i for i in range(joint_num) if p.getJointInfo(tiago, i)[2] == p.JOINT_FIXED]
p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 1)
# move with two wheels
ii = 0
while 1:
    # wheel = [12,14,16,18] # control the turnings of side wheels
    # wheel = [13,15,17,19] # side wheels forward
    wheel = [9,11]
    radius = 0.0985
    cylinder = 0.04
    currentIndex = wheel
    rdm = [random.uniform(0, 1)] * 3
    # for ii in range(len(wheel)):
    # p.resetBaseVelocity(tiago,linearVelocity=[1,0,0])
    # p.changeConstraint()
    # p.resetBasePositionAndOrientation(tiago, (0+ii, 0+ii, 0.65), (0.0, 0.0, 0.0, 1.0))
    p.stepSimulation()

    state_wheel_l = p.getLinkState(tiago, 11, computeLinkVelocity=1)
    state_wheel_r = p.getLinkState(tiago, 9, computeLinkVelocity=1)
    position_wheel_left = np.array(state_wheel_l[0])
    position_wheel_right = np.array(state_wheel_r[0])
    linVelo_wheel = 5
    angVelo_wheel = 0
    # linVelo_wheel = (np.array(state_wheel_l[6]) + np.array(state_wheel_r[6])) / 2
    # angVelo_wheel = (np.array(state_wheel_l[7]) + np.array(state_wheel_r[7])) / 2
    wheel_axle_half = 0.5 * (np.linalg.norm(position_wheel_right - position_wheel_left) - cylinder)
    left_wheel_wish = (linVelo_wheel - angVelo_wheel * wheel_axle_half) / radius
    right_wheel_wish = (linVelo_wheel + angVelo_wheel * wheel_axle_half) / radius
    scaling = 1.113
    interval = 1/240
    p.setJointMotorControlArray(bodyIndex=tiago, jointIndices=currentIndex,
                                 controlMode=p.VELOCITY_CONTROL,
                                 targetVelocities= #[linVelo_wheel]*len(wheel),
                                [scaling * right_wheel_wish, scaling * left_wheel_wish],
                                 velocityGains=[1]*len(wheel), forces=[500]*len(wheel))
    # ii += interval * 0.2
    print(currentIndex, 'is moving')
    print('linear velocity is',np.array(state_wheel_l[6]))
    sleep(interval)


import numpy as np
import pandas as pd
import soundfile as sf
import torch
from torch.utils.data.dataset import Dataset
import random as random
import os

class LibriMix(Dataset):
    """" Dataset class for Librimix source separation tasks.

    Args:
        csv_dir (str): The path to the metatdata file
        task (str): One of ``'enh_single'``, ``'enh_both'``, ``'sep_clean'`` or
            ``'sep_noisy'``.

            * ``'enh_single'`` for single speaker speech enhancement.
            * ``'enh_both'`` for multi speaker speech enhancement.
            * ``'sep_clean'`` for two-speaker clean source separation.
            * ``'sep_noisy'`` for two-speaker noisy source separation.

        sample_rate (int) : The sample rate of the sources and mixtures
        n_src (int) : The number of sources in the mixture
        segment (int) : The desired sources and mixtures length in s
    """
    dataset_name = 'LibriMix'

    def __init__(self, csv_dir, task, sample_rate=16000,
                 n_src=2, segment=4):
        self.csv_dir = csv_dir
        self.task = task
        # Get the csv corresponding to the task
        if task == 'enh_single':
            md_file = [f for f in os.listdir(csv_dir) if 'single' in f][0]
            self.csv_path = os.path.join(self.csv_dir, md_file)
        elif task == 'enh_both':
            md_file = [f for f in os.listdir(csv_dir) if 'both' in f][0]
            self.csv_path = os.path.join(self.csv_dir, md_file)
            md_clean_file = [f for f in os.listdir(csv_dir) if 'clean' in f][0]
            self.df_clean = pd.read_csv(os.path.join(csv_dir, md_clean_file))
        elif task == 'sep_clean':
            md_file = [f for f in os.listdir(csv_dir) if 'clean' in f][0]
            self.csv_path = os.path.join(self.csv_dir, md_file)
        elif task == 'sep_noisy':
            md_file = [f for f in os.listdir(csv_dir) if 'both' in f][0]
            self.csv_path = os.path.join(self.csv_dir, md_file)
        self.segment = segment
        self.sample_rate = sample_rate
        # Open csv file
        self.df = pd.read_csv(self.csv_path)
        # Get rid of the utterances too short
        if self.segment is not None:
            max_len = len(self.df)
            self.seg_len = int(self.segment * self.sample_rate)
            # Ignore the file shorter than the desired_length
            self.df = self.df[self.df['length'] >= self.seg_len]
            print(f"Drop {max_len - len(self.df)} utterances from {max_len} "
                  f"(shorter than {segment} seconds)")
        else:
            self.seg_len = None
        self.n_src = n_src

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        # Get the row in dataframe
        row = self.df.iloc[idx]
        # Get mixture path
        self.mixture_path = row['mixture_path']
        sources_list = []
        # If there is a seg start point is set randomly
        if self.seg_len is not None:
            start = random.randint(0, row['length'] - self.seg_len)
            stop = start + self.seg_len
        else:
            start = 0
            stop = None
        # If task is enh_both then the source is the clean mixture
        if 'enh_both' in self.task:
            mix_clean_path = self.df_clean.iloc[idx]['mixture_path']
            s, _ = sf.read(mix_clean_path, dtype='float32', start=start,
                           stop=stop)
            sources_list.append(s)

        else:
            # Read sources
            for i in range(self.n_src):
                source_path = row[f'source_{i + 1}_path']
                s, _ = sf.read(source_path, dtype='float32', start=start,
                               stop=stop)
                sources_list.append(s)
        # Read the mixture
        mixture, _ = sf.read(self.mixture_path, dtype='float32', start=start,
                             stop=stop)
        # Convert to torch tensor
        mixture = torch.from_numpy(mixture)
        # Stack sources
        sources = np.vstack(sources_list)
        # Convert sources to tensor
        sources = torch.from_numpy(sources)
        return mixture, sources


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import pandas
import numpy as np
import warnings, multiprocessing
from .diffdata import DiffData

def run_forcephotometry(
    target,
    download_files=True,
    nprocess=4,
    verbose=True,
    filecheck=True,
    ignore_warnings=False,
    **kwargs):
    """ """
    fp = ForcePhotometry.from_name(target)
    fp.load_metadata()
    if download_files:
        fp.io.download_data(nprocess=nprocess)
    fp.load_filepathes(filecheck=filecheck, ignore_warnings=ignore_warnings)
    fp.run_forcefit(update_diffdata=False, verbose=verbose, **kwargs)

class ForcePhotometry:
    """ """

    # =============== #
    #    Init         #
    # =============== #
    def __init__(self, ztftarget):
        """ """
        self.io = ztftarget

    @classmethod
    def from_name(cls, name):
        """ Load the ForcePhotometry from Target Name """
        from .io import ZTFTarget

        return cls(ZTFTarget.from_name(name))

    @classmethod
    def from_coords(cls, ra, dec, jdmin=None, jdmax=None, name="unknown"):
        """ Load the ForcePhotometry from Target Name """
        from .io import ZTFTarget

        return cls(ZTFTarget.from_coords(ra, dec, jdmin=jdmin, jdmax=jdmax, name=name))

    # =============== #
    #   Method        #
    # =============== #
    # -------- #
    # LOADDER  #
    # -------- #
    def load_metadata(self, **kwargs):
        """ """
        if not hasattr(self.io, "marshal"):
            self.io.load_marshal(**kwargs)
        self.io.load_metadata()

    def load_filepathes(self, **kwargs):
        """
        **kwargs eventually goes to ztfquery.Query.get_local_data()
                 -> exists=True, indexes=None, filecheck=True, ignore_warnings=False, etc.

        """
        self._filepathes = self.io.get_diffimg_forcepsf_filepath(**kwargs)
        self._diffdata = None

    def store(self, filename=None):
        """ """
        if filename is None:
            from .io import LOCALDATA

            filename = LOCALDATA + "/%s.csv" % self.io.name
        self._data_forcefit.to_csv(filename, index=False)

    # -------- #
    # FITTER   #
    # -------- #
    def run_forcefit(self, indexes=None, update_diffdata=False,
                         store=False, verbose=True, no_badsub=False,
                         force_refit=False, nprocess=4):
        """ """
        import gc

        if indexes is None:
            indexes = np.arange(self.nsources)

        from .io import LOCALDATA

        # Try to load previous fit results
        try:
            path_previous_results = os.path.join(
                LOCALDATA, "{}.csv".format(self.io.name)
            )
            previous_results = pandas.read_csv(
                path_previous_results, float_precision="round_trip"
            )
        except FileNotFoundError:
            previous_results = pandas.DataFrame(columns=["filename"])

        dataout = {}
        if verbose:
            print("Starting run_forcefit() for %d image differences" % len(indexes))

        from astropy.utils.console import ProgressBar

        index_count = len(indexes)
        _coords = [self.io.get_coordinate()] * index_count
        _update_diffdata = [update_diffdata] * index_count
        _no_badsub = [no_badsub] * index_count
        _verbose = [verbose] * index_count
        _previous_results = [previous_results] * index_count
        _force_refit = [force_refit] * index_count

        if nprocess > 1:
            dataout_df = pandas.DataFrame(
                columns=[
                    "sigma",
                    "sigma.err",
                    "ampl",
                    "ampl.err",
                    "fval",
                    "chi2",
                    "chi2dof",
                    "filename",
                    "humidity",
                    "filter",
                    "obsmjd",
                    "ccdid",
                    "amp_id",
                    "gain",
                    "readnoi",
                    "darkcur",
                    "magzp",
                    "magzpunc",
                    "magzprms",
                    "clrcoeff",
                    "clrcounc",
                    "zpclrcov",
                    "zpmed",
                    "zpavg",
                    "zprmsall",
                    "clrmed",
                    "clravg",
                    "clrrms",
                    "qid",
                    "rcid",
                    "seeing",
                    "maglim",
                    "status",
                    "filterid",
                    "fieldid",
                    "moonalt",
                    "moonillf",
                    "target_x",
                    "target_y",
                    "data_hasnan",
                ]
            )
            bar = ProgressBar(index_count)
            dataout_list = []

            with multiprocessing.Pool(nprocess) as p:
                for j, dataout_i in enumerate(
                    p.imap_unordered(
                        self.get_ith_diffdata_multiprocess,
                        zip(
                            indexes,
                            self.filepathes,
                            _coords,
                            _update_diffdata,
                            _no_badsub,
                            _verbose,
                            _previous_results,
                            _force_refit,
                        ),
                    )
                ):
                    if bar is not None:
                        bar.update(j)
                    dataout_list.append(dataout_i)
                if bar is not None:
                    bar.update(index_count)

            for dataout_i in dataout_list:
                if dataout_i is not None:
                    dataout_i = pandas.Series(dataout_i, index=dataout_df.columns)
                    dataout_df = dataout_df.append(dataout_i, ignore_index=True)
            self._data_forcefit = dataout_df
            if store:
                self.store()

        else:
            for i in indexes:
                if verbose:
                    print("running %d " % i)
                    print(self.filepathes[i][0].split("/")[-1])

                filename_modified = self.filepathes[i][0].split("/")[-1][:-26] + ".fits"
                query = previous_results.query("filename == @filename_modified")

                # Check query of dataframe. If it is 40 characters long, this datapoint has already been fitted
                # and will not be fitted again (but it WILL be fitted again if force_refit is True)

                # Parse the query to a dictionary to check if it is 40 characters long or not
                if len(query) == 1:
                    query = query.to_dict("r")[0]

                # Now do the actual check
                if len(query) == 40 and not force_refit:
                    if verbose:
                        print("not fitting %d " % i)
                    dataout[i] = query

                else:
                    if verbose:
                        print("fitting %d " % i)
                    diffdata = self.get_ith_diffdata(i, update=update_diffdata)
                    has_nan = np.any(np.isnan(diffdata.diffimg))
                    if has_nan and no_badsub:
                        print("NaNs in the image, skipped")
                    else:
                        try:
                            fitresults = diffdata.fit_flux()
                            datainfo = diffdata.get_main_info()
                            dataout[i] = {**fitresults, **datainfo}
                            dataout[i]["data_hasnan"] = has_nan
                        except ValueError:
                            warnings.warn(
                                "Shape of diffimg and psfimg do not correspond (index: %d). Skipping."
                                % (i)
                            )
                            pass
                    del diffdata
                gc.collect()

            self._data_forcefit = pandas.DataFrame(dataout).T
            if store:
                self.store()

    def get_ith_diffdata(self, index, update=False, rebuild=False, **kwargs):
        """ loads and returns a DiffData object corresponding
        at the ith-entry of the self.filepathes

        Parameters
        ----------
        index: [int]
            which entry iof self.filepathes/self.diffdata do you want ?

        update: [bool] -optional-
            if the DiffData is derived (because unknown or rebuild), should this
            be updated to self.diffdata

        rebuild: [bool] -optional-
            Even if self.DiffData[index] exists, shall this remeasure it ?

        Returns
        -------
        DiffData
        """
        if self.diffdata is None or self.diffdata[index] is None or rebuild:
            diffdata = DiffData(
                *self.filepathes[index], self.io.get_coordinate(), **kwargs
            )
            if update:
                self.diffdata[index] = diffdata
            return diffdata

        return self.diffdata[index]

    @staticmethod
    def get_ith_diffdata_multiprocess(args):
        import gc

        (
            index,
            filepath,
            coords,
            update_diffdata,
            no_badsub,
            verbose,
            previous_results,
            force_refit,
        ) = args

        filename_modified = filepath[0].split("/")[-1][:-26] + ".fits"
        query = previous_results.query("filename == @filename_modified")

        # Check query of dataframe. If it is 40 characters long, this datapoint has already been fitted
        # and will not be fitted again (but it WILL be fitted again if force_refit is True)

        # Parse the query to a dictionary to check if it is 40 characters long or not
        if len(query) == 1:
            query = query.to_dict("r")[0]

        # Now do the actual check
        if len(query) == 40 and not force_refit:
            if verbose:
                print("not fitting %d " % index)
            dataout = query

        else:
            if verbose:
                print("fitting %d " % index)
                print(filepath[0].split("/")[-1])
            diffdata = DiffData(*filepath, coords)
            has_nan = np.any(np.isnan(diffdata.diffimg))
            if has_nan and no_badsub:
                print("NaNs in the image, skipped")
                return None
            else:
                try:
                    fitresults = diffdata.fit_flux()
                    datainfo = diffdata.get_main_info()
                    dataout = {**fitresults, **datainfo}
                    dataout["data_hasnan"] = has_nan
                except ValueError:
                    warnings.warn(
                        "Shape of diffimg and psfimg do not correspond (index: %d). Skipping."
                        % (index)
                    )
                    return None
            del diffdata

        gc.collect()
        return dataout

    # --------- #
    #  PLOTTER  #
    # --------- #
    def show_lc(self, ax=None, scalezp=25, reference=False, **kwargs):
        """ """
        # - Figure
        if ax is None:
            import matplotlib.pyplot as mpl

            fig = mpl.figure()
            ax = fig.add_subplot(111)
        else:
            fig = ax.figure

        x, y, err = self.data_forcefit[["obsmjd", "ampl", "ampl.err"]].values.T
        f0coef = (
            10 ** (-(self.data_forcefit["magzp"] - scalezp) / 2.5)
            if scalezp is not None
            else 1
        )

        for i, band_ in [[1, "C0"], [2, "C2"], [3, "C1"]]:
            if i not in self.data_forcefit["filterid"]:
                continue
            flag = np.asarray(self.data_forcefit["filterid"] == i, dtype="bool")

            prop = {**dict(marker="o", color=band_, zorder=5), **kwargs}
            ax.scatter(x[flag], (y * f0coef)[flag], **prop)
            prop["zorder"] = prop.get("zorder", 5) - 1
            ax.errorbar(
                x[flag],
                (y * f0coef)[flag],
                yerr=err[flag],
                ls="None",
                ecolor="0.6",
                **prop
            )

        ax.axhline(0, ls="--", color="0.5")
        return ax

    # =============== #
    #  Properties     #
    # =============== #
    @property
    def data_forcefit(self):
        """ DataFrame containing the forcefit results and data info """
        if not hasattr(self, "_data_forcefit"):
            raise AttributeError("data_forcefit not loaded. run self.run_forcefit()")
        return self._data_forcefit

    @property
    def nsources(self):
        """ """
        return len(self.filepathes)

    @property
    def metadata(self):
        """ """
        return self.io.zquery.metatable

    @property
    def filepathes(self):
        """ List of DiffImage and PSFDiffImage file path """
        if not hasattr(self, "_filepathes"):
            self.load_filepathes()
        return self._filepathes

    def had_filepathes(self):
        """ """
        return hasattr(self, "_filepathes") and self._filepathes is not None

    @property
    def diffdata(self):
        """ DiffData object associated with the filepathes"""
        if not hasattr(self, "_diffdata") or self._diffdata is None:
            if self.had_filepathes():
                self._diffdata = [None for i in range(self.nsources)]
            else:
                self._diffdata = None

        return self._diffdata