import numpy as np
import tensorflow as tf

from Mybase import layers
from Mybase.layers import *
from Mybase.layers_utils import *
from Mybase.losses import *

from Mybase.deep_metric_utils.bbox import *
from Mybase.deep_metric_utils.embcls_layer import *
from Mybase.deep_metric_utils.embcls_target_layer import *

class Resnet101_Deep_Metric(object):

    def __init__(self, cls_num=81, reg=1e-4, drp=0.5, typ=tf.float32):

        self.cls_num = cls_num #class number
        self.reg     = reg     #regularization
        self.drp     = drp     #dropout
        self.typ     = typ     #dtype

        self.mod_tra = True    #mode training
        self.glb_pol = False   #global pooling
        self.inc_btm = True    #include bottom block

        #resnet block setting
        self.res_set = [( 256,  64, [1, 1], [1, 1],  3, True ),  #conv2x 256 /4
                        ( 512, 128, [1, 1], [2, 2],  4, True ),  #conv3x 128 /8  #use
                        (1024, 256, [1, 1], [2, 2], 23, True ),  #conv4x 64  /16 #23--->101 #6--->50
                        (2048, 512, [1, 1], [2, 2],  3, True )]  #conv5x 64  /16 #use
        self.out_srd = 4 #output stride

    def forward(self, imgs=None, lbls=None, gbxs=None, gmks=None, gbx_nums=None, mtra=None, scp=None):
        #lbl=label #scp=scope #mod_tra=mode train

        img_shp = imgs.get_shape().as_list()
        img_num, img_hgt, img_wdh = img_shp[0], img_shp[1], img_shp[2]
        fet_hgt = img_hgt // (self.out_srd*2)
        fet_wdh = img_wdh // (self.out_srd*2)
        img_shp = np.stack([img_hgt, img_wdh], axis=0)
        fet_shp = np.stack([fet_hgt, fet_wdh], axis=0)
        #common parameters
        com_pams = {
            'com':    {'reg': self.reg, 'wscale': 0.01, 'dtype': self.typ, 'reuse': False, 'is_train': False, 'trainable': True},
            'bn':     {'eps': 1e-5, 'decay': 0.9997}, #0.9997
            'relu':   {'alpha': -0.1},
            'conv':   {'number': 64,'shape':[7, 7],'rate':1,'stride':[2, 2],'padding':'SAME','use_bias':True},
            'deconv': {'number':256,'shape':[2, 2],'rate':1,'stride':[2, 2],'padding':'SAME','out_shape':[128, 128],'use_bias':False},
            'max_pool':     {'shape': [3, 3], 'stride': [2, 2], 'padding': 'VALID'},
            'resnet_block': {'block_setting': self.res_set, 'output_stride': self.out_srd},
            'resnet_unit':  {'depth_output':1024, 'depth_bottle':256, 'use_branch':True, 'shape':[1, 1], 'stride':[1, 1], 'rate':1},
            'pyramid':   {'depth': 256},
            'glb_pool':  {'axis':  [1, 2]},
            'reshape':   {'shape': []},
            'squeeze':   {'axis':  [1, 2]},
            'transpose': {'perm':  [0, 3, 1, 2, 4]},
            'affine':    {'dim': 1024, 'use_bias': True},
            #'bilstm':   {'num_h': self.fet_dep//2, 'num_o': None, 'fbias': 1.0, 'tmajr': False},
            #'concat':   {'axis': 0},
            #'split':    {'axis': 0, 'number': img_num},
            #'dropout':  {'keep_p': self.dropout},
        }
        #####################Get the first feature map!####################
        if self.inc_btm:
            print('Get the first feature map!')
            opas = {'op': [{'op': 'conv_bn_relu1', 'loop': 1, 'params':{'com':{'trainable': True }}}, #(None, 512, 512, 64)
                           {'op': 'max_pool1',     'loop': 1, 'params':{}}, #(None, 256, 256, 64) #pool2
                          ], 'loop': 1}
            tsr_out = layers_module1(imgs, 0, com_pams, opas, mtra)
            print('')
        #####################Get the resnet blocks!#########################
        print('Get the resnet block!')
        fet_lst = []
        opas = {'op': [{'op': 'resnet_block2', 'loop': 1, 'params':{}}], 'loop': 1}
        fet_lst.extend(layers_module1(tsr_out, 1, com_pams, opas, mtra))
        assert len(fet_lst) == 4, 'The first resnet block is worng!'
        tsr_out = fet_lst[-1]
        print('')
        ################Get the image classification results!###############
        if self.glb_pol: # Global average pooling.
            print('Get the image classification results!')
            com_pams['conv'] = {'number':self.cls_num,'shape':[1, 1],'rate':1,'stride':[1, 1],'padding':'SAME','use_bias':True}
            opas = {'op': [{'op': 'glb_pool1', 'loop': 1, 'params': {}},
                           {'op': 'conv1',     'loop': 1, 'params': {}},
                           {'op': 'squeeze1',  'loop': 1, 'params': {}},
                          ], 'loop': 1}
            scrs     = layers_module1(tsr_out, 99, com_pams, opas, mtra) #class scores
            prbs     = tf.nn.softmax(scrs) #class probabilities
            los_dat  = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lbls, logits=scrs))
            los_reg  = tf.add_n(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))
            los      = los_dat + los_reg
            loss     = tf.stack([los, los_dat, los_reg], axis=0)
            return loss, scrs
            print('')
        ###############Get embedding vectors for each pixel!################
        print('Get embedding vectors for each pixel!')
        '''
        We start by learning an embedding space, so that pixels that correspond to the same object instance
        are close, and pixels that correspond to different objects (including the background) are far.
        That is to say, embedding vectors which are similar are more likely to belong to the same object instance.
        '''
        com_pams['com']['trainable'] = True
        com_pams['conv'] = {'number': 256, 'shape': [1, 1], 'rate': 1, 'stride': [1, 1], 'padding': 'SAME', 'use_bias': True}
        opas = {'op': [{'op': 'conv_bn_relu1',   'loop': 4, 'params':{'conv':{'shape': [3, 3]}}},
                       {"op": 'deconv_bn_relu1', 'loop': 1, 'params':{}},
                       {'op': 'conv1',           'loop': 1, 'params':{'conv':{'number': 64}}},
                      ], 'loop': 1}
        fet_emb = layers_module1(tsr_out, 2, com_pams, opas, mtra) #(N, H, W, 64)
        print('')
        ####Get Mask classification and seedness scores for each pixel!#####
        print('Get Mask classification and seedness scores for each pixel!')
        '''
        The model predicts a class label for the mask centered at each pixel, as well as a confidence score
        that this pixel would make a good “seed” for creating a mask.
        That is to say, the model predicts the classification score of the mask each pixel will generate if
        picked as a seed. We derive the seediness scores from the classification scores and use them to choose
        which seed points in the image to sample. Each seed point generates a mask based on the embedding
        vectors; each mask is then associated with a class label and confidence score.
        '''
        com_pams['com']['trainable'] = True
        com_pams['conv'] = {'number': 256, 'shape': [1, 1], 'rate': 1, 'stride': [1, 1], 'padding': 'SAME', 'use_bias': True}
        opas = {'op': [{'op': 'conv_bn_relu1',   'loop': 4, 'params':{'conv':{'shape': [3, 3]}}},
                       {"op": 'deconv_bn_relu1', 'loop': 1, 'params':{}},
                       {'op': 'conv1',           'loop': 1, 'params':{'conv':{'number': self.cls_num*4}}},
                       {'op': 'reshape1',        'loop': 1, 'params':{'reshape':{'shape': [img_num, fet_hgt, fet_wdh, 4, -1]}}},
                      ], 'loop': 1}
        fet_cls = layers_module1(tsr_out, 3, com_pams, opas, mtra) #(N, H, W, 4, C+1)
        print('')
        #####################Get the mask prediction!#######################
        print('Get the mask prediction!')
        EL = EmbclsLayer(self.cls_num, img_shp, fet_shp)
        boxs, msks, msk_clss, msk_prbs, msk_nums = EL.generate_msks(fet_emb, fet_cls)
        print('')

        if self.mod_tra:
            #######################Get the mask losses!##########################
            print('Get the mask losses!')
            ET = EmbClsTargetLayer(self.cls_num, img_shp, fet_shp)
            sims_los, prbs_los = ET.generate_embcls_loss(fet_emb, fet_cls, gbxs, gmks, gbx_nums)
            los_dat  = sims_los * 10.0 + prbs_los * 2.0
            los_reg  = tf.add_n(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))
            los      = los_dat  + los_reg
            loss     = tf.stack([los, sims_los, prbs_los, los_reg], axis=0)
            print('')
            return loss, boxs, msks, msk_clss, msk_prbs, msk_nums
        else:
            return boxs, msks, msk_clss, msk_prbs, msk_nums


# --------------------------------------------------------
# Tensorflow Faster R-CNN
# Licensed under The MIT License [see LICENSE for details]
# Written by Xinlei Chen
# --------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import tensorflow.contrib.slim as slim
from tensorflow.contrib.slim import losses
from tensorflow.contrib.slim import arg_scope
import numpy as np

from nets.network import Network
from model.config import cfg

