#!/usr/bin/env python3

import math
import multiprocessing
import os
import time
import datetime
import sys

import cv2
import matplotlib
matplotlib.use('Agg')

import numpy as np
import tensorflow as tf
import tensorlayer as tl
from pycocotools.coco import maskUtils

import _pickle as cPickle

sys.path.append('.')

from train_config import config
from openpose_plus.models import model
from openpose_plus.utils import PoseInfo, draw_results, get_heatmap, get_vectormap, load_mscoco_dataset, tf_repeat

tf.logging.set_verbosity(tf.logging.DEBUG)
tl.logging.set_verbosity(tl.logging.DEBUG)

tl.files.exists_or_mkdir(config.LOG.vis_path, verbose=False)  # to save visualization results
tl.files.exists_or_mkdir(config.MODEL.model_path, verbose=False)  # to save model files

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

# FIXME: Don't use global variables.
# define hyper-parameters for training
batch_size = config.TRAIN.batch_size
lr_decay_every_step = config.TRAIN.lr_decay_every_step
n_step = config.TRAIN.n_step
save_interval = config.TRAIN.save_interval
weight_decay_factor = config.TRAIN.weight_decay_factor
lr_init = config.TRAIN.lr_init
lr_decay_factor = config.TRAIN.lr_decay_factor

# FIXME: Don't use global variables.
# define hyper-parameters for model
model_path = config.MODEL.model_path
n_pos = config.MODEL.n_pos
hin = config.MODEL.hin
win = config.MODEL.win
hout = config.MODEL.hout
wout = config.MODEL.wout

tensorboard = True

def get_pose_data_list(im_path, ann_path):
    """
    train_im_path : image folder name
    train_ann_path : coco json file name
    """
    print("[x] Get pose data from {}".format(im_path))
    data = PoseInfo(im_path, ann_path, False)
    imgs_file_list = data.get_image_list()
    objs_info_list = data.get_joint_list()
    mask_list = data.get_mask()
    targets = list(zip(objs_info_list, mask_list))
    if len(imgs_file_list) != len(objs_info_list):
        raise Exception("number of images and annotations do not match")
    else:
        print("{} has {} images".format(im_path, len(imgs_file_list)))
    return imgs_file_list, objs_info_list, mask_list, targets

def make_model(img, results, mask, is_train=True, reuse=False):
    confs = results[:, :, :, :n_pos]
    pafs = results[:, :, :, n_pos:]
    m1 = tf_repeat(mask, [1, 1, 1, n_pos])
    m2 = tf_repeat(mask, [1, 1, 1, n_pos * 2])

    cnn, b1_list, b2_list, net = model(img, n_pos, m1, m2, is_train, reuse)

    # define loss
    losses = []
    last_losses_l1 = []
    last_losses_l2 = []
    stage_losses = []

    for idx, (l1, l2) in enumerate(zip(b1_list, b2_list)):
        loss_l1 = tf.nn.l2_loss((l1.outputs - confs) * m1)
        loss_l2 = tf.nn.l2_loss((l2.outputs - pafs) * m2)

        losses.append(tf.reduce_mean([loss_l1, loss_l2]))
        stage_losses.append(loss_l1 / config.TRAIN.batch_size)
        stage_losses.append(loss_l2 / config.TRAIN.batch_size)

    last_conf = b1_list[-1].outputs
    last_paf = b2_list[-1].outputs
    last_losses_l1.append(loss_l1)
    last_losses_l2.append(loss_l2)
    l2_loss = 0.0

    for p in tl.layers.get_variables_with_name('kernel', True, True):
        l2_loss += tf.contrib.layers.l2_regularizer(weight_decay_factor)(p)
    total_loss = tf.reduce_sum(losses) / config.TRAIN.batch_size + l2_loss

    log_tensors = {'total_loss': total_loss, 'stage_losses': stage_losses, 'l2_loss': l2_loss}
    net.cnn = cnn
    net.img = img  # net input
    net.last_conf = last_conf  # net output
    net.last_paf = last_paf  # net output
    net.confs = confs  # GT
    net.pafs = pafs  # GT
    net.m1 = m1  # mask1, GT
    net.m2 = m2  # mask2, GT
    net.stage_losses = stage_losses
    net.l2_loss = l2_loss
    return net, total_loss, log_tensors
    # return total_loss, last_conf, stage_losses, l2_loss, cnn, last_paf, img, confs, pafs, m1, net

def _data_aug_fn(image, ground_truth):
    """Data augmentation function."""
    ground_truth = cPickle.loads(ground_truth)
    ground_truth = list(ground_truth)

    annos = ground_truth[0]
    mask = ground_truth[1]
    h_mask, w_mask, _ = np.shape(image)
    # mask
    mask_miss = np.ones((h_mask, w_mask), dtype=np.uint8)

    for seg in mask:
        bin_mask = maskUtils.decode(seg)
        bin_mask = np.logical_not(bin_mask)
        mask_miss = np.bitwise_and(mask_miss, bin_mask)

    ## image data augmentation
    # # randomly resize height and width independently, scale is changed
    # image, annos, mask_miss = tl.prepro.keypoint_random_resize(image, annos, mask_miss, zoom_range=(0.8, 1.2))# removed hao
    # # random rotate
    # image, annos, mask_miss = tl.prepro.keypoint_random_rotate(image, annos, mask_miss, rg=15.0)# removed hao
    # # random left-right flipping
    # image, annos, mask_miss = tl.prepro.keypoint_random_flip(image, annos, mask_miss, prob=0.5)# removed hao

    M_rotate = tl.prepro.affine_rotation_matrix(angle=(-180, 180))  # original paper: -40~40
    # M_flip = tl.prepro.affine_horizontal_flip_matrix(prob=0.5) # hao removed: bug, keypoints will have error
    M_zoom = tl.prepro.affine_zoom_matrix(zoom_range=(0.5, 1.1))  # original paper: 0.5~1.1
    # M_shear = tl.prepro.affine_shear_matrix(x_shear=(-0.1, 0.1), y_shear=(-0.1, 0.1))
    M_combined = M_rotate.dot(M_zoom)
    # M_combined = M_rotate.dot(M_flip).dot(M_zoom)#.dot(M_shear)
    # M_combined = tl.prepro.affine_zoom_matrix(zoom_range=0.9) # for debug
    h, w, _ = image.shape
    transform_matrix = tl.prepro.transform_matrix_offset_center(M_combined, x=w, y=h)
    image = tl.prepro.affine_transform_cv2(image, transform_matrix)
    mask_miss = tl.prepro.affine_transform_cv2(mask_miss, transform_matrix, border_mode='replicate')
    annos = tl.prepro.affine_transform_keypoints(annos, transform_matrix)

    # random resize height and width together
    # image, annos, mask_miss = tl.prepro.keypoint_random_resize_shortestedge(
    #     image, annos, mask_miss, min_size=(hin, win), zoom_range=(0.95, 1.6)) # removed hao
    # random crop
    # image, annos, mask_miss = tl.prepro.keypoint_random_crop(image, annos, mask_miss, size=(hin, win))  # with padding # removed hao

    image, annos, mask_miss = tl.prepro.keypoint_random_flip(image, annos, mask_miss, prob=0.5)
    image, annos, mask_miss = tl.prepro.keypoint_resize_random_crop(image, annos, mask_miss, size=(hin, win)) # hao add

    # generate result maps including keypoints heatmap, pafs and mask
    h, w, _ = np.shape(image)
    height, width, _ = np.shape(image)
    heatmap = get_heatmap(annos, height, width)
    vectormap = get_vectormap(annos, height, width)
    resultmap = np.concatenate((heatmap, vectormap), axis=2)

    image = np.array(image, dtype=np.float32)

    img_mask = mask_miss.reshape(hin, win, 1)
    image = image * np.repeat(img_mask, 3, 2)

    resultmap = np.array(resultmap, dtype=np.float32)
    mask_miss = cv2.resize(mask_miss, (hout, wout), interpolation=cv2.INTER_AREA)
    mask_miss = np.array(mask_miss, dtype=np.float32)
    return image, resultmap, mask_miss

def _map_fn(img_list, annos):
    """TF Dataset pipeline."""
    image = tf.read_file(img_list)
    image = tf.image.decode_jpeg(image, channels=3)  # get RGB with 0~1
    image = tf.image.convert_image_dtype(image, dtype=tf.float32)
    # Affine transform and get paf maps
    image, resultmap, mask = tf.py_func(_data_aug_fn, [image, annos], [tf.float32, tf.float32, tf.float32])

    image = tf.reshape(image, [hin, win, 3])
    resultmap = tf.reshape(resultmap, [hout, wout, n_pos * 3])
    mask = tf.reshape(mask, [hout, wout, 1])

    image = tf.image.random_brightness(image, max_delta=45./255.)   # 64./255. 32./255.)  caffe -30~50
    image = tf.image.random_contrast(image, lower=0.5, upper=1.5)   # lower=0.2, upper=1.8)  caffe 0.3~1.5
    # image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
    # image = tf.image.random_hue(image, max_delta=0.1)
    image = tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0)

    return image, resultmap, mask

