# coding:utf-8

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable,Function
import numpy as np
import torchvision.models as models
from model.utils.config import cfg
from model.rpn.bbox_transform import bbox_overlaps_batch, bbox_overlaps
from model.roi_layers import nms
#from model.roi_crop.functions.roi_crop import RoICropFunction
import cv2
import pdb
import random
from torch.utils.data.sampler import Sampler
from sklearn.cluster import KMeans
import shutil
import os
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
plt.style.use('seaborn-whitegrid')
plt.rc('font',family='Times New Roman',size=12)
from sklearn.manifold import TSNE
def to_np_array(x):
    if x is None:
        return None
    # if isinstance(x, Variable): x = x.data
    return x.cpu().data.numpy() if torch.is_tensor(x) else np.array(x)

class sampler(Sampler):
  def __init__(self, train_size, batch_size):
    self.num_data = train_size
    self.num_per_batch = int(train_size / batch_size)
    self.batch_size = batch_size
    self.range = torch.arange(0,batch_size).view(1, batch_size).long()
    self.leftover_flag = False
    if train_size % batch_size:
      self.leftover = torch.arange(self.num_per_batch*batch_size, train_size).long()
      self.leftover_flag = True

  def __iter__(self):
    rand_num = torch.randperm(self.num_per_batch).view(-1,1) * self.batch_size
    self.rand_num = rand_num.expand(self.num_per_batch, self.batch_size) + self.range

    self.rand_num_view = self.rand_num.view(-1)

    if self.leftover_flag:
      self.rand_num_view = torch.cat((self.rand_num_view, self.leftover),0)

    return iter(self.rand_num_view)

  def __len__(self):
    return self.num_data

class EFocalLoss(nn.Module):
    r"""
        This criterion is a implemenation of Focal Loss, which is proposed in
        Focal Loss for Dense Object Detection.

            Loss(x, class) = - \alpha (1-softmax(x)[class])^gamma \log(softmax(x)[class])

        The losses are averaged across observations for each minibatch.
        Args:
            alpha(1D Tensor, Variable) : the scalar factor for this criterion
            gamma(float, double) : gamma > 0; reduces the relative loss for well-classiﬁed examples (p > .5),
                                   putting more focus on hard, misclassiﬁed examples
            size_average(bool): size_average(bool): By default, the losses are averaged over observations for each minibatch.
                                However, if the field size_average is set to False, the losses are
                                instead summed for each minibatch.
    """

    def __init__(self, class_num, alpha=None, gamma=2, size_average=True):
        super(EFocalLoss, self).__init__()
        if alpha is None:
            self.alpha = Variable(torch.ones(class_num, 1) * 1.0)
        else:
            if isinstance(alpha, Variable):
                self.alpha = alpha
            else:
                self.alpha = Variable(alpha)
        self.gamma = gamma
        self.class_num = class_num
        self.size_average = size_average

    def forward(self, inputs, targets):
        N = inputs.size(0)
        # print(N)
        C = inputs.size(1)
        # inputs = F.sigmoid(inputs)
        P = torch.softmax(inputs, dim=-1)
        class_mask = inputs.data.new(N, C).fill_(0)
        class_mask = Variable(class_mask)
        ids = targets.view(-1, 1)
        class_mask.scatter_(1, ids.data, 1.)
        # print(class_mask)

        if inputs.is_cuda and not self.alpha.is_cuda:
            self.alpha = self.alpha.cuda()
        alpha = self.alpha[ids.data.view(-1)]

        probs = (P * class_mask).sum(1).view(-1, 1)
        log_p = probs.log()
        # print('probs size= {}'.format(probs.size()))
        # print(probs)
        batch_loss = -alpha * torch.exp(-self.gamma * probs) * log_p
        # print('-----bacth_loss------')
        # print(batch_loss)

        if self.size_average:
            loss = batch_loss.mean()
        else:
            loss = batch_loss.sum()
        return loss
class FocalLoss(nn.Module):
    r"""
        This criterion is a implemenation of Focal Loss, which is proposed in
        Focal Loss for Dense Object Detection.

            Loss(x, class) = - \alpha (1-softmax(x)[class])^gamma \log(softmax(x)[class])

        The losses are averaged across observations for each minibatch.
        Args:
            alpha(1D Tensor, Variable) : the scalar factor for this criterion
            gamma(float, double) : gamma > 0; reduces the relative loss for well-classiﬁed examples (p > .5),
                                   putting more focus on hard, misclassiﬁed examples
            size_average(bool): size_average(bool): By default, the losses are averaged over observations for each minibatch.
                                However, if the field size_average is set to False, the losses are
                                instead summed for each minibatch.
    """

    def __init__(self, class_num, alpha=None, gamma=2, size_average=True,sigmoid=False,reduce=True):
        super(FocalLoss, self).__init__()
        if alpha is None:
            self.alpha = Variable(torch.ones(class_num, 1) * 1.0)
        else:
            if isinstance(alpha, Variable):
                self.alpha = alpha
            else:
                self.alpha = Variable(alpha)
        self.gamma = gamma
        self.class_num = class_num
        self.size_average = size_average
        self.sigmoid = sigmoid
        self.reduce = reduce
    def forward(self, inputs, targets):
        N = inputs.size(0)
        # print(N)
        C = inputs.size(1)
        if self.sigmoid:
            P = torch.sigmoid(inputs)
            #F.softmax(inputs)
            if targets == 0:
                probs = 1 - P#(P * class_mask).sum(1).view(-1, 1)
                log_p = probs.log()
                batch_loss = - (torch.pow((1 - probs), self.gamma)) * log_p
            if targets == 1:
                probs = P  # (P * class_mask).sum(1).view(-1, 1)
                log_p = probs.log()
                batch_loss = - (torch.pow((1 - probs), self.gamma)) * log_p
        else:
            #inputs = F.sigmoid(inputs)
            P = torch.softmax(inputs, dim=-1)

            class_mask = inputs.data.new(N, C).fill_(0)
            class_mask = Variable(class_mask)
            ids = targets.view(-1, 1)
            class_mask.scatter_(1, ids.data, 1.)
            # print(class_mask)

            if inputs.is_cuda and not self.alpha.is_cuda:
                self.alpha = self.alpha.cuda()
            alpha = self.alpha[ids.data.view(-1)]

            probs = (P * class_mask).sum(1).view(-1, 1)

            log_p = probs.log()
            # print('probs size= {}'.format(probs.size()))
            # print(probs)

            batch_loss = -alpha * (torch.pow((1 - probs), self.gamma)) * log_p
            # print('-----bacth_loss------')
            # print(batch_loss)

        if not self.reduce:
            return batch_loss
        if self.size_average:
            loss = batch_loss.mean()
        else:
            loss = batch_loss.sum()
        return loss
class FocalPseudo(nn.Module):
    r"""
        This criterion is a implemenation of Focal Loss, which is proposed in
        Focal Loss for Dense Object Detection.

            Loss(x, class) = - \alpha (1-softmax(x)[class])^gamma \log(softmax(x)[class])

        The losses are averaged across observations for each minibatch.
        Args:
            alpha(1D Tensor, Variable) : the scalar factor for this criterion
            gamma(float, double) : gamma > 0; reduces the relative loss for well-classiﬁed examples (p > .5),
                                   putting more focus on hard, misclassiﬁed examples
            size_average(bool): size_average(bool): By default, the losses are averaged over observations for each minibatch.
                                However, if the field size_average is set to False, the losses are
                                instead summed for each minibatch.
    """
    def __init__(self, class_num, alpha=None, gamma=2, size_average=True,threshold=0.8):
        super(FocalPseudo, self).__init__()
        if alpha is None:
            self.alpha = Variable(torch.ones(class_num, 1)*1.0)
        else:
            if isinstance(alpha, Variable):
                self.alpha = alpha
            else:
                self.alpha = Variable(alpha)
        self.gamma = gamma
        self.class_num = class_num
        self.size_average = size_average
        self.threshold = threshold

    def forward(self, inputs):
        N = inputs.size(0)
        C = inputs.size(1)
        inputs = inputs[0,:,:]
        #print(inputs)
        #pdb.set_trace()
        inputs,ind = torch.max(inputs,1)
        ones = torch.ones(inputs.size()).cuda()
        value = torch.where(inputs>self.threshold,inputs,ones)
        #
        #pdb.set_trace()
        #ind
        #print(value)
        try:
            ind = value.ne(1)
            indexes = torch.nonzero(ind)
            #value2 = inputs[indexes]
            inputs = inputs[indexes]
            log_p = inputs.log()
            # print('probs size= {}'.format(probs.size()))
            # print(probs)
            if not self.gamma == 0:
                batch_loss = - (torch.pow((1 - inputs), self.gamma)) * log_p
            else:
                batch_loss = - log_p
        except:
            #inputs = inputs#[indexes]
            log_p = value.log()
            # print('probs size= {}'.format(probs.size()))
            # print(probs)
            if not self.gamma == 0:
                batch_loss = - (torch.pow((1 - inputs), self.gamma)) * log_p
            else:
                batch_loss = - log_p
        # print('-----bacth_loss------')
        # print(batch_loss)
        #batch_loss = batch_loss #* weight
        if self.size_average:
            try:
                loss = batch_loss.mean() #+ 0.1*balance
            except:
                pdb.set_trace()
        else:
            loss = batch_loss.sum()
        return loss
def CrossEntropy(output, label):
        criteria = torch.nn.CrossEntropyLoss()
        loss = criteria(output, label)
        return loss
class GradReverse(Function):
    def __init__(self, lambd):
        self.lambd = lambd

    def forward(self, x):
        return x.view_as(x)

    def backward(self, grad_output):
        #pdb.set_trace()
        return (grad_output * -self.lambd)

def grad_reverse(x, lambd=1.0):
    return GradReverse(lambd)(x)

def save_net(fname, net):
    import h5py
    h5f = h5py.File(fname, mode='w')
    for k, v in net.state_dict().items():
        h5f.create_dataset(k, data=v.cpu().numpy())

def load_net(fname, net):
    import h5py
    h5f = h5py.File(fname, mode='r')
    for k, v in net.state_dict().items():
        param = torch.from_numpy(np.asarray(h5f[k]))
        v.copy_(param)

def weights_normal_init(model, dev=0.01):
    if isinstance(model, list):
        for m in model:
            weights_normal_init(m, dev)
    else:
        for m in model.modules():
            if isinstance(m, nn.Conv2d):
                m.weight.data.normal_(0.0, dev)
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0.0, dev)

def clip_gradient(model, clip_norm):
    """Computes a gradient clipping coefficient based on gradient norm."""
    totalnorm = 0
    for p in model.parameters():
        if p.requires_grad:
            modulenorm = p.grad.data.norm()
            totalnorm += modulenorm ** 2
    totalnorm = torch.sqrt(totalnorm).item()

    norm = (clip_norm / max(totalnorm, clip_norm))
    #print(norm)
    for p in model.parameters():
        if p.requires_grad:
            p.grad.mul_(norm)

def vis_detections(im, class_name, dets, thresh=0.8):
    """Visual debugging of detections."""
    for i in range(np.minimum(10, dets.shape[0])):
        bbox = tuple(int(np.round(x)) for x in dets[i, :4])
        score = dets[i, -1]
        if score > thresh:
            cv2.rectangle(im, bbox[0:2], bbox[2:4], (0, 0, 204), 3)
            # cv2.rectangle(im, (bbox[0], bbox[1] - 45), (bbox[0]+250, bbox[1] + 5), (255, 0, 0), thickness=-1)
            cv2.putText(im, '%s: %.2f' % (class_name, score), (bbox[0], bbox[1] - 6), cv2.FONT_HERSHEY_PLAIN,
                        2, (100, 100, 100), thickness=3)
        # if score > thresh:
        #     cv2.rectangle(im, bbox[0:2], bbox[2:4], (0, 204, 0), 2)
        #     cv2.putText(im, '%s: %.3f' % (class_name, score), (bbox[0], bbox[1] + 15), cv2.FONT_HERSHEY_PLAIN,
        #                 1.0, (0, 0, 255), thickness=1)
    return im