class vgg16(Network):
  def __init__(self, batch_size=1):
    Network.__init__(self, batch_size=batch_size)
    self._arch = 'vgg16'

  def build_network(self, sess, is_training=True):
    with tf.variable_scope('vgg_16', 'vgg_16',
                           regularizer=tf.contrib.layers.l2_regularizer(cfg.TRAIN.WEIGHT_DECAY)):
      # select initializers
      if cfg.TRAIN.TRUNCATED:
        initializer = tf.truncated_normal_initializer(mean=0.0, stddev=0.01)
        initializer_bbox = tf.truncated_normal_initializer(mean=0.0, stddev=0.001)
      else:
        initializer = tf.random_normal_initializer(mean=0.0, stddev=0.01)
        initializer_bbox = tf.random_normal_initializer(mean=0.0, stddev=0.001)

      net = slim.repeat(self._image, 2, slim.conv2d, 64, [3, 3],
                        trainable=False, scope='conv1')
      net = slim.max_pool2d(net, [2, 2], padding='SAME', scope='pool1')
      net = slim.repeat(net, 2, slim.conv2d, 128, [3, 3],
                        trainable=False, scope='conv2')
      net = slim.max_pool2d(net, [2, 2], padding='SAME', scope='pool2')
      net = slim.repeat(net, 3, slim.conv2d, 256, [3, 3],
                        trainable=is_training, scope='conv3')
      net = slim.max_pool2d(net, [2, 2], padding='SAME', scope='pool3')
      net = slim.repeat(net, 3, slim.conv2d, 512, [3, 3],
                        trainable=is_training, scope='conv4')
      net = slim.max_pool2d(net, [2, 2], padding='SAME', scope='pool4')
      net = slim.repeat(net, 3, slim.conv2d, 512, [3, 3],
                        trainable=is_training, scope='conv5')
      self._act_summaries.append(net)
      self._layers['conv5_3'] = net
      # build the anchors for the image
      self._anchor_component()

      # rpn
      rpn = slim.conv2d(net, 512, [3, 3], trainable=is_training, weights_initializer=initializer, scope="rpn_conv/3x3")
      self._act_summaries.append(rpn)
      rpn_cls_score = slim.conv2d(rpn, self._num_anchors * 2, [1, 1], trainable=is_training,
                                  weights_initializer=initializer,
                                  padding='VALID', activation_fn=None, scope='rpn_cls_score')
      # change it so that the score has 2 as its channel size
      rpn_cls_score_reshape = self._reshape_layer(rpn_cls_score, 2, 'rpn_cls_score_reshape')
      rpn_cls_prob_reshape = self._softmax_layer(rpn_cls_score_reshape, "rpn_cls_prob_reshape")
      rpn_cls_prob = self._reshape_layer(rpn_cls_prob_reshape, self._num_anchors * 2, "rpn_cls_prob")
      rpn_bbox_pred = slim.conv2d(rpn, self._num_anchors * 4, [1, 1], trainable=is_training,
                                  weights_initializer=initializer,
                                  padding='VALID', activation_fn=None, scope='rpn_bbox_pred')
      if is_training:
        rois, roi_scores = self._proposal_layer(rpn_cls_prob, rpn_bbox_pred, "rois")
        rpn_labels = self._anchor_target_layer(rpn_cls_score, "anchor")
        # Try to have a determinestic order for the computing graph, for reproducibility
        with tf.control_dependencies([rpn_labels]):
          rois, _ = self._proposal_target_layer(rois, roi_scores, "rpn_rois")
      else:
        if cfg.TEST.MODE == 'nms':
          rois, _ = self._proposal_layer(rpn_cls_prob, rpn_bbox_pred, "rois")
        elif cfg.TEST.MODE == 'top':
          rois, _ = self._proposal_top_layer(rpn_cls_prob, rpn_bbox_pred, "rois")
        else:
          raise NotImplementedError

      # rcnn
      if cfg.POOLING_MODE == 'crop':
        pool5 = self._crop_pool_layer(net, rois, "pool5")
      else:
        raise NotImplementedError

      pool5_flat = slim.flatten(pool5, scope='flatten')
      fc6 = slim.fully_connected(pool5_flat, 4096, scope='fc6')
      if is_training:
        fc6 = slim.dropout(fc6, scope='dropout6')
      fc7 = slim.fully_connected(fc6, 4096, scope='fc7')
      if is_training:
        fc7 = slim.dropout(fc7, scope='dropout7')
      cls_score = slim.fully_connected(fc7, self._num_classes, weights_initializer=initializer,
                                       trainable=is_training,
                                       activation_fn=None, scope='cls_score')
      cls_prob = self._softmax_layer(cls_score, "cls_prob")
      bbox_pred = slim.fully_connected(fc7, self._num_classes * 4, weights_initializer=initializer_bbox,
                                       trainable=is_training,
                                       activation_fn=None, scope='bbox_pred')

      self._predictions["rpn_cls_score"] = rpn_cls_score
      self._predictions["rpn_cls_score_reshape"] = rpn_cls_score_reshape
      self._predictions["rpn_cls_prob"] = rpn_cls_prob
      self._predictions["rpn_bbox_pred"] = rpn_bbox_pred
      self._predictions["cls_score"] = cls_score
      self._predictions["cls_prob"] = cls_prob
      self._predictions["bbox_pred"] = bbox_pred
      self._predictions["rois"] = rois

      self._score_summaries.update(self._predictions)

      return rois, cls_prob, bbox_pred


import numpy as np
import os
import pdb
from PIL import Image, ImageDraw
import cv2
import glob as gb
# import av
# from skvideo.io import ffprobe
import pandas as pd
import sys
import pdb

"""
    convert hmdb-51 detection pickle file to kinetics detection format
"""

src_det_file_path = '/home/jinchoi/src/rehab/dataset/action/HMDB51/detectron_results/detection_merged_missing_dets.pkl'
tgt_det_file_path = '/home/jinchoi/src/rehab/dataset/action/HMDB51/detectron_results/hmdb51_detections_height_256pixels.npy'

dets = np.load(src_det_file_path, allow_pickle=True)

new_fmt_dets = dict()

for i,human_boxes in enumerate(dets['all_boxes'][1]):
    parsed_filename = dets['roidb'][i][0].split('/')
    cur_cls = parsed_filename[-3]
    vid = parsed_filename[-2]
    idx = int(parsed_filename[-1].split('.')[0].split('_')[1])

    if cur_cls not in new_fmt_dets:
        new_fmt_dets[cur_cls] = dict()

    if vid not in new_fmt_dets[cur_cls]:
        new_fmt_dets[cur_cls][vid] = dict()

    new_fmt_dets[cur_cls][vid][int(idx)] = dict(
                                            human_boxes = human_boxes,
                                            frame = str(idx)
                                               )

wrapped_dets = dict(
                     video_res_info=(-1, 256),
                     dets = new_fmt_dets
)

np.save(tgt_det_file_path, wrapped_dets)


import tensorflow as tf
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)
from absl import app, flags, logging
from absl.flags import FLAGS
import core.utils as utils
from core.yolov4 import filter_boxes
from tensorflow.python.saved_model import tag_constants
from PIL import Image
import cv2
import numpy as np
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession

flags.DEFINE_string('framework', 'tf', '(tf, tflite, trt')
flags.DEFINE_string('weights', './checkpoints/yolov4-416',
                    'path to weights file')
flags.DEFINE_integer('size', 416, 'resize images to')
flags.DEFINE_boolean('tiny', False, 'yolo or yolo-tiny')
flags.DEFINE_string('model', 'yolov4', 'yolov3 or yolov4')
flags.DEFINE_string('image', './data/kite.jpg', 'path to input image')
flags.DEFINE_string('output', 'result.png', 'path to output image')
flags.DEFINE_float('iou', 0.45, 'iou threshold')
flags.DEFINE_float('score', 0.25, 'score threshold')