def single_train(training_dataset):
    ds = training_dataset.shuffle(buffer_size=4096)  # shuffle before loading images
    ds = ds.repeat(n_epoch)
    ds = ds.map(_map_fn, num_parallel_calls=multiprocessing.cpu_count() // 2)  # decouple the heavy map_fn
    ds = ds.batch(config.TRAIN.batch_size)  # TODO: consider using tf.contrib.map_and_batch
    ds = ds.prefetch(2)
    iterator = ds.make_one_shot_iterator()
    one_element = iterator.get_next()
    net, total_loss, log_tensors = make_model(*one_element, is_train=True, reuse=False)
    x_ = net.img  # net input
    last_conf = net.last_conf  # net output
    last_paf = net.last_paf  # net output
    confs_ = net.confs  # GT
    pafs_ = net.pafs  # GT
    mask = net.m1  # mask1, GT
    # net.m2 = m2                 # mask2, GT
    stage_losses = net.stage_losses
    l2_loss = net.l2_loss
    temp_tot_loss, temp_l2_loss = (float("inf"),float("inf"))

    global_step = tf.Variable(1, trainable=False)
    print('Start - n_step: {} batch_size: {} lr_init: {} lr_decay_every_step: {}'.format(
        n_step, config.TRAIN.batch_size, lr_init, lr_decay_every_step))
    with tf.variable_scope('learning_rate'):
        lr_v = tf.Variable(lr_init, trainable=False)

    opt = tf.train.MomentumOptimizer(lr_v, 0.9)
    train_op = opt.minimize(total_loss, global_step=global_step)
    tfconfig = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)

    # start training
    with tf.Session(config=tfconfig) as sess:
        sess.run(tf.global_variables_initializer())

        if tensorboard:

            if not os.path.exists('summaries'):
                os.mkdir('summaries')
            if not os.path.exists(os.path.join('summaries', 'first')):
                os.mkdir(os.path.join('summaries', 'first'))

            summ_writer = tf.summary.FileWriter(os.path.join('summaries', 'run'+str(datetime.datetime.now())), sess.graph)

            tf.summary.scalar('total_loss', total_loss)
            tf.summary.scalar('l2_loss', l2_loss)

            for ix, ll in enumerate(stage_losses):
                tf.summary.scalar('stage{}_loss'.format(ix), ll)
            merge = tf.summary.merge_all()

        # restore pre-trained weights
        try:
            # tl.files.load_and_assign_npz(sess, os.path.join(model_path, 'pose.npz'), net)
            tl.files.load_and_assign_npz_dict(sess=sess, name=os.path.join(model_path, config.MODEL.model_file))
        except:
            print("no pre-trained model")

        if config.MODEL.initial_weights:
            # restore pre-trained mobilnet weights
            try:
                tl.files.load_and_assign_npz_dict(sess=sess, name=os.path.join(model_path, config.MODEL.initial_weights_file))
            except:
                print("no pre-trained mobilnet model")

        # train until the end
        sess.run(tf.assign(lr_v, lr_init))
        while True:
            tic = time.time()
            step = sess.run(global_step)
            if step != 0 and (step % lr_decay_every_step == 0):
                new_lr_decay = lr_decay_factor**(step // lr_decay_every_step)
                sess.run(tf.assign(lr_v, lr_init * new_lr_decay))

            [_, _loss, _stage_losses, _l2, conf_result, paf_result] = \
                sess.run([train_op, total_loss, stage_losses, l2_loss, last_conf, last_paf])

            # tstring = time.strftime('%d-%m %H:%M:%S', time.localtime(time.time()))
            lr = sess.run(lr_v)
            print('Total Loss at iteration {} / {} is: {} Learning rate {:10e} l2_loss {:10e} Took: {}s'.format(
                step, n_step, _loss, lr, _l2,
                time.time() - tic))
            for ix, ll in enumerate(_stage_losses):
                print('Network#', ix, 'For Branch', ix % 2 + 1, 'Loss:', ll)

            if temp_tot_loss > _loss:
                temp_tot_loss = _loss
                # save some results
                [img_out, confs_ground, pafs_ground, conf_result, paf_result,
                    mask_out] = sess.run([x_, confs_, pafs_, last_conf, last_paf, mask])
                draw_results(img_out, confs_ground, conf_result, pafs_ground, paf_result, mask_out,
                    'train_best_%d_' % step)

                # save model
                # tl.files.save_npz(
                #    net.all_params, os.path.join(model_path, 'pose' + str(step) + '.npz'), sess=sess)
                # tl.files.save_npz(net.all_params, os.path.join(model_path, 'pose.npz'), sess=sess)
                tl.files.save_npz_dict(
                    net.all_params, os.path.join(model_path, 'pose_best.npz'), sess=sess)

            if tensorboard:
                summ = sess.run(merge)
                summ_writer.add_summary(summ, step)

            # save intermediate results and model
            if (step != 0) and (step % save_interval == 0):
                # save some results
                [img_out, confs_ground, pafs_ground, conf_result, paf_result,
                 mask_out] = sess.run([x_, confs_, pafs_, last_conf, last_paf, mask])
                draw_results(img_out, confs_ground, conf_result, pafs_ground, paf_result, mask_out, 'train_%d_' % step)
                # save model
                # tl.files.save_npz(
                #    net.all_params, os.path.join(model_path, 'pose' + str(step) + '.npz'), sess=sess)
                # tl.files.save_npz(net.all_params, os.path.join(model_path, 'pose.npz'), sess=sess)
                tl.files.save_npz_dict(net.all_params, os.path.join(model_path, 'pose' + str(step) + '.npz'), sess=sess)
                tl.files.save_npz_dict(net.all_params, os.path.join(model_path, config.MODEL.model_file), sess=sess)
            if step == n_step:  # training finished
                break

def _mock_map_fn(img_list, annos):
    """TF Dataset pipeline."""
    image = tf.read_file(img_list)
    image = tf.image.decode_jpeg(image, channels=3)  # get RGB with 0~1
    image = tf.image.convert_image_dtype(image, dtype=tf.float32)

    image = np.ones((hin, win, 3), dtype=np.float32)
    resultmap = np.ones((hout, wout, 57), dtype=np.float32)
    mask = np.ones((hout, wout, 1), dtype=np.float32)

    return image, resultmap, mask

def _parallel_train_model(img, results, mask):
    net, total_loss, log_tensors = make_model(img, results, mask, is_train=True, reuse=False)
    return net, total_loss, log_tensors

def parallel_train(training_dataset):
    import horovod.tensorflow as hvd

    hvd.init()  # Horovod

    ds = training_dataset.shuffle(buffer_size=4096)
    ds = ds.shard(num_shards=hvd.size(), index=hvd.rank())
    ds = ds.repeat(n_epoch)
    ds = ds.map(_map_fn, num_parallel_calls=4)
    ds = ds.batch(config.TRAIN.batch_size)
    ds = ds.prefetch(buffer_size=1)

    iterator = ds.make_one_shot_iterator()
    one_element = iterator.get_next()
    net, total_loss, log_tensors = make_model(*one_element, is_train=True, reuse=False)
    x_ = net.img  # net input
    last_conf = net.last_conf  # net output
    last_paf = net.last_paf  # net output
    confs_ = net.confs  # GT
    pafs_ = net.pafs  # GT
    mask = net.m1  # mask1, GT
    # net.m2 = m2                 # mask2, GT
    stage_losses = net.stage_losses
    l2_loss = net.l2_loss
    temp_tot_loss, temp_l2_loss = (float("inf"),float("inf"))

    global_step = tf.Variable(1, trainable=False)
    # TODO
    # scaled_lr = lr_init * hvd.size()  # Horovod: scale the learning rate linearly
    scaled_lr = lr_init  # Linear scaling rule is not working in openpose training.
    with tf.variable_scope('learning_rate'):
        lr_v = tf.Variable(scaled_lr, trainable=False)

    opt = tf.train.MomentumOptimizer(lr_v, 0.9)
    opt = hvd.DistributedOptimizer(opt)  # Horovod
    train_op = opt.minimize(total_loss, global_step=global_step)
    tfconfig = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)

    tfconfig.gpu_options.allow_growth = True  # Horovod
    tfconfig.gpu_options.visible_device_list = str(hvd.local_rank())  # Horovod

    # Add variable initializer.
    init = tf.global_variables_initializer()

    # Horovod: broadcast initial variable states from rank 0 to all other processes.
    # This is necessary to ensure consistent initialization of all workers when
    # training is started with random weights or restored from a checkpoint.
    bcast = hvd.broadcast_global_variables(0)  # Horovod

    # Horovod: adjust number of steps based on number of GPUs.
    global n_step, lr_decay_every_step
    n_step = n_step // hvd.size() + 1  # Horovod
    lr_decay_every_step = lr_decay_every_step // hvd.size() + 1  # Horovod

    # Start training
    with tf.Session(config=tfconfig) as sess:
        if tensorboard:
            if hvd.rank() == 0:

                if not os.path.exists('summaries'):
                    os.mkdir('summaries')

                summ_writer = tf.summary.FileWriter(os.path.join('summaries', 'run'+str(datetime.datetime.now())), sess.graph)

            tf.summary.scalar('total_loss', total_loss)
            tf.summary.scalar('l2_loss', l2_loss)

            for ix, ll in enumerate(stage_losses):
                tf.summary.scalar('stage{}_loss'.format(ix), ll)
            merge = tf.summary.merge_all()

        #Init model and load weights
        init.run()

        #TODO
        if hvd.rank() == 0:  # Horovod
            # restore pre-trained weights
            try:
                # tl.files.load_and_assign_npz(sess, os.path.join(model_path, 'pose.npz'), net)
                tl.files.load_and_assign_npz_dict(sess=sess, name=os.path.join(model_path, config.MODEL.model_file))
            except:
                print("no pre-trained model")

            if config.MODEL.initial_weights:
                # restore pre-trained mobilnet weights
                try:
                    tl.files.load_and_assign_npz_dict(sess=sess, name=os.path.join(model_path, config.MODEL.initial_weights_file))
                except:
                    print("no pre-trained mobilnet model")

        bcast.run()  # Horovod bcast vars across workers
        print('Worker{}: Initialized'.format(hvd.rank()))
        print('Worker{}: Start - n_step: {} batch_size: {} lr_init: {} lr_decay_every_step: {}'.format(
            hvd.rank(), n_step, config.TRAIN.batch_size, lr_init, lr_decay_every_step))

        # train until the end
        while True:
            step = sess.run(global_step)
            if step == n_step:
                break

            tic = time.time()
            if step != 0 and (step % lr_decay_every_step == 0):
                new_lr_decay = lr_decay_factor**(step // lr_decay_every_step)
                sess.run(tf.assign(lr_v, scaled_lr * new_lr_decay))

            [_, _loss, _stage_losses, _l2, conf_result, paf_result] = \
                sess.run([train_op, total_loss, stage_losses, l2_loss, last_conf, last_paf])

            # tstring = time.strftime('%d-%m %H:%M:%S', time.localtime(time.time()))
            lr = sess.run(lr_v)
            print(
                'Worker{}: Total Loss at iteration {} / {} is: {} Learning rate {:10e} l2_loss {:10e} Took: {}s'.format(
                    hvd.rank(), step, n_step, _loss, lr, _l2,
                    time.time() - tic))
            for ix, ll in enumerate(_stage_losses):
                print('Worker{}:', hvd.rank(), 'Network#', ix, 'For Branch', ix % 2 + 1, 'Loss:', ll)

            # save intermediate results and model
            if hvd.rank() == 0:  # Horovod
                if temp_tot_loss > _loss:
                    temp_tot_loss = _loss
                    # save some results
                    [img_out, confs_ground, pafs_ground, conf_result, paf_result,
                     mask_out] = sess.run([x_, confs_, pafs_, last_conf, last_paf, mask])
                    draw_results(img_out, confs_ground, conf_result, pafs_ground, paf_result, mask_out,
                                 'train_best_%d_' % step)

                    # save model
                    # tl.files.save_npz(
                    #    net.all_params, os.path.join(model_path, 'pose' + str(step) + '.npz'), sess=sess)
                    # tl.files.save_npz(net.all_params, os.path.join(model_path, 'pose.npz'), sess=sess)
                    tl.files.save_npz_dict(
                        net.all_params, os.path.join(model_path, 'pose_best.npz'), sess=sess)

                if tensorboard:
                    summ = sess.run(merge)
                    summ_writer.add_summary(summ, step)

                if (step != 0) and (step % save_interval == 0):
                    # save some results
                    [img_out, confs_ground, pafs_ground, conf_result, paf_result,
                     mask_out] = sess.run([x_, confs_, pafs_, last_conf, last_paf, mask])
                    draw_results(img_out, confs_ground, conf_result, pafs_ground, paf_result, mask_out,
                                 'train_%d_' % step)

                    # save model
                    # tl.files.save_npz(
                    #    net.all_params, os.path.join(model_path, 'pose' + str(step) + '.npz'), sess=sess)
                    # tl.files.save_npz(net.all_params, os.path.join(model_path, 'pose.npz'), sess=sess)
                    tl.files.save_npz_dict(
                        net.all_params, os.path.join(model_path, 'pose' + str(step) + '.npz'), sess=sess)
                    tl.files.save_npz_dict(net.all_params, os.path.join(model_path, config.MODEL.model_file), sess=sess)

if __name__ == '__main__':

    if 'coco' in config.DATA.train_data:
        # automatically download MSCOCO data to "data/mscoco..."" folder
        train_im_path, train_ann_path, val_im_path, val_ann_path, _, _ = \
            load_mscoco_dataset(config.DATA.data_path, config.DATA.coco_version, task='person')

        # read coco training images contains valid people
        train_imgs_file_list, train_objs_info_list, train_mask_list, train_targets = \
            get_pose_data_list(train_im_path, train_ann_path)

        # read coco validating images contains valid people (you can use it for training as well)
        val_imgs_file_list, val_objs_info_list, val_mask_list, val_targets = \
            get_pose_data_list(val_im_path, val_ann_path)

    if 'custom' in config.DATA.train_data:
        ## read your own images contains valid people
        ## 1. if you only have one folder as follow:
        ##   data/your_data
        ##           /images
        ##               0001.jpeg
        ##               0002.jpeg
        ##           /coco.json
        # your_imgs_file_list, your_objs_info_list, your_mask_list, your_targets = \
        #     get_pose_data_list(config.DATA.your_images_path, config.DATA.your_annos_path)
        ## 2. if you have a folder with many folders: (which is common in industry)
        folder_list = tl.files.load_folder_list(path='data/your_data')
        your_imgs_file_list, your_objs_info_list, your_mask_list = [], [], []
        for folder in folder_list:
            _imgs_file_list, _objs_info_list, _mask_list, _targets = \
                get_pose_data_list(os.path.join(folder, 'images'), os.path.join(folder, 'coco.json'))
            print(len(_imgs_file_list))
            your_imgs_file_list.extend(_imgs_file_list)
            your_objs_info_list.extend(_objs_info_list)
            your_mask_list.extend(_mask_list)
        print("number of own images found:", len(your_imgs_file_list))

    # choose dataset for training
    if config.DATA.train_data == 'coco':
        # 1. only coco training set
        imgs_file_list = train_imgs_file_list
        train_targets = list(zip(train_objs_info_list, train_mask_list))
    elif config.DATA.train_data == 'custom':
        # 2. only your own data
        imgs_file_list = your_imgs_file_list
        train_targets = list(zip(your_objs_info_list, your_mask_list))
    elif config.DATA.train_data == 'coco_and_custom':
        # 3. your own data and coco training set
        imgs_file_list = train_imgs_file_list + your_imgs_file_list
        train_targets = list(zip(train_objs_info_list + your_objs_info_list, train_mask_list + your_mask_list))
    else:
        raise Exception('please choose a valid config.DATA.train_data setting.')

    # define data augmentation
    def generator():
        """TF Dataset generator."""
        assert len(imgs_file_list) == len(train_targets)
        for _input, _target in zip(imgs_file_list, train_targets):
            yield _input.encode('utf-8'), cPickle.dumps(_target)

    n_epoch = math.ceil(n_step / (len(imgs_file_list) / config.TRAIN.batch_size))
    dataset = tf.data.Dataset().from_generator(generator, output_types=(tf.string, tf.string))

    if config.TRAIN.train_mode == 'single':
        single_train(dataset)
    elif config.TRAIN.train_mode == 'parallel':
        parallel_train(dataset)
    else:
        raise Exception('Unknown training mode')


import os
import math
try:
    from rio_hist.match import histogram_match as rio_match
    has_rio = True
except ImportError:
    has_rio = False

import mercantile
from shapely.geometry import box, shape, mapping, asShape

import numpy as np
try:
    from matplotlib import pyplot as plt
    has_pyplot = True
except:
    has_pyplot = False

class PlotMixin(object):

    def rgb(self, **kwargs):
        ''' Convert the image to a 3 band RGB for plotting

        This method shares the same arguments as plot(). It will perform visual adjustment on the
        image and prepare the data for plotting in MatplotLib. Values are converted to an
        appropriate precision and the axis order is changed to put the band axis last.
        '''

        if "bands" in kwargs:
            use_bands = kwargs["bands"]
            assert len(use_bands) == 3, 'Plot method only supports single or 3-band outputs'
            del kwargs["bands"]
        else:
            use_bands = self._rgb_bands
        if kwargs.get('blm') == True:
            return self.histogram_match(use_bands, **kwargs)
        if "histogram" not in kwargs:
            if "stretch" not in kwargs:
                kwargs['stretch'] = [2,98]
            return self.histogram_stretch(use_bands, **kwargs)
        elif kwargs["histogram"] == "equalize":
            return self.histogram_equalize(use_bands, **kwargs)
        elif kwargs["histogram"] == "match":
            return self.histogram_match(use_bands, **kwargs)
        elif kwargs["histogram"] == "minmax":
            return self.histogram_stretch(use_bands, stretch=[0, 100], **kwargs)
        elif kwargs["histogram"] == "ignore":
            data = self._read(self[use_bands,...], **kwargs)
            return np.rollaxis(data, 0, 3)
        else:
            raise KeyError('Unknown histogram parameter, use "equalize", "match", "minmax", or "ignore"')

    def histogram_equalize(self, use_bands, **kwargs):
        ''' Equalize and the histogram and normalize value range
            Equalization is on all three bands, not per-band'''
        data = self._read(self[use_bands,...], **kwargs)
        data = np.rollaxis(data.astype(np.float32), 0, 3)
        flattened = data.flatten()
        if 0 in data:
            masked = np.ma.masked_values(data, 0).compressed()
            image_histogram, bin_edges = np.histogram(masked, 256)
        else:
            image_histogram, bin_edges = np.histogram(flattened, 256)
        bins = (bin_edges[:-1] + bin_edges[1:]) / 2.0
        cdf = image_histogram.cumsum()
        cdf = cdf / float(cdf[-1])
        image_equalized = np.interp(flattened, bins, cdf).reshape(data.shape)
        if 'stretch' in kwargs or 'gamma' in kwargs:
            return self._histogram_stretch(image_equalized, **kwargs)
        else:
            return image_equalized

    def histogram_match(self, use_bands, blm_source=None, **kwargs):
        ''' Match the histogram to existing imagery '''
        assert has_rio, "To match image histograms please install rio_hist"
        data = self._read(self[use_bands,...], **kwargs)
        data = np.rollaxis(data.astype(np.float32), 0, 3)
        if 0 in data:
            data = np.ma.masked_values(data, 0)
        bounds = self._reproject(box(*self.bounds), from_proj=self.proj, to_proj="EPSG:4326").bounds
        if blm_source == 'browse':
            from gbdxtools.images.browse_image import BrowseImage
            ref = BrowseImage(self.cat_id, bbox=bounds).read()
        else:
            from gbdxtools.images.tms_image import TmsImage
            tms = TmsImage(zoom=self._calc_tms_zoom(self.affine[0]), bbox=bounds, **kwargs)
            ref = np.rollaxis(tms.read(), 0, 3)
        out = np.dstack([rio_match(data[:,:,idx], ref[:,:,idx].astype(np.double)/255.0)
                        for idx in range(data.shape[-1])])
        if 'stretch' in kwargs or 'gamma' in kwargs:
            return self._histogram_stretch(out, **kwargs)
        else:
            return out

    def histogram_stretch(self, use_bands, **kwargs):
        ''' entry point for contrast stretching '''
        data = self._read(self[use_bands,...], **kwargs)
        data = np.rollaxis(data.astype(np.float32), 0, 3)
        return self._histogram_stretch(data, **kwargs)

    def _histogram_stretch(self, data, **kwargs):
        ''' perform a contrast stretch and/or gamma adjustment '''
        limits = {}
        # get the image min-max statistics
        for x in range(3):
            band = data[:,:,x]
            try:
                limits[x] = np.percentile(band, kwargs.get("stretch", [0,100]))
            except IndexError:
                # this band has no dynamic range and cannot be stretched
                return data
        # compute the stretch
        for x in range(3):
            band = data[:,:,x]
            if 0 in band:
                band = np.ma.masked_values(band, 0).compressed()
            top = limits[x][1]
            bottom = limits[x][0]
            if top != bottom: # catch divide by zero
                data[:,:,x] = (data[:,:,x] - bottom) / float(top - bottom) * 255.0
        data = np.clip(data, 0, 255).astype("uint8")
        # gamma adjust
        if "gamma" in kwargs:
            invGamma = 1.0 / kwargs['gamma']
            lut = np.array([((i / 255.0) ** invGamma) * 255
		            for i in np.arange(0, 256)]).astype("uint8")
            data = np.take(lut, data)
        return data

    def ndvi(self, **kwargs):
        """
        Calculates Normalized Difference Vegetation Index using NIR and Red of an image.

        Returns: numpy array with ndvi values
        """
        data = self._read(self[self._ndvi_bands,...]).astype(np.float32)
        return (data[0,:,:] - data[1,:,:]) / (data[0,:,:] + data[1,:,:])

    def ndwi(self):
        """
        Calculates Normalized Difference Water Index using Coastal and NIR2 bands for WV02, WV03.
        For Landsat8 and sentinel2 calculated by using Green and NIR bands.

        Returns: numpy array of ndwi values
        """
        data = self._read(self[self._ndwi_bands,...]).astype(np.float32)
        return (data[1,:,:] - data[0,:,:]) / (data[0,:,:] + data[1,:,:])

    def plot(self, spec="rgb", **kwargs):
        ''' Plot the image with MatplotLib

        Plot sizing includes default borders and spacing. If the image is shown in Jupyter the outside whitespace will be automatically cropped to save size, resulting in a smaller sized image than expected.

        Histogram options:
            * 'equalize': performs histogram equalization on the image.
            * 'minmax': stretch the pixel range to the minimum and maximum input pixel values. Equivalent to stretch=[0,100].
            * 'match': match the histogram to the Maps API imagery. Pass the additional keyword blm_source='browse' to match to the Browse Service (image thumbnail) instead.
            * 'ignore': Skip dynamic range adjustment, in the event the image is already correctly balanced and the values are in the correct range.

        Gamma values greater than 1 will brighten the image midtones, values less than 1 will darken the midtones.

        Plots generated with the histogram options of 'match' and 'equalize' can be combined with the stretch and gamma options. The stretch and gamma adjustments will be applied after the histogram adjustments.

        Args:
            w (float or int): width of plot in inches at 72 dpi, default is 10
            h (float or int): height of plot in inches at 72 dpi, default is 10
            bands (list): bands to use for plotting, such as bands=[4,2,1]. Defaults to the image's natural RGB bands. This option is useful for generating pseudocolor images when passed a list of three bands. If only a single band is provided, a colormapped plot will be generated instead.
            colormap (str): MatPlotLib colormap name to use for single band images. Default is colormap='Grey_R'.
            histogram (str): either 'equalize', 'minmax', 'match', or ignore
            stretch (list): stretch the histogram between two percentile values, default is [2,98]
            gamma (float): adjust image gamma, default is 1.0
        '''

        if self.shape[0] == 1 or ("bands" in kwargs and len(kwargs["bands"]) == 1):
            if "cmap" in kwargs:
                cmap = kwargs["cmap"]
                del kwargs["cmap"]
            else:
                cmap = "Greys_r"
            self._plot(tfm=self._single_band, cmap=cmap, **kwargs)
        else:
            if spec == "rgb" and self._has_token(**kwargs):
                self._plot(tfm=self.rgb, **kwargs)
            else:
                self._plot(tfm=getattr(self, spec), **kwargs)

    def _has_token(self, **kwargs):
        if "access_token" in kwargs or "MAPBOX_API_KEY" in os.environ:
            return True
        else:
            return False

    def _plot(self, tfm=lambda x: x, **kwargs):
        assert has_pyplot, "To plot images please install matplotlib"
        assert self.shape[1] and self.shape[-1], "No data to plot, dimensions are invalid {}".format(str(self.shape))

        f, ax1 = plt.subplots(1, figsize=(kwargs.get("w", 10), kwargs.get("h", 10)))
        ax1.axis('off')
        ax1.set_title(kwargs.get('title', ''), fontsize=kwargs.get('fontsize', 22))
        plt.imshow(tfm(**kwargs), interpolation='nearest', cmap=kwargs.get("cmap", None))
        plt.show(block=False)

    def _read(self, data, **kwargs):
        if hasattr(data, 'read'):
            return data.read(**kwargs)
        else:
            return data.compute(get=threaded_get)

    def _single_band(self, **kwargs):
        arr = self._read(self, **kwargs)
        return arr[0,:,:]

    def _calc_tms_zoom(self, scale):
        for z in range(15,20):
            b = mercantile.bounds(0,0,z)
            if scale > math.sqrt((b.north - b.south)*(b.east - b.west) / (256*256)):
                return z

class BandMethodsTemplate(object):
    @property
    def _rgb_bands(self):
        raise NotImplementedError("RGB bands undefined for image type {}".format(self.__class__.__name__))

    @property
    def _ndvi_bands(self):
        raise NotImplementedError("NDVI bands undefined for image type {}".format(self.__class__.__name__))

    @property
    def _ndwi_bands(self):
        raise NotImplementedError("NDWI bands undefined for image type {}".format(self.__class__.__name__))


import torch
import os
import numpy as np
import cv2
from PIL import Image
import matplotlib.pyplot as plt
from .config import Config
from .utils.to_sqlite import insert_vector_db, insert_human_db, insert_infer_db, load_gallery_from_db, convertToBinaryData, load_human_db, convertImgtoBlob, convertBlobtoIMG
from .utils.reranking import re_ranking

from .model import make_model
from torch.backends import cudnn
import torchvision.transforms as T
from .utils.metrics import cosine_similarity, euclidean_distance
import pickle

class reid_inference:
    """Reid Inference class.
    """

    def __init__(self):
        cudnn.benchmark = True
        self.Cfg = Config()
        self.model = make_model(self.Cfg, 255)
        self.model.load_param(self.Cfg.TEST_WEIGHT)
        self.model = self.model.to('cuda')
        self.transform = T.Compose([
                T.Resize(self.Cfg.INPUT_SIZE),
                T.ToTensor(),
                T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
            ])
        print(f'Model loaded with weight from {self.Cfg.TEST_WEIGHT}')
        self.model.eval()
        print('Ready to Eval')
        print('Loading from DB...')
        self.all_img_id, self.all_patch_img, self.all_gal_feat = load_gallery_from_db() #load from vectorkb_table
        self.human_dict = load_human_db()
        self._tmp_img = ""
        self._tmp_galfeat = ""
        print('Data loaded. You can start infer an image using to_gallery_feat --> query_feat --> infer')

    def to_gallery_feat(self, image_patch_or_path, flip=True, norm=True):
        """
        Use to build gallery feat on images picked from Deep Sort.
        This is different from normal query feature extraction as this has flipped & normed feature,
        to improve the matching precision.
        Takes image path or PIL image directly.
        To be combined with INFER function at the end of troubleshooting
        """
        if type(image_patch_or_path) is str:
            query_img = Image.open(image_patch_or_path)
        else:
            query_img = image_patch_or_path

        input = torch.unsqueeze(self.transform(query_img), 0)
        input = input.to('cuda')
        with torch.no_grad():
            if flip:
                gal_feat = torch.FloatTensor(input.size(0), 2048).zero_().cuda()
                for i in range(2):
                    if i == 1:
                        inv_idx = torch.arange(input.size(3) - 1, -1, -1).long().cuda()
                        input = input.index_select(3, inv_idx)
                    f = self.model(input)
                    gal_feat = gal_feat + f
            else:
                gal_feat = self.model(input)

            if norm:
                gal_feat = torch.nn.functional.normalize(gal_feat, dim=1, p=2)

            self._tmp_img = query_img  #temp save PIL image here
            self._tmp_galfeat = gal_feat #temp save gal_feat here
            return gal_feat

    def to_query_feat(self, image_patch_or_path):
        """
        image - input image path.
        for finding query feature, no flipping and normalization is done.
        This function returns feature (1,2048) tensor.

        """
        if type(image_patch_or_path) is str:
            query_img = Image.open(image_patch_or_path)
        else:
            query_img = image_patch_or_path

        input = torch.unsqueeze(self.transform(query_img), 0)
        input = input.to('cuda')
        with torch.no_grad():
            query_feat = self.model(input)
        return query_feat

    def infer(self, query_feat, query_img_id, top_k= 3, reranking=True):
        if len(self.all_gal_feat)>0:
            # if reranking:
            #     dist_mat = 1 - re_ranking(query_feat, self.all_gal_feat, k1=30, k2=10, lambda_value=0.2)[0]
            #     indices = np.argsort(dist_mat)[::-1]

            # else:
            dist_mat = torch.nn.functional.cosine_similarity(query_feat, self.all_gal_feat).cpu().numpy()
            indices = np.argsort(dist_mat)[::-1][:50] #to test if use 50 or use all better

            #do reranking
            if reranking:
                candidate_gal_feat = torch.index_select(self.all_gal_feat, 0, torch.tensor([indices]).cuda()[0])
                rerank_dist = re_ranking(query_feat, candidate_gal_feat, k1=30, k2=6, lambda_value=0.3)[0]
                rerank_idx = np.argsort(1-rerank_dist)[::-1]
                indices = np.array([indices[i] for i in rerank_idx])

            #if match found --> insert to human_table, need a human list too. make it into class
            #if no match found --> insert new identity to human_table.
            if dist_mat[indices[0]] >= self.Cfg.THRESHOLD:

                #match found
                matched_img_id = self.all_img_id[indices[0]]
                identity = self.human_dict[matched_img_id]
                print(f"Match found! Identity is {identity}")

                #insert to human_table & dict
                insert_human_db(query_img_id, identity, "Matched")
                self.human_dict[query_img_id] = identity

                #insert query image to gallery table & list. Recalling we have self.tmp variables
                query_img_blob = convertImgtoBlob(self._tmp_img)
                insert_vector_db(query_img_id, query_img_blob, pickle.dumps(self._tmp_galfeat) )

                #insert to record table
                record = [query_img_id , query_img_blob]
                for k in range(top_k):
                    try:
                        record.append(self.all_img_id[indices[k]])
                        record.append(convertImgtoBlob(self.all_patch_img[indices[k]]))
                        record.append(dist_mat.item(indices[k]))
                    except:
                        record.append(None)
                        record.append(None)
                        record.append(None)
                insert_infer_db(record)

            elif (len(indices)>= 2) and (dist_mat[indices[1]] >= 0.75):
                #match found
                matched_img_id = self.all_img_id[indices[1]]
                identity = self.human_dict[matched_img_id]
                print(f"Match found! Identity is {identity} --> SECOND MATCH")

                #insert to human_table & dict
                insert_human_db(query_img_id, identity, "Matched")
                self.human_dict[query_img_id] = identity

                #insert query image to gallery table & list. Recalling we have self.tmp variables
                query_img_blob = convertImgtoBlob(self._tmp_img)
                insert_vector_db(query_img_id, query_img_blob, pickle.dumps(self._tmp_galfeat) )

                #insert to record table
                record = [query_img_id , query_img_blob]
                for k in range(top_k):
                    try:
                        record.append(self.all_img_id[indices[k+1]])
                        record.append(convertImgtoBlob(self.all_patch_img[indices[k+1]]))
                        record.append(dist_mat.item(indices[k+1]))
                    except:
                        record.append(None)
                        record.append(None)
                        record.append(None)
                insert_infer_db(record)

            else:
                #no match found
                new_identity = str(max(map(int, self.human_dict.values()))+1)
                print(f"No match found! Creating new identity -- {new_identity}")

                #insert to human_table & dict
                insert_human_db(query_img_id, new_identity, "New")
                self.human_dict[query_img_id] = new_identity

                #insert query image to gallery table & list. Recalling we have self.tmp variables
                query_img_blob = convertImgtoBlob(self._tmp_img)
                insert_vector_db(query_img_id, query_img_blob, pickle.dumps(self._tmp_galfeat) )

                #insert to record table
                record = [query_img_id , query_img_blob]
                for k in range(top_k):
                    record.append(None)
                    try:
                        record.append(convertImgtoBlob(self.all_patch_img[indices[k]]))
                        record.append(dist_mat.item(indices[k]))
                    except:
                        record.append(None)
                        record.append(None)
                insert_infer_db(record)

            #Putting these records into memory database
            self.all_img_id.append(query_img_id)
            self.all_patch_img.append(self._tmp_img)
            self.all_gal_feat = torch.cat([self.all_gal_feat, self._tmp_galfeat])

        else:
            #new record
            new_identity = str(1)
            print(f"No match found! Creating new identity -- {new_identity}")

            #insert to human_table & dict
            insert_human_db(query_img_id, new_identity, "New")
            self.human_dict[query_img_id] = new_identity

            #insert query image to gallery table & list. Recalling we have self.tmp variables
            query_img_blob = convertImgtoBlob(self._tmp_img)
            insert_vector_db(query_img_id, query_img_blob, pickle.dumps(self._tmp_galfeat) )

            #insert to record table
            record = [query_img_id , query_img_blob]
            for k in range(top_k):
                record.append(None)
                record.append(None)
                record.append(None)
            insert_infer_db(record)

            #Putting these records into memory database
            self.all_img_id.append(query_img_id)
            self.all_patch_img.append(self._tmp_img)
            self.all_gal_feat = torch.cat([self.all_gal_feat, self._tmp_galfeat])

    # def build_all_gallery(dir_to_gal_folder = self.Cfg.GALLERY_DIR, to_db = False):
    #     """
    #     TAKE NOTEE!! TO BE MODIFIED AS WE NO LONGER NEED TO MASS UPLOAD FROM
    #     IMG FILE TO GALLERY DB.
    #     """
    #     all_gal_feat = []
    #     all_img_id = os.listdir(dir_to_gal_folder) #this is id rather than path

    #     db_feat = []
    #     db_img = []

    #     print(f'Building gallery from {dir_to_gal_folder}...')
    #     for img in all_img_id:
    #         gal_feat = to_gallery_feat(dir_to_gal_folder + "/" + img)
    #         all_gal_feat.append(gal_feat)
    #         db_feat.append(pickle.dumps(gal_feat))
    #         db_img.append(convertToBinaryData(dir_to_gal_folder + "/" + img))

    #     all_gal_feat = torch.cat(all_gal_feat, dim=0)

    #     if to_db:
    #         db_img_path = [dir_to_gal_folder + "/" + img for img in all_img_id]
    #         db_humam_id = [img.split('_')[0] for img in all_img_id]
    #         insert_vector_db(all_img_id, db_img_path, db_img, db_feat)
    #         insert_human_db(all_img_id, db_humam_id)
    #         print('All gallery uploaded to DB.')
    #     else:
    #         return all_gal_feat, all_img_id

    # UNSUPPORTED
                    # plt.subplot(1, top_k+2, 1)
                    # plt.title('Query')
                    # plt.axis('off')
                    # query_img = Image.open(query_img_path)
                    # plt.imshow(np.asarray(query_img))

                    # for k in range(top_k):
                    #     plt.subplot(1, top_k+2, k+3)
                    #     name = str(indices[k]) + '\n' + '{0:.2f}'.format(dist_mat[indices[k]])
                    #     img = np.asarray(Image.open(self.all_img_path[indices[k]]))
                    #     plt.title(name)
                    #     plt.axis('off')
                    #     plt.imshow(img)
                    # plt.show()


import os
import cv2
import csv
import time
import numpy as np

from PIL import Image
from mtcnn.mtcnn import MTCNN

from models.util import utils

from sklearn.preprocessing import Normalizer
from scipy.spatial.distance import cosine

import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array

in_encoder = Normalizer('l2')
print("Using gpu: {0}".format(tf.test.is_gpu_available(
                            cuda_only=False,
                            min_cuda_compute_capability=None)))

class FaceNet:
    def __init__(self):
        self.CONFIDENCE       = 0.5
        self.THRESHOLD        = 0.3
        self.LABELS           = self.init_label()
        self.COLORS           = self.init_color()
        self.net              = self.init_weight()
        self.model_Face       = self.init_model()
        self.database         = self.init_database()

    # function to read database
    def init_database(self):
        print("read data from database...")
        database = {}

        dbPath = utils.get_file_path("webApp/data", "dict.csv")
        reader = csv.reader(open(dbPath), delimiter='\n')

        for row in reader:
            data = row[0].split(",", 1)
            encode = data[1].replace('"','')
            encode = encode.replace('[','')
            encode = encode.replace(']','')
            encode = np.fromstring(encode, dtype=float, sep=',')
            database[data[0]] = encode
        return database

    def init_model(self):
        print("initial facenet model...")
        facePath = utils.get_file_path("webApp/data", "facenet_keras.h5")
       # cvPath = os.path.sep.join(["data", "haarcascade_frontalface_alt2.xml"])
        #faceCascade = cv2.CascadeClassifier(cvPath)
        return load_model(facePath, custom_objects={ 'loss': self.triplet_loss }, compile=False)

    # initial labels
    def init_label(self):
        print("initial yolo label...")
        labelsPath = utils.get_file_path("webApp/cfg", "classes.names")
        return open(labelsPath).read().strip().split("\n")

    # initial colors
    def init_color(self):
        print("initial yolo colors...")
        np.random.seed(42)
        return np.random.randint(0, 255, size=(len(self.LABELS), 3), dtype="uint8")

    def init_weight(self):
        print("[INFO] loading YOLO from disk...")
        # derive the paths to the YOLO weights and model configuration
        weightsPath = utils.get_file_path("webApp/data", "yolov4.weights")
        configPath = utils.get_file_path("webApp/cfg", "yolov4.cfg")
        return cv2.dnn.readNetFromDarknet(configPath, weightsPath)

    def triplet_loss(y_true, y_pred, alpha = 0.2):
        print("using triplet_loss function...")
        anchor, positive, negative = y_pred[0], y_pred[1], y_pred[2]

        # Step 1: Compute the (encoding) distance between the anchor and the positive
        pos_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, positive)), axis=-1)
        # Step 2: Compute the (encoding) distance between the anchor and the negative
        neg_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, negative)), axis=-1)
        # Step 3: subtract the two previous distances and add alpha.
        basic_loss = tf.add(tf.subtract(pos_dist, neg_dist), alpha)
        # Step 4: Take the maximum of basic_loss and 0.0. Sum over the training examples.
        loss = tf.reduce_sum(tf.maximum(basic_loss, 0.0))

        return loss

    # get face embedding and perform face recognition
    def get_embedding(self, image):
        print("get enbedding code function begin...")
        # scale pixel values
        face = image.astype('float32')
        # standardization
        mean, std = face.mean(), face.std()
        face = (face - mean) / std
        face = cv2.resize(face,(160,160))
        face = np.expand_dims(face, axis=0)
        encode = self.model_Face.predict(face)[0]
        return encode

    def find_person(self, encoding, min_dist=1):
        print("find person function begin...")
        min_dist = float("inf")
        encoding = in_encoder.transform(np.expand_dims(encoding, axis=0))[0]
        for (name, db_enc) in self.database.items():
            dist = cosine(db_enc, encoding)
            if dist < 0.5 and dist < min_dist:
                min_dist = dist
                identity = name

        if min_dist > 0.5:
            return "None"
        else:
            return identity
        return "None"

    def detect(self, frame, W, H):
        # construct a blob from the input frame and then perform a forward
        # pass of the YOLO object detector, giving us our bounding boxes
        # and associated probabilities
        blob = cv2.dnn.blobFromImage(frame, 1 / 255.0, (416, 416), swapRB=True, crop=False)
        self.net.setInput(blob)

        ln = self.net.getLayerNames()
        ln = [ln[i[0] - 1] for i in self.net.getUnconnectedOutLayers()]

        # calculate execute time
        start = time.time()
        layerOutputs = self.net.forward(ln)
        end = time.time()

        print("[INFO] YOLO took {:.6f} seconds".format(end - start))

        # initialize our lists of detected bounding boxes, confidences,
        # and class IDs, respectively
        boxes = []
        confidences = []
        classIDs = []

        #faces_list = []
        #encodes = []
        names = []

        # loop over each of the layer outputs
        for output in layerOutputs:
            # loop over each of the detections
            for detection in output:
                # extract the class ID and confidence of the current object detection
                scores = detection[5:]
                classID = np.argmax(scores)
                confidence = scores[classID]

                # filter out weak predictions by ensuring the detected
                # probability is greater than the minimum probability
                if confidence > self.CONFIDENCE:
                    # scale the bounding box coordinates back relative to
                    # the size of the image
                    box = detection[0:4] * np.array([W, H, W, H])
                    (centerX, centerY, width, height) = box.astype("int")

                    # use the center (x, y)-coordinates to derive the top
                    # and and left corner of the bounding box
                    x = int(centerX - (width / 2))
                    y = int(centerY - (height / 2))

                    # update bounding box coordinates, confidences and class IDs
                    boxes.append([x, y, int(width), int(height)])
                    confidences.append(float(confidence))
                    if classID == 0 :
                        classIDs.append(classID)
                        names.append('')
                    elif classID == 1:
                        #openCV
                        #convert to greyscale
                        #faces_list=[]
                        #encodes=[]
                        crop = frame[y:y+int(height), x:x+int(width)]

                        #gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
                        #detect face
                        #faces = faceCascade.detectMultiScale(crop,
                                                            #scaleFactor=1.1,
                                                            #minNeighbors=5,
                                                            ##minSize=(60, 60),
                                                            #flags=cv2.CASCADE_SCALE_IMAGE)
                        #to draw rectangle
                        label = ""
                        #for (x, y, w, h) in faces:
                        #face_frame = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
                        face_frame = img_to_array(crop)
                        name = "None"
                        if face_frame.size != 0 :
                            face_frame = cv2.resize(face_frame, (160, 160))
                            encode = self.get_embedding(face_frame)
                            name = self.find_person(encode)
                        if name == "None":
                            label = "Not found"
                        else :
                            label = name

                        classIDs.append(classID)
                        names.append(label)

        # apply non-maximal suppression to suppress weak, overlapping bounding boxes
        idxs = cv2.dnn.NMSBoxes(boxes, confidences, self.CONFIDENCE, self.THRESHOLD)

        # ensure at least one detection exists
        if len(idxs) > 0:
            # loop over the indexes we are keeping
            for i in idxs.flatten():
                # extract the bounding box coordinates
                (x, y) = (boxes[i][0], boxes[i][1])
                (w, h) = (boxes[i][2], boxes[i][3])

                # draw a bounding box rectangle and label on the frame
                color = [int(c) for c in self.COLORS[classIDs[i]]]
                cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
                text = "{}: {:.4f}".format(self.LABELS[classIDs[i]]+":"+names[i], confidences[i])
                cv2.putText(frame, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 5)

    # use MTCNN to detect faces and return face array
    def extract_mtcnn_face(self, filename, required_size=(160, 160)):
        print("extracting face from image")
        detector = MTCNN()

        image = Image.open(filename)
        # convert to RGB, if needed
        image = image.convert('RGB')
        # convert to array
        pixels = np.asarray(image)
        # detect faces in the image
        results = detector.detect_faces(pixels)
        # extract the bounding box from the first face
        x1, y1, width, height = results[0]['box']
        # deal with negative pixel index
        x1, y1 = abs(x1), abs(y1)
        x2, y2 = x1 + width, y1 + height
        # extract the face
        face = pixels[y1:y2, x1:x2]
        # resize pixels to the model size
        image = Image.fromarray(face)
        image = image.resize(required_size)
        face_array = np.asarray(image)

        return face_array

    # encode person and save into db
    def save_encode_db(self, label, filename):
        print("encoding was begining for: " + label)

        # extract face
        imagePath = os.path.sep.join(["uploads", filename])
        face_frame = self.extract_mtcnn_face(imagePath)

        # get enbedding code
        self.database[label] = self.get_embedding(face_frame)

        # write into db
        dbPath = os.path.sep.join(["data", "dict.csv"])
        with open(dbPath, 'w', newline='') as csv_file:
            writer = csv.writer(csv_file)
            for key, value in self.database.items():
               value = list(value)
               writer.writerow([key, value])