def show_results(img, image_file, results, num_classes, threshold=0.6, save_fig=True):
    plt.clf()
    plt.imshow(img)
    plt.axis('off')
    ax = plt.gca()

    colors = plt.cm.hsv(np.linspace(0, 1, num_classes)).tolist()

    for i in range(0, results.shape[0]):
        score = results[i, 4]
        if threshold and score < threshold:
            continue

        label = int(results[i, 5])
        iou_before = results[i, 6]
        iou_after = results[i, 7]
        color = colors[label % num_classes]

        xmin = int(round(results[i, 0]))
        ymin = int(round(results[i, 1]))
        xmax = int(round(results[i, 2]))
        ymax = int(round(results[i, 3]))
        coords = (xmin, ymin), xmax - xmin, ymax - ymin
        ax.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=3))
        display_text = '%.2f, %.2f, %.2f' % (score,iou_before, iou_after)
        ax.text(xmin, ymin, display_text, bbox={'facecolor':color, 'alpha':0.5})
    if save_fig:
        plt.savefig(image_file, bbox_inches="tight")
    #plt.show()

def show_results_score(img, image_file, results, num_classes, threshold=0.6, save_fig=True):
    plt.clf()
    plt.imshow(img)
    plt.axis('off')
    ax = plt.gca()

    colors = plt.cm.hsv(np.linspace(0, 1, num_classes)).tolist()

    for i in range(0, results.shape[0]):
        score = results[i, 4]
        #norm = results[i, -1]
        if threshold and score < threshold:
            continue

        label = int(results[i, 5])
        color = colors[label % num_classes]

        xmin = int(round(results[i, 0]))
        ymin = int(round(results[i, 1]))
        xmax = int(round(results[i, 2]))
        ymax = int(round(results[i, 3]))
        coords = (xmin, ymin), xmax - xmin, ymax - ymin
        ax.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=1))
        #display_text = '%.2f %.2f' % (score, norm)
        ax.text(xmin, ymin, '', bbox={'facecolor':color, 'alpha':0.5})
    if save_fig:
        plt.savefig(image_file, bbox_inches="tight")
    #plt.show()

def show_rois(img, image_file, results, save_fig=True):
    plt.clf()
    plt.imshow(img)
    plt.axis('off')
    ax = plt.gca()

    colors = plt.cm.hsv(np.linspace(0, 1, 2)).tolist()

    for i in range(0, results.shape[0]):
        color = colors[0]

        xmin = int(round(results[i, 0]))
        ymin = int(round(results[i, 1]))
        xmax = int(round(results[i, 2]))
        ymax = int(round(results[i, 3]))
        coords = (xmin, ymin), xmax - xmin, ymax - ymin
        ax.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=1))
    if save_fig:
        plt.savefig(image_file, bbox_inches="tight")

def adjust_learning_rate(optimizer, decay=0.1):
    """Sets the learning rate to the initial LR decayed by 0.5 every 20 epochs"""
    for param_group in optimizer.param_groups:
        param_group['lr'] = decay * param_group['lr']
import math
def calc_supp(iter,iter_total=80000):
    p = float(iter) / iter_total
    #print(math.exp(-10*p))
    return 2 / (1 + math.exp(-10*p)) - 1
# def adjust_learning_rate(optimizer, decay=0.1,lr_init = 0.001):
#     """Sets the learning rate to the initial LR decayed by 0.5 every 20 epochs"""
#     for param_group in optimizer.param_groups:
#         param_group['lr'] = decay * lr_init#param_group['lr']

def save_checkpoint(state, filename, is_best, best_score=0):
    output_dir = os.path.dirname(filename)
    new_filename = os.path.join(output_dir, 'checkpoint.pth')
    torch.save(state, new_filename)
    if is_best:
        output_dir = os.path.dirname(filename)
        filename_best = os.path.join(output_dir, 'model_best'+str(best_score)+'.pth')
        shutil.copyfile(new_filename, filename_best)

def _smooth_l1_loss(bbox_pred, bbox_targets, bbox_inside_weights, bbox_outside_weights,\
     sigma=1.0, dim=[1], reduction='mean'):

    sigma_2 = sigma ** 2
    box_diff = bbox_pred - bbox_targets
    in_box_diff = bbox_inside_weights * box_diff
    abs_in_box_diff = torch.abs(in_box_diff)
    smoothL1_sign = (abs_in_box_diff < 1. / sigma_2).detach().float()
    in_loss_box = torch.pow(in_box_diff, 2) * (sigma_2 / 2.) * smoothL1_sign \
                  + (abs_in_box_diff - (0.5 / sigma_2)) * (1. - smoothL1_sign)
    out_loss_box = bbox_outside_weights * in_loss_box
    loss_box = out_loss_box
    for i in sorted(dim, reverse=True):
      loss_box = loss_box.sum(i)
    if reduction == 'mean':
        loss_box = loss_box.mean()
    return loss_box

def _crop_pool_layer(bottom, rois, max_pool=True):
    # code modified from
    # https://github.com/ruotianluo/pytorch-faster-rcnn
    # implement it using stn
    # box to affine
    # input (x1,y1,x2,y2)
    """
    [  x2-x1             x1 + x2 - W + 1  ]
    [  -----      0      ---------------  ]
    [  W - 1                  W - 1       ]
    [                                     ]
    [           y2-y1    y1 + y2 - H + 1  ]
    [    0      -----    ---------------  ]
    [           H - 1         H - 1      ]
    """
    rois = rois.detach()
    batch_size = bottom.size(0)
    D = bottom.size(1)
    H = bottom.size(2)
    W = bottom.size(3)
    roi_per_batch = rois.size(0) / batch_size
    x1 = rois[:, 1::4] / 16.0
    y1 = rois[:, 2::4] / 16.0
    x2 = rois[:, 3::4] / 16.0
    y2 = rois[:, 4::4] / 16.0

    height = bottom.size(2)
    width = bottom.size(3)

    # affine theta
    zero = Variable(rois.data.new(rois.size(0), 1).zero_())
    theta = torch.cat([\
      (x2 - x1) / (width - 1),
      zero,
      (x1 + x2 - width + 1) / (width - 1),
      zero,
      (y2 - y1) / (height - 1),
      (y1 + y2 - height + 1) / (height - 1)], 1).view(-1, 2, 3)

    if max_pool:
      pre_pool_size = cfg.POOLING_SIZE * 2
      grid = F.affine_grid(theta, torch.Size((rois.size(0), 1, pre_pool_size, pre_pool_size)))
      bottom = bottom.view(1, batch_size, D, H, W).contiguous().expand(roi_per_batch, batch_size, D, H, W)\
                                                                .contiguous().view(-1, D, H, W)
      crops = F.grid_sample(bottom, grid)
      crops = F.max_pool2d(crops, 2, 2)
    else:
      grid = F.affine_grid(theta, torch.Size((rois.size(0), 1, cfg.POOLING_SIZE, cfg.POOLING_SIZE)))
      bottom = bottom.view(1, batch_size, D, H, W).contiguous().expand(roi_per_batch, batch_size, D, H, W)\
                                                                .contiguous().view(-1, D, H, W)
      crops = F.grid_sample(bottom, grid)

    return crops, grid

def _affine_grid_gen(rois, input_size, grid_size):

    rois = rois.detach()
    x1 = rois[:, 1::4] / 16.0
    y1 = rois[:, 2::4] / 16.0
    x2 = rois[:, 3::4] / 16.0
    y2 = rois[:, 4::4] / 16.0

    height = input_size[0]
    width = input_size[1]

    zero = Variable(rois.data.new(rois.size(0), 1).zero_())
    theta = torch.cat([\
      (x2 - x1) / (width - 1),
      zero,
      (x1 + x2 - width + 1) / (width - 1),
      zero,
      (y2 - y1) / (height - 1),
      (y1 + y2 - height + 1) / (height - 1)], 1).view(-1, 2, 3)

    grid = F.affine_grid(theta, torch.Size((rois.size(0), 1, grid_size, grid_size)))

    return grid

def _affine_theta(rois, input_size):

    rois = rois.detach()
    x1 = rois[:, 1::4] / 16.0
    y1 = rois[:, 2::4] / 16.0
    x2 = rois[:, 3::4] / 16.0
    y2 = rois[:, 4::4] / 16.0

    height = input_size[0]
    width = input_size[1]

    zero = Variable(rois.data.new(rois.size(0), 1).zero_())

    # theta = torch.cat([\
    #   (x2 - x1) / (width - 1),
    #   zero,
    #   (x1 + x2 - width + 1) / (width - 1),
    #   zero,
    #   (y2 - y1) / (height - 1),
    #   (y1 + y2 - height + 1) / (height - 1)], 1).view(-1, 2, 3)

    theta = torch.cat([\
      (y2 - y1) / (height - 1),
      zero,
      (y1 + y2 - height + 1) / (height - 1),
      zero,
      (x2 - x1) / (width - 1),
      (x1 + x2 - width + 1) / (width - 1)], 1).view(-1, 2, 3)

    return theta
'''
def compare_grid_sample():
    # do gradcheck
    N = random.randint(1, 8)
    C = 2 # random.randint(1, 8)
    H = 5 # random.randint(1, 8)
    W = 4 # random.randint(1, 8)
    input = Variable(torch.randn(N, C, H, W).cuda(), requires_grad=True)
    input_p = input.clone().data.contiguous()

    grid = Variable(torch.randn(N, H, W, 2).cuda(), requires_grad=True)
    grid_clone = grid.clone().contiguous()

    out_offcial = F.grid_sample(input, grid)
    grad_outputs = Variable(torch.rand(out_offcial.size()).cuda())
    grad_outputs_clone = grad_outputs.clone().contiguous()
    grad_inputs = torch.autograd.grad(out_offcial, (input, grid), grad_outputs.contiguous())
    grad_input_off = grad_inputs[0]

    crf = RoICropFunction()
    grid_yx = torch.stack([grid_clone.data[:,:,:,1], grid_clone.data[:,:,:,0]], 3).contiguous().cuda()
    out_stn = crf.forward(input_p, grid_yx)
    grad_inputs = crf.backward(grad_outputs_clone.data)
    grad_input_stn = grad_inputs[0]
    pdb.set_trace()

    delta = (grad_input_off.data - grad_input_stn).sum()
'''

def proposals_to_centers(proposals):
    """
    :param proposals: [N, 5], (b_ix, x1, y1, x2, y2)
    :return: centers [N, 2], (b_ix, center_x, center_y)
    """
    cx = (proposals[:, 3] + proposals[:, 1]) / 2.0
    cy = (proposals[:, 4] + proposals[:, 2]) / 2.0
    center = np.vstack([cx, cy]).transpose()
    return center