def main(_argv):
    config = ConfigProto()
    config.gpu_options.allow_growth = True
    session = InteractiveSession(config=config)
    STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
    input_size = FLAGS.size
    #image_path = FLAGS.image
    images = ['/content/dataframe/frame'+ str(cf) + '.jpg' for cf in range(0, 2010,10)]
    bb = np.empty([len(images), 1, 50, 4])
    cc = np.empty([len(images), 1, 50])
    ii = 0
    for image_path in images:
      original_image = cv2.imread(str(image_path))
      original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)

      # image_data = utils.image_preprocess(np.copy(original_image), [input_size, input_size])
      image_data = cv2.resize(original_image, (input_size, input_size))
      image_data = image_data / 255.
      # image_data = image_data[np.newaxis, ...].astype(np.float32)

      images_data = []
      for i in range(1):
          images_data.append(image_data)
      images_data = np.asarray(images_data).astype(np.float32)

      if FLAGS.framework == 'tflite':
          interpreter = tf.lite.Interpreter(model_path=FLAGS.weights)
          interpreter.allocate_tensors()
          input_details = interpreter.get_input_details()
          output_details = interpreter.get_output_details()
          print(input_details)
          print(output_details)
          interpreter.set_tensor(input_details[0]['index'], images_data)
          interpreter.invoke()
          pred = [interpreter.get_tensor(output_details[i]['index']) for i in range(len(output_details))]
          if FLAGS.model == 'yolov3' and FLAGS.tiny == True:
              boxes, pred_conf = filter_boxes(pred[1], pred[0], score_threshold=0.25, input_shape=tf.constant([input_size, input_size]))
          else:
              boxes, pred_conf = filter_boxes(pred[0], pred[1], score_threshold=0.25, input_shape=tf.constant([input_size, input_size]))
      else:
          saved_model_loaded = tf.saved_model.load(FLAGS.weights, tags=[tag_constants.SERVING])
          infer = saved_model_loaded.signatures['serving_default']
          batch_data = tf.constant(images_data)
          pred_bbox = infer(batch_data)
          for key, value in pred_bbox.items():
              boxes = value[:, :, 0:4]
              pred_conf = value[:, :, 4:]

      boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression(
          boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)),
          scores=tf.reshape(
              pred_conf, (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])),
          max_output_size_per_class=50,
          max_total_size=50,
          iou_threshold=FLAGS.iou,
          score_threshold=FLAGS.score
      )
      pred_bbox = [boxes.numpy(), scores.numpy(), classes.numpy(), valid_detections.numpy()]
      """out_boxes, out_scores, out_classes, num_boxes = pred_bbox
      image_h, image_w, _ = original_image.shape
      for i in range(num_boxes[0]):
          if int(out_classes[0][i]) < 0 : continue
          coor = out_boxes[0][i]
          coor[0] = int(coor[0] * image_h)
          coor[2] = int(coor[2] * image_h)
          coor[1] = int(coor[1] * image_w)
          coor[3] = int(coor[3] * image_w)

          fontScale = 0.5
          score = out_scores[0][i]
          class_ind = int(out_classes[0][i])
          c1, c2 = (coor[1], coor[0]), (coor[3], coor[2])"""
      bb[ii]= boxes
      cc[ii]= classes
      ii = ii+1
    np.save('/content/framebb.npy', bb)
    np.save('/content/framecc.npy',cc)

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


import os
import random

import numpy as np
import torch
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
import apex.amp as amp

from data_processing import data_loader
from ResNet import DKU_SEResNet
from myexp import ex
from utils import train,test,evaluate
from loss import FocalLoss
@ex.config
def my_config():
    access_type = 'PA'
    # feature_type = 'LFCC'

    # set paths to the wave files and protocols
    feature_type = 'GDgram'  # [GDgram, spec, logspec]
    print('training with feature: {}'.format(feature_type))
    server = 1
    if server == 1:  # local test
        pathToASVspoof2019Data = '/data/ASVspoof2019/'
        _dir_dataset = os.path.join(pathToASVspoof2019Data, access_type)
        train_dir = os.path.join(_dir_dataset,'ASVspoof2019_PA_train/{}_magnitude_1024_400_240/'.format(feature_type))
        dev_dir = os.path.join(_dir_dataset,'ASVspoof2019_PA_dev/{}_magnitude_1024_400_240/'.format(feature_type))
        eval_dir = os.path.join(_dir_dataset, 'ASVspoof2019_PA_eval/{}_magnitude_1024_400_240/'.format(feature_type))
    elif server == 0:
        pathToASVspoof2019Data = '/data/dyq/anti-spoofing/ASVspoof2019/'
        _dir_dataset = os.path.join(pathToASVspoof2019Data, access_type)
        train_dir = os.path.join('/fast/dyq/ASVspoof2019_PA_train', '{}_magnitude_1024_400_240/'.format(feature_type))
        dev_dir = os.path.join('/fast/dyq/ASVspoof2019_PA_dev', '{}_magnitude_1024_400_240/'.format(feature_type))
        eval_dir = os.path.join(_dir_dataset, 'ASVspoof2019_PA_eval/{}_magnitude_1024_400_240/'.format(feature_type))

    # train_dir = os.path.join('/fast/dyq/ASVspoof2019_PA_train', 'spec_magnitude_1024_400_240/')
    # dev_dir = os.path.join('/fast/dyq/ASVspoof2019_PA_dev', 'spec_magnitude_1024_400_240/')

    trainProtocolFile = os.path.join(_dir_dataset, 'ASVspoof2019_' + access_type + '_cm_protocols',
                                     'ASVspoof2019.' + access_type + '.cm.train.trn.txt')
    devProtocolFile = os.path.join(_dir_dataset, 'ASVspoof2019_' + access_type + '_cm_protocols',
                                   'ASVspoof2019.' + access_type + '.cm.dev.trl.txt')
    evalProtocolFile = os.path.join(_dir_dataset, 'ASVspoof2019_' + access_type + '_cm_protocols',
                                    'ASVspoof2019.' + access_type + '.cm.eval.trl.txt')
    save_dir = 'results/'
    batch_size = 16
    train_batch = 16
    dev_batch = 48
    nb_time = 500
    epoch = 30
    save_interval = 1000
    log_interval = 200
    lr = 0.0005
    load_model = False
    load_file = save_dir + 'LCNN-1000.pth'
    test_first = False
    eval_mode = False
    num_workers = 1
    model_params = {
        'FOCAL_GAMMA' : 2,  # gamma parameter for focal loss; if obj is not focal loss, set this to None
        'LOSS_WEIGHT': [1., 9.]
    }

    # Writer will output to ./runs/ directory by default
    writer = SummaryWriter()

def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True

@ex.automain
def main(model_params,writer, eval_mode, load_file,seed):
    setup_seed(seed)
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    print("Use", device)
    model = DKU_SEResNet(2).to(device)
    #define loss function
    if model_params['FOCAL_GAMMA']:
        print('Training with focal loss')
        focal_obj = FocalLoss(gamma=model_params['FOCAL_GAMMA'], alpha=model_params['LOSS_WEIGHT'])
    else:
        focal_obj = None
    optimizer = optim.AdamW(model.parameters(), weight_decay = 5e-5)

    # Allow Amp to perform casts as required by the opt_level
    model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

    train_loader, dev_loader, eval_loader = data_loader()
    if eval_mode is False:
        train(model, device, train_loader, dev_loader, optimizer,focal_obj=focal_obj)
    else:
        evaluate(model, device, eval_loader, load_file,focal_obj=focal_obj)

    writer.close()

'''
command:
python main.py with 'epoch=50' 'lr=0.001'  'load_model=False' 'load_file=results/Model-epoch-25.pth' 'test_first=False' 'num_workers=1' 'eval_mode=False'
'''


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torchvision import datasets, transforms
from torch.utils import data
import pretrainedmodels
import pandas as pd
import numpy as np
from sklearn.utils.class_weight import compute_class_weight
import os
import cv2
from skimage.io import imread
from torch.utils.data.sampler import WeightedRandomSampler, BatchSampler
from albumentations import (
    HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90, Normalize, RandomGamma, RandomBrightnessContrast, HueSaturationValue, CLAHE, ChannelShuffle,
    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,
    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine,
    IAASharpen, IAAEmboss, RandomContrast, RandomBrightness, Flip, OneOf, Compose, PadIfNeeded, RandomCrop, Resize
)
from tqdm import tqdm_notebook
from sklearn.metrics import f1_score, roc_auc_score
from sklearn.model_selection import train_test_split, StratifiedKFold, GroupKFold
import matplotlib.pyplot as plt
import pretrainedmodels.utils as utils
from sklearn.metrics import roc_auc_score
use_cuda = torch.cuda.is_available()
import pickle

'''
Get predictions for each fold (OOF preds) and test data using TTA
'''

device = torch.device("cuda" if use_cuda else "cpu")
DIR = './'
train_im_dir = DIR+'/train'
test_im_dir = DIR+'/test'
model_dir = os.path.join(DIR,'resnet34')
model_name = 'resnet34'
train_data = pd.read_csv(os.path.join(DIR,'train_labels.csv'))
test_data = pd.read_csv(os.path.join(DIR,'sample_submission.csv'))
patch_ids = pd.read_csv(os.path.join(DIR,'patch_id_wsi.csv'))
train_data = pd.merge(train_data, patch_ids, on='id')
n_groups = 15
skf = GroupKFold(n_splits=n_groups)
folds_id_train = []
folds_label_train = []
folds_id_val = []
folds_label_val = []
for train_index, test_index in skf.split(train_data['id'].values, train_data['label'].values, train_data['wsi'].values):
    folds_id_train.append(train_data['id'].values[train_index])
    folds_id_val.append(train_data['id'].values[test_index])
    folds_label_train.append(train_data['label'].values[train_index])
    folds_label_val.append(train_data['label'].values[test_index])
test_id = test_data['id'].values
test_label = test_data['label'].values