#!/usr/bin/python
# -*- coding: utf-8 -*-
__author__ = 'ar'

import os
import numpy as np
import lmdb

from keras.utils import np_utils

from app.backend.core.datasets.dbconfig import DBImage2DConfig
from app.backend.core.datasets.dbimageinfo import DatasetImage2dInfo
from app.backend.core.datasets.imgproc2d import ImageTransformer2D

#########################
class BatcherImage2DLMDB:
    cfg         = None
    sizeBatch   = 64
    numTrain    = -1
    numVal      = -1
    numLbl      = -1
    lbl         = None
    keysTrain   = None
    keysVal     = None
    shapeImg    = None
    meanImg     = None
    meanChannel = None
    meanChImage = None
    dbTrain     = None
    dbVal       = None
    isRemoveMean= True
    scaleFactor = 1./255.
    def __init__(self, parPathDB=None, parSizeBatch=-1, scaleFactor=-1.):
        if parPathDB is None:
            #FIXME: check this point, LMDBBatcher is not initialized correctly
            return
        try:
            self.cfg     = DatasetImage2dInfo(parPathDB)
            self.cfg.loadDBInfo(isBuildSearchIndex=False)
            tpathTrainDB = self.cfg.pathDbTrain
            tpathValDB   = self.cfg.pathDbVal
            self.dbTrain = lmdb.open(tpathTrainDB, readonly=True)
            self.dbVal   = lmdb.open(tpathValDB,   readonly=True)
            with self.dbTrain.begin() as txnTrain, self.dbVal.begin() as txnVal:
                self.lbl    = self.cfg.labels
                self.numLbl = len(self.lbl)
                self.numTrain = self.dbTrain.stat()['entries']
                self.numVal   = self.dbVal.stat()['entries']
                with txnTrain.cursor() as cursTrain, txnVal.cursor() as cursVal:
                    self.keysTrain = np.array([key for key, _ in cursTrain])
                    self.keysVal   = np.array([key for key, _ in cursVal])
                    timg,_ = ImageTransformer2D.decodeLmdbItem2NNSampple(txnTrain.get(self.keysTrain[0]))
                    self.shapeImg = timg.shape
                if parSizeBatch > 1:
                    self.sizeBatch = parSizeBatch
                if scaleFactor > 0:
                    self.scaleFactor = scaleFactor
                self.loadMeanProto()
        except lmdb.Error as err:
            self.pathDataDir = None
            print 'LMDBReader.Error() : [%s]' % err
    def loadFromTrainDir(self, pathTrainDir, parImgShape=None):
        if parImgShape is not None:
            self.shapeImg = parImgShape
        self.pathDataDir = pathTrainDir
        tpathLabels = self.getPathLabels()
        #FIXME: potential bug
        if os.path.isfile(tpathLabels):
            with open(tpathLabels,'r') as f:
                self.lbl = f.read().splitlines()
        self.loadMeanProto()
    def close(self):
        if self.isOk():
            self.numTrain   = -1
            self.numVal     = -1
            self.numLbl     = -1
            self.dbTrain.close()
            self.dbVal.close()
    def getPath(self,localPath):
        if self.pathDataDir is not None:
            return os.path.join(self.pathDataDir, localPath)
        else:
            return localPath
    def getPathLabels(self):
        return self.cfg.pathLabels
    def getPathTrainDB(self):
        return self.cfg.pathDbTrain
    def getPathValDB(self):
        return self.cfg.pathDbVal
    def getPathMeanProto(self):
        return self.cfg.pathMeanData
    def loadMeanProto(self):
        self.meanImg = ImageTransformer2D.loadImageFromBinaryBlog(self.cfg.pathMeanData)*self.scaleFactor
        self.meanChannel = np.mean(self.meanImg, axis=(1, 2))
        self.meanChImage = np.zeros(self.meanImg.shape, dtype=self.meanImg.dtype)
        for ii in range(self.meanImg.shape[0]):
            self.meanChImage[ii,:,:] = self.meanChannel[ii]
    def isOk(self):
        return ((self.numTrain > 0) and (self.numVal > 0) and (self.numLbl>0))
    def getBatch(self, isTrainData=True, isShuffle=True, batchSize=-1, reshape2Shape=None):
        if batchSize<1:
            batchSize = self.sizeBatch
        if self.isOk():
            if isTrainData:
                tptrkDB  = self.dbTrain
                tptrKeys = self.keysTrain
            else:
                tptrkDB  = self.dbVal
                tptrKeys = self.keysVal
            with tptrkDB.begin() as txn:
                if isShuffle:
                    np.random.shuffle(tptrKeys)
                tbatchKeys = tptrKeys[:batchSize]
                tsizeX = [batchSize, self.shapeImg[0], self.shapeImg[1], self.shapeImg[2]]
                dataX = np.zeros(tsizeX, np.float32)  # FIXME: [1] check data type! [float32/float64]
                dataY = []
                for ii in xrange(batchSize):
                    timg,tlabel=ImageTransformer2D.decodeLmdbItem2NNSampple(txn.get(b'%s' % tbatchKeys[ii]))
                    timg = timg.astype(np.float)*self.scaleFactor
                    if self.isRemoveMean:
                        timg-=self.meanChImage
                    dataY.append(tlabel)
                    dataX[ii] = timg  # FIXME: [1] check data type! [float32/float64]
                dataY = np_utils.to_categorical(np.array(dataY), self.numLbl)
                if ((reshape2Shape is None) or (reshape2Shape[1:] == self.shapeImg)):
                    return (dataX, dataY)
                else:
                    #FIXME: reshaping can raise Exception "shape mismatch"
                    dataX=dataX.reshape([self.sizeBatch] + reshape2Shape[1:])
                    return (dataX, dataY)
        else:
            return None
    def loadAllDataTrain(self):
        if self.isOk():
            return self.getBatch(isTrainData=True, isShuffle=False, batchSize=self.numTrain)
        else:
            return None
    def loadAllDataVal(self):
        if self.isOk():
            return self.getBatch(isTrainData=True, isShuffle=False, batchSize=self.numVal)
        else:
            return None
    def getBatchTrain(self, reshape2Shape=None):
        return self.getBatch(isTrainData=True, reshape2Shape=reshape2Shape)
    def getBatchVal(self, reshape2Shape=None):
        return self.getBatch(isTrainData=False, reshape2Shape=reshape2Shape)