def compute_cluster_targets(proposals, features, N_cluster=4, threshold=128):
    '''
    Args:
        proposals:[N, k], k>=5(b_ix, x1,y1,x2,y2, ...), N = 512
        features: [N, 4096],
    Return:
        batch_rois: [N_cluster, 128, 4096]
        batch_cluster_center: [N_cluster, 2], (center_x, center_y)
    '''
    proposals_np = to_np_array(proposals)
    features_np = to_np_array(features)
    centers = proposals_to_centers(proposals_np)

    """
    KMeans part
    """
    kmeans = KMeans(n_clusters=N_cluster, random_state=0).fit(centers)

    cluster_center = kmeans.cluster_centers_
    cluster_labels = kmeans.labels_

    batch_rois_cluster = []
    for cluster_idx in range(0, N_cluster):
        keep_ix = np.where(cluster_labels[:] == cluster_idx)[0]

        if keep_ix.shape[0] < threshold:
            keep_ix_new = np.random.choice(keep_ix.shape[0], threshold, replace=True)
            keep_ix2 = keep_ix[keep_ix_new]
            batch_rois_tmp = features_np[keep_ix2]
        else:
            keep_ix2 = keep_ix[0:threshold]
            batch_rois_tmp = features_np[keep_ix2]

        # batch_rois_tmp = features[keep_ix]
        batch_rois_cluster.append(batch_rois_tmp)

    batch_rois_cluster = np.stack(batch_rois_cluster, axis=0) # (N_cluster, threshold, 4096)
    batch_rois_cluster = np.mean(batch_rois_cluster, axis=1)

    f = lambda x: (torch.from_numpy(x)).float().cuda().contiguous()
    batch_rois_cluster = f(batch_rois_cluster)
    # batch_mask_labels = f(batch_mask_labels).float()
    return batch_rois_cluster, cluster_center

def compute_mean_cluster_targets(proposals, features, N_cluster=4, threshold=128):
    '''
    Args:
        proposals:[N, k], k>=5(b_ix, x1,y1,x2,y2, ...), N = 512
        features: [N, 4096],
    Return:
        batch_rois: [N_cluster, 128, 4096]
        batch_cluster_center: [N_cluster, 2], (center_x, center_y)
    '''
    proposals_np = to_np_array(proposals)
    features_np = to_np_array(features)
    centers = proposals_to_centers(proposals_np)

    """
    KMeans part
    """
    kmeans = KMeans(n_clusters=N_cluster, random_state=0).fit(centers)

    cluster_center = kmeans.cluster_centers_
    cluster_labels = kmeans.labels_

    batch_rois_cluster = []
    for cluster_idx in range(0, N_cluster):
        keep_ix = np.where(cluster_labels[:] == cluster_idx)[0]

        batch_rois_tmp = features_np[keep_ix].mean(0)

        # batch_rois_tmp = features[keep_ix]
        batch_rois_cluster.append(batch_rois_tmp)

    batch_rois_cluster = np.stack(batch_rois_cluster, axis=0) # (N_cluster, threshold, 4096)

    f = lambda x: (torch.from_numpy(x)).float().cuda().contiguous()
    batch_rois_cluster = f(batch_rois_cluster)
    # batch_mask_labels = f(batch_mask_labels).float()
    return batch_rois_cluster, cluster_center

def compute_weighted_mean_cluster_targets(proposals, features, prob, N_cluster=4, threshold=128):
    '''
    Args:
        proposals:[N, k], k>=5(b_ix, x1,y1,x2,y2, ...), N = 512
        features: [N, 4096],
    Return:
        batch_rois: [N_cluster, 128, 4096]
        batch_cluster_center: [N_cluster, 2], (center_x, center_y)
    '''
    proposals_np = to_np_array(proposals)
    features_np = to_np_array(features)
    centers = proposals_to_centers(proposals_np)
    prob_np = to_np_array(prob)
    obj_prob_np = 1 - prob_np[:,0]
    """
    KMeans part
    """
    kmeans = KMeans(n_clusters=N_cluster, random_state=0).fit(centers)

    cluster_center = kmeans.cluster_centers_
    cluster_labels = kmeans.labels_

    batch_rois_cluster = []
    for cluster_idx in range(0, N_cluster):
        keep_ix = np.where(cluster_labels[:] == cluster_idx)[0]

        batch_rois_tmp = features_np[keep_ix]
        batch_rois_weight = obj_prob_np[keep_ix] / obj_prob_np[keep_ix].sum()
        batch_rois_tmp = (batch_rois_tmp * batch_rois_weight.reshape(-1,1)).mean(0)

        batch_rois_cluster.append(batch_rois_tmp)

    batch_rois_cluster = np.stack(batch_rois_cluster, axis=0) # (N_cluster, threshold, 4096)

    f = lambda x: (torch.from_numpy(x)).float().cuda().contiguous()
    batch_rois_cluster = f(batch_rois_cluster)
    # batch_mask_labels = f(batch_mask_labels).float()
    return batch_rois_cluster, cluster_center
def compute_entropy(predict, softmax=True):
    n, k = predict.size()
    if softmax:
        predict_prob = torch.softmax(predict, dim=-1)
    else:
        predict_prob = predict
    loss_ent = -1 * torch.mean(torch.sum(predict_prob *
                                             (torch.log(predict_prob + 1e-5)), 1))

    return loss_ent

def accuracy(output, target, topk=(1, ), ignore_index=-1):
    """Computes the precision@k for the specified values of k"""
    keep = torch.nonzero(target != ignore_index).squeeze()
    #logger.info('target.shape:{0}, keep.shape:{1}'.format(target.shape, keep.shape))
    assert (keep.dim() == 1)
    target = target[keep]
    output = output[keep]
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res

def select_roi(all_rois, gt_boxes=None, num_boxes=None):
    overlaps = bbox_overlaps_batch(all_rois, gt_boxes)

    max_overlaps, gt_assignment = torch.max(overlaps, 2)

    batch_size = overlaps.size(0)
    num_proposal = overlaps.size(1)
    num_boxes_per_img = overlaps.size(2)

    offset = torch.arange(0, batch_size)*gt_boxes.size(1)
    offset = offset.view(-1, 1).type_as(gt_assignment) + gt_assignment

    # changed indexing way for pytorch 1.0
    labels = gt_boxes[:,:,4].contiguous().view(-1)[(offset.view(-1),)].view(batch_size, -1)

    labels_batch = labels.new(batch_size, num_proposal).zero_()

    # Guard against the case when an image has fewer than max_fg_rois_per_image
    # foreground RoIs
    for i in range(batch_size):

        fg_inds = torch.nonzero(max_overlaps[i] >= cfg.TRAIN.FG_THRESH).view(-1)
        labels_batch[i][fg_inds] = labels[i][fg_inds]

    return labels_batch

def cluster_rois(batch_size, rois, cls_prob, pooled_feat):
        rois = rois[:,:,1:]
        for i in range(batch_size):
            boxes = rois[i]
            scores= cls_prob.view(batch_size, -1, self.n_classes)[i]
            feats = pooled_feat.view(batch_size, -1, 4096)[i]
            thresh = 0.01
            class_feats_list = []
            #pdb.set_trace()
            for j in range(self.n_classes):
                inds = torch.nonzero(scores[:,j]>thresh).view(-1)
                # if there is det
                if inds.numel() > 0:
                    cls_scores = scores[:,j][inds]
                    _, order = torch.sort(cls_scores, 0, True)
                    cls_boxes = boxes[inds, :]
                    j_feats = feats[inds,:]
                    cls_dets = torch.cat((cls_boxes, cls_scores.unsqueeze(1)), 1)
                    cls_dets = cls_dets[order]
                    j_feats = j_feats[order]

                    try:
                        keep = nms(cls_dets[:,:4], cls_dets[:,-1], 0.3)
                    except:
                        print(cls_dets.size())

                    keep = keep.view(-1).long()
                    num_clusters = keep.size(0)
                    cluster_feats_list = []
                    for k in range(num_clusters):
                        start = keep[k]
                        if k != num_clusters-1:
                            end = keep[k+1]
                        else:
                            end = cls_dets.size(0)
                        try:
                            c_cls_dets = cls_dets[start:end]
                        except:
                            print(order, start, end)
                        c_feats = j_feats[start:end]
                        cluster_feats = (c_feats * (c_cls_dets[:,-1] / c_cls_dets[:,-1].sum()).unsqueeze(1)).sum(0)
                        cluster_feats_list.append(cluster_feats)
                    class_feats = torch.stack(cluster_feats_list, 0).mean(0)
                else:
                    class_feats = (feats * (scores[:,j] / scores[:,j].sum()).unsqueeze(1)).sum(0)
                class_feats_list.append(class_feats)
            return torch.stack(class_feats_list,0)

def cluster_rois_v1(batch_size, rois, roi_feats):
    rois = rois.view(batch_size, -1, 5)[:,:,1:]
    roi_feats = roi_feats.view(batch_size, -1, roi_feats.size(-1))
    n_rois = rois.size(1)
    score = torch.arange(n_rois, 0, -1).cuda().float() / n_rois
    cluster_roi_feats = list()

    for i in range(batch_size):
        overlaps = bbox_overlaps(rois[i], rois[i])
        overlaps[overlaps<0.5] = 0
        overlaps /= overlaps.sum(1, keepdim=True)
        fused_roi_feat = torch.matmul(overlaps, roi_feats[i])
        keep = nms(rois[i], score, 0.5)
        cluster_roi_feats.append(fused_roi_feat[keep])

    return torch.cat(cluster_roi_feats, 0)

def cluster_rois_v3(batch_size, rois, roi_feats):
    rois = rois.view(batch_size, -1, 5)[:,:,1:]
    roi_feats = roi_feats.view(batch_size, -1, roi_feats.size(-1))
    n_rois = rois.size(1)
    cluster_roi_feats = list()

    for i in range(batch_size):
        overlaps = bbox_overlaps(rois[i], rois[i])
        overlaps[overlaps<0.5] = 0
        overlaps /= overlaps.sum(1, keepdim=True)
        fused_roi_feat = torch.matmul(overlaps, roi_feats[i])
        cluster_roi_feats.append(fused_roi_feat)

    return torch.cat(cluster_roi_feats, 0)

def cluster_rois_v2(batch_size, rois, roi_feats):
    rois = rois.view(batch_size, -1, 5)[:,:,1:]
    roi_feats = roi_feats.view(batch_size, -1, roi_feats.size(-1))
    n_rois = rois.size(1)
    cluster_roi_feats = list()
    for i in range(batch_size):
        overlaps = bbox_overlaps(rois[i], rois[i])
        overlaps[overlaps<0.5] = 0
        overlaps /= overlaps.sum(1, keepdim=True)
        fused_roi_feat = torch.matmul(overlaps, roi_feats[i])
        cluster_roi_feats.append(fused_roi_feat)
    return torch.cat(cluster_roi_feats, 0)

def assign_pseudo_label(cls_prob, bg_threshold=0.75, max_n_bgs=64, fg_threshold=0.5):

    n_classes = cls_prob.size(1)
    pseudo_labels = torch.zeros(cls_prob.size(0)).cuda().long() - 1
    n_fgs = 0
    for j in range(1,n_classes):
        inds = torch.nonzero(cls_prob[:,j] >= fg_threshold).view(-1)
        pseudo_labels[inds] = j
        n_fgs += inds.size(0)
    fg_inds = (cls_prob[:,0] >= bg_threshold).nonzero().view(-1)
    if fg_inds.size(0) > max_n_bgs:
        pseudo_labels[fg_inds[:max_n_bgs]] = 0
    return pseudo_labels

def assign_real_label(all_rois, gt_boxes):
        """Generate a random sample of RoIs comprising foreground and background
        examples.
        """
        # overlaps: (rois x gt_boxes)

        overlaps = bbox_overlaps_batch(all_rois, gt_boxes)

        max_overlaps, gt_assignment = torch.max(overlaps, 2)

        batch_size = overlaps.size(0)
        num_proposal = overlaps.size(1)
        num_boxes_per_img = overlaps.size(2)

        offset = torch.arange(0, batch_size)*gt_boxes.size(1)
        offset = offset.view(-1, 1).type_as(gt_assignment) + gt_assignment

        # changed indexing way for pytorch 1.0
        labels = gt_boxes[:,:,4].contiguous().view(-1)[(offset.view(-1),)].view(batch_size, -1)

        labels_batch = labels.new(batch_size, num_proposal).zero_() -1

        # Guard against the case when an image has fewer than max_fg_rois_per_image
        # foreground RoIs
        for i in range(batch_size):

            fg_inds = torch.nonzero(max_overlaps[i] >= cfg.TRAIN.FG_THRESH).view(-1)
            fg_num_rois = fg_inds.numel()

            # Select background RoIs as those within [BG_THRESH_LO, BG_THRESH_HI)
            bg_inds = torch.nonzero((max_overlaps[i] < cfg.TRAIN.BG_THRESH_HI) &
                                    (max_overlaps[i] >= cfg.TRAIN.BG_THRESH_LO)).view(-1)
            bg_num_rois = bg_inds.numel()

            # Select sampled values from various arrays:
            labels_batch[i][fg_inds] = labels[i][fg_inds]
            labels_batch[i][bg_inds] = 0
        return labels_batch