val_preds = []
val_labels = []
test_preds = []
scores_CV = []
for valid_idx in range(n_groups):
    base_model = pretrainedmodels.resnet34(num_classes=1000,pretrained='imagenet').to(device) #load pretrained as base
    model = Net(base_model, 512).to(device) # create model
    model.load_state_dict(torch.load(model_dir+'/resnet34.fold{}.best.pt'.format(valid_idx))) #loading weights
    model.eval()
    valid_preds_idx = np.zeros((len(folds_id_val[valid_idx])))
    valid_target_idx = np.zeros((len(folds_id_val[valid_idx])))
    test_preds_idx = np.zeros((len(test_label)))
    val_loader = torch.utils.data.DataLoader(DataGenerator(folds_id_val[valid_idx], folds_label_val[valid_idx], validation_aug, train_im_dir),
                                             shuffle=False, pin_memory=False, num_workers=1,batch_size=1)
    test_loader = torch.utils.data.DataLoader(DataGenerator(test_id, test_label, validation_aug, test_im_dir),
                                              shuffle=False, pin_memory=False, num_workers=1,batch_size=1)
    #predction for validation data
    with torch.no_grad():
        for batch_idx, (x, target) in enumerate(tqdm_notebook(val_loader)):
            #output = protein_model(x.to(device, dtype=torch.float))
            image = np.rollaxis(x.numpy()[0], 0, 3)
            images = make_tta_heavy(image,n_images=8) #create 8 images for random augmentations to take mean prediciton of each
            output = model(torch.from_numpy(images).to(device, dtype=torch.float))
            output = output.mean()
            valid_preds_idx[batch_idx] = output
            valid_target_idx[batch_idx] = target
    val_preds.append(valid_preds_idx)
    val_labels.append(valid_target_idx)
    score_CV_idx = roc_auc_score(valid_target_idx, valid_preds_idx)
    scores_CV.append(score_CV_idx)
    print('fold {}, score {}'.format(valid_idx, score_CV_idx))
    #predction for test data
    with torch.no_grad():
        for batch_idx, (x, target) in enumerate(tqdm_notebook(test_loader)):
            image = np.rollaxis(x.numpy()[0], 0, 3)
            images = make_tta_heavy(image,n_images=8)
            output = model(torch.from_numpy(images).to(device, dtype=torch.float))
            output = output.mean()
            #output = protein_model(x.to(device, dtype=torch.float))
            test_preds_idx[batch_idx] = output
    test_preds.append(test_preds_idx)

val_preds_combined = np.hstack(val_preds)
val_labels_combined = np.hstack(val_labels)
#average test predctions over each fold
test_preds_combined = np.vstack(test_preds)
test_preds_combined = np.mean(test_preds_combined, axis=0)
cv_rocauc = roc_auc_score(val_labels_combined, val_preds_combined
print('Total roc auc'.format(cv_rocauc))
d = dict({'oof_preds':val_preds_combined,'oof_labels':val_labels_combined,'test_preds':test_preds_combined})
with open(os.path.join(model_dir, "resnet34.tta.Preds.pickle"), "wb") as output_file:
    pickle.dump(d, output_file)
#create sample submission
test_data['label'] = d['test_preds']
test_data.to_csv(os.path.join(model_dir, 'resnet34.prediction.tta.csv') ,sep=',',index=False)


import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torchvision.models as models
from torch.autograd import Variable
import numpy as np
import sys

from Enhancement_models.target_loss.model.utils.config import cfg
from Enhancement_models.target_loss.model.rpn.rpn import _RPN
from Enhancement_models.target_loss.model.roi_layers import ROIAlign, ROIPool

# from model.roi_pooling.modules.roi_pool import _RoIPooling
# from model.roi_align.modules.roi_align import RoIAlignAvg
from Enhancement_models.target_loss.model.rpn.proposal_target_layer_cascade import _ProposalTargetLayer
import time
import pdb
from Enhancement_models.target_loss.model.utils.net_utils import _smooth_l1_loss, _crop_pool_layer, _affine_grid_gen, _affine_theta

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): 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=1, size_average=True):
        super(FocalLoss, self).__init__()
        if alpha is None:
            self.alpha = Variable(torch.ones(class_num, 1))
            #self.alpha[7] = 0.75
            #self.alpha[15] = 0.75
        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)
        C = inputs.size(1)
        P = F.softmax(inputs)

        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 self.size_average:
            loss = batch_loss.mean()
        else:
            loss = batch_loss.sum()
        return loss
class _fasterRCNN(nn.Module):
    """ faster RCNN """
    def __init__(self, classes, module_type, class_agnostic):
        super(_fasterRCNN, self).__init__()

        self.module_type = module_type
        self.classes = classes
        self.n_classes = len(classes)
        self.class_agnostic = class_agnostic
        # loss
        self.RCNN_loss_cls = 0
        self.RCNN_loss_bbox = 0

        # define rpn
        self.RCNN_rpn = _RPN(self.dout_base_model)
        self.RCNN_proposal_target = _ProposalTargetLayer(self.n_classes)

        # self.RCNN_roi_pool = _RoIPooling(cfg.POOLING_SIZE, cfg.POOLING_SIZE, 1.0/16.0)
        # self.RCNN_roi_align = RoIAlignAvg(cfg.POOLING_SIZE, cfg.POOLING_SIZE, 1.0/16.0)

        self.RCNN_roi_pool = ROIPool((cfg.POOLING_SIZE, cfg.POOLING_SIZE), 1.0/16.0)
        self.RCNN_roi_align = ROIAlign((cfg.POOLING_SIZE, cfg.POOLING_SIZE), 1.0/16.0, 0)
        #if self.module_type == 'vgg16':
        #    self.feature = torch.zeros(4096).cuda()
        #else:
        #    self.feature = torch.zeros(2048).cuda()
        self.focalloss = FocalLoss(21)

    def forward(self, im_data, im_info, gt_boxes, num_boxes, person_feat, car_feat, is_detect = True, detection_loss=True, culsum_loss=False):
        batch_size = im_data.size(0)

        # feed image data to base model to obtain base feature map
        if self.module_type == 'vgg16':
            base_feat_layer1 = self.RCNN_base[:7](im_data)
            #base_feat_layer2 = self.RCNN_base[7:19](base_feat_layer1)
            base_feat = self.RCNN_base[7:](base_feat_layer1)
        else:
            #print(self.RCNN_base[:5])
            base_feat_layer1 = self.RCNN_base[:5](im_data)
            base_feat = self.RCNN_base[5:](base_feat_layer1)
        if is_detect is False:
            return base_feat_layer1, None, None, None
        else:
            self.training = self.training and num_boxes
            # feed base feature map tp RPN to obtain rois
            #print('im_info',im_info)
            #print(self.training)
            #print('boxes',gt_boxes)
            rois, rpn_loss_cls, rpn_loss_bbox = self.RCNN_rpn(base_feat, im_info, gt_boxes, num_boxes)

            # if it is training phrase, then use ground trubut bboxes for refining
            if self.training:
                roi_data = self.RCNN_proposal_target(rois, gt_boxes, num_boxes)
                rois, rois_label, rois_target, rois_inside_ws, rois_outside_ws = roi_data

                rois_label = Variable(rois_label.view(-1).long())
                rois_target = Variable(rois_target.view(-1, rois_target.size(2)))
                rois_inside_ws = Variable(rois_inside_ws.view(-1, rois_inside_ws.size(2)))
                rois_outside_ws = Variable(rois_outside_ws.view(-1, rois_outside_ws.size(2)))
            else:
                rois_label = None
                rois_target = None
                rois_inside_ws = None
                rois_outside_ws = None
                rpn_loss_cls = 0
                rpn_loss_bbox = 0

            rois = Variable(rois)
            # do roi pooling based on predicted rois
            if cfg.POOLING_MODE == 'align':
                pooled_feat = self.RCNN_roi_align(base_feat, rois.view(-1, 5))
            elif cfg.POOLING_MODE == 'pool':
                pooled_feat = self.RCNN_roi_pool(base_feat, rois.view(-1,5))

            # feed pooled features to top model
            pooled_feat = self._head_to_tail(pooled_feat)

            # compute bbox offset
            bbox_pred = self.RCNN_bbox_pred(pooled_feat)
            if self.training and not self.class_agnostic:
                # select the corresponding columns according to roi labels
                bbox_pred_view = bbox_pred.view(bbox_pred.size(0), int(bbox_pred.size(1) / 4), 4)
                bbox_pred_select = torch.gather(bbox_pred_view, 1, rois_label.view(rois_label.size(0), 1, 1).expand(rois_label.size(0), 1, 4))
                bbox_pred = bbox_pred_select.squeeze(1)

            # compute object classification probability
            cls_score = self.RCNN_cls_score(pooled_feat)
            cls_prob = F.softmax(cls_score, 1)

            car_feature = 0
            person_feature = 0
            if culsum_loss:
              if self.training is False:
                  pooled_feat = pooled_feat[:128,:]
                  cls_prob = cls_prob[:128,:]

              if self.module_type == 'vgg16':
                  car_feature = torch.zeros(4096).cuda()  #
                  person_feature = torch.zeros(4096).cuda()
              else:
                  car_feature = torch.zeros(2048).cuda()  #
                  person_feature = torch.zeros(2048).cuda()
              if isinstance(car_feat,int):
                  car_feature = 0
              else:
                  car_feature.copy_(car_feat.detach_())
              if isinstance(person_feat,int):
                  person_feature = 0
              else:
                  person_feature.copy_(person_feat.detach_())
              delta = 0
              if self.training:
                  max_index = rois_label
              else:
                  max_index = cls_prob.argmax(dim = 1)
              #print(max_index)
              select_feat_car = pooled_feat[max_index == 7]
              select_feat_person = pooled_feat[max_index == 15]
              if(len(select_feat_car)):
                  select_feat_car=select_feat_car.mean(dim=0)
                  #print(select_feat_car)
                  if(isinstance(car_feat,int)):
                      car_feature = select_feat_car
                      print('===> Car_feature inital')
                  else:
                      delta = torch.cosine_similarity(select_feat_car, car_feat, dim = 0).item()
                      car_feature = delta * select_feat_car + (1 - delta) * car_feat
              if(len(select_feat_person)):
                  select_feat_person=select_feat_person.mean(dim=0)
                  if(isinstance(person_feat,int)):
                      person_feature = select_feat_person
                      print('===> Person_feature inital')
                  else:
                      delta = torch.cosine_similarity(select_feat_person, person_feat, dim = 0).item()
                      person_feature = delta * select_feat_person + (1 - delta) * person_feat

            RCNN_loss_cls = 0
            RCNN_loss_bbox = 0
            target_loss = 0

            if self.training and detection_loss:
                # classification loss
                RCNN_loss_cls = F.cross_entropy(cls_score, rois_label)
                #print(cls_score)
                #print(rois_label)
                #RCNN_loss_cls = self.focalloss(cls_score, rois_label)

                # bounding box regression L1 loss
                RCNN_loss_bbox = _smooth_l1_loss(bbox_pred, rois_target, rois_inside_ws, rois_outside_ws)

                # the detection loss
                target_loss = rpn_loss_cls + rpn_loss_bbox + RCNN_loss_cls + RCNN_loss_bbox

            return base_feat_layer1, target_loss, person_feature, car_feature

    def _init_weights(self):
        def normal_init(m, mean, stddev, truncated=False):
            """
            weight initalizer: truncated normal and random normal.
            """
            # x is a parameter
            if truncated:
                m.weight.data.normal_().fmod_(2).mul_(stddev).add_(mean) # not a perfect approximation
            else:
                m.weight.data.normal_(mean, stddev)
                m.bias.data.zero_()

        normal_init(self.RCNN_rpn.RPN_Conv, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_rpn.RPN_cls_score, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_rpn.RPN_bbox_pred, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_cls_score, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_bbox_pred, 0, 0.001, cfg.TRAIN.TRUNCATED)

    def create_architecture(self):
        self._init_modules()
        self._init_weights()