if __name__ == '__main__':
    pass


import os
import argparse
from numpy.lib.function_base import append
from numpy.lib.npyio import load
from tqdm import tqdm
from abc import ABCMeta, abstractmethod
from bvh import Bvh

import numpy as np
from scipy.spatial.transform import Rotation as R
import torch
import torch.nn as nn
from torch.utils.data import Dataset
from torch.utils.data import DataLoader, RandomSampler
from torch.utils.data.sampler import BatchSampler
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch_optimizer import RAdam
from torch.optim import SGD

from default_param import _C as cfg
from tools import *

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

    parser.add_argument(
        "--debiaser_type",
        default="Linear",
        choices=["Linear", "NN"])
    parser.add_argument(
        "--config_file",
        default="")
    parser.add_argument(
        "--orig_data_npz",
        type=str,
        default="../dataset/dataset_split.npz")
    parser.add_argument(
        "--phys_data_npz",
        type=str,
        default="../dataset/dataset_Fixed_phys_kick.npz")
    parser.add_argument(
        "--aug_data_npz",
        type=str,
        default="../dataset/dataset_IK_phys_kick.npz")
    parser.add_argument(
        "--joint_subset",
        type=str,
        choices=["all", "ankle"])
    parser.add_argument(
        "--rep_in",
        type=str,
        default='expmap')
    parser.add_argument(
        "--rep_out",
        type=str,
        default='expmap')
    parser.add_argument(
        "--show_diff",
        action="store_true")
    parser.add_argument(
        "--act_class",
        default="kick")
    parser.add_argument(
        "--debug",
        action="store_true")
    return parser.parse_args()