'''
def vis_roi_feat(dataset_name, i, s_boxes, s_gt_boxes, s_roi_feat, t_boxes, t_gt_boxes, t_roi_feat,version='default'):
    os.makedirs('./tsne_images/'+str(dataset_name)+'/'+version, exist_ok=True)
    path = './tsne_images/'+str(dataset_name)+'/'+version+'/'+'vis.jpg'

    iou_before = bbox_overlaps(s_boxes.squeeze(), s_gt_boxes.view(-1,5)[:,:4])
    max_iou_before = iou_before.max(1, keepdim=True)[0]
    s_pos_inds = (max_iou_before >= 0.5).nonzero().view(-1).cpu().numpy()
    s_neg_inds = (max_iou_before < 0.4).nonzero().view(-1).cpu().numpy()
    n_s = s_roi_feat.size(0)

    iou_before = bbox_overlaps(t_boxes.squeeze(), t_gt_boxes.view(-1,5)[:,:4])
    max_iou_before = iou_before.max(1, keepdim=True)[0]
    t_pos_inds = (max_iou_before >= 0.5).nonzero().view(-1).cpu().numpy() + n_s
    t_neg_inds = (max_iou_before < 0.4).nonzero().view(-1).cpu().numpy() + n_s

    roi_feat = torch.cat([s_roi_feat, t_roi_feat], 0)
    roi_embed = TSNE(n_components=2).fit_transform(roi_feat.detach().cpu().numpy())

    fig = plt.figure()
    ax = plt.subplot()
    ax.scatter(roi_embed[s_pos_inds, 0], roi_embed[s_pos_inds, 1], c='green')
    ax.scatter(roi_embed[s_neg_inds, 0], roi_embed[s_neg_inds, 1], c='blue')
    ax.scatter(roi_embed[t_pos_inds, 0], roi_embed[t_pos_inds, 1], c='cyan')
    ax.scatter(roi_embed[t_neg_inds, 0], roi_embed[t_neg_inds, 1], c='red')
    plt.savefig(path)
    print('saved  '+path)
    plt.close()
'''
def vis_roi_feat(dataset_name, step, s_all_roi_feat, s_all_overlaps, t_all_roi_feat, t_all_overlaps, version='default'):
    os.makedirs('./tsne_images/'+str(dataset_name)+'/'+version, exist_ok=True)
    path = './tsne_images/'+str(dataset_name)+'/'+version+'/'+ str(step) + '.jpg'

    s_pos_inds = (s_all_overlaps >= 0.5).nonzero().view(-1).cpu().numpy()
    s_neg_inds = (s_all_overlaps < 0.4).nonzero().view(-1).cpu().numpy()
    n_s = s_all_roi_feat.size(0)

    t_pos_inds = (t_all_overlaps >= 0.5).nonzero().view(-1).cpu().numpy() + n_s
    t_neg_inds = (t_all_overlaps < 0.4).nonzero().view(-1).cpu().numpy() + n_s

    roi_feat = torch.cat([s_all_roi_feat, t_all_roi_feat], 0)
    roi_embed = TSNE(n_components=2, perplexity=15, learning_rate=400).fit_transform(roi_feat.detach().cpu().numpy())

    fig = plt.figure()
    ax = plt.subplot()
    ax.scatter(roi_embed[s_pos_inds, 0], roi_embed[s_pos_inds, 1], c='green')
    ax.scatter(roi_embed[s_neg_inds, 0], roi_embed[s_neg_inds, 1], c='blue')
    ax.scatter(roi_embed[t_pos_inds, 0], roi_embed[t_pos_inds, 1], c='cyan')
    ax.scatter(roi_embed[t_neg_inds, 0], roi_embed[t_neg_inds, 1], c='red')
    plt.savefig(path)
    print('saved  '+path)
    plt.close()


# MoCo-related code is modified from https://github.com/facebookresearch/moco
import sys
import math
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
sys.path.append('../')
from backbone.select_backbone import select_backbone

# utils
@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor)
        for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output

class InfoNCE(nn.Module):
    '''
    Basically, it's a MoCo for video input: https://arxiv.org/abs/1911.05722
    '''
    def __init__(self, network='s3d', dim=128, K=2048, m=0.999, T=0.07):
        '''
        dim: feature dimension (default: 128)
        K: queue size; number of negative keys (default: 2048)
        m: moco momentum of updating key encoder (default: 0.999)
        T: softmax temperature (default: 0.07)
        '''
        super(InfoNCE, self).__init__()

        self.dim = dim
        self.K = K
        self.m = m
        self.T = T

        # create the encoders (including non-linear projection head: 2 FC layers)
        backbone, self.param = select_backbone(network)
        feature_size = self.param['feature_size']
        self.encoder_q = nn.Sequential(
                            backbone,
                            nn.AdaptiveAvgPool3d((1,1,1)),
                            nn.Conv3d(feature_size, feature_size, kernel_size=1, bias=True),
                            nn.ReLU(),
                            nn.Conv3d(feature_size, dim, kernel_size=1, bias=True))

        backbone, _ = select_backbone(network)
        self.encoder_k = nn.Sequential(
                            backbone,
                            nn.AdaptiveAvgPool3d((1,1,1)),
                            nn.Conv3d(feature_size, feature_size, kernel_size=1, bias=True),
                            nn.ReLU(),
                            nn.Conv3d(feature_size, dim, kernel_size=1, bias=True))

        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            param_k.data.copy_(param_q.data)  # initialize
            param_k.requires_grad = False  # not update by gradient

        # create the queue
        self.register_buffer("queue", torch.randn(dim, K))
        self.queue = nn.functional.normalize(self.queue, dim=0)
        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))

        # Notes: for handling sibling videos, e.g. for UCF101 dataset

    @torch.no_grad()
    def _momentum_update_key_encoder(self):
        '''Momentum update of the key encoder'''
        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys):
        # gather keys before updating queue
        keys = concat_all_gather(keys)

        batch_size = keys.shape[0]

        ptr = int(self.queue_ptr)
        assert self.K % batch_size == 0  # for simplicity

        # replace the keys at ptr (dequeue and enqueue)
        self.queue[:, ptr:ptr + batch_size] = keys.T
        ptr = (ptr + batch_size) % self.K  # move pointer

        self.queue_ptr[0] = ptr

    @torch.no_grad()
    def _batch_shuffle_ddp(self, x):
        '''
        Batch shuffle, for making use of BatchNorm.
        *** Only support DistributedDataParallel (DDP) model. ***
        '''
        # gather from all gpus
        batch_size_this = x.shape[0]
        x_gather = concat_all_gather(x)
        batch_size_all = x_gather.shape[0]

        num_gpus = batch_size_all // batch_size_this

        # random shuffle index
        idx_shuffle = torch.randperm(batch_size_all).cuda()

        # broadcast to all gpus
        torch.distributed.broadcast(idx_shuffle, src=0)

        # index for restoring
        idx_unshuffle = torch.argsort(idx_shuffle)

        # shuffled index for this gpu
        gpu_idx = torch.distributed.get_rank()
        idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]

        return x_gather[idx_this], idx_unshuffle

    @torch.no_grad()
    def _batch_unshuffle_ddp(self, x, idx_unshuffle):
        '''
        Undo batch shuffle.
        *** Only support DistributedDataParallel (DDP) model. ***
        '''
        # gather from all gpus
        batch_size_this = x.shape[0]
        x_gather = concat_all_gather(x)
        batch_size_all = x_gather.shape[0]

        num_gpus = batch_size_all // batch_size_this

        # restored index for this gpu
        gpu_idx = torch.distributed.get_rank()
        idx_this = idx_unshuffle.view(num_gpus, -1)[gpu_idx]

        return x_gather[idx_this]

    def forward(self, block):
        '''Output: logits, targets'''
        (B, N, *_) = block.shape # [B,N,C,T,H,W]
        assert N == 2
        x1 = block[:,0,:].contiguous()
        x2 = block[:,1,:].contiguous()

        # compute query features
        q = self.encoder_q(x1)  # queries: B,C,1,1,1
        q = nn.functional.normalize(q, dim=1)
        q = q.view(B, self.dim)

        in_train_mode = q.requires_grad

        # compute key features
        with torch.no_grad():  # no gradient to keys
            if in_train_mode: self._momentum_update_key_encoder()  # update the key encoder

            # shuffle for making use of BN
            x2, idx_unshuffle = self._batch_shuffle_ddp(x2)

            k = self.encoder_k(x2)  # keys: B,C,1,1,1
            k = nn.functional.normalize(k, dim=1)

            # undo shuffle
            k = self._batch_unshuffle_ddp(k, idx_unshuffle)

        k = k.view(B, self.dim)

        # compute logits
        l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
        l_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])

        # logits: B,(1+K)
        logits = torch.cat([l_pos, l_neg], dim=1)

        # apply temperature
        logits /= self.T

        # labels: positive key indicators
        labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()

        # dequeue and enqueue
        if in_train_mode: self._dequeue_and_enqueue(k)

        return logits, labels

class UberNCE(InfoNCE):
    '''
    UberNCE is a supervised version of InfoNCE,
    it uses labels to define positive and negative pair
    Still, use MoCo to enlarge the negative pool
    '''
    def __init__(self, network='s3d', dim=128, K=2048, m=0.999, T=0.07):
        """
        dim: feature dimension (default: 128)
        K: queue size; number of negative keys (default: 2048)
        m: moco momentum of updating key encoder (default: 0.999)
        T: softmax temperature (default: 0.07)
        """
        super(UberNCE, self).__init__(network, dim, K, m, T)
        # extra queue to store label
        self.register_buffer("queue_label", torch.ones(K, dtype=torch.long) * -1)

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys, labels):
        # gather keys before updating queue
        keys = concat_all_gather(keys)
        labels = concat_all_gather(labels)

        batch_size = keys.shape[0]

        ptr = int(self.queue_ptr)
        assert self.K % batch_size == 0  # for simplicity

        # replace the keys at ptr (dequeue and enqueue)
        self.queue[:, ptr:ptr + batch_size] = keys.T
        self.queue_label[ptr:ptr + batch_size] = labels
        ptr = (ptr + batch_size) % self.K  # move pointer

        self.queue_ptr[0] = ptr

    def forward(self, block, k_label):
        '''Output: logits, binary mask for positive pairs
        '''
        (B, N, *_) = block.shape # [B,N,C,T,H,W]
        assert N == 2
        x1 = block[:,0,:].contiguous()
        x2 = block[:,1,:].contiguous()

        # compute query features
        q = self.encoder_q(x1)  # queries: B,C,1,1,1
        q = nn.functional.normalize(q, dim=1)
        q = q.view(B, self.dim)

        in_train_mode = q.requires_grad

        # compute key features
        with torch.no_grad():  # no gradient to keys
            if in_train_mode: self._momentum_update_key_encoder()  # update the key encoder

            # shuffle for making use of BN
            x2, idx_unshuffle = self._batch_shuffle_ddp(x2)

            k = self.encoder_k(x2)  # keys: B,C,1,1,1
            k = nn.functional.normalize(k, dim=1)

            # undo shuffle
            k = self._batch_unshuffle_ddp(k, idx_unshuffle)

        k = k.view(B, self.dim)

        # compute logits
        l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
        l_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])

        # logits: B,(1+K)
        logits = torch.cat([l_pos, l_neg], dim=1)

        # apply temperature
        logits /= self.T

        # mask: binary mask for positive keys
        mask = k_label.unsqueeze(1) == self.queue_label.unsqueeze(0) # B,K
        mask = torch.cat([torch.ones((mask.shape[0],1), dtype=torch.long, device=mask.device).bool(),
                          mask], dim=1) # B,(1+K)

        # dequeue and enqueue
        if in_train_mode: self._dequeue_and_enqueue(k, k_label)

        return logits, mask