'''
class _fasterRCNN(nn.Module):
    """ faster RCNN """
    def __init__(self, classes, module_type, class_agnostic):
        super(_fasterRCNN, self).__init__()

        self.module_type = module_type
        self.classes = classes
        self.n_classes = len(classes)
        self.class_agnostic = class_agnostic
        # loss
        self.RCNN_loss_cls = 0
        self.RCNN_loss_bbox = 0

        # define rpn
        self.RCNN_rpn = _RPN(self.dout_base_model)
        self.RCNN_proposal_target = _ProposalTargetLayer(self.n_classes)

        # self.RCNN_roi_pool = _RoIPooling(cfg.POOLING_SIZE, cfg.POOLING_SIZE, 1.0/16.0)
        # self.RCNN_roi_align = RoIAlignAvg(cfg.POOLING_SIZE, cfg.POOLING_SIZE, 1.0/16.0)

        self.RCNN_roi_pool = ROIPool((cfg.POOLING_SIZE, cfg.POOLING_SIZE), 1.0/16.0)
        self.RCNN_roi_align = ROIAlign((cfg.POOLING_SIZE, cfg.POOLING_SIZE), 1.0/16.0, 0)
        #if self.module_type == 'vgg16':
        #    self.feature = torch.zeros(4096).cuda()
        #else:
        #    self.feature = torch.zeros(2048).cuda()
        self.focalloss = FocalLoss(21)

    def forward(self, im_data, im_info, gt_boxes, num_boxes, person_feat, car_feat, detection_loss=True, culsum_loss=False):
        batch_size = im_data.size(0)

        # feed image data to base model to obtain base feature map
        base_feat_layer1 = self.RCNN_base[:6](im_data)
        base_feat = self.RCNN_base[6:](base_feat_layer1)
        self.training = self.training and num_boxes
        # feed base feature map tp RPN to obtain rois
        #print('im_info',im_info)
        #print(self.training)
        #print('boxes',gt_boxes)
        rois, rpn_loss_cls, rpn_loss_bbox = self.RCNN_rpn(base_feat, im_info, gt_boxes, num_boxes)

        # if it is training phrase, then use ground trubut bboxes for refining
        if self.training:
            roi_data = self.RCNN_proposal_target(rois, gt_boxes, num_boxes)
            rois, rois_label, rois_target, rois_inside_ws, rois_outside_ws = roi_data

            rois_label = Variable(rois_label.view(-1).long())
            rois_target = Variable(rois_target.view(-1, rois_target.size(2)))
            rois_inside_ws = Variable(rois_inside_ws.view(-1, rois_inside_ws.size(2)))
            rois_outside_ws = Variable(rois_outside_ws.view(-1, rois_outside_ws.size(2)))
        else:
            rois_label = None
            rois_target = None
            rois_inside_ws = None
            rois_outside_ws = None
            rpn_loss_cls = 0
            rpn_loss_bbox = 0

        rois = Variable(rois)
        # do roi pooling based on predicted rois
        if cfg.POOLING_MODE == 'align':
            pooled_feat = self.RCNN_roi_align(base_feat, rois.view(-1, 5))
        elif cfg.POOLING_MODE == 'pool':
            pooled_feat = self.RCNN_roi_pool(base_feat, rois.view(-1,5))

        # feed pooled features to top model
        pooled_feat = self._head_to_tail(pooled_feat)

        # compute bbox offset
        bbox_pred = self.RCNN_bbox_pred(pooled_feat)
        if self.training and not self.class_agnostic:
            # select the corresponding columns according to roi labels
            bbox_pred_view = bbox_pred.view(bbox_pred.size(0), int(bbox_pred.size(1) / 4), 4)
            bbox_pred_select = torch.gather(bbox_pred_view, 1, rois_label.view(rois_label.size(0), 1, 1).expand(rois_label.size(0), 1, 4))
            bbox_pred = bbox_pred_select.squeeze(1)

        # compute object classification probability
        cls_score = self.RCNN_cls_score(pooled_feat)
        cls_prob = F.softmax(cls_score, 1)

        car_feature = 0
        person_feature = 0
        if culsum_loss:
          if self.training is False:
              pooled_feat = pooled_feat[:128,:]
              cls_prob = cls_prob[:128,:]

          if self.module_type == 'vgg16':
              car_feature = torch.zeros(4096).cuda()  #
              person_feature = torch.zeros(4096).cuda()
          else:
              car_feature = torch.zeros(2048).cuda()  #
              person_feature = torch.zeros(2048).cuda()
          if isinstance(car_feat,int):
              car_feature = 0
          else:
              car_feature.copy_(car_feat.detach_())
          if isinstance(person_feat,int):
              person_feature = 0
          else:
              person_feature.copy_(person_feat.detach_())
          delta = 0
          if self.training:
              max_index = rois_label
          else:
              max_index = cls_prob.argmax(dim = 1)
          #print(max_index)
          select_feat_car = pooled_feat[max_index == 7]
          select_feat_person = pooled_feat[max_index == 15]
          if(len(select_feat_car)):
              select_feat_car=select_feat_car.mean(dim=0)
              #print(select_feat_car)
              if(isinstance(car_feat,int)):
                  car_feature = select_feat_car
                  print('===> Car_feature inital')
              else:
                  delta = torch.cosine_similarity(select_feat_car, car_feat, dim = 0).item()
                  car_feature = delta * select_feat_car + (1 - delta) * car_feat
          if(len(select_feat_person)):
              select_feat_person=select_feat_person.mean(dim=0)
              if(isinstance(person_feat,int)):
                  person_feature = select_feat_person
                  print('===> Person_feature inital')
              else:
                  delta = torch.cosine_similarity(select_feat_person, person_feat, dim = 0).item()
                  person_feature = delta * select_feat_person + (1 - delta) * person_feat

        RCNN_loss_cls = 0
        RCNN_loss_bbox = 0
        target_loss = 0

        if self.training and detection_loss:
            # classification loss
            RCNN_loss_cls = F.cross_entropy(cls_score, rois_label)

            # bounding box regression L1 loss
            RCNN_loss_bbox = _smooth_l1_loss(bbox_pred, rois_target, rois_inside_ws, rois_outside_ws)

            # the detection loss
            target_loss = rpn_loss_cls + rpn_loss_bbox + RCNN_loss_cls + RCNN_loss_bbox

        return target_loss, person_feature, car_feature

    def _init_weights(self):
        def normal_init(m, mean, stddev, truncated=False):
            """
            weight initalizer: truncated normal and random normal.
            """
            # x is a parameter
            if truncated:
                m.weight.data.normal_().fmod_(2).mul_(stddev).add_(mean) # not a perfect approximation
            else:
                m.weight.data.normal_(mean, stddev)
                m.bias.data.zero_()

        normal_init(self.RCNN_rpn.RPN_Conv, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_rpn.RPN_cls_score, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_rpn.RPN_bbox_pred, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_cls_score, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_bbox_pred, 0, 0.001, cfg.TRAIN.TRUNCATED)

    def create_architecture(self):
        self._init_modules()
        self._init_weights()
'''


import keras
import tensorflow as tf