def load_npz(path):
    with np.load(path, allow_pickle=True) as dataset:
        sort = np.argsort(dataset["names"])
        motions = dataset["motions"][sort]
        names = dataset["names"][sort]
        return motions, names

def search_by_str(query, key):
    idx = []
    for q in query:
        idx.append(np.char.find(key, q) != -1)

    return np.array(idx)

def offset_motion(orig_motion, phys_motion, aug_motion, **kwargs):
    len_orig = len(orig_motion)
    len_phys = len(phys_motion)
    len_aug = len(aug_motion)
    len_min = min(len_orig, len_phys)

    diff = []
    for i, aug_pose in enumerate(aug_motion):
        idx = int(i / len_aug * len_min)
        aug_motion[i], diff_pose = offset_pose(
            orig_motion[idx], phys_motion[idx], aug_pose, **kwargs)
        diff.append(diff_pose)

    return aug_motion, np.mean(diff, axis=0)

def offset_pose(
        orig_pose,
        phys_pose,
        aug_pose,
        joint_subset="all",
        rep_in='euler',
        rep_out='euler',
        ret_diff=False):
    aug_pose[:3] += orig_pose[:3] - phys_pose[:3]
    rot_orig = pose2rot(orig_pose[3:], rep=rep_in)
    rot_phys = pose2rot(phys_pose[3:], rep=rep_in)
    rot_aug = pose2rot(aug_pose[3:], rep=rep_in)
    rot_diff = rot_orig * rot_phys.inv()

    if joint_subset == "ankle":
        joint_id = [9, 15]
        rot_diff = rot_mask(rot_diff, joint_id)

    rot_aug = rot_diff * rot_aug
    aug_pose[3:] = rot2pose(rot_aug, rep=rep_out)

    if ret_diff:
        return aug_pose, rot2pose(rot_diff, rep=rep_out)
    else:
        return aug_pose, np.full(np.size(aug_pose), np.inf)