class CoCLR(InfoNCE):
    '''
    CoCLR: using another view of the data to define positive and negative pair
    Still, use MoCo to enlarge the negative pool
    '''
    def __init__(self, network='s3d', dim=128, K=2048, m=0.999, T=0.07, topk=5, reverse=False):
        '''
        dim: feature dimension (default: 128)
        K: queue size; number of negative keys (default: 2048)
        m: moco momentum of updating key encoder (default: 0.999)
        T: softmax temperature (default: 0.07)
        '''
        super(CoCLR, self).__init__(network, dim, K, m, T)

        self.topk = topk

        # create another encoder, for the second view of the data
        backbone, _ = select_backbone(network)
        feature_size = self.param['feature_size']
        self.sampler = nn.Sequential(
                            backbone,
                            nn.AdaptiveAvgPool3d((1,1,1)),
                            nn.Conv3d(feature_size, feature_size, kernel_size=1, bias=True),
                            nn.ReLU(),
                            nn.Conv3d(feature_size, dim, kernel_size=1, bias=True))
        for param_s in self.sampler.parameters():
            param_s.requires_grad = False  # not update by gradient

        # create another queue, for the second view of the data
        self.register_buffer("queue_second", torch.randn(dim, K))
        self.queue_second = nn.functional.normalize(self.queue_second, dim=0)

        # for handling sibling videos, e.g. for UCF101 dataset
        self.register_buffer("queue_vname", torch.ones(K, dtype=torch.long) * -1)
        # for monitoring purpose only
        self.register_buffer("queue_label", torch.ones(K, dtype=torch.long) * -1)

        self.queue_is_full = False
        self.reverse = reverse

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys, keys_second, vnames):
        # gather keys before updating queue
        keys = concat_all_gather(keys)
        keys_second = concat_all_gather(keys_second)
        vnames = concat_all_gather(vnames)
        # labels = concat_all_gather(labels)

        batch_size = keys.shape[0]

        ptr = int(self.queue_ptr)
        assert self.K % batch_size == 0  # for simplicity

        # replace the keys at ptr (dequeue and enqueue)
        self.queue[:, ptr:ptr + batch_size] = keys.T
        self.queue_second[:, ptr:ptr + batch_size] = keys_second.T
        self.queue_vname[ptr:ptr + batch_size] = vnames
        self.queue_label[ptr:ptr + batch_size] = torch.ones_like(vnames)
        ptr = (ptr + batch_size) % self.K  # move pointer

        self.queue_ptr[0] = ptr

    def forward(self, block1, block2, k_vsource):
        '''Output: logits, targets'''
        (B, N, *_) = block1.shape # B,N,C,T,H,W
        assert N == 2
        x1 = block1[:,0,:].contiguous()
        f1 = block1[:,1,:].contiguous()
        x2 = block2[:,0,:].contiguous()
        f2 = block2[:,1,:].contiguous()

        if self.reverse:
            x1, f1 = f1, x1
            x2, f2 = f2, x2

        # compute query features
        q = self.encoder_q(x1)  # queries: B,C,1,1,1
        q = nn.functional.normalize(q, dim=1)
        q = q.view(B, self.dim)

        in_train_mode = q.requires_grad

        # compute key features
        with torch.no_grad():  # no gradient to keys
            if in_train_mode: self._momentum_update_key_encoder()  # update the key encoder

            # shuffle for making use of BN
            x2, idx_unshuffle = self._batch_shuffle_ddp(x2)

            k = self.encoder_k(x2)  # keys: B,C,1,1,1
            k = nn.functional.normalize(k, dim=1)

            # undo shuffle
            k = self._batch_unshuffle_ddp(k, idx_unshuffle)
            k = k.view(B, self.dim)

            # compute key feature for second view
            kf = self.sampler(f2) # keys: B,C,1,1,1
            kf = nn.functional.normalize(kf, dim=1)
            kf = kf.view(B, self.dim)

        # if queue_second is full: compute mask & train CoCLR, else: train InfoNCE

        # compute logits
        l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
        l_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])

        # logits: N,(1+K)
        logits = torch.cat([l_pos, l_neg], dim=1)

        # apply temperature
        logits /= self.T

        # mask: binary mask for positive keys
        # handle sibling videos, e.g. for UCF101. It has no effect on K400
        mask_source = k_vsource.unsqueeze(1) == self.queue_vname.unsqueeze(0) # B,K
        mask = mask_source.clone()

        if not self.queue_is_full:
            self.queue_is_full = torch.all(self.queue_label != -1)
            if self.queue_is_full: print('\n===== queue is full now =====')

        if self.queue_is_full and (self.topk != 0):
            mask_sim = kf.matmul(self.queue_second.clone().detach())
            mask_sim[mask_source] = - np.inf # mask out self (and sibling videos)
            _, topkidx = torch.topk(mask_sim, self.topk, dim=1)
            topk_onehot = torch.zeros_like(mask_sim)
            topk_onehot.scatter_(1, topkidx, 1)
            mask[topk_onehot.bool()] = True

        mask = torch.cat([torch.ones((mask.shape[0],1), dtype=torch.long, device=mask.device).bool(),
                          mask], dim=1)

        # dequeue and enqueue
        if in_train_mode: self._dequeue_and_enqueue(k, kf, k_vsource)

        return logits, mask.detach()


#-*- coding:utf-8 -*-
"""
 @author: scorpio.lu
 @datetime:2020-06-11 15:22
 @software: PyCharm
 @contact: luyi@zhejianglab.com

            ----------
             路有敬亭山
            ----------

"""
import os
import argparse
import sys
from datetime import datetime
import numpy as np
import math
import cv2
from scipy.spatial.distance import cdist
import oneflow as flow
from reid_model import resreid, HS_reid
from data_loader import Market1501

parser = argparse.ArgumentParser(description="flags for person re-identification")
parser.add_argument("--gpu_num_per_node", type=int, default=1, required=False)
parser.add_argument("--model", type=str, default="resreid", required=False, help="resreid or pcbreid")
parser.add_argument("--batch_size", type=int, default=300, required=False)
parser.add_argument("--data_dir", type=str, default='/home/oneflow_reid/person_reid/dataset', required=False, help="dataset directory")
parser.add_argument("-image_height", "--image_height", type=int, default=256, required=False)
parser.add_argument("-image_width", "--image_width", type=int, default=128, required=False)
parser.add_argument("--use_tensorrt", dest="use_tensorrt", action="store_true", default=False, required=False, help="inference with tensorrt")
parser.add_argument("--model_load_dir", type=str, default='/home/oneflow_reid/person_reid/model', required=False, help="model load directory")
parser.add_argument("--log_dir", type=str, default="./output", required=False, help="log info save directory")
args = parser.parse_args()

model={'resreid': resreid, 'HS-reid': HS_reid}
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
flow.config.gpu_device_num(args.gpu_num_per_node)
if args.use_tensorrt:
    func_config.use_tensorrt(True)

input_blob = flow.FixedTensorDef((args.batch_size, 3, args.image_height, args.image_width), dtype=flow.float)
#input_blob = flow.MirroredTensorDef((args.batch_size, 3, args.image_height, args.image_width), dtype=flow.float)

def resize_image(img, origin_h, origin_w, image_height, image_width):
    w = image_width
    h = image_height
    resized=np.zeros((3, image_height, image_width), dtype=np.float32)
    part=np.zeros((3, origin_h, image_width), dtype = np.float32)
    w_scale = (float)(origin_w - 1) / (w - 1)
    h_scale = (float)(origin_h - 1) / (h - 1)

    for c in range(w):
        if c == w-1 or origin_w == 1:
            val = img[:, :, origin_w-1]
        else:
            sx = c * w_scale
            ix = int(sx)
            dx = sx - ix
            val = (1 - dx) * img[:, :, ix] + dx * img[:, :, ix+1]
        part[:, :, c] = val
    for r in range(h):
        sy = r * h_scale
        iy = int(sy)
        dy = sy - iy
        val = (1-dy)*part[:, iy, :]
        resized[:, r, :] = val
        if r==h-1 or origin_h==1:
            continue
        resized[:, r, :] = resized[:, r, :] + dy * part[:, iy+1, :]
    return resized

def batch_image_preprocess(img_paths, img_height, img_weidth):
    result_list = []
    base = np.ones([args.image_height, args.image_width])
    norm_mean = [base * 0.485, base * 0.456, base * 0.406]  # imagenet mean
    norm_std = [0.229, 0.224, 0.225]  # imagenet std
    for img_path in img_paths:
        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
        img = img.transpose(2, 0, 1).astype(np.float32)  # hwc->chw
        img = img / 255  # /255  # to tensor
        img[[0, 1, 2], :, :] = img[[2, 1, 0], :, :]  # bgr2rgb

        w = img_weidth
        h = img_height
        origin_h = img.shape[1]
        origin_w = img.shape[2]

        resize_img = resize_image(img, origin_h, origin_w, h, w)
        # normalize

        resize_img[0] = (resize_img[0] - norm_mean[0])/ norm_std[0]
        resize_img[1] = (resize_img[1] - norm_mean[1]) / norm_std[1]
        resize_img[2] = (resize_img[2] - norm_mean[2]) / norm_std[2]
        result_list.append(resize_img)
    results = np.asarray(result_list).astype(np.float32)

    return results

def evaluate(qf, q_pids, q_camids, gf, g_pids, g_camids, max_rank=50):
    num_g = len(gf)
    num_q = len(qf)
    print('Computing distance matrix  ...')
    dist = cdist(qf, gf).astype(np.float16)
    dist = np.power(dist, 2).astype(np.float16)
    print('Computing CMC and mAP ...')
    if num_g < max_rank:
        max_rank = num_g
        print('Note: number of gallery samples is quite small, got {}'.format(num_g))
    indices = np.argsort(dist, axis=1)
    matches = (g_pids[indices] == q_pids[:, np.newaxis]).astype(np.int32)

    all_cmc = []
    all_AP = []
    num_valid_q = 0.

    for q_idx in range(num_q):

        q_pid = q_pids[q_idx]
        q_camid = q_camids[q_idx]

        order = indices[q_idx]
        remove = (g_pids[order] == q_pid) & (g_camids[order] == q_camid)
        keep = np.invert(remove)

        raw_cmc = matches[q_idx][keep]
        if not np.any(raw_cmc):
            continue

        cmc = raw_cmc.cumsum()
        cmc[cmc > 1] = 1

        all_cmc.append(cmc[:max_rank])
        num_valid_q += 1.

        num_rel = raw_cmc.sum()
        tmp_cmc = raw_cmc.cumsum()
        tmp_cmc = [x / (i + 1.) for i, x in enumerate(tmp_cmc)]
        tmp_cmc = np.asarray(tmp_cmc) * raw_cmc
        AP = tmp_cmc.sum() / num_rel
        all_AP.append(AP)

    assert num_valid_q > 0, 'Error: all query identities do not appear in gallery'

    all_cmc = np.asarray(all_cmc).astype(np.float32)
    all_cmc = all_cmc.sum(0) / num_valid_q
    mAP = np.mean(all_AP)

    return all_cmc, mAP

@flow.function(func_config)
def reid_eval_job(image=input_blob):
    features = resreid(image, trainable=False)
    return features