from keras_retinanet import models
from keras_retinanet.utils.image import read_image_bgr, preprocess_image, resize_image
from keras_retinanet.utils.visualization import draw_box, draw_caption
from keras_retinanet.utils.colors import label_color

# import miscellaneous modules
# import matplotlib.pyplot as plt
import cv2
import os
import numpy as np
import time

def get_session():
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    return tf.Session(config=config)

def main():
    keras.backend.tensorflow_backend.set_session(get_session())
    # model_path = './snapshots/1/resnet50_pins_n_solder_12_inference.h5'
    model_path = './snapshots/resnet50_pins_n_solder_14_inference.h5'
    model = models.load_model(model_path, backbone_name='resnet50')

    # load label to names mapping for visualization purposes
    labels_to_names = {0: 'pin', 1: 'solder', 2: 'anomaly_bent_solder'}

    video = cv2.VideoCapture('/media/trevol/HDD/DiskE/Computer_Vision_Task/Video_6.mp4')
    video.set(cv2.CAP_PROP_POS_FRAMES, 269)
    while True:
        ret, frame = video.read()
        if not ret:
            break
        image = predict_on_image(model, labels_to_names, frame, thresh=0.96)
        image = putFramePos(image, video.get(cv2.CAP_PROP_POS_FRAMES))
        image = resize(image, .5)
        cv2.imshow('Video', image)
        if cv2.waitKey() == 27:
            break
    video.release()

def predict_on_image(model, labels_to_names, image, thresh):
    # copy to draw on
    draw = image.copy()
    # draw = cv2.cvtColor(draw, cv2.COLOR_BGR2RGB, dst=draw)

    # preprocess image for network
    image = preprocess_image(image)
    image, scale = resize_image(image)

    # process image
    start = time.time()
    boxes, scores, labels = model.predict_on_batch(np.expand_dims(image, axis=0))
    print("processing time: ", time.time() - start)

    # correct for image scale
    boxes /= scale

    # visualize detections
    for box, score, label in zip(boxes[0], scores[0], labels[0]):
        # scores are sorted so we can break
        if score < thresh:
            break

        # color = label_color(label)
        if score < 0.99:
            color = (0, 0, 255)
        elif label == 1:
            color = (0, 255, 0)
        else:
            color = (0, 255, 255)
        b = np.round(box, 0).astype(int)
        draw_box(draw, b, color=color, thickness=1)

        caption = f"{labels_to_names[label]} {score:.2f}"
        draw_caption(draw, b, caption)

    # return cv2.cvtColor(draw, cv2.COLOR_RGB2BGR, dst=draw)
    return draw

def putFramePos(frame, pos):
    cv2.putText(frame, str(int(pos)), (10, 40), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255))
    return frame

def resize(img, factor):
    dsize = tuple(np.multiply(img.shape[1::-1], factor).astype(int))
    return cv2.resize(img, dsize, interpolation=cv2.INTER_AREA)

if __name__ == '__main__':
    main()


#!/usr/bin/env python3
"""
The purpose of this script is to:
1. reset and repopulate "Shops" table in database
2. choose a winning lottery number, reward winners, and reset "Lottery" table in database
3. determine tournament winners, reward winners, and reset "Battles" table in database

PS: FOR FULL AUTOMATION, SCHEDULE THIS FILE TO RUN AUTOMATICALLY AT AN INTERVAL.
PS2: MAKE SURE TO CHANGE PATHS TO FULL FILE PATHS IF THIS IS AUTOMATED
"""

import configparser
import sys
import random
import discord
import numpy
import os
# add parent folder to module path. can comment this out if using virtual environment
sys.path.append('..')

from numpy import random
from Users import Users
from Database import Database
from discord.ext import commands
from pathlib import Path
from datetime import date

# change working directory to parent to simplify file paths
os.chdir("..")

# startup discord client
client = commands.Bot(command_prefix=["=", "%"])

# set up parser to config through our .ini file with our bot's token
config = configparser.ConfigParser()
bot_token_path = Path("tokens/tokenbot.ini") # use forward slash "/" for path directories
# confirm the token is located in the above path
if bot_token_path.is_file():
    config.read(bot_token_path)
    # we now have the bot's token
    TOKEN = config.get('BOT1', 'token')
else:
    print("\n","Discord bot token not found at: ",bot_token_path,"... Please correct file path in daily_maintenance.py file.")
    sys.exit()

@client.event
async def on_ready():
    print('Logged in as\n'
          + client.user.name + '\n'
          + client.user.id + '\n'
          + '---------')

    # open database
    db = Database(0)
    db.connect()

    """ PERFORM DAILY SHOP MAINTENANCE NOW! """
    db.reset_shop()

    # make sure to change this path to the full file path if you plan to use crontab, or else it will not work
    config = configparser.ConfigParser()
    items_path = Path("db_and_words/all_items.ini")  # use forward slash "/" for path directories
    if items_path.is_file():
        config.read(items_path)
    else:
        print("\n", "Shop items file not found at: ", items_path, "... Please correct file path in daily_maintenance.py file.")
        sys.exit()

    # parse through each section in the items in the .ini file
    # each item has a 17% chance to be inserted into the daily shop
    for item in config.sections():
        if 17 >= random.randint(1, 100) >= 1:
            item_name = config.get(item, 'name')
            item_type = config.get(item, 'type')
            item_level = config.get(item, 'level')
            item_price = config.get(item, 'price')

            db.insert_shop_item(item_name, item_type, item_level, item_price)

    """ PERFORM DAILY LOTTERY MAINTENANCE NOW """
    # generate a random winning number 1-5
    win_number = random.randint(1,5)
    # get python list of winner ticket id's who match today's winning number
    std_winners, prem_winners = db.get_lottery_winners(win_number)
    # we have today's winners now, so reset lottery
    db.reset_lottery()
    std_winners_string = ''
    for winner in std_winners:
        # create instance of each basic ticket user who won, and update their money
        user = Users(winner)
        user.update_user_money(user.get_user_level(0) * 100)
        # alter each item on list to discord @ format and concatenate into 1 string to ping winners below
        std_winners_string += ('\n**TICKET ID:** ' + winner + ' <@' + winner + '>')

    prem_winners_string = ''
    for winner in prem_winners:
        # create instance of each premium ticket user who won, and update their
        user = Users(winner)
        user.update_user_money(user.get_user_level(0) * 170)
        # alter each item on list to discord @ format and concatenate into 1 string to ping winners below
        prem_winners_string += ('\n**TICKET ID:** ' + winner + ' <@' + winner + '>')

    # prepare string of shop reset notifcation & lottery results to send to every discord server
    if not std_winners:
        std_winners_string = '\n<a:worrycry:525209793405648896> no basic winners...  <a:worrycry:525209793405648896>'
    if not prem_winners:
        prem_winners_string = '\n<a:worrycry:525209793405648896> no premium winners...  <a:worrycry:525209793405648896>'

    # split shop & lottery announcements into 2 strings
    # have to insert encode \u200B for spaces when using discord encoding
    global_announcement1 = "__**SHOP ANNOUNCEMENT**__ \u200B \u200B" \
                           + "_" + str(date.today()) + "_" \
                           + "\n:shopping_cart: \u200B Daily shop just reset... Check out **=shop**!\n"
    global_announcement2 = "__**LOTTERY ANNOUNCEMENT**__ \u200B \u200B_" + str(date.today()) + "_" \
                           + "\nToday's winning number is... **" \
                           + str(win_number) + "**\n\n__The lucky **basic** winners:__  " \
                           + std_winners_string + "\n__The lucky **premium** winners:__  " \
                           + prem_winners_string
    global_announcement3 = "__**PATCH ANNOUNCEMENT**__\nPatch **2.7** is here! Use **=adopt** command for your own pet!"

    # embed shop announcement and lottery announcement, and set thumbnails of a shopping cart and "money rain frog" gif
    em = discord.Embed(description=global_announcement1, colour=0x607d4a)
    em.set_thumbnail(url="https://i.imgur.com/rS6tXmD.gif")
    em2 = discord.Embed(description=global_announcement2, colour=0x607d4a)
    em2.set_thumbnail(url="https://cdn.discordapp.com/emojis/525200274340577290.gif?size=64")
    em3 = discord.Embed(description=global_announcement3, colour=0x607d4a)
    em3.set_thumbnail(url="http://i66.tinypic.com/25g7akg.jpg")

    # for each server the bot is in, post the lottery results in the lottery channel
    for server in client.servers:
        # try to get the server's announcement status
        # if they turned off announcements, skip to next loop iteration
        try:
            if db.get_server_announcements_status(server.id) == 0:
                continue
        except:
            pass
        # create boolean for each server, to dictate whether or not the 'lottery' channel could be located
        channel_found = 0
        for channel in server.channels:
            try:
                if channel.name == 'lottery':
                    channel_found = 1
                    await client.send_message(channel, embed=em)
                    await client.send_message(channel, embed=em2)
                    await client.send_message(channel, embed=em3)
            except:
                pass

        # if there were no channels found with the name 'lottery'.....
        # make the channel, then send the results
        if channel_found == 0:
            try:
                channel = await client.create_channel(server, 'lottery', type=discord.ChannelType.text)
                await client.send_message(channel, embed=em)
                await client.send_message(channel, embed=em2)
            except:
                # if the bot failed to make the channel, simply move on
                pass

    """ PERFORM DAILY TOURNAMENT MAINTENANCE NOW! """
    for server in client.servers:
        server_fighters_ids = db.get_server_tourney_members(server.id)
        # only do anything else if the server has more than 1 entry
        # need more than 1 fighter entry to have a tournament
        if len(server_fighters_ids) > 1:
            total_stats_pool = 0
            server_fighters_weights = []
            # 2 loops necessary here for the math:
            # we need the sum pool of stats of all fighters in the server
            for fighter_id in server_fighters_ids:
                fighter = Users(fighter_id)
                # algorithm for calculating a fighter's stats in tourneys: (item score + user level*2 + 22)
                fighter_stats = (fighter.get_user_item_score() + (fighter.get_user_level(0) * 2)) + 22
                total_stats_pool += fighter_stats

            # now we need to divide each fighter's stats by the sum stat pool in order to get specific win chances
            for fighter_id in server_fighters_ids:
                fighter = Users(fighter_id)
                # algorithm for calculating a fighter's stats in tourneys: (item score + user level*2 + 22)
                fighter_stats = (fighter.get_user_item_score() + (fighter.get_user_level(0) * 2)) + 22
                fighter_weight = fighter_stats / total_stats_pool
                server_fighters_weights.append(fighter_weight)

            # print("\nPool: " + str(total_stats_pool) + " All fighters: " + str(
            #     server_fighters_ids) + " All weights: " + str(server_fighters_weights))

            # https://docs.scipy.org/doc/numpy/reference/generated/numpy.random.choice.html
            # random's choices package did not offer a "weighted shuffle"
            # decided to use numpy here for the weighted choice, because it can do a a "weighted shuffle"
            # a weighted shuffle was required to randomly generate an ordered list of winners influenced by weights
            server_winners = numpy.random.choice(server_fighters_ids, size=len(server_fighters_ids),
                                                 replace=False, p=server_fighters_weights)

            # make object of first and second place users
            first_place = Users(server_winners[0])
            second_place = Users(server_winners[1])

            # update their money accordingly
            prize1 = first_place.get_user_level(0) * 170
            prize2 = second_place.get_user_level(0) * 100
            first_place.update_user_money(prize1)
            second_place.update_user_money(prize2)
            # first place will count as a win in records
            first_place.update_user_records(0, 1, prize1)
            # second place will count as a win in records
            second_place.update_user_records(0, 1, prize2)

            # prepare string for local server tourney announcement
            # have to insert encode \u200B for spaces when using discord encoding
            local_server_announcement = '__**TOURNEY ANNOUNCEMENT**__ \u200B \u200B_' + str(date.today())\
                                        + '_\n\n**:trophy: 1st place: ** ' + '<@' \
                                        + server_winners[0] + '> :trophy:  __Prize__: **$' + str(prize1) \
                                        + '**\n' + '**:trophy: 2nd place:** ' + '<@' + server_winners[1] \
                                        + '> :trophy:  __Prize__: **$' + str(prize2) + '**\n'

            # find the channel in the server and state the results
            for channel in server.channels:
                # try to get the server's announcement status
                # if they turned off announcements, skip to next loop iteration
                try:
                    if db.get_server_announcements_status(server.id) == 0:
                        continue
                except:
                    pass
                if channel.name == 'lottery':
                    # if there were more than 2 fighters, make an "honorable mentions" string to append to announcement
                    if len(server_fighters_ids) > 2:
                        # counter to represent placings- losers list will start at 3rd place
                        counter = 3
                        loser_string = '\n__Honorable mentions__\n'
                        for loser in server_winners[2:]:
                            user = Users(server_winners[counter-1])
                            user.update_user_records(1, 0, 0)
                            if counter == 3:
                                loser_string += ('**' + str(counter) + 'rd place: <@' + loser + '>**\n')
                            else:
                                loser_string += ('**' + str(counter) + 'th place: <@' + loser + '>**\n')
                            counter += 1
                        local_server_announcement += loser_string

                        # embed combined announcement, with emoji of 64x64 "nunchuck frog", then send it
                        em = discord.Embed(description=local_server_announcement, colour=0x607d4a)
                        em.set_thumbnail(url="https://cdn.discordapp.com/emojis/493220414206509056.gif?size=64")
                        try:
                            await client.send_message(channel, embed=em)
                        except:
                            pass

                    # else if only 2 entries in the server's tournament, just send first and second place results
                    elif len(server_fighters_ids) == 2:
                        # embed announcement, with emoji of 64x64 "nunchuck frog", then send it
                        em = discord.Embed(description=local_server_announcement, colour=0x607d4a)
                        em.set_thumbnail(url="https://cdn.discordapp.com/emojis/493220414206509056.gif?size=64")
                        try:
                            await client.send_message(channel, embed=em)
                        except:
                            pass

    # reset tournament entries after every server's tournament is processed
    db.reset_tourney()

    # end this daily maintenance program
    sys.exit(0)