def rot_mask(rot, ids):
    expmap = rot.as_rotvec()
    expmap_mask = [v if i in ids else [0, 0, 0] for i, v in enumerate(expmap)]
    return R.from_rotvec(expmap_mask)

def pose2rot(pose, rep='euler'):
    pose = np.reshape(pose, [-1, 3])
    if rep == 'euler':
        rot = R.from_euler('ZYX', pose, degrees=True)
    elif rep == 'expmap':
        rot = R.from_rotvec(pose)
    else:
        raise ValueError

    return rot

def rot2pose(rot, rep='euler'):
    if rep == 'euler':
        pose = rot.as_euler('ZYX', degrees=True)
    elif rep == 'expmap':
        pose = rot.as_rotvec()
    else:
        raise ValueError

    return pose.flatten()

def select_debiaser(debiaser_type):
    if debiaser_type == "Linear":
        return Linear_Debiaser
    if debiaser_type == "NN":
        return NN_Debiaser

class Debiaser(metaclass=ABCMeta):
    @abstractmethod
    def __init__(self, orig_data_npz, phys_data_npz, **kwargs):
        pass

    @abstractmethod
    def type(self):
        return ""

    @abstractmethod
    def process(self, aug_data_npz, **kwargs):
        pass

class Linear_Debiaser(Debiaser):
    def __init__(self, orig_data_npz, phys_data_npz, **kwargs):
        self.orig_motions, self.orig_names = load_npz(orig_data_npz)
        self.phys_motions, self.phys_names = load_npz(phys_data_npz)

        idx_matched_motions = search_by_str(self.orig_names, self.phys_names)
        self.orig_motions = self.orig_motions[np.any(
            idx_matched_motions, axis=1)]
        self.orig_names = self.orig_names[np.any(idx_matched_motions, axis=1)]

        self.kwargs = kwargs

    def type(self):
        return "Linear"

    def process(self, aug_data_npz):
        aug_motions, aug_names = load_npz(aug_data_npz)

        idx_matched_motions = search_by_str(self.orig_names, aug_names)
        diffs = []
        for idxs, orig_motion, phys_motion in tqdm(zip(
                idx_matched_motions, self.orig_motions, self.phys_motions)):
            for i, aug_motion in enumerate(aug_motions[idxs]):
                # assert orig_names[j] in aug_names[idxs][i]
                aug_motions[idxs][i], diff = offset_motion(
                    orig_motion, phys_motion, aug_motion, **self.kwargs)
                diffs.append(diff)

        self.save_npz(aug_data_npz, aug_motions)

        return diffs

    def save_npz(self, aug_data_npz, aug_motions):
        with np.load(aug_data_npz, allow_pickle=True) as dataset:
            np.savez(self.get_outpath(aug_data_npz),
                     motions=aug_motions,
                     names=dataset["names"],
                     subjects=dataset["subjects"],
                     classes=dataset["classes"],
                     max_len=dataset["max_len"],
                     min_len=dataset["min_len"],
                     dim_pose=dataset["dim_pose"],
                     frame_times=dataset["frame_times"],
                     omit_root_pos=dataset["omit_root_pos"],
                     rep=dataset["rep"],
                     abs_angle=dataset["abs_angle"],
                     allow_pickle=True)

    def get_outpath(self, aug_data_npz):
        split_name = os.path.splitext(aug_data_npz)
        save_path = f"{split_name[0]}_offset_{self.type()}{split_name[1]}"
        return save_path