class ReIDInference(object):
    def __init__(self):
        check_point = flow.train.CheckPoint()
        if args.model_load_dir:
            assert os.path.isdir(args.model_load_dir)
            print("Restoring model from {}.".format(args.model_load_dir))
            check_point.load(args.model_load_dir)
        else:
            print("Init model on demand.")
            check_point.init()
            snapshot_save_path = os.path.join(args.model_save_dir, "last_snapshot")
            if not os.path.exists(snapshot_save_path):
                os.makedirs(snapshot_save_path)
            print("Saving model to {}.".format(snapshot_save_path))
            check_point.save(snapshot_save_path)

    def inference(self, imgs):
        query_images = batch_image_preprocess(imgs, args.image_height, args.image_width)
        batch_times = math.ceil(len(imgs)/args.batch_size)
        features = []
        for i in range(batch_times):
            start = max(0, i*args.batch_size)
            end = min((i+1)*args.batch_size, len(query_images))
            array = query_images[start:end]
            feature = reid_eval_job([array]).get()
            features.extend(feature.ndarray_list_[0])
        return features

def main():
    print("=".ljust(66, "="))
    print("Running {}: num_gpu = {}.".format(args.model, args.gpu_num_per_node))
    print("=".ljust(66, "="))
    for arg in vars(args):
        print("{} = {}".format(arg, getattr(args, arg)))
    print("-".ljust(66, "-"))
    print("Time stamp: {}".format(str(datetime.now().strftime("%Y-%m-%d-%H:%M:%S"))))
    flow.env.grpc_use_no_signal()
    flow.env.log_dir(args.log_dir)
    obj = ReIDInference()

    print("Loading data from {}".format(args.data_dir))
    dataset = Market1501(root=args.data_dir)
    query_img, query_id, query_cam_id = zip(*dataset.query)
    gallery_img, gallery_id, gallery_cam_id = zip(*dataset.gallery)

    print('extracting query features...')
    query_features = obj.inference(query_img)
    print('extracting query features done...')
    print('extracting gallery features...')
    gallery_features = obj.inference(gallery_img)
    print('extracting gallery features done...')
    cmc, mAP = evaluate(query_features, np.array(query_id), np.array(query_cam_id), gallery_features, np.array(gallery_id), np.array(gallery_cam_id))
    print("=".ljust(30, "=")+" Result "+ "=".ljust(30, "="))
    print('mAP: {:.1%}'.format(mAP))
    print('CMC curve')
    for r in [1, 5, 10]:
        print('Rank-{:<3}: {:.1%}'.format(r, cmc[r - 1]))
    print("=".ljust(66, "="))

if __name__ == "__main__":
    main()


#!/usr/bin/python3

import  RPi.GPIO as GPIO
import time
import numpy

import cv2
#from CocoDrink import *

import multiprocessing as mp

HIGH = 1
LOW = 0

def LoadColorImage():
	path = '/home/pi/Tapomatic/v2020.0/static/images/BacardiLogoV0.png'
	originalImage = cv2.imread(path)

	print('Dimensions: ', originalImage.shape)
	imgHeight, imgWidth, colorChannels = originalImage.shape
	print('Height: ', imgHeight)
	print('Width: ', imgWidth)
	print('Color Channels: ', colorChannels)

	cv2.imshow('Color Image', originalImage)
	cv2.waitKey(0)
	cv2.destroyAllWindows()

	#return originalImage

def LoadGrayScaleImage():
	path = '/home/pi/Tapomatic/v2020.0/static/images/BacardiLogoV0.png'
	grayScaleImage = cv2.imread(path, cv2.IMREAD_GRAYSCALE)

	print('Dimensions: ', grayScaleImage.shape)
	imgHeight, imgWidth =  grayScaleImage.shape
	print('Height: ', imgHeight)
	print('Width: ', imgWidth)

	#cv2.imshow('Gray Image', grayScaleImage)
	#cv2.waitKey(0)
	#cv2.destroyAllWindows()

	return grayScaleImage

def LoadBlackWhiteImage():
	path = '/home/pi/Tapomatic/v2020.0/static/images/BacardiLogoV1.png'
	originalImage = cv2.imread(path, cv2.IMREAD_UNCHANGED)

	# Convert to gray scale first to apply better thresholding which will create $                        grayImg = cv2.cvtColor(originalImage, cv2.COLOR_BGR2GRAY)
	grayImg = cv2.cvtColor(originalImage, cv2.COLOR_BGR2GRAY)

	# Anything above 127 on a scale from 0 to 255 is WHITE
	(thresh, bwImage) = cv2.threshold(grayImg, 127, 255, cv2.THRESH_BINARY)

	print('Dimensions: ', bwImage.shape)
	imgHeight, imgWidth =  bwImage.shape
	print('Height: ', imgHeight)
	print('Width: ', imgWidth)

	cv2.imshow('B & W Image', bwImage)
	cv2.waitKey(0)
	cv2.destroyAllWindows()

	return bwImage

def BurnImage(img, laserConstant):
	"""
	Toogle GPIO pin possibly connected to a high power relay HIGH to turn ON a LASER
	Puts CPU to sleep so NOT a threadable function yet

	Key arguments:
	laserConfig -- TODO REMOVE?

	Return value:
	NOTHING
	"""
	pixelBurnDuration = 0.015

	dutyCycle = 100
	frequency = 100                                         # Desired LASER pulse in Hz

	imageBurnComplete = False
	pixelNum = 1

	path = '/home/pi/Tapomatic/v2020.0/static/images/BacardiLogoV0.png'
	grayScaleImage = cv2.imread(path, cv2.IMREAD_GRAYSCALE)

	imgHeight, imgWidth =  grayScaleImage.shape
	totalPixels = imgHeight * imgWidth
	print('Total Number Of Pixels: ', totalPixels)

	highTime = 1/frequency  * dutyCycle * laserConstant
	print('High Time: ', highTime)
	print('Low Time: ', highTime - (1/frequency))

	while(not imageBurnComplete):

		#https://www.machinelearningplus.com/python/parallel-processing-python/
		pool = mp.Pool(2)
		results = pool.apply(MoveLaserStepperMotor, args=(totalPixels, pixelDwellDuration, frequency))

	pool.close

		time.sleep(highTime)
		GPIO.output(7, HIGH)
		time.sleep(highTime - (1/frequency))
		GPIO.output(7, LOW)

		#imageBurnComplete = MoveLaserStepperMotor(pixelDwellDuration, frequency)
		pixelNum = pixelNum + 1
		if(pixelNum >= totalPixel):
			imageBurnComplete = True

def MoveLaserStepperMotor(pixelDwellDuration, frequency):
	print('NOT DONE')

if __name__ == "__main__":

	print("Number of processors: ", mp.cpu_count())

	GPIO.setmode(GPIO.BOARD)
	GPIO.setup(7, GPIO.OUT, initial=GPIO.LOW)

	LoadColorImage()
	LoadGrayScaleImage()
	LoadBlackWhiteImage()
	BurnImage(LoadGrayScaleImage, 1.0)

	# Relays
	#GPIO.setup(15, GPIO.OUT, initial=GPIO.LOW)
	#GPIO.setup(31, GPIO.OUT, initial=GPIO.LOW)
	#GPIO.setup(37, GPIO.OUT, initial=GPIO.LOW)

	GPIO.output(7, HIGH)
	#GPIO.output(15, HIGH)
	#GPIO.output(31, HIGH)
	GPIO.output(37, HIGH)

	time.sleep(3)

	GPIO.cleanup()
	##GPIO.cleanup(15)
	##GPIO.cleanup(31)
	##GPIO.cleanup(37)


import os
import threading
import time
from collections import defaultdict
import atexit

import numpy
from gi.repository import Gtk, GObject
from matplotlib import rcParams
from matplotlib.animation import FuncAnimation
from matplotlib.backends.backend_gtk3cairo import FigureCanvasGTK3Cairo as FigureCanvas
from matplotlib.figure import Figure
from matplotlib.lines import Line2D
from mxdc.utils import misc
from mxdc.widgets import dialogs

rcParams['font.family'] = 'Cantarell'
rcParams['font.size'] = 10

COLORS = [
    '#1f77b4',
    '#ff7f0e',
    '#2ca02c',
    '#d62728',
    '#9467bd',
    '#8c564b',
    '#e377c2',
    '#7f7f7f',
    '#bcbd22',
    '#17becf'
]

class TickerChart(Gtk.Box):
    __gsignals__ = {
        'cursor-time': (GObject.SignalFlags.RUN_FIRST, None, (object,)),
    }

    def __init__(self, interval=100, view=20, keep=None, linewidth=1):
        super().__init__()
        self.fig = Figure(dpi=72)
        self.canvas = FigureCanvas(self.fig)
        self.pack_start(self.canvas, True, True, 0)

        self.data = {}
        self.plots = {}
        self.info = {}
        self.alternates = set()
        self.active = None

        self.axes = []
        self.axes.append(self.fig.add_subplot(111))
        self.fig.subplots_adjust(left=0.12, right=0.88)

        self.axes[0].set_xlabel('seconds ago')
        self.interval = interval  # milliseconds
        self.view_range = view  # seconds

        self.keep_range = keep or (view * 4)  # seconds
        self.linewidth = linewidth

        self.keep_size = int(self.keep_range * 1000 / self.interval)
        self.view_step = view // 2
        self.deviation = 20

        self.view_time = time.time()
        self.add_data('time')
        self.paused = False
        self.show_all()

    def pause(self):
        self.paused = True

    def resume(self):
        self.paused = False

    def zoom_out(self):
        self.view_range = max(self.view_range - self.view_step, self.view_step)
        self.update()

    def zoom_in(self):
        self.view_range = min(self.view_range + self.view_step, self.keep_range)
        self.update()

    def incr_margin(self):
        self.deviation = min(self.deviation + 5, 50)
        self.update()

    def decr_margin(self):
        self.deviation = max(self.deviationi - 5, 5)
        self.update()

    def add_data(self, name):
        if name in self.data:
            return
        self.data[name] = numpy.empty(self.keep_size)
        self.data[name][:] = numpy.nan

    def resize_data(self):
        for name, data in list(self.data.items()):
            self.data[name] = numpy.empty(self.keep_size)
            if self.max_samples > len(data):
                self.data[name][-len(data):] = data
            else:
                self.data[name] = data[-self.keep_size:]

    def select_active(self, name):
        if name in self.alternates:
            self.active = name
            self.axes[1].set_ylabel(name)

    def add_alternate(self, name):
        self.alternates.add(name)

    def shift_data(self):
        for name, data in list(self.data.items()):
            data[:-1] = data[1:]

    def add_plot(self, name, color=None, linestyle='-', axis=0, alternate=False):
        assert axis in [0, 1], 'axis must be 0 or 1'
        if axis == 1 and len(self.axes) == 1:
            self.axes.append(self.axes[0].twinx())
        if not color:
            color = COLORS[len(self.plots)]
        self.plots[name] = Line2D([], [], color=color, linewidth=self.linewidth, linestyle=linestyle)
        self.axes[axis].add_line(self.plots[name])
        self.axes[axis].set_ylabel(name, color=color)
        self.info[name] = {'color': color, 'linestyle': linestyle, 'axis': axis}
        if alternate:
            self.add_alternate(name)
            self.select_active(name)
        self.add_data(name)

    def clear(self):
        for name, line in list(self.plots.items()):
            self.data[name][:] = numpy.nan

    def update(self):
        if self.paused:
            return
        selector = ~numpy.isnan(self.data['time'])
        selector[selector] = (
            (self.data['time'][selector] > (self.view_time - self.view_range))
            & (self.data['time'][selector] <= self.view_time)
        )

        if selector.sum() < 2:
            return
        now = self.data['time'][selector][-1]
        x_data = self.data['time'][selector] - now
        xmin, xmax = min(x_data.min(), -self.view_range), x_data.max()

        extrema = defaultdict(lambda: (numpy.nan, numpy.nan))

        for name, line in list(self.plots.items()):
            if name in self.alternates and name != self.active: continue
            axis = self.info[name]['axis']
            ymin, ymax = extrema[axis]
            y_data = self.data[name][selector]
            mn, mx = misc.get_min_max(y_data, ldev=self.deviation, rdev=self.deviation)
            ymin, ymax = numpy.nanmin([ymin, mn]), numpy.nanmax([ymax, mx])
            extrema[axis] = (ymin, ymax)
            line.set_data(x_data, y_data)

        for i, (ymin, ymax) in list(extrema.items()):
            if ymin != ymax:
                self.axes[i].set_ylim(ymin, ymax)
            if xmin != xmax:
                self.axes[i].set_xlim(xmin, xmax)

    def redraw(self):
        self.canvas.draw_idle()

    def animate(self, i):
        self.update()
        return list(self.plots.values())

    def save(self):
        dialog = Gtk.FileChooserDialog(
            "Save Chart ...", dialogs.MAIN_WINDOW, Gtk.FileChooserAction.SAVE,
            (Gtk.STOCK_CANCEL, Gtk.ResponseType.CANCEL, Gtk.STOCK_SAVE, Gtk.ResponseType.OK)
        )
        dialog.set_size_request(600, 300)
        response = dialog.run()
        if response == Gtk.ResponseType.OK:
            img_filename = dialog.get_filename()
            if os.access(os.path.dirname(img_filename), os.W_OK):
                self.fig.savefig(img_filename)
        dialog.destroy()