client.run(TOKEN)


# -*- coding: utf-8 -*-
import cv2
import numpy as np
import time
import autopy
import pyautogui

from enum import IntEnum

from utils import get_config_param
from handController.gestureKeyboard import gestureKeyboard
from handController.handTrackingModule.HandTrackingModule import HandDetector

# Mode
class Mode(IntEnum):
    NEUTRAL = 0
    CURSOR = 1

class GestureController():
    def __init__(self):
        # Load config params
        self.frameR = int(get_config_param("GestureController", "frameReduction"))
        self.smoothening = int(get_config_param("GestureController", "pointerSmootheningFactor"))
        self.scroll_scale_factor = int(get_config_param("GestureController", "scrollScaleDownFactor"))

        self.cap = cv2.VideoCapture(0)
        self.wCam = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.hCam = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        self.detector = HandDetector(maxHands=1)
        self.wScr, self.hScr = autopy.screen.size()

        self.stopped = False
        self.mode = Mode.NEUTRAL
        self.gestureKeyboard = gestureKeyboard.GestureKeyboard()

    def cleanup(self):
        self.cap.release()
        cv2.destroyAllWindows()

    def bringup_keyboard(self, vid_src):
        self.gestureKeyboard.start(vid_src)

    def main(self):
        pTime = 0
        plocX, plocY = 0, 0
        clocX, clocY = 0, 0

        while not self.stopped:
            # 1. Find hand Landmarks
            success, img = self.cap.read()
            hands, img = self.detector.findHands(img)

            cv2.rectangle(img, (self.frameR, self.frameR), (self.wCam - self.frameR, self.hCam - self.frameR),
                (255, 0, 255), 2)

            # 2. Get the tip of the index and middle fingers
            if hands:
                # Hand 1
                hand1 = hands[0]
                lmList = hand1["lmList"]  # List of 21 Landmark points

                if len(lmList) != 0:
                    x1, y1 = lmList[8]
                    x2, y2 = lmList[12]
                    # print(x1, y1, x2, y2)

                # 3. Check which fingers are up
                fingers = self.detector.fingersUp(hand1)
                print(fingers)

                # 4. Only Index Finger : Moving Mode
                if (len(fingers) > 0):
                    if fingers[1] == 1 and fingers[2] == 0:
                        self.mode = Mode.CURSOR
                        # 4a. Convert Coordinates
                        x3 = np.interp(x1, (self.frameR, self.wCam - self.frameR), (0, self.wScr))
                        y3 = np.interp(y1, (self.frameR, self.hCam - self.frameR), (0, self.hScr))

                        # 4b. Smoothen Values
                        clocX = plocX + (x3 - plocX) / self.smoothening
                        clocY = plocY + (y3 - plocY) / self.smoothening

                        # 4c. Move Mouse
                        try:
                            autopy.mouse.move(self.wScr - clocX, clocY)
                            cv2.circle(img, (x1, y1), 15, (255, 0, 255), cv2.FILLED)
                            plocX, plocY = clocX, clocY
                        except ValueError:
                            pass

                    # 5. Both Index and middle fingers are up : Clicking Mode
                    elif fingers[1] == 1 and fingers[2] == 1 and fingers[3] == 0:
                        # 5a. Find distance between fingers
                        length, lineInfo, img = self.detector.findDistance(lmList[8], lmList[12], img)
                        # 5b. Click mouse if distance short
                        if length < 30:
                            cv2.circle(img, (lineInfo[4], lineInfo[5]),
                            15, (0, 255, 0), cv2.FILLED)
                            autopy.mouse.click()
                            time.sleep(0.15)

                    # 6. Index, middle and ring fingers are up : scroll mode
                    elif fingers[1] == 1 and fingers[2] == 1 and fingers[3] == 1 and fingers[4] == 0:
                        # 6a. Convert Coordinates
                        x3 = np.interp(x1, (self.frameR, self.wCam - self.frameR), (0, self.wScr))
                        y3 = np.interp(y1, (self.frameR, self.hCam - self.frameR), (0, self.hScr))
                        # 6b. Smoothen Values
                        clocX = plocX + (x3 - plocX) / self.smoothening
                        clocY = plocY + (y3 - plocY) / self.smoothening

                        scroll_amt = int((clocY - (self.hScr / 2)) / self.scroll_scale_factor)
                        pyautogui.scroll(scroll_amt)

                        cv2.circle(img, (x1, y1), 15, (255, 86, 255), cv2.FILLED)
                        plocX, plocY = clocX, clocY

                    # 7. Index, middle, ring and pinky fingers are up : keyboard
                    elif fingers[0] == 0 and fingers[1] == 1 and fingers[2] == 1 and fingers[3] == 1 and fingers[4] == 1:
                        self.bringup_keyboard(self.cap)

                # 7. Frame Rate
                cTime = time.time()
                fps = 1 / (cTime - pTime)
                pTime = cTime
                cv2.putText(img, str(int(fps)), (20, 50), cv2.FONT_HERSHEY_PLAIN, 3,
                (255, 0, 0), 3)

            else:
                if self.mode == Mode.CURSOR:
                    autopy.mouse.click()
                    time.sleep(0.15)

                    self.mode = Mode.NEUTRAL

            # 8. Display
            cv2.imshow("Image", img)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    def stop(self):
        self.stopped = True