class NN_Debiaser(Linear_Debiaser):
    def __init__(self, orig_data_npz, phys_data_npz, **kwargs):
        super().__init__(orig_data_npz, phys_data_npz, **kwargs)

        self.cfg = cfg

        load_config(self.cfg, kwargs['config_file'], kwargs['act_class'], "phys" if "phys" in phys_data_npz else "kin")
        angle_dim = 3 if self.cfg.INPUT.REP == "expmap" else 3
        input_size = (3 + self.cfg.INPUT.NUM_JOINT * angle_dim) * len(self.cfg.INPUT.TYPE)
        output_size = 3 + 6 * self.cfg.INPUT.NUM_JOINT
        self.model = ModelWithLoss(self.cfg, input_size, output_size).cuda()
        path = os.path.join(self.cfg.OUTPUT_DIR, "model.pth")
        if os.path.exists(path):
            self.model.load_state_dict(torch.load(path))
        else:
            self.train()

    def train(self):
        dataloader = self.build_dataloader()
        if cfg.SOLVER.OPTIMIZER == "radam":
            optimizer = RAdam(
                filter(
                    lambda p: p.requires_grad,
                    self.model.parameters()),
                lr=cfg.SOLVER.LR,
                weight_decay=cfg.MODEL.REG_WEIGHT)
        elif cfg.SOLVER.OPTIMIZER == "sgd":
            optimizer = SGD(
                filter(
                    lambda p: p.requires_grad,
                    self.model.parameters()),
                lr=cfg.SOLVER.LR,
                momentum=0.9)
        scheduler = ReduceLROnPlateau(optimizer,
                                      factor=cfg.SOLVER.DECAY_FACTOR,
                                      patience=10)

        with tqdm(range(cfg.SOLVER.ITER), ncols=100) as pbar:
            for i in pbar:
                self.model.train()
                logging_loss = 0
                for data_dict in dataloader:
                    optimizer.zero_grad()
                    loss = self.model(data_dict)
                    loss = loss.mean()

                    loss.backward()

                    optimizer.step()
                    logging_loss += loss.item()
                scheduler.step(logging_loss)
                logging_loss /= len(dataloader)

                np.set_printoptions(
                    precision=4, floatmode='maxprec', suppress=True)

                pbar.set_postfix(
                    dict(
                        loss=np.asarray(
                            [logging_loss]),
                        lr=[optimizer.param_groups[0]['lr']]
                    ))

                if optimizer.param_groups[0]['lr'] < cfg.SOLVER.LR * 1e-3:
                    print("early stopping")
                    break

        torch.save(
            self.model.state_dict(),
            os.path.join(
                self.cfg.OUTPUT_DIR,
                "model.pth"))

    def build_dataset(self):
        x, y = [], []
        for i in range(len(self.orig_motions)):
            orig_motion = self.orig_motions[i]
            phys_motion = self.phys_motions[i]
            length = min(len(orig_motion), len(phys_motion))
            phys_motion = phys_motion[:length]
            orig_motion = orig_motion[:length]

            phys_motion = self.preprocess(phys_motion)

            x.append(phys_motion)
            y.append(orig_motion)
        x = np.concatenate(x, axis=0)
        y = np.concatenate(y, axis=0)
        assert len(x) == len(y)
        return SimpleDataset(x=x, y=y)

    def preprocess(self, motion):
        if "vel" in self.cfg.INPUT.TYPE:
            vel = motion[1:, :] - motion[:-1, :]
            try:
                vel = np.pad(vel, [[0,1], [0,0]], mode='edge')
                motion = np.concatenate([motion, vel], axis=1)
            except ValueError:
                motion = np.concatenate([motion, np.zeros_like(motion)], axis=1)
        return motion

    def build_dataloader(self):
        dataset = self.build_dataset()
        sampler = sampler = RandomSampler(dataset)
        batch_sampler = BatchSampler(sampler=sampler,
                                     batch_size=cfg.SOLVER.BATCH_SIZE,
                                     drop_last=False)
        loader = DataLoader(
            dataset,
            num_workers=cfg.SOLVER.NUM_WORKERS,
            batch_sampler=batch_sampler)
        return loader

    def type(self):
        return "NN"

    def process(self, aug_data_npz, **kwargs):
        self.model.eval()
        aug_motions, aug_names = load_npz(aug_data_npz)

        diffs = []
        with torch.no_grad():
            for i in tqdm(range(len(aug_motions))):
                aug_motions[i], diff = self.model.predict(
                    {'input': torch.Tensor(self.preprocess(aug_motions[i]))})
                diffs.append(diff)
        self.save_npz(aug_data_npz, aug_motions)

        return np.concatenate(diffs, axis=0)