class ChartManager(GObject.GObject):
    def __init__(self, interval=100, view=20):
        GObject.GObject.__init__(self)
        self.chart = TickerChart(interval=interval, view=view)
        self.sources = {}
        self.values = {}
        self.interval = interval / 1000.  # convert from milliseconds to seconds
        self.animation = FuncAnimation(self.chart.fig, self.chart.animate, None, interval=interval, blit=False)
        self.start()
        atexit.register(self.stop)

    def add_plot(self, dev, name, signal='changed', color=None, linestyle='-', axis=0, alternate=False):
        self.chart.add_plot(name, color=color, linestyle=linestyle, axis=axis, alternate=alternate)
        self.values[name] = numpy.nan
        self.sources[name] = dev.connect(signal, self.collect_data, name)

    def select_active(self, name):
        self.chart.select_active(name)

    def zoom_in(self, *args, **kwargs):
        self.chart.zoom_in()

    def zoom_out(self, *args, **kwargs):
        self.chart.zoom_out()

    def clear(self, *args, **kwargs):
        self.chart.clear()

    def save(self, *args, **kwargs):
        self.chart.save()

    def collect_data(self, dev, value, name):
        self.values[name] = value

    def start(self):
        """Start the Data monitor thread """
        self._stopped = False
        worker_thread = threading.Thread(name="TickerSampler", target=self.update_data)
        worker_thread.setDaemon(True)
        worker_thread.start()

    def stop(self):
        self._stopped = True

    def pause(self, *args, **kwargs):
        self.chart.pause()

    def resume(self, *args, **kwargs):
        self.chart.resume()

    def is_paused(self):
        return self.chart.paused

    def update_data(self):
        # update the values of the array every interval seconds, shift left
        while not self._stopped:
            self.chart.shift_data()
            if not self.is_paused():
                for name, value in list(self.values.items()):
                    if name in self.chart.data:
                        self.chart.data[name][-1] = value
            self.chart.data['time'][-1] = time.time()
            self.chart.view_time = time.time()
            time.sleep(self.interval)

    def cleanup(self):
        self.stop()


from bokeh.plotting import Figure
from bokeh.layouts import row, column

from bokeh.models import HoverTool, ColumnDataSource, Range1d
from bokeh.models import LinearColorMapper, ColorBar

from qlf_models import QLFModels
from dashboard.bokeh.helper import get_palette

import numpy as np
import logging
from bokeh.resources import CDN
from bokeh.embed import file_html

import os
from dashboard.models import Job, Process, Fibermap

spectro_data = os.environ.get('DESI_SPECTRO_DATA')

logger = logging.getLogger(__name__)

class GlobalSnr:
    def __init__(self, process_id, arm):
        self.selected_process_id = process_id
        self.selected_arm = arm

    def data_source_arm(self, fmap, arm):
        """ Creating data source for plots
        """
        data_model = {
            'resids': [],
            'hover': [],
            'OBJ_TYPE': [],
            'ra':  [],
            'dec': [],
            'cam': [],
        }

        process_id = self.selected_process_id
        process = Process.objects.get(pk=process_id)
        joblist = [entry.camera.camera for entry in Job.objects.filter(
            process_id=process_id)]

        ra_tile = fmap.fiber_ra
        dec_tile = fmap.fiber_dec
        otype_tile = fmap.objtype

        objlist = sorted(set(otype_tile))
        if 'SKY' in objlist:
            objlist.remove('SKY')

        ra_snr = []
        dec_snr = []
        resids_snr = []
        ot_snr = []
        cam_snr = []

        for spect in list(range(10)):
            cam = arm+str(spect)

            ra_petal = []
            dec_petal = []
            resids_petal = []
            ot_petal = []
            cam_petal = []
            if cam in joblist:
                mergedqa = QLFModels().get_output(
                    self.selected_process_id, cam)

                med_snr = np.array(
                    mergedqa['TASKS']['CHECK_SPECTRA']['METRICS']["MEDIAN_SNR"])

                # color values:
                resids_petal = (
                    mergedqa['TASKS']['CHECK_SPECTRA']['METRICS']['SNR_RESID'])

                ra_petal = []
                dec_petal = []
                fibers_snr = []

                obj = np.arange(len(objlist))

                for t in range(len(obj)):
                    otype = list(objlist)[t]
                    oid = np.where(np.array(list(objlist)) == otype)[0][0]
                    if otype == 'STD':
                        fibers = mergedqa['GENERAL_INFO']['STD_FIBERID']
                    elif otype == 'TGT':
                        fibers = mergedqa['GENERAL_INFO']['STAR_FIBERID']

                        # fibers = mergedqa['GENERAL_INFO']['%s_FIBERID'%otype]

                    fibers_snr = fibers_snr + fibers

                    for j in range(len(fibers)):
                        ras = mergedqa['GENERAL_INFO']['RA'][fibers[j]]
                        decs = mergedqa['GENERAL_INFO']['DEC'][fibers[j]]
                        ra_petal.append(ras)
                        dec_petal.append(decs)
                        ot_petal.append(otype)

                ra_snr = ra_snr+ra_petal
                dec_snr = dec_snr+dec_petal
                resids_snr = resids_snr + \
                    [i if i > -999 else np.nan for i in resids_petal]
                ot_snr = ot_snr + ot_petal
                cam_snr = cam_snr + [cam]*len(ra_petal)

            else:
                ra_snr = ra_snr + list(ra_tile[500*spect: 500*(spect + 1)])
                dec_snr = dec_snr + list(dec_tile[500*spect: 500*(spect + 1)])
                resids_snr = resids_snr + 500*[np.nan]
                ot_snr = ot_snr + list(otype_tile[500*spect: 500*(spect + 1)])
                cam_snr = cam_snr + [str(cam)]*500

        data_model['ra'] = ra_snr  # ra_tile
        data_model['dec'] = dec_snr  # dec_tile
        data_model['resids'] = resids_snr
        data_model['hover'] = ['%4.3f' % (ires) for ires in resids_snr]
        data_model['OBJ_TYPE'] = ot_snr
        data_model['cam'] = cam_snr

        source = ColumnDataSource(data=data_model)
        return source

    def data_source(self, fmap):
        """ Creating data source for plots
        """
        data_model = {
            'x_b': [],
            'x_r': [],
            'x_z': [],
            'OBJ_TYPE': [],
            'ra':  [],
            'dec': [],
            'cam_b': [],
            'cam_r': [],
            'cam_z': []
        }

        process_id = self.selected_process_id
        joblist = [entry.camera.camera for entry in Job.objects.filter(
            process_id=process_id)]

        ra_tile = fmap.fiber_ra
        dec_tile = fmap.fiber_dec
        otype_tile = fmap.objtype

        objlist = sorted(set(otype_tile))
        if 'SKY' in objlist:
            objlist.remove('SKY')

        fibaux = []

        y = []
        cam_inst = []
        ra_all = []
        dec_all = []

        for spect in list(range(1)):
            cam = self.selected_arm+str(spect)
            if cam in joblist:
                mergedqa = QLFModels().get_output(
                    self.selected_process_id, cam)

                # Assign available objects
                objtype = otype_tile[500*spect: 500*(spect + 1)]

                med_snr = np.array(
                    mergedqa['TASKS']['CHECK_SPECTRA']['METRICS']["MEDIAN_SNR"])
                resids = mergedqa['TASKS']['CHECK_SPECTRA']['METRICS']['SNR_RESID']

                obj = np.arange(len(objlist))

                rayes = []
                decyes = []

                for t in range(len(obj)):
                    otype = list(objlist)[t]
                    oid = np.where(np.array(list(objlist)) == otype)[0][0]

                    if otype == 'STD':
                        fibers = mergedqa['GENERAL_INFO']['STAR_FIBERID']
                    else:
                        fibers = mergedqa['GENERAL_INFO']['%s_FIBERID' % otype]

                    fibaux = fibaux + [500*spect + i for i in fibers]

                    for i in range(len(fibers)):
                        ras = mergedqa['GENERAL_INFO']['RA'][fibers[i]]
                        decs = mergedqa['GENERAL_INFO']['DEC'][fibers[i]]
                        rayes.append(ras)
                        decyes.append(decs)

                ra_all = ra_all + rayes
                dec_all = dec_all + decyes

                nanresids = [i if i > -9999. else np.nan for i in resids]

                y = y + nanresids

            else:
                y = y + 500*[np.nan]
                ra_all = ra_all + \
                    [i for i in ra_tile[500*spect:500*(spect + 1)]]
                dec_all = dec_all + \
                    [i for i in dec_tile[500*spect:500*(spect + 1)]]
                fibaux = fibaux + list(range(500*spect, 500*(spect+1)))

            cam_inst = cam_inst + [cam]*500

            data_model['x_' + cam[0]] = y
            data_model['cam_'+cam[0]] = cam_inst

        data_model['OBJ_TYPE'] = [otype_tile[ii] for ii in fibaux]
        data_model['ra'] = ra_all
        data_model['dec'] = dec_all

        source = ColumnDataSource(data=data_model)

        return source

    def wedge_plot(self, wedge_arm, fmap, common_source=None, sigma_kind='x'):
        ra_center = fmap.exposure.telra
        dec_center = fmap.exposure.teldec

        fiber_tooltip = """

                    RESIDUAL SNR:
                    @y

                    RA:
                    @ra

                    DEC:
                    @dec

                    Obj Type:
                    @OBJ_TYPE

                    CAM:
                    @cam

        """
        fiber_tooltip = fiber_tooltip.replace('@y', '@hover')

        hover = HoverTool(tooltips=fiber_tooltip)

        my_palette = get_palette("bwr")  # "seismic")#"RdYlBu_r")#"viridis")
        source = common_source

        sigma = source.data['resids']  # ['{}_'.format(sigma_kind) +wedge_arm]
        rng_min, rng_max = np.nanmin(sigma), np.nanmax(sigma)
        rng = rng_max-rng_min

        if np.isnan(rng_min) or np.isnan(rng_max):
            fill_color = 'lightgray'
        else:
            mapper = LinearColorMapper(palette=my_palette,  nan_color='darkgrey',
                                       low=-0.2,
                                       high=0.2)

            fill_color = {'field': 'resids', 'transform': mapper}

        radius = 0.017
        radius_hover = 0.018

        xrange = Range1d(start=ra_center + 2, end=ra_center-2)
        yrange = Range1d(start=dec_center+1.8, end=dec_center-1.8)

        p = Figure(title='SNR (ARM {})'.format(wedge_arm), x_axis_label='RA', y_axis_label='DEC', plot_width=600, plot_height=600, tools=[hover, "pan,wheel_zoom,reset,box_zoom,crosshair"], active_drag="box_zoom", x_range=xrange, y_range=yrange
                   )

        p.title.align = 'center'
        p.circle('ra', 'dec', source=source, name="data", radius=radius,
                 fill_color=fill_color,
                 line_color='black', line_width=0.4,
                 hover_line_color='red')

        p.circle('ra', 'dec', source=source, name="data", radius=radius_hover,
                 hover_fill_color=fill_color,
                 fill_color=None,
                 line_color=None, line_width=3, hover_line_color='orange')

        if 'mapper' in locals():
            cbar = Figure(height=p.plot_height,
                          width=140,
                          toolbar_location=None,
                          min_border=0,
                          outline_line_color=None,
                          )

            color_bar = ColorBar(color_mapper=mapper, label_standoff=14,
                                 major_label_text_font_style="bold", padding=26,
                                 major_label_text_align='right',
                                 major_label_text_font_size="10pt",
                                 location=(0, 0))
            cbar.title.align = 'center'
            cbar.title.text_font_size = '10pt'
            cbar.add_layout(color_bar, 'left')
            p_list = [cbar, p]
        else:
            p_list = [p]

        return p_list

    def load_qa(self):
        process_id = self.selected_process_id
        process = Process.objects.get(pk=process_id)
        exposure = process.exposure
        fmap = Fibermap.objects.filter(exposure=exposure).last('pk')

        src_arm = self.data_source_arm(fmap, self.selected_arm)

        # , common_source=source)
        p = self.wedge_plot(self.selected_arm, fmap, common_source=src_arm)
        layout = row(column(row(p)))

        return file_html(layout, CDN, "Global SNR")