if __name__ == "__main__":
    gestureController = GestureController()
    gestureController.main()

    gestureController.stop()
    gestureController.cleanup()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import time
import argparse

import cv2 as cv
import numpy as np
import tensorflow as tf

def get_args():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--image",
        type=str,
        default='image/sample.jpg',
    )
    parser.add_argument(
        "--model",
        type=str,
        default='saved_model_nsfw/model_float16_quant.tflite',
    )
    parser.add_argument(
        "--input_size",
        type=str,
        default='224,224',
    )

    args = parser.parse_args()

    return args

def run_inference(
    interpreter,
    input_size,
    image,
):
    # Pre process:Resize, RGB->BGR, float32 cast
    input_image = cv.resize(image, dsize=(input_size[1], input_size[0]))
    input_image = cv.cvtColor(input_image, cv.COLOR_BGR2RGB)

    vgg_mean = [104, 117, 123]
    input_image = input_image - vgg_mean

    input_image = np.expand_dims(input_image, axis=0)
    input_image = input_image.astype('float32')

    # Inference
    input_details = interpreter.get_input_details()
    interpreter.set_tensor(input_details[0]['index'], input_image)
    interpreter.invoke()

    output_details = interpreter.get_output_details()
    result = interpreter.get_tensor(output_details[0]['index'])

    # Post process
    result = np.squeeze(result).astype(np.float32)

    return result

def main():
    args = get_args()

    model_path = args.model
    input_size = [int(i) for i in args.input_size.split(',')]

    # Load model
    interpreter = tf.lite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()

    # read image
    image_path = args.image
    image = cv.imread(image_path)

    # Inference execution
    start_time = time.time()
    result = run_inference(
        interpreter,
        input_size,
        image,
    )
    elapsed_time = time.time() - start_time

    print('Elapsed Time :', '{:.1f}'.format(elapsed_time * 1000) + "ms")
    print('sfw :', '{:.3f}'.format(result[0]))
    print('nsfw:', '{:.3f}'.format(result[1]))

if __name__ == '__main__':
    main()


from __future__ import division
import tensorflow as tf
import numpy as np
from xml.etree.ElementTree import Element, SubElement, tostring

def _parse_input_text(input_text, tokenizer, max_seq_length):
    tokens = tokenizer.tokenize(input_text)
    if len(tokens) > max_seq_length - 2:
        tokens = tokens[0:(max_seq_length - 2)]
    tokens.insert(0, '[CLS]')
    tokens.append('[SEP]')

    return tokens

def average_last_layer_by_head(attentions):
    """
    Computes the attention attention weights by:
        1. Take the attention weights from the last multi-head attention layer assigned
           to the CLS token
        2. Average each token across attention heads
        3. Normalize across tokens

    :param attentions: list of dictionaries of the form
        {'layer_name': (batch_size, num_multihead_attn, sequence_length, sequence_length)}
    :return: a tensor of weights
    """
    last_multihead_attn = list(attentions[-1].values())[0]
    tf.logging.info("Attention look! " + str(last_multihead_attn.shape))
    last_multihead_attn = tf.reshape(last_multihead_attn, [1,12,128,128])#todo important! these are hard coded. fix them
    tf.logging.info("Attention look! " + str(last_multihead_attn.shape))

    # For each multihead attention, get the attention weights going into the CLS token
    cls_attn = last_multihead_attn[:, :, 0, :]

    # Average across attention heads
    cls_attn = tf.reduce_mean(cls_attn, axis=1)

    # Normalize across tokens
    total_weights = tf.reduce_sum(cls_attn, axis=-1, keepdims=True)
    norm_cls_attn = cls_attn / total_weights
    tf.logging.info("Attention look! " + str(norm_cls_attn.shape))

    return norm_cls_attn

def average_first_layer_by_head(attentions):
    """
    Computes the attention attention weights by:
        1. Take the attention weights from the first multi-head attention layer assigned
           to the CLS token
        2. Average each token across attention heads
        3. Normalize across tokens

    :param attentions: list of dictionaries of the form
        {'layer_name': (batch_size, num_multihead_attn, sequence_length, sequence_length)}
    :return: a tensor of weights
    """
    last_multihead_attn = list(attentions[0].values())[0]

    # For each multihead attention, get the attention weights going into the CLS token
    cls_attn = last_multihead_attn[:, :, 0, :]

    # Average across attention heads
    cls_attn = tf.reduce_mean(cls_attn, axis=1)

    # Normalize across tokens
    total_weights = tf.reduce_sum(cls_attn, axis=-1, keepdims=True)
    norm_cls_attn = cls_attn / total_weights

    return norm_cls_attn

def average_layer_i_on_token_j_by_head(layer_index, token_index, attentions):
    """
    General function to average attention weights by heads then across tokens in layer layer_index for token_index
    :param layer_index: The layer we want to extract the attention weights from
    :param token_index: Which token of the layer we want to extract the attention e.g. [CLS], [SEP], etc
    :param attentions: list of dictionaries of the form
        {'layer_name': (batch_size, num_multihead_attn, sequence_length, sequence_length)}
    :return: a tensor of weights
    """
    target_attention_layer = list(attentions[layer_index].values())[0]
    token_attn = target_attention_layer[:, :, token_index, :]

    token_attn = tf.reduce_mean(token_attn, axis=1)

    total_weights = tf.reduce_sum(token_attn, axis=-1, keepdims=True)
    norm_token_attn = token_attn / total_weights

    return norm_token_attn

def viz_attention(tokens, token_weights, target_label, pred_label, pred_probs, review_id, viz_relative=False):
    """
    Returns a html string with the tokens highlighted according to its weights
    :param tokens: A list of strings
    :param token_weights: A rank 1 numpy array where each element is the weight to attach to the corresponding token
    :param target_label: The input's ground truth label
    :param pred_label: The predicted label
    :param pred_probs: The array of predicted probabilities
    :param review_id: The input's id
    :param viz_relative: boolean indicating whether to normalize token_weights by dividing it by the max weight
    :return: A html string
    """
    prob_0, prob_1 = pred_probs
    token_weights = token_weights / np.max(token_weights) if viz_relative else token_weights

    top = Element('span')
    top.set('style', 'white-space: pre-wrap;')

    title = SubElement(top, 'h3')
    title.text = review_id

    header = SubElement(top, 'p')
    header.set('style', 'white-space: pre-line')
    header.text = """
    Target label: {}
    Predicted label: {}
    Probabilities: [{:.4f}, {:.4f}]
    """.format(target_label, pred_label, prob_0, prob_1)

    input_text = SubElement(top, 'p')
    for token, weight in zip(tokens, token_weights):
        child = SubElement(input_text, 'span')
        child.set('style', 'background-color: rgba(0, 0, 256, {})'.format(weight))
        child.text = token
        child.tail = ' '

    return tostring(top)

def merge_wordpiece_tokens(paired_tokens):
    """
    Combine tokens that have been broken up during the wordpiece tokenization process
    The weights of the merged token is the average of their wordpieces
    :param paired_tokens: A list of tuples with the form (wordpiece_token, attention_weight)
    :return: A list of tuples with the form (word_token, attention_weight)
    """
    new_paired_tokens = []
    n_tokens = len(paired_tokens)

    i = 0

    while i < n_tokens:
        current_token, current_weight = paired_tokens[i]
        if current_token.startswith('##'):
            previous_token, previous_weight = new_paired_tokens.pop()
            merged_token = previous_token
            merged_weight = [previous_weight]
            while current_token.startswith('##'):
                merged_token = merged_token + current_token.replace('##', '')
                merged_weight.append(current_weight)
                i = i + 1
                current_token, current_weight = paired_tokens[i]
            merged_weight = np.mean(merged_weight)
            new_paired_tokens.append((merged_token, merged_weight))

        else:
            new_paired_tokens.append((current_token, current_weight))
            i = i + 1
    return new_paired_tokens

def ignore_token_weights(paired_tokens, stopwords):
    """
    Set attention weights of certain tokens to 0 given a list of stopwords
    :param paired_tokens: A list of tuple with the form (token, attention_weight)
    :param stopwords: A list of words whose attention weight we want to set to 0
    :return: paired_tokens with the some weights set to 0 based on stopwords
    """
    new_paired_tokens = []
    for token, weight in paired_tokens:
        weight = 0 if token in stopwords else weight
        new_paired_tokens.append((token, weight))

    return new_paired_tokens