class ModelWithLoss(nn.Module):
    def __init__(self, cfg, input_size, output_size):
        super(ModelWithLoss, self).__init__()

        if cfg.MODEL.TYPE == "mlp":
            self.model = MLP(
                cfg,
                input_size=input_size,
                output_size=output_size)

        self.build_loss(cfg)

    def forward(self, data_dict):
        ret = self.pred(data_dict)
        loss = self.loss(ret)

        return loss

    def pred(self, data_dict):
        x = data_dict["input"].float().cuda()
        pred = self.model(x)

        pred_pos = x[:, :3] + pred[:, :3]
        res = compute_rotation_matrix_from_Rodriguez(x[:, 3:3+16*3].reshape(-1, 3))
        pred_diff = compute_rotation_matrix_from_ortho6d(
            pred[:, 3:].reshape(-1, 6))
        pred_rot = torch.matmul(res, pred_diff)
        data_dict["pred_pos"] = pred_pos
        data_dict["pred_rot"] = pred_rot
        data_dict["pred_diff"] = pred_diff

        return data_dict

    def predict(self, data_dict):
        shape = data_dict["input"].shape
        ret_dict = self.pred(data_dict)
        pred_pos = ret_dict['pred_pos'].cpu().numpy()
        pred_rot = R.from_matrix(ret_dict['pred_rot'].cpu(
        ).numpy()).as_rotvec().reshape([shape[0], -1])
        pred_diff = R.from_matrix(ret_dict['pred_diff'].cpu(
        ).numpy()).as_rotvec().reshape([shape[0], -1])
        return np.concatenate([pred_pos, pred_rot], axis=1), pred_diff

    def loss(self, data_dict):
        loss = 0
        gt_rot = compute_rotation_matrix_from_Rodriguez(
            data_dict["gt"][:, 3:].float().cuda().reshape(-1, 3))
        loss += self.pred_loss(data_dict["pred_rot"].squeeze(),
                               gt_rot)
        loss += self.pred_loss(data_dict["pred_pos"].squeeze(),
                               data_dict["gt"][:, :3].float().cuda())

        for name, param in self.model.named_parameters():
            if 'weight' in name:
                loss += torch.norm(param, self.reg_norm) * self.reg_weight

        return loss

    def build_loss(self, cfg):
        self.pred_loss = choose_loss(cfg.MODEL.PRED_LOSS)
        self.reg_norm = 1 if cfg.MODEL.REG_TYPE == "l1" else 2
        self.reg_weight = cfg.MODEL.REG_WEIGHT

class MLP(nn.Module):
    def __init__(self, cfg, input_size, output_size):
        super(MLP, self).__init__()

        dim_h = cfg.MODEL.DIM_HIDDEN
        num_l = cfg.MODEL.NUM_HIDDEN
        dim_o = output_size
        coef = cfg.MODEL.DIM_HIDDEN_FACT

        self.h_size = [input_size]
        for _ in range(num_l):
            self.h_size.append(dim_h)
            dim_h = int(dim_h * coef)
        self.h_size.append(dim_o)

        self.layers = nn.ModuleList([nn.Linear(
            self.h_size[k], self.h_size[k + 1]) for k in range(len(self.h_size) - 1)])
        self.norm = nn.BatchNorm1d(self.h_size[0])
        self.drop = nn.Dropout(p=cfg.MODEL.P_DROPOUT)
        self.act = nn.LeakyReLU(0.01)

    def forward(self, x):
        x = self.drop(x)

        for i in range(0, len(self.layers) - 1):
            x = self.act(self.layers[i](x))

        x = self.layers[-1](x)
        return x

class SimpleDataset(Dataset):
    def __init__(self, x, y=[]):

        self.x = x
        self.y = y
        self.test = False
        if len(y) == 0:
            self.test = True

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

    def __getitem__(self, idx):
        x = self.x[idx]
        y = self.y[idx] if not self.test else -1.0
        return {"input": x, "gt": y}

def choose_loss(loss_name):
    if loss_name == "l1":
        return nn.L1Loss()
    elif loss_name == "mse":
        return nn.MSELoss()
    elif loss_name == "bce":
        return nn.BCELoss()
    else:
        raise ValueError(f"unknown loss function {loss_name}")

def load_config(cfg, conf_path, act_class, phys_or_kin):
    if os.path.isfile(conf_path):
        print(f"Configuration file loaded from {conf_path}.")
        cfg.merge_from_file(conf_path)
        cfg.OUTPUT_DIR = os.path.join(
            cfg.OUTPUT_DIR, os.path.splitext(
                os.path.basename(conf_path))[0], act_class, phys_or_kin)
    else:
        print("Use default configuration.")
        cfg.OUTPUT_DIR = os.path.join(cfg.OUTPUT_DIR, "default", act_class, phys_or_kin)

    cfg.freeze()

    print(f"output dirname: {cfg.OUTPUT_DIR}")
    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
    if os.path.isfile(conf_path):
        import shutil
        shutil.copy2(conf_path, os.path.join(cfg.OUTPUT_DIR, 'config.yaml'))

def visualize_diff(diffs, type):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots()

    diffs = np.array(diffs)
    shape = diffs.shape
    diffs = np.mean(np.reshape(diffs, [-1, int(shape[1] / 3), 3]), axis=2)
    bp = ax.boxplot(diffs)

    plt.title('offset diff on each joint')
    plt.grid()
    plt.savefig(f"diff_{type}.png")

def main():
    args = parse_args()

    kwargs = {'joint_subset': args.joint_subset,
              'rep_in': args.rep_in,
              'rep_out': args.rep_out,
              'ret_diff': args.show_diff}

    if args.debiaser_type != "Linear":
        kwargs['config_file'] = args.config_file
        kwargs['act_class'] = args.act_class

    debiaser = select_debiaser(
        args.debiaser_type)(
        args.orig_data_npz,
        args.phys_data_npz,
        **kwargs)
    diffs = debiaser.process(args.aug_data_npz)
    if args.show_diff:
        visualize_diff(diffs, args.debiaser_type)

    if args.debug and args.rep_out == "expmap":
        data_path = '../../data/bvh/hdm05/'
        bvhpath = os.path.join(
            data_path, 'HDM_bd_cartwheelLHandStart1Reps_001_120.bvh')

        outpath = "debug"
        if os.path.exists(outpath):
            import shutil
            shutil.rmtree(outpath)
        os.makedirs(outpath, exist_ok=True)

        dataset = np.load(debiaser.get_outpath(args.aug_data_npz), allow_pickle=True)
        #dataset = np.load('../dataset/dataset_Fixed_phys_kick.npz', allow_pickle=True)
        motions = dataset['motions']
        names = dataset['names']
        frame_times = dataset['frame_times']

        from joblib import Parallel, delayed
        Parallel(n_jobs=12)(delayed(saveBvh)(motion, name, frame_time, outpath) for motion, name, frame_time in tqdm(zip(motions, names, frame_times)))

def loadBvh(bvhpath):
    with open(bvhpath) as f:
        return Bvh(f.read())

def repTransfer_reverse(rep, rot_joint):
    if rep == 'quat':
        return [R.from_quat(quat) for quat in rot_joint]
    elif rep == 'expmap':
        return [R.from_rotvec(rotvec) for rotvec in rot_joint]
    elif rep == 'ortho6d':
        return [ortho6dToR(ortho6d) for ortho6d in rot_joint]
    elif rep == 'euler':
        rot_joint = [r.from_euler('ZYX', degrees=True) for r in rot_joint]
    else:
        print("unknown 3D rotation representation")

def ortho6dToR(ortho6d):
    assert len(ortho6d) == 6

    x_raw = ortho6d[0:3]
    y_raw = ortho6d[3:6]

    x = x_raw / np.linalg.norm(x_raw)
    z = np.cross(x, y_raw)
    z = z / np.linalg.norm(z)
    y = np.cross(z, x)

    return R.from_matrix([x, y, z])

def rot_dict_to_frames(mocap, rot_dict):
    frames = []
    for bvhjoint in mocap.get_joints_names():
        euler_angle = [rotationToBvhEuler(r) for r in rot_dict[bvhjoint]]
        frames.append(euler_angle)
    return frames

def rotationToBvhEuler(rotation):
    euler = rotation.as_euler('ZYX', degrees=True)  # extrinsic euler
    return [euler[0], euler[1], euler[2]]

def add_bvh_joint(used_joints, mocap, rot_dict):

    for bvhjoint in mocap.get_joints_names():
        if bvhjoint in used_joints:
            pass
        else:
            parent_joint = mocap.joint_parent(bvhjoint).name
            rot_dict[bvhjoint] = [rp for rp in rot_dict[parent_joint]]

    return rot_dict

def saveBvh(motion, name, frame_time, outpath):
    data_path = '../../data/bvh/hdm05/'
    bvhpath = os.path.join(
        data_path, 'HDM_bd_cartwheelLHandStart1Reps_001_120.bvh')

    position = motion[:, :3]
    motion = np.reshape(motion[:, 3:], [-1, 16, 3])

    mocap = loadBvh(bvhpath)

    used_joints = ["Hips",
                    "Spine", "Spine1", "Neck",
                    "RightUpLeg", "RightLeg", "RightFoot",
                    "RightArm", "RightForeArm", "RightHand",
                    "LeftUpLeg", "LeftLeg", "LeftFoot",
                    "LeftArm", "LeftForeArm", "LeftHand"]

    rot_dict = {}
    for i, joint_name in enumerate(used_joints):
        rot_dict[joint_name] = repTransfer_reverse("expmap", motion[:, i])

    rot_dict = add_bvh_joint(used_joints, mocap, rot_dict)
    frames = rot_dict_to_frames(mocap, rot_dict)
    frameNum = len(motion)
    frames = np.concatenate(frames, axis=1)
    frames = np.reshape(frames, [frameNum, -1])
    frames = np.concatenate([np.squeeze(position), frames], axis=1)
    writeBvh(
        bvhpath,
        os.path.join(
            outpath,
            name + '_offset.bvh'),
        frames,
        frameTime=frame_time,
        frameNum=frameNum)

def writeBvh(bvhpath, outpath, frames, frameTime=None, frameNum=None):
    print("saving bvh file to {}".format(outpath))
    with open(outpath, 'w') as fw:
        with open(bvhpath, 'r') as fr:
            while True:
                line = fr.readline()

                if 'Frame Time' in line:
                    if frameTime is None:
                        pass
                    else:
                        line = "Frame Time: {}\n".format(frameTime)
                    fw.write(line)
                    break
                elif 'Frames' in line:
                    if frameNum is None:
                        frameNum = len(frames)
                    else:
                        pass
                    line = "Frames: {}\n".format(frameNum)
                    fw.write(line)
                else:
                    fw.write(line)
            for vector in frames:
                line_aligned = ' '.join(vector.astype(str)) + '\n'
                fw.write(line_aligned)

if __name__ == '__main__':
    main()