import random
import argparse
import requests
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import matplotlib.patheffects as path_effects
from cartopy.io.img_tiles import OSM
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
import cartopy.feature as cpf
from geopy import Point
from geopy.distance import vincenty

# Create global empty variables for update_flights function
flight_list = {}
coords_list = []
color_list = []
annotation_list = []

def create_map(long, lat, ext):
    # Set projection type
    projection = ccrs.PlateCarree()
    # Fetch map tiles
    tiles = OSM()

    # Create figure and ax with gridlines and formated axes
    fig, ax = plt.subplots(figsize=(10, 10),
                           subplot_kw=dict(projection=projection))
    gl = ax.gridlines(draw_labels=True)
    gl.xlabels_top = False
    gl.ylabels_right = False
    gl.xformatter = LONGITUDE_FORMATTER
    gl.yformatter = LATITUDE_FORMATTER

    # Take part of projection calculated by ext
    ax.set_extent(ext, projection)
    # Add tiles to ax with zoom 7 and sline36 interpolation
    ax.add_image(tiles, 7, interpolation='spline36')
    # Plot observing point
    ax.plot([long], [lat], 'bs')
    # Add title to plit
    fig.suptitle('Live Flight Tracker', fontsize=16)

    # Create empty scatter plot
    track_flights = ax.scatter([], [], marker='o', c=[], s=14, alpha=.85, edgecolors = 'k')
    return fig, ax, track_flights

def update_flights(self, long, lat, dist, fig, ax, track_flights):

    global coords_list
    global color_list
    global annotation_list
    global flight_list

    print('---------------------')

    # Request to AdsExchange API
    url = 'http://public-api.adsbexchange.com/VirtualRadar/AircraftList.json'
    payload = {'lat': lat, 'lng': long, 'fDstL': 0, 'fDstU': dist}
    r = requests.get(url, params=payload, headers={'Connection':'close'})

    # If call was not correct, return with no change
    if not r.status_code == 200:
        return track_flights, annotation_list
    js_str = r.json()
    # If response is empty, return with no change
    if js_str['lastDv'] == str(-1):
        return track_flights, annotation_list

    # Clean annotation list
    for anot in annotation_list:
        anot.remove()
    annotation_list[:] = []

    # Extract latitude, longitude and name of all flights
    for flight in js_str['acList']:
        # If light not in dict, add it
        if not flight['Icao'] in flight_list:
            flight_list[flight['Icao']] = {'coords':[], 'postime':[], 'color':None, 'inimage':10}
            flight_list[flight['Icao']]['color'] = (random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1))

        coords = (flight['Long'], flight['Lat'])
        icao = flight['Icao']
        postime = flight['PosTime']

        # Set limit for not found in repsone to 10
        flight_list[icao]['inimage'] = 10

        # Check if flight data is not from history, return the lastest if so
        if (len(flight_list[icao]['postime']) > 1 and
                int(flight_list[icao]['postime'][-1]) >= int(postime)):
                coords = flight_list[icao]['coords'][-1]
                postime = flight_list[icao]['postime'][-1]

        # Add coords and postime to lists in flight dictionary
        flight_list[icao]['coords'].append(coords)
        flight_list[icao]['postime'].append(postime)

        # Append coords and color to list for scatter plot input
        coords_list.append(flight_list[icao]['coords'][-1])
        color_list.append(flight_list[icao]['color'])

        # Create annotation of flight each flight
        anonnotation = ax.text(coords[0], coords[1], icao, fontsize=12,
                                   fontweight='bold', size=9,
                                   color=flight_list[icao]['color'],
                                   path_effects=[path_effects.withStroke(linewidth=4, foreground="black")])
        annotation_list.append(anonnotation)

    # Error handling when first call is unsucessful
    if not flight_list:
        return

    # Unzip coodinates tuples to lat and long
    uzipped_coords = [[*x] for x in zip(*coords_list)]
    # Create np array of lats and longs
    data = np.array(uzipped_coords, dtype=object)
    data = np.transpose(data)

    # Put lat and logs into scatter plot
    track_flights.set_offsets(data)
    track_flights.set_facecolors(color_list)

    # Check if flight is still in image if not subtract 1 from limit of not found
    # If in list restore to 10 or if 0 delete flight from flights dictionary
    for key in flight_list.copy():
        flight_list[key]['inimage'] = flight_list[key]['inimage'] - 1
        if flight_list[key]['inimage'] < 1:
            for i in flight_list[key]['coords']:
                coords_list.remove(i)
            color_list = [col for col in color_list if col != flight_list[key]['color']]
            del flight_list[key]

    return track_flights, annotation_list

def create_extent(long, lat, dist):
    # Calculate gps coordinates for dist square around point of interest using Vincenty's formula
    extent_north = vincenty(kilometers=dist).destination(Point(lat, long), 0).format_decimal()
    extent_east = vincenty(kilometers=dist).destination(Point(lat, long), 90).format_decimal()
    extent_south = vincenty(kilometers=dist).destination(Point(lat, long), 180).format_decimal()
    extent_west = vincenty(kilometers=dist).destination(Point(lat, long), 270).format_decimal()
    return [float(extent_west.split(',')[1]), float(extent_east.split(',')[1]),
            float(extent_south.split(',')[0]), float(extent_north.split(',')[0])]

def parse_args():
    # Initialize the parser
    parser = argparse.ArgumentParser(description="Live Flight Tracker")
    parser.add_argument("latitude", type=float, help="Latitude of a observing point")
    parser.add_argument("longitude", type=float, help="Longitude of a observing point")
    parser.add_argument("-d", "--distance", type=int, default=150,
                        help="destination in km to cover from observing point")
    args = parser.parse_args()

    return args.distance, args.latitude, args.longitude


import numpy as np
import struct
from deca.file import ArchiveFile
from deca.ff_aaf import load_aaf_header
from deca.ff_avtx import DdImageHeader
from deca.ff_types import *

# guess type
raw_image_size = {
    0x1000000: '',
    0x800000: '',
    0x400000: '',
    0x280000: '',
    0x200000: '',
    0x140000: '',
    0x100000: '',
    0xa0000: '',
    0x80000: '',
    0x40000: '',
    0x20000: '',
}

def file_stats(f, file_size):
    # TODO huge files will take a WHILE
    counts = np.zeros((256,), dtype=np.uint32)
    while True:
        buf = f.read(1024*1024)
        if len(buf) == 0:
            break
        buf = bytearray(buf)
        cnt = np.bincount(buf, minlength=256)
        counts = counts + cnt
    return counts

def determine_file_type_and_size(f, file_size0):
    file_type = None
    file_sub_type = None
    file_size = file_size0

    start_pos = f.tell()
    magic = f.read(256)
    magic_int = None
    if len(magic) >= 20:
        magic_int = struct.unpack('I', magic[0:4])[0]

        if b' FDA' == magic[0:4]:
            file_type = FTYPE_ADF
        elif b'\x00FDA' == magic[0:4]:
            file_type = FTYPE_ADF0
        elif b'AVTX' == magic[0:4]:
            file_type = FTYPE_AVTX
            header = DdImageHeader()
            header.deserialize_ddsc(magic)
            file_sub_type = header.dds_header_dxt10.dxgiFormat
        elif b'DDS ' == magic[0:4]:
            file_type = FTYPE_DDS
            header = DdImageHeader()
            header.deserialize_dds(magic)
            file_sub_type = header.dds_header_dxt10.dxgiFormat
        elif b'AAF' == magic[0:3].upper():
            file_type = FTYPE_AAF
            f.seek(start_pos)
            aafh = load_aaf_header(f)
            file_size = aafh.size_u
        elif b'RTPC' == magic[0:4]:
            file_type = FTYPE_RTPC
        elif b'CFX' == magic[0:3]:
            file_type = FTYPE_GFX
        elif b'GFX' == magic[0:3]:
            file_type = FTYPE_GFX
        elif b'RIFF' == magic[0:4]:
            file_type = FTYPE_RIFF
        elif b'OggS' == magic[0:4]:
            file_type = FTYPE_OGG
        elif b'BM6' == magic[0:3]:
            file_type = FTYPE_BMP
        elif b'BM8' == magic[0:3]:
            file_type = FTYPE_BMP
        elif b'MDI\x00' == magic[0:4]:
            file_type = FTYPE_MDI
        elif b'PFX\x00' == magic[0:4]:
            file_type = FTYPE_PFX
        elif b'SARC' == magic[4:8]:
            file_type = FTYPE_SARC
        elif b'TAG0' == magic[4:8]:
            file_type = FTYPE_TAG0
        elif b'FSB5' == magic[16:20]:
            file_type = FTYPE_FSB5C
        elif b'\x57\xE0\xE0\x57\x10\xC0\xC0\x10' == magic[0:8]:
            file_type = FTYPE_H2014
        elif b'\x05\x00\x00\x00RBMDL' == magic[0:9]:
            file_type = FTYPE_RBMDL
        elif b'KB2' == magic[0:3]:
            file_type = FTYPE_BINK_KB2
        elif b'BIK' == magic[0:3]:
            file_type = FTYPE_BINK_BIK
        elif b'GT0C' == magic[0:4]:
            file_type = FTYPE_GT0C

    # need to inspect file structure

    fm = ArchiveFile(f)

    if file_type is None:
        # OBC files with (u32)4, (u32)count , 80 * count bytes, something to do with areas on the map? object placement?
        fm.seek(start_pos)
        ver = fm.read_u32()
        cnt = fm.read_u32()
        if ver == 4 and cnt * 80 + 8 == file_size0:
            file_type = FTYPE_OBC

    if file_type is None:  # text file only contains text bytes, json, xml, ...
        fm.seek(start_pos)
        counts = file_stats(fm, file_size0)
        all_sum = np.sum(counts)
        pri_sum = np.sum(counts[[9, 10, 13] + list(range(20, 128))])
        if all_sum == pri_sum:
            file_type = FTYPE_TXT

    return file_type, file_size, magic_int, file_sub_type