from __future__ import print_function

import numpy as np
import pdb
from multiprocessing import Pool
import functools

import numpy as np
import os
import sys
import argparse
import time
import math
from PIL import Image

import torchvision.utils as vutils
import tensorboard_logger as tb_logger
import torch
import torch.backends.cudnn as cudnn
from torchvision import transforms, datasets

from util import TwoCropTransform, AverageMeter, GansetDataset, GansteerDataset, OnlineDataset
from util import adjust_learning_rate, warmup_learning_rate, accuracy
from util import set_optimizer, save_model
from networks.resnet_big import SupConResNet, SupCEResNet
from losses import SupConLoss
import oyaml as yaml
import pbar as pbar

try:
    import apex
    from apex import amp, optimizers
except ImportError:
    pass

def getitem(idx, root_dir, transform, class_to_idx, numcontrast=0):
    if torch.is_tensor(idx):
        idx = idx.tolist()

    img_name = '{}/{}_anchor'.format(root_dir, idx)

    if numcontrast > 0:
        img_name_neighbor = img_name.replace('anchor','neighbor')
    else:
        img_name_neighbor = img_name

    while not os.path.isfile(img_name) or not os.path.isfile(img_name_neighbor):
        print("Image {} missing ".format(img_name))
        time.sleep(2)

    image = Image.open(img_name)
    image_neighbor = Image.open(img_name_neighbor)
    label = img_name.split('/')[-2]
    label = class_to_idx[label]
    if transform:
        image = transform(image)
        image_neighbor = transform(image_neighbor)

    return image, image_neighbor, label

=======
>>>>>>> 0c8c9108f080fd28b5bcef9c7e0c39e06fedb243
def parse_option():
    parser = argparse.ArgumentParser('argument for training')
    parser.add_argument('--encoding_type', type=str, default='contrastive',
                        choices=['contrastive', 'crossentropy', 'autoencoding'])
    parser.add_argument('--print_freq', type=int, default=10,
                        help='print frequency')

    parser.add_argument('--numiter', type=int, default=130000,
                        help='save frequency')
    parser.add_argument('--save_freq', type=int, default=20,
                        help='save frequency')
    parser.add_argument('--batch_size', type=int, default=256,
                        help='batch_size')
    parser.add_argument('--num_workers', type=int, default=16,
                        help='num of workers to use')
    parser.add_argument('--epochs', type=int, default=200,
                        help='number of training epochs')
    parser.add_argument('--showimg', action='store_true', help='display image in tensorboard')
    parser.add_argument('--resume', default='', type=str, help='whether to resume training')

    # optimization
    parser.add_argument('--learning_rate', type=float, default=0.03,
                        help='learning rate')
    parser.add_argument('--lr_decay_epochs', type=str, default='120,160',
                        help='where to decay lr, can be a list')
    parser.add_argument('--lr_decay_rate', type=float, default=0.1,
                        help='decay rate for learning rate')
    parser.add_argument('--weight_decay', type=float, default=1e-4,
                        help='weight decay')
    parser.add_argument('--momentum', type=float, default=0.9,
                        help='momentum')

    # model dataset
    parser.add_argument('--model', type=str, default='resnet50')
    parser.add_argument('--dataset', type=str, default='biggan',
                        choices=['biggan', 'cifar10', 'cifar100', 'imagenet100', 'imagenet100K', 'imagenet'], help='dataset')

    ## Ali: todo: this should be based on opt.encoding type and remove the default (revisit every default) and name of the model for saving
    # method
    parser.add_argument('--numcontrast', type=int, default=20,
                        help='num of workers to use')
    parser.add_argument('--method', type=str, default='SimCLR',
                        choices=['SupCon', 'SimCLR'], help='choose method')
    parser.add_argument('--walk_method', type=str, help='choose method')

    # temperature
    parser.add_argument('--temp', type=float, default=0.1,
                        help='temperature for loss function')

    # other setting
    parser.add_argument('--cosine', action='store_true', help='using cosine annealing')
    parser.add_argument('--syncBN', action='store_true',
                        help='using synchronized batch normalization')
    parser.add_argument('--warm', action='store_true',
                        help='warm-up for large batch training')
    parser.add_argument('--trial', type=str, default='0',
                        help='id for recording multiple runs')

    # specifying folders
    parser.add_argument('-d', '--data_folder', type=str,
                        default='/data/scratch-oc40/jahanian/ganclr_results/ImageNet100',
                        help='the data folder')
    parser.add_argument('-s', '--cache_folder', type=str,
                        default='/data/scratch-oc40/jahanian/ganclr_results/',
                        help='the saving folder')

    opt = parser.parse_args()

    # set the path according to the environment
    opt.data_folder = opt.data_folder
    if opt.encoding_type == 'crossentropy':
        opt.method = 'SupCE'
        opt.model_path = os.path.join(opt.cache_folder, 'SupCE/{}_models'.format(opt.dataset))
        opt.tb_path = os.path.join(opt.cache_folder, 'SupCE/{}_tensorboard'.format(opt.dataset))
    else:
        opt.model_path = os.path.join(opt.cache_folder, 'SupCon/{}_models'.format(opt.dataset))
        opt.tb_path = os.path.join(opt.cache_folder, 'SupCon/{}_tensorboard'.format(opt.dataset))

    iterations = opt.lr_decay_epochs.split(',')
    opt.lr_decay_epochs = list([])
    for it in iterations:
        opt.lr_decay_epochs.append(int(it))

    opt.model_name = '{}_{}online_{}_{}_ncontrast.{}_lr_{}_decay_{}_bsz_{}_temp_{}_trial_{}'.\
            format(opt.method, opt.dataset, opt.walk_method, opt.model, opt.numcontrast, opt.learning_rate,
            opt.weight_decay, opt.batch_size, opt.temp, opt.trial)

    if opt.cosine:
        opt.model_name = '{}_cosine'.format(opt.model_name)

    opt.model_name = '{}_{}'.format(opt.model_name, os.path.basename(opt.data_folder))
    # warm-up for large-batch training,
    if opt.batch_size > 256:
        opt.warm = True
    if opt.warm:
        opt.model_name = '{}_warm'.format(opt.model_name)
        opt.warmup_from = 0.01
        opt.warm_epochs = 10
        if opt.cosine:
            eta_min = opt.learning_rate * (opt.lr_decay_rate ** 3)
            opt.warmup_to = eta_min + (opt.learning_rate - eta_min) * (
                    1 + math.cos(math.pi * opt.warm_epochs / opt.epochs)) / 2
        else:
            opt.warmup_to = opt.learning_rate

    opt.tb_folder = os.path.join(opt.tb_path, opt.model_name)
    if not os.path.isdir(opt.tb_folder):
        os.makedirs(opt.tb_folder)

    opt.save_folder = os.path.join(opt.model_path, opt.model_name)
    if not os.path.isdir(opt.save_folder):
        os.makedirs(opt.save_folder)

    if opt.dataset == 'biggan' or opt.dataset == 'imagenet100' or opt.dataset == 'imagenet100K' or opt.dataset == 'imagenet':
        # or 256 as you like
        opt.img_size = 256
        opt.n_cls = 1000
    elif opt.dataset == 'cifar10' or opt.dataset == 'cifar100':
        opt.img_size = 32

    return opt

def set_loader(opt):
    # construct data loader
    if opt.dataset == 'cifar10':
        mean = (0.4914, 0.4822, 0.4465)
        std = (0.2023, 0.1994, 0.2010)
    elif opt.dataset == 'cifar100':
        mean = (0.5071, 0.4867, 0.4408)
        std = (0.2675, 0.2565, 0.2761)
    elif opt.dataset == 'biggan' or opt.dataset == 'imagenet100' or opt.dataset == 'imagenet100K' or opt.dataset == 'imagenet':
        mean = (0.485, 0.456, 0.406)
        std = (0.229, 0.224, 0.225)
    else:
        raise ValueError('dataset not supported: {}'.format(opt.dataset))
    normalize = transforms.Normalize(mean=mean, std=std)
    opt.mean = mean
    opt.std = std

    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(size=int(opt.img_size*0.875), scale=(0.2, 1.)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomApply([
            transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)
        ], p=0.8),
        transforms.RandomGrayscale(p=0.2),
        transforms.ToTensor(),
        normalize,
    ])

    train_dataset = OnlineDataset(root_dir=os.path.join(opt.data_folder, 'train'), transform=train_transform, numcontrast=opt.numcontrast)
    return train_dataset

def set_model(opt):
    if opt.encoding_type == 'contrastive':
        model = SupConResNet(name=opt.model, img_size=int(opt.img_size*0.875))
        criterion = SupConLoss(temperature=opt.temp)

    elif opt.encoding_type == 'crossentropy':
        model = SupCEResNet(name=opt.model, num_classes=opt.n_cls, img_size=int(opt.img_size*0.875))
        criterion = torch.nn.CrossEntropyLoss()

    elif opt.encoding_type == 'autoencoding':
        print("TODO(ali): Implement here")
        raise NotImplementedError

    # enable synchronized Batch Normalization
    if opt.syncBN:
        model = apex.parallel.convert_syncbn_model(model)

    if torch.cuda.is_available():
        if torch.cuda.device_count() > 1:
            model.encoder = torch.nn.DataParallel(model.encoder)
        model = model.cuda()
        criterion = criterion.cuda()
        cudnn.benchmark = True

    return model, criterion

def train(train_loader, model, criterion, optimizer, epoch, opt, pool):
    """one epoch training"""
    model.train()

    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()

    top1 = AverageMeter()
    end = time.time()

    ## Ali: Todo: this data loading depends on if we generate positive on the fly or load them. if loading then we have data[0] and data[1]
    ## so we need to check for opt.encoding_type (currently it's good for SupCon with gan data, i.e., positive pairs are from gan)

    # for idx, (images, labels) in enumerate(train_loader):
    #     data_time.update(time.time() - end)
    print("Start train")
    index_total = (epoch-1) * opt.numiter
    index_batch = 0

    func_pool = functools.partial(getitem,
                root_dir=train_loader.root_dir,
                transform=train_loader.transform,
                class_to_idx=train_loader.class_to_idx)

    while index_batch < opt.numiter:
        im_indices = [i+index_total for i in range(opt.batch_size)]
        data = pool.map(func_pool, im_indices)
        index_batch += opt.batch_size
        index_total += opt.batch_size
        pdb.set_trace()
=======
    start_ep = 0
    # if loading from ckpt: Ali todo
    #     start_ep = checkpoint['epoch'] + 1

    epoch_batches = 1600 // opt.batch_size
    for epoch, epoch_loader in enumerate(pbar(
        epoch_grouper(train_loader, epoch_batches),
        total=(opt.niter-start_ep)), start_ep):

        # stopping condition
        if epoch > opt.niter:
            break

        # run a train epoch of epoch_batches batches
        for step, (z_batch,) in enumerate(pbar(
            epoch_loader, total=epoch_batches), 1):

        # now get z_batch as if it's your data

    for idx, data in enumerate(train_loader):
>>>>>>> 0c8c9108f080fd28b5bcef9c7e0c39e06fedb243

        if len(data) == 2:
            images = data[0]
            labels = data[1]
        elif len(data) == 3:
            images = data[:2]
            labels = data[2]
        else:
            raise NotImplementedError

        data_time.update(time.time() - end)
        if opt.encoding_type != 'contrastive':
            # We only pick one of images
            images = images[1]
        else:
            ims = images[0]
            anchors = images[1]
            images = torch.cat([images[0].unsqueeze(1), images[1].unsqueeze(1)],
                               dim=1)
            # print('2) images shape', images.shape)

            images = images.view(-1, 3, int(opt.img_size*0.875), int(opt.img_size*0.875)).cuda(non_blocking=True)
            # print('3) images shape', images.shape)

        labels = labels.cuda(non_blocking=True)
        bsz = labels.shape[0]
        # warm-up learning rate
        warmup_learning_rate(opt, epoch, idx, len(train_loader), optimizer)
        # compute loss

        if opt.encoding_type == 'contrastive':
            features = model(images)
            features = features.view(bsz, 2, -1)
            if opt.method == 'SupCon':
                loss = criterion(features, labels)
            elif opt.method == 'SimCLR':
                loss = criterion(features)
            else:
                raise ValueError('contrastive method not supported: {}'.
                                 format(opt.method))
        elif opt.encoding_type == 'crossentropy':
            output = model(images)
            loss = criterion(output, labels)
            acc1, acc5 = accuracy(output, labels, topk=(1, 5))
            top1.update(acc1[0], bsz)
        else:
            raise NotImplementedError

        # update metric
        losses.update(loss.item(), bsz)

        # SGD
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        # print info
        if (idx + 1) % opt.print_freq == 0:
            if opt.encoding_type == 'crossentropy':
                print('Train: [{0}][{1}/{2}]\t'
                      'BT {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                      'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
                      'loss {loss.val:.3f} ({loss.avg:.3f})\t'
                      'Acc@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
                       epoch, idx + 1, len(train_loader), batch_time=batch_time,
                       data_time=data_time, loss=losses, top1=top1))
            else:
                print('Train: [{0}][{1}/{2}]\t'
                      'BT {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                      'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
                      'loss {loss.val:.3f} ({loss.avg:.3f})'.format(
                       epoch, idx + 1, len(train_loader), batch_time=batch_time,
                       data_time=data_time, loss=losses))
            sys.stdout.flush()
    other_metrics = {}

    if opt.encoding_type == 'crossentropy':
        other_metrics['top1_acc'] = top1.avg
    else:
        if opt.showimg:
            other_metrics['image'] = [ims[:8], anchors[:8]]

    return losses.avg, other_metrics

def training_loader(truncation, batch_size, global_seed=0):
    '''
    Returns an infinite generator that runs through randomized z
    batches, forever.
    '''
    g_epoch = 1
    while True:
        z_data = biggan_networks.truncated_noise_dataset(truncation=truncation,
                                                         batch_size=10000,
                                                         seed=g_epoch + global_seed)
        dataloader = torch.utils.data.DataLoader(
                z_data,
                shuffle=False,
                batch_size=batch_size,
                num_workers=10,
                pin_memory=True)
        for batch in dataloader:
            yield batch
        g_epoch += 1

def epoch_grouper(loader, epoch_size, num_epochs=None):
    '''
    To use with the infinite training loader: groups the training data
    batches into epochs of the given size.
    '''
    it = iter(loader)
    epoch = 0
    while True:
        chunk_it = itertools.islice(it, epoch_size)
        try:
            first_el = next(chunk_it)
        except StopIteration:
            return
        yield itertools.chain((first_el,), chunk_it)
        epoch += 1
        if num_epochs is not None and epoch >= num_epochs:
            return
>>>>>>> 0c8c9108f080fd28b5bcef9c7e0c39e06fedb243

def main():
    opt = parse_option()

    with open(os.path.join(opt.save_folder, 'optE.yml'), 'w') as f:
        yaml.dump(vars(opt), f, default_flow_style=False)

    # build data loader
    # opt.encoding_type tells us how to get training data
    train_loader = set_loader(opt)

    # build model and criterion
    # opt.encoding_type tells us what to put as the head; choices are:
    # contrastive -> mlp or linear
    # crossentropy -> one linear for pred_y
    # autoencoding -> one linear for pred_z and one linear for pred_y
    model, criterion = set_model(opt)

    # build optimizer
    optimizer = set_optimizer(opt, model)

    # tensorboard
    logger = tb_logger.Logger(logdir=opt.tb_folder, flush_secs=2)

    # training routine
    init_epoch = 1
    if len(opt.resume) > 0:
        model_ckp = torch.load(opt.resume)
        init_epoch = model_ckp['epoch'] + 1
        model.load_state_dict(model_ckp['model'])
        optimizer.load_state_dict(model_ckp['optimizer'])
    pool = Pool(processes=opt.num_workers)

    for epoch in range(init_epoch, opt.epochs + 1):
        adjust_learning_rate(opt, optimizer, epoch)

        # train for one epoch
        time1 = time.time()
        loss, other_metrics = train(train_loader, model, criterion, optimizer, epoch, opt, pool)
        time2 = time.time()
        print('epoch {}, total time {:.2f}'.format(epoch, time2 - time1))

        # tensorboard logger
        logger.log_value('loss', loss, epoch)
        logger.log_value('learning_rate', optimizer.param_groups[0]['lr'], epoch)
        for metric_name, metric_value in other_metrics.items():
            if metric_name == 'image':
                images = metric_value
                anchors = images[0]
                otherims = images[1]
                bs = anchors.shape[0]
                grid_images = vutils.make_grid(
                        torch.cat((anchors, otherims)), nrow=bs)
                grid_images *= np.array(opt.std)[:, None, None]
                grid_images += np.array(opt.mean)[:, None, None]
                grid_images = (255*grid_images.cpu().numpy()).astype(np.uint8)
                grid_images = grid_images[None, :].transpose(0,2,3,1)
                logger.log_images(metric_name, grid_images, epoch)
            else:
                logger.log_value(metric_name, metric_value, epoch)

        if epoch % opt.save_freq == 0:
            save_file = os.path.join(
                opt.save_folder, 'ckpt_epoch_{epoch}.pth'.format(epoch=epoch))
            save_model(model, optimizer, opt, epoch, save_file)

    # save the last model
    save_file = os.path.join(
        opt.save_folder, 'last.pth')
    save_model(model, optimizer, opt, opt.epochs, save_file)

if __name__ == '__main__':
    main()


"""
This file is part of the repo: https://github.com/tencent-ailab/hifi3dface

If you find the code useful, please cite our paper:

"High-Fidelity 3D Digital Human Creation from RGB-D Selfies."
Xiangkai Lin*, Yajing Chen*, Linchao Bao*, Haoxian Zhang, Sheng Wang, Xuefei Zhe, Xinwei Jiang, Jue Wang, Dong Yu, and Zhengyou Zhang.
arXiv: https://arxiv.org/abs/2010.05562

Copyright (c) [2020] [Tencent AI Lab]

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

# from __future__ import absolute_import
# from __future__ import division
# from __future__ import print_function

import cv2, os, importlib, math
import os.path as osp
import numpy as np
import scipy.io as scio
import tensorflow as tf

def create_mtcnn_pb(sess):

    pnet_fun = lambda img: sess.run(
        ("pnet/conv4-2/BiasAdd:0", "pnet/prob1:0"), feed_dict={"pnet/input:0": img}
    )
    rnet_fun = lambda img: sess.run(
        ("rnet/conv5-2/conv5-2:0", "rnet/prob1:0"), feed_dict={"rnet/input:0": img}
    )
    onet_fun = lambda img: sess.run(
        ("onet/conv6-2/conv6-2:0", "onet/conv6-3/conv6-3:0", "onet/prob1:0"),
        feed_dict={"onet/input:0": img},
    )

    return pnet_fun, rnet_fun, onet_fun

def detect_face(img, minsize, pnet, rnet, onet, threshold, factor):
    """Detects faces in an image, and returns bounding boxes and points for them.
    img: input image
    minsize: minimum faces' size
    pnet, rnet, onet: caffemodel
    threshold: threshold=[th1, th2, th3], th1-3 are three steps's threshold
    factor: the factor used to create a scaling pyramid of face sizes to detect in the image.
    """
    factor_count = 0
    total_boxes = np.empty((0, 9))
    points = np.empty(0)
    h = img.shape[0]
    w = img.shape[1]
    minl = np.amin([h, w])
    m = 12.0 / minsize
    minl = minl * m
    # create scale pyramid
    scales = []
    while minl >= 12:
        scales += [m * np.power(factor, factor_count)]
        minl = minl * factor
        factor_count += 1

    # first stage
    for scale in scales:
        hs = int(np.ceil(h * scale))
        ws = int(np.ceil(w * scale))
        im_data = imresample(img, (hs, ws))
        im_data = (im_data - 127.5) * 0.0078125
        img_x = np.expand_dims(im_data, 0)
        img_y = np.transpose(img_x, (0, 2, 1, 3))
        out = pnet(img_y)
        out0 = np.transpose(out[0], (0, 2, 1, 3))
        out1 = np.transpose(out[1], (0, 2, 1, 3))

        boxes, _ = generateBoundingBox(
            out1[0, :, :, 1].copy(), out0[0, :, :, :].copy(), scale, threshold[0]
        )

        # inter-scale nms
        pick = nms(boxes.copy(), 0.5, "Union")
        if boxes.size > 0 and pick.size > 0:
            boxes = boxes[pick, :]
            total_boxes = np.append(total_boxes, boxes, axis=0)

    numbox = total_boxes.shape[0]
    if numbox > 0:
        pick = nms(total_boxes.copy(), 0.7, "Union")
        total_boxes = total_boxes[pick, :]
        regw = total_boxes[:, 2] - total_boxes[:, 0]
        regh = total_boxes[:, 3] - total_boxes[:, 1]
        qq1 = total_boxes[:, 0] + total_boxes[:, 5] * regw
        qq2 = total_boxes[:, 1] + total_boxes[:, 6] * regh
        qq3 = total_boxes[:, 2] + total_boxes[:, 7] * regw
        qq4 = total_boxes[:, 3] + total_boxes[:, 8] * regh
        total_boxes = np.transpose(np.vstack([qq1, qq2, qq3, qq4, total_boxes[:, 4]]))
        total_boxes = rerec(total_boxes.copy())
        total_boxes[:, 0:4] = np.fix(total_boxes[:, 0:4]).astype(np.int32)
        dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph = pad(total_boxes.copy(), w, h)

    numbox = total_boxes.shape[0]
    if numbox > 0:
        # second stage
        tempimg = np.zeros((24, 24, 3, numbox))
        for k in range(0, numbox):
            tmp = np.zeros((int(tmph[k]), int(tmpw[k]), 3))
            tmp[dy[k] - 1 : edy[k], dx[k] - 1 : edx[k], :] = img[
                y[k] - 1 : ey[k], x[k] - 1 : ex[k], :
            ]
            if (
                tmp.shape[0] > 0
                and tmp.shape[1] > 0
                or tmp.shape[0] == 0
                and tmp.shape[1] == 0
            ):
                tempimg[:, :, :, k] = imresample(tmp, (24, 24))
            else:
                return np.empty()
        tempimg = (tempimg - 127.5) * 0.0078125
        tempimg1 = np.transpose(tempimg, (3, 1, 0, 2))
        out = rnet(tempimg1)
        out0 = np.transpose(out[0])
        out1 = np.transpose(out[1])
        score = out1[1, :]
        ipass = np.where(score > threshold[1])
        total_boxes = np.hstack(
            [total_boxes[ipass[0], 0:4].copy(), np.expand_dims(score[ipass].copy(), 1)]
        )
        mv = out0[:, ipass[0]]
        if total_boxes.shape[0] > 0:
            pick = nms(total_boxes, 0.7, "Union")
            total_boxes = total_boxes[pick, :]
            total_boxes = bbreg(total_boxes.copy(), np.transpose(mv[:, pick]))
            total_boxes = rerec(total_boxes.copy())

    numbox = total_boxes.shape[0]
    if numbox > 0:
        # third stage
        total_boxes = np.fix(total_boxes).astype(np.int32)
        dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph = pad(total_boxes.copy(), w, h)
        tempimg = np.zeros((48, 48, 3, numbox))
        for k in range(0, numbox):
            tmp = np.zeros((int(tmph[k]), int(tmpw[k]), 3))
            tmp[dy[k] - 1 : edy[k], dx[k] - 1 : edx[k], :] = img[
                y[k] - 1 : ey[k], x[k] - 1 : ex[k], :
            ]
            if (
                tmp.shape[0] > 0
                and tmp.shape[1] > 0
                or tmp.shape[0] == 0
                and tmp.shape[1] == 0
            ):
                tempimg[:, :, :, k] = imresample(tmp, (48, 48))
            else:
                return np.empty()
        tempimg = (tempimg - 127.5) * 0.0078125
        tempimg1 = np.transpose(tempimg, (3, 1, 0, 2))
        out = onet(tempimg1)
        out0 = np.transpose(out[0])
        out1 = np.transpose(out[1])
        out2 = np.transpose(out[2])
        score = out2[1, :]
        points = out1
        ipass = np.where(score > threshold[2])
        points = points[:, ipass[0]]
        total_boxes = np.hstack(
            [total_boxes[ipass[0], 0:4].copy(), np.expand_dims(score[ipass].copy(), 1)]
        )
        mv = out0[:, ipass[0]]

        w = total_boxes[:, 2] - total_boxes[:, 0] + 1
        h = total_boxes[:, 3] - total_boxes[:, 1] + 1
        points[0:5, :] = (
            np.tile(w, (5, 1)) * points[0:5, :] + np.tile(total_boxes[:, 0], (5, 1)) - 1
        )
        points[5:10, :] = (
            np.tile(h, (5, 1)) * points[5:10, :]
            + np.tile(total_boxes[:, 1], (5, 1))
            - 1
        )
        if total_boxes.shape[0] > 0:
            total_boxes = bbreg(total_boxes.copy(), np.transpose(mv))
            pick = nms(total_boxes.copy(), 0.7, "Min")
            total_boxes = total_boxes[pick, :]
            points = points[:, pick]

    return total_boxes, points

def generateBoundingBox(imap, reg, scale, t):
    """Use heatmap to generate bounding boxes"""
    stride = 2
    cellsize = 12

    imap = np.transpose(imap)
    dx1 = np.transpose(reg[:, :, 0])
    dy1 = np.transpose(reg[:, :, 1])
    dx2 = np.transpose(reg[:, :, 2])
    dy2 = np.transpose(reg[:, :, 3])
    y, x = np.where(imap >= t)
    if y.shape[0] == 1:
        dx1 = np.flipud(dx1)
        dy1 = np.flipud(dy1)
        dx2 = np.flipud(dx2)
        dy2 = np.flipud(dy2)
    score = imap[(y, x)]
    reg = np.transpose(np.vstack([dx1[(y, x)], dy1[(y, x)], dx2[(y, x)], dy2[(y, x)]]))
    if reg.size == 0:
        reg = np.empty((0, 3))
    bb = np.transpose(np.vstack([y, x]))
    q1 = np.fix((stride * bb + 1) / scale)
    q2 = np.fix((stride * bb + cellsize - 1 + 1) / scale)
    boundingbox = np.hstack([q1, q2, np.expand_dims(score, 1), reg])
    return boundingbox, reg

def imresample(img, sz):
    im_data = cv2.resize(
        img, (sz[1], sz[0]), interpolation=cv2.INTER_AREA
    )  # @UndefinedVariable
    return im_data

def nms(boxes, threshold, method):
    if boxes.size == 0:
        return np.empty((0, 3))
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    s = boxes[:, 4]
    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    I = np.argsort(s)
    pick = np.zeros_like(s, dtype=np.int16)
    counter = 0
    while I.size > 0:
        i = I[-1]
        pick[counter] = i
        counter += 1
        idx = I[0:-1]
        xx1 = np.maximum(x1[i], x1[idx])
        yy1 = np.maximum(y1[i], y1[idx])
        xx2 = np.minimum(x2[i], x2[idx])
        yy2 = np.minimum(y2[i], y2[idx])
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        if method is "Min":
            o = inter / np.minimum(area[i], area[idx])
        else:
            o = inter / (area[i] + area[idx] - inter)
        I = I[np.where(o <= threshold)]
    pick = pick[0:counter]
    return pick

# function [bboxA] = rerec(bboxA)
def rerec(bboxA):
    """Convert bboxA to square."""
    h = bboxA[:, 3] - bboxA[:, 1]
    w = bboxA[:, 2] - bboxA[:, 0]
    l = np.maximum(w, h)
    bboxA[:, 0] = bboxA[:, 0] + w * 0.5 - l * 0.5
    bboxA[:, 1] = bboxA[:, 1] + h * 0.5 - l * 0.5
    bboxA[:, 2:4] = bboxA[:, 0:2] + np.transpose(np.tile(l, (2, 1)))
    return bboxA

# function [boundingbox] = bbreg(boundingbox,reg)
def bbreg(boundingbox, reg):
    """Calibrate bounding boxes"""
    if reg.shape[1] == 1:
        reg = np.reshape(reg, (reg.shape[2], reg.shape[3]))

    w = boundingbox[:, 2] - boundingbox[:, 0] + 1
    h = boundingbox[:, 3] - boundingbox[:, 1] + 1
    b1 = boundingbox[:, 0] + reg[:, 0] * w
    b2 = boundingbox[:, 1] + reg[:, 1] * h
    b3 = boundingbox[:, 2] + reg[:, 2] * w
    b4 = boundingbox[:, 3] + reg[:, 3] * h
    boundingbox[:, 0:4] = np.transpose(np.vstack([b1, b2, b3, b4]))
    return boundingbox

# function [dy edy dx edx y ey x ex tmpw tmph] = pad(total_boxes,w,h)
def pad(total_boxes, w, h):
    """Compute the padding coordinates (pad the bounding boxes to square)"""
    tmpw = (total_boxes[:, 2] - total_boxes[:, 0] + 1).astype(np.int32)
    tmph = (total_boxes[:, 3] - total_boxes[:, 1] + 1).astype(np.int32)
    numbox = total_boxes.shape[0]

    dx = np.ones((numbox), dtype=np.int32)
    dy = np.ones((numbox), dtype=np.int32)
    edx = tmpw.copy().astype(np.int32)
    edy = tmph.copy().astype(np.int32)

    x = total_boxes[:, 0].copy().astype(np.int32)
    y = total_boxes[:, 1].copy().astype(np.int32)
    ex = total_boxes[:, 2].copy().astype(np.int32)
    ey = total_boxes[:, 3].copy().astype(np.int32)

    tmp = np.where(ex > w)
    edx.flat[tmp] = np.expand_dims(-ex[tmp] + w + tmpw[tmp], 1)
    ex[tmp] = w

    tmp = np.where(ey > h)
    edy.flat[tmp] = np.expand_dims(-ey[tmp] + h + tmph[tmp], 1)
    ey[tmp] = h

    tmp = np.where(x < 1)
    dx.flat[tmp] = np.expand_dims(2 - x[tmp], 1)
    x[tmp] = 1

    tmp = np.where(y < 1)
    dy.flat[tmp] = np.expand_dims(2 - y[tmp], 1)
    y[tmp] = 1

    return dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph

def detect_with_MTCNN(origin_images_dir, out_dir, pb_path, mode="no_depth"):
    print("MTCNN detect")

    if os.path.exists(out_dir) is False:
        os.makedirs(out_dir)

    minsize = 20  # minimum size of face
    threshold = [0.5, 0.6, 0.6]  # three steps's threshold
    factor = 0.709  # scale factor

    with tf.Graph().as_default():
        graph_def = tf.GraphDef()
        graph_file = pb_path
        with open(graph_file, "rb") as f:
            print("hello")
            graph_def.ParseFromString(f.read())
            tf.import_graph_def(graph_def, name="")

        sess = tf.Session()
        with sess.as_default():
            tf.global_variables_initializer().run()
            pnet, rnet, onet = create_mtcnn_pb(sess)

    # find files
    import glob

    if mode == "depth":
        files = glob.glob(osp.join(origin_images_dir, "*.jpg"))
        files.extend(glob.glob(osp.join(origin_images_dir, "*.JPG")))
        dep_files = glob.glob(osp.join(origin_images_dir, "*.png"))
        dep_files.extend(glob.glob(osp.join(origin_images_dir, "*.PNG")))
        files.sort()
        dep_files.sort()
    else:
        files = glob.glob(osp.join(origin_images_dir, "*.jpg"))
        files.extend(glob.glob(osp.join(origin_images_dir, "*.png")))
        files.extend(glob.glob(osp.join(origin_images_dir, "*.JPG")))
        files.extend(glob.glob(osp.join(origin_images_dir, "*.PNG")))
        files.sort()
    print("=========================")
    # print("img:", files)
    # print("Dep:", dep_files)

    # detect face bbox
    count = 0
    names_list = []
    dep_name_list = []
    for index in range(0, len(files)):
        img = cv2.imread(files[index])
        bounding_boxes, points = detect_face(
            img, minsize, pnet, rnet, onet, threshold, factor
        )  # bounding_boxes.shape: (n, 5) points.shape: (10, n)
        if len(bounding_boxes) == 1:
            points = np.transpose(points)
            batch_imgs = img
            batch_bboxes = bounding_boxes
            batch_points = points
            batch_names = files[index].split("/")[-1]
            names_list.append(batch_names)

            scio.savemat(
                os.path.join(out_dir, batch_names[:-4] + ".mat"),
                {
                    "batch_bboxes": batch_bboxes.astype(np.float64),
                    "batch_points": batch_points.astype(np.float64),
                },
            )
            if mode == "depth":
                dep_name_list.append(dep_files[index].split("/")[-1])

        elif len(bounding_boxes) > 1:
            print("too much face to detect by MTCNN, only select first person")
            points = np.transpose(points[:, 0:1])
            batch_imgs = img
            batch_bboxes = bounding_boxes[0:1, :]
            batch_points = points
            batch_names = files[index].split("/")[-1]
            names_list.append(batch_names)

            scio.savemat(
                os.path.join(out_dir, batch_names[:-4] + ".mat"),
                {
                    "batch_bboxes": batch_bboxes.astype(np.float64),
                    "batch_points": batch_points.astype(np.float64),
                },
            )
            if mode == "depth":
                dep_name_list.append(dep_files[index].split("/")[-1])

        else:
            print("no face to detect by MTCNN, please input single person photo")
            pass
            # raise Exception("no face or much face to detect, please input single person photo")

        count = count + 1
        if (count % 100 == 0) | (count == len(files)):
            print("has run MTCNN: " + str(count) + " / " + str(len(files)))
    sess.close()

    if mode == "depth":
        return names_list, dep_name_list
    else:
        return names_list


import sys, time
import glob
import pygame
import random as rnd
import numpy.random
poisson = numpy.random.poisson

from util import *
from graphics import Screen
from game import Game, parseCardFile

from AI import AI
from story import Story
from cardheap import Deck
import getDecision

# Report seed (for debugging)
if len(sys.argv) > 1:
    SEED = int(sys.argv[1])
else:
    SEED = int(time.time())
rnd.seed(SEED)
numpy.random.seed(SEED)
print 'SEED:', SEED

resolutions = ((750,520),    #0 Eee-PC Netbook Window
               (1024,600),   #1 Eee-PC Netbook Fullscreen
               (1280,800),   #2 LowRes (WideScreen)
               (1280,1024),  #3 Standard 5:4 Display
               (1268,970),   #4 Window filling a standard 1280x1024 Display
               (1440,900),   #5 Macbook Pro 15" (WideScreen)
               (1920,1200))  #6 Standard WideScreen HD Display, Macbook Pro 17"

enableFullscreen  = (False, #0
                     True)  #1

backgrounds = ("cthulhu_1440x900.jpg",               #0
               "cthulhu_1920x1200.jpg",              #1 Cool Cthulhu Illustration
               "dark_backgr_1920x1200.jpg",          #2
               "minimalistic_Symbol_1920x1200.jpg",  #3
               "old_inn_1920x1200.jpg",              #4
               "tentacles_1920x1200.jpg",            #5 Cool Tentacle Illustration
               "whiteWood_1920x1200.jpg",            #6 White
               "wood_1440x1050.jpg",                 #7 Good Table
               "woodB_1920x1200.jpg",                #8
               "woodC_1920x1200.jpg",                #9
               "woodD_1920x1200.jpg",                #10
               "woodE_2560x1600.jpg",                #11
               "woodF_1280x800.jpg")                 #12

###############################
##### CHOOSE DISPLAY MODE #####

resolution = 4    # 4  5  6
fullscreen = 0    # 0  0  1
background = 6    # 7  0  0

###############################
###############################

pygame.init()

# Initialize screen
resolution = resolutions[resolution]
fullscreen = enableFullscreen[fullscreen]
background = "Backgrounds/" + backgrounds[background]

screen = Screen(resolution, background, fullscreen,
                caption = "Call of Cthulhu LCG")

try:

    # MAKE UP SOME PLAYERS AND CARDS

    # Players
    P1 = AI('Frrmack')
    P2 = AI('Boris')
    # P1 = Player('Frrmack')
    # P2 = Player('Boris')

    # Decks
    Deck1 = Deck('Random Deck 1')
    Deck2 = Deck('Random Deck 2')
    for n in range(50):
        # Cardtype = [Character, Event, Support][poisson(0.5) % 3]
        # redcard = Cardtype( 'Red Card %i'%(n+1), randomize=True )
        # bluecard = Cardtype( 'Blue Card %i'%(n+1), randomize=True )
        cardfile = rnd.choice(glob.glob('Cards/*.card'))
        card = parseCardFile(cardfile)
        Deck1.add(card)
        card = parseCardFile(cardfile)
        Deck2.add(card)

    P1.useDeck(Deck1)
    P2.useDeck(Deck2)

    # Stories
    storydeck = Deck('Random Story Deck')
    for n in range(10):
        # name = rnd.choice('ABCDEFGHIJKLXYZW') + \
        #     rnd.choice('ABCDEFGHIJKLXYZW')
        # storycard = Story( 'Story %s' % name  )
        cardfile = rnd.choice(glob.glob('Stories/*.card'))
        storycard = parseCardFile(cardfile)
        storydeck.add(storycard)

    # NEW GAME SETUP
    game = Game(screen, P1,P2)
    P1.screen = screen
    P2.screen = screen

    # Draw domains
    P1.domainPanel.draw()
    P2.domainPanel.draw()

    # Shuffle decks
    rnd.shuffle(P1.deck)
    rnd.shuffle(P2.deck)
    rnd.shuffle(storydeck)

    # Arrange stories
    game.storydeck = storydeck
    game.initiateStories()
    game.drawStories()
    screen.update()

    # Draw 8 cards each from the deck
    print '%s DRAWS 8 cards' % P1.name
    print '%s DRAWS 8 cards\n' % P2.name
    P1.drawCard(8)
    P2.drawCard(8)

    #screen.msgBox("%s x %s" % (screen.width, screen.height), OKBox=True)

    screen.readClick()

    # Attach 1 card to each domain
    for ply in [P1, P2]:
        for i in range(3):
            card, domain = getDecision.attachOneCardToADomain(ply)
            domain = ply.domains[i]
            ply.attach2Domain(card, domain)

    screen.readClick()

    # P1 starts
    game.ActivePlayer = P1
    game.DefendingPlayer = P2

    # PLAY FOR A LIMITED NUMBER OF TURNS (for now, coding and debugging)

    for turn in range(20):
        printTurnHeader("__________ START %s's TURN ___________\n" % game.ActivePlayer.name)
        game.screen.msgBox("%s's turn starts" % game.ActivePlayer.name)

        # REFRESH PHASE
        # Ready all exhausted cards
        # Restore 1 insane character
        # Refresh Domains
        printPhaseHeader('--- REFRESH PHASE -------------------------------\n')

        for card in game.ActivePlayer.board.cardsOfState('exhausted'):
            print game.ActivePlayer.name,'READIES',card
            card.ready()

        for domain in game.ActivePlayer.domains:
            if domain.isDrained():
                print game.ActivePlayer.name,'REFRESHES',domain
                domain.refresh()

        if len(game.ActivePlayer.board.cardsOfState('insane')) > 0:
            card = getDecision.restoreOneInsane(game.ActivePlayer)
            if card is not None:
                print game.ActivePlayer.name,'RESTORES',card
                card.restore()

        game.ActivePlayer.board.redraw()

        screen.readClick()
        print

        # ACTIONS

        # DRAW PHASE
        # Draw 2 cards
        printPhaseHeader('--- DRAW PHASE ----------------------------------\n')
        if turn == 0:
            nCards, grammar = 1, 'card'
        else:
            nCards, grammar = 2, 'cards'
        game.ActivePlayer.drawCard(nCards)
        print game.ActivePlayer.name, 'draws %i %s' % (nCards,grammar),'\n'
        # print game.ActivePlayer.handReport(), '\n'

        screen.readClick()

        # ACTIONS

        # RESOURCE PHASE
        # Attach one card to a domain
        printPhaseHeader('--- RESOURCE PHASE ------------------------------\n')

        # get decision
        card, domain = getDecision.attachOneCardToADomain(game.ActivePlayer)

        print game.ActivePlayer.name, 'ATTACHES', card, '\n\t TO', domain, '\n'
        game.ActivePlayer.attach2Domain(card, domain)

        # Report
        # print game.ActivePlayer.domainReport(), '\n'

        # ACTIONS

        # OPERATIONS PHASE
        # Play cards from your hand (also actions may take place)
        printPhaseHeader('--- OPERATIONS PHASE -----------------------------\n')

        while True:
            # keep playing until you choose no card to play
            card, target, domain = getDecision.playCardFromHand(game.ActivePlayer)

            if card is None:
                break

            # target stuff is not coded yet

            if domain is None:
                print game.ActivePlayer.name, 'PLAYS', card, '\n'
            else:
                print game.ActivePlayer.name, 'PLAYS %s USING %s \n' % (card, domain)

            game.ActivePlayer.play(card, target, domain)
            screen.readClick()

        # Report
        # print game.report(), '\n'
        # print game.ActivePlayer.report(), '\n'

        # STORY PHASE
        printPhaseHeader('--- STORY PHASE ----------------------------------\n')
        clear(4)

        # STORY PHASE I
        # Active Player commits characters to stories
        printPhaseHeader('________ Active Player Commits ________\n')

        if turn == 0:
            print "NO COMMITS ON THE FIRST TURN\n"
        else:

            while True:
                # keep committing until you choose no character to commit

                character, story = getDecision.commitCharacterToStoryWhenActive(game.ActivePlayer)

                if character is None:
                    break

                print game.ActivePlayer.name,'COMMITS',character,'\n\t TO',story,'\n'
                game.ActivePlayer.commit(character, story)

        # Report
        # print game.report(showCommitted=True), '\n'
        # print game.ActivePlayer.report(), '\n'

        screen.readClick()

        # ACTIONS

        # STORY PHASE II
        # Defending Player commits characters to stories
        printPhaseHeader('________ Defending Player Commits ________\n')

        while True:
            # keep committing until you choose no character to commit

            character, story = getDecision.commitCharacterToStoryWhenDefending(game.DefendingPlayer)

            if character is None:
                break

            print game.DefendingPlayer.name,'COMMITS',character,'\n\t TO',story,'\n'
            game.DefendingPlayer.commit(character, story)

        screen.readClick()

        # ACTIONS

        # STORY PHASE III
        # Resolve stories
        printPhaseHeader('________ Stories are Resolved ________\n')

        activeStories = filter(lambda st: st.isAnyCommitted(), game.stories)
        if len(activeStories) == 0:
            print 'NO STORIES TO RESOLVE\n'
            clear(15)
        else:
            for story in activeStories:
                print storyColor(story.name), 'is resolved:'
                story.resolve()
                print

        screen.readClick()

        # Response to struggles / success

        # Uncommit characters
        for story in activeStories:
            story.uncommitAll()

        # ACTIONS

        # END OF TURN

        # End of turn, swap active player
        game.ActivePlayer, game.DefendingPlayer = game.DefendingPlayer, game.ActivePlayer

except:
    raise

finally:
    print 'SEED:', SEED


import cv2
import json
import numpy as np
import os
import time
import glob

from model import efficientdet
from utils import preprocess_image, postprocess_boxes
from utils.draw_boxes import draw_boxes

def main():
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'

    phi = 1
    weighted_bifpn = True
    model_path = 'efficientdet-d1.h5'
    image_sizes = (512, 640, 768, 896, 1024, 1280, 1408)
    image_size = image_sizes[phi]
    # coco classes
    classes = {value['id'] - 1: value['name'] for value in json.load(open('coco_90.json', 'r')).values()}
    num_classes = 90
    score_threshold = 0.3
    colors = [np.random.randint(0, 256, 3).tolist() for _ in range(num_classes)]
    _, model = efficientdet(phi=phi,
                            weighted_bifpn=weighted_bifpn,
                            num_classes=num_classes,
                            score_threshold=score_threshold)
    model.load_weights(model_path, by_name=True)

    for image_path in glob.glob('datasets/VOC2007/JPEGImages/*.jpg'):
        image = cv2.imread(image_path)
        src_image = image.copy()
        # BGR -> RGB
        image = image[:, :, ::-1]
        h, w = image.shape[:2]

        image, scale = preprocess_image(image, image_size=image_size)
        # run network
        start = time.time()
        boxes, scores, labels = model.predict_on_batch([np.expand_dims(image, axis=0)])
        boxes, scores, labels = np.squeeze(boxes), np.squeeze(scores), np.squeeze(labels)
        print(time.time() - start)
        boxes = postprocess_boxes(boxes=boxes, scale=scale, height=h, width=w)

        # select indices which have a score above the threshold
        indices = np.where(scores[:] > score_threshold)[0]

        # select those detections
        boxes = boxes[indices]
        labels = labels[indices]

        draw_boxes(src_image, boxes, scores, labels, colors, classes)

        cv2.namedWindow('image', cv2.WINDOW_NORMAL)
        cv2.imshow('image', src_image)
        cv2.waitKey(0)

if __name__ == '__main__':
    main()


'''
300D Model - Train / Test (epochs)
=-=-=
Batch size = 512
Fixed GloVe
- 300D SumRNN + Translate + 3 MLP (1.2 million parameters) - 0.8315 / 0.8235 / 0.8249 (22 epochs)
- 300D GRU + Translate + 3 MLP (1.7 million parameters) - 0.8431 / 0.8303 / 0.8233 (17 epochs)
- 300D LSTM + Translate + 3 MLP (1.9 million parameters) - 0.8551 / 0.8286 / 0.8229 (23 epochs)

Following Liu et al. 2016, I don't update the GloVe embeddings during training.
Unlike Liu et al. 2016, I don't initialize out of vocabulary embeddings randomly and instead leave them zeroed.

The jokingly named SumRNN (summation of word embeddings) is 10-11x faster than the GRU or LSTM.

Original numbers for sum / LSTM from Bowman et al. '15 and Bowman et al. '16
=-=-=
100D Sum + GloVe - 0.793 / 0.753
100D LSTM + GloVe - 0.848 / 0.776
300D LSTM + GloVe - 0.839 / 0.806
'''

from __future__ import print_function

import json
import os
import pickle
import tempfile
from argparse import ArgumentParser

import numpy as np

import keras
import keras.backend as K
from keras.callbacks import EarlyStopping, History, ModelCheckpoint
from keras.layers import Dense, Dropout, Input, TimeDistributed, recurrent
from keras.layers.noise import AlphaDropout
from keras.layers.embeddings import Embedding
from keras.layers.merge import concatenate
from keras.layers.normalization import BatchNormalization
from keras.layers.wrappers import Bidirectional
from keras.models import Model
from keras.preprocessing.sequence import pad_sequences
from keras.preprocessing.text import Tokenizer
from keras.regularizers import l2
from keras.utils import np_utils

np.random.seed(1337)  # for reproducibility

def extract_tokens_from_binary_parse(parse):
    return parse.replace('(', ' ').replace(')', ' ').replace('-LRB-', '(').replace('-RRB-', ')').split()

def yield_examples(fn, skip_no_majority=True, limit=None):
    for i, line in enumerate(open(fn)):
        if limit and i > limit:
            break
        data = json.loads(line)
        label = data['gold_label']
        s1 = ' '.join(extract_tokens_from_binary_parse(data['sentence1_binary_parse']))
        s2 = ' '.join(extract_tokens_from_binary_parse(data['sentence2_binary_parse']))
        if skip_no_majority and label == '-':
            continue
        yield (label, s1, s2)

def get_data(fn, limit=None):
    raw_data = list(yield_examples(fn=fn, limit=limit))
    left = [s1 for _, s1, s2 in raw_data]
    right = [s2 for _, s1, s2 in raw_data]
    print(max(len(x.split()) for x in left))
    print(max(len(x.split()) for x in right))

    LABELS = {'contradiction': 0, 'neutral': 1, 'entailment': 2}
    Y = np.array([LABELS[l] for l, s1, s2 in raw_data])
    Y = np_utils.to_categorical(Y, len(LABELS))

    return left, right, Y

def main(config):
    training = get_data('snli_1.0_train.jsonl')
    validation = get_data('snli_1.0_dev.jsonl')
    test = get_data('snli_1.0_test.jsonl')

    tokenizer = Tokenizer(lower=False, filters='')
    tokenizer.fit_on_texts(training[0] + training[1])

    # Lowest index from the tokenizer is 1 - we need to include 0 in our vocab count
    VOCAB = len(tokenizer.word_counts) + 1
    LABELS = {'contradiction': 0, 'neutral': 1, 'entailment': 2}
    # RNN = recurrent.LSTM
    # RNN = lambda *args, **kwargs: Bidirectional(recurrent.LSTM(*args, **kwargs))
    # RNN = recurrent.GRU
    # RNN = lambda *args, **kwargs: Bidirectional(recurrent.GRU(*args, **kwargs))
    # Summation of word embeddings
    RNN = None
    LAYERS = 1
    USE_GLOVE = True
    TRAIN_EMBED = False
    EMBED_HIDDEN_SIZE = 300
    SENT_HIDDEN_SIZE = 300
    BATCH_SIZE = 512
    PATIENCE = 4  # 8
    MAX_EPOCHS = config.max_epochs
    MAX_LEN = 42
    DP = 0.2
    L2 = 4e-6
    ACTIVATION = config.activation
    OPTIMIZER = config.optimizer

    if config.dropout_type == 'AlphaDropout':
        DROPOUT = AlphaDropout
    else:
        DROPOUT = Dropout

    print('RNN / Embed / Sent = {}, {}, {}'.format(RNN, EMBED_HIDDEN_SIZE, SENT_HIDDEN_SIZE))
    print('GloVe / Trainable Word Embeddings = {}, {}'.format(USE_GLOVE, TRAIN_EMBED))

    to_seq = lambda X: pad_sequences(tokenizer.texts_to_sequences(X), maxlen=MAX_LEN)
    prepare_data = lambda data: (to_seq(data[0]), to_seq(data[1]), data[2])

    training = prepare_data(training)
    validation = prepare_data(validation)
    test = prepare_data(test)

    print('Build model...')
    print('Vocab size =', VOCAB)

    GLOVE_STORE = 'precomputed_glove.weights'
    if USE_GLOVE:
        if not os.path.exists(GLOVE_STORE + '.npy'):
            print('Computing GloVe')

            embeddings_index = {}
            f = open('glove.840B.300d.txt')
            for line in f:
                values = line.split(' ')
                word = values[0]
                coefs = np.asarray(values[1:], dtype='float32')
                embeddings_index[word] = coefs
            f.close()

            # prepare embedding matrix
            embedding_matrix = np.zeros((VOCAB, EMBED_HIDDEN_SIZE))
            for word, i in tokenizer.word_index.items():
                embedding_vector = embeddings_index.get(word)
                if embedding_vector is not None:
                    # words not found in embedding index will be all-zeros.
                    embedding_matrix[i] = embedding_vector
                else:
                    print('Missing from GloVe: {}'.format(word))

            np.save(GLOVE_STORE, embedding_matrix)

        print('Loading GloVe')
        embedding_matrix = np.load(GLOVE_STORE + '.npy')

        print('Total number of null word embeddings:')
        print(np.sum(np.sum(embedding_matrix, axis=1) == 0))

        embed = Embedding(
            VOCAB, EMBED_HIDDEN_SIZE, weights=[embedding_matrix], input_length=MAX_LEN, trainable=TRAIN_EMBED)
    else:
        embed = Embedding(VOCAB, EMBED_HIDDEN_SIZE, input_length=MAX_LEN)

    rnn_kwargs = dict(output_dim=SENT_HIDDEN_SIZE, dropout_W=DP, dropout_U=DP)
    SumEmbeddings = keras.layers.core.Lambda(lambda x: K.sum(x, axis=1), output_shape=(SENT_HIDDEN_SIZE, ))

    translate = TimeDistributed(Dense(SENT_HIDDEN_SIZE, activation=ACTIVATION))

    premise = Input(shape=(MAX_LEN, ), dtype='int32')
    hypothesis = Input(shape=(MAX_LEN, ), dtype='int32')

    prem = embed(premise)
    hypo = embed(hypothesis)

    prem = translate(prem)
    hypo = translate(hypo)

    if RNN and LAYERS > 1:
        for l in range(LAYERS - 1):
            rnn = RNN(return_sequences=True, **rnn_kwargs)
            if config.add_batchnorm:
                prem = BatchNormalization()(rnn(prem))
                hypo = BatchNormalization()(rnn(hypo))

    rnn = SumEmbeddings if not RNN else RNN(return_sequences=False, **rnn_kwargs)

    prem = rnn(prem)
    hypo = rnn(hypo)

    if config.add_batchnorm:
        prem = BatchNormalization()(prem)
        hypo = BatchNormalization()(hypo)

    joint = concatenate([prem, hypo])
    joint = DROPOUT(DP)(joint)
    for i in range(3):
        joint = Dense(2 * SENT_HIDDEN_SIZE, activation=ACTIVATION, kernel_regularizer=l2(L2) if L2 else None)(joint)
        joint = DROPOUT(DP)(joint)

        if config.add_batchnorm:
            joint = BatchNormalization()(joint)

    pred = Dense(len(LABELS), activation='softmax')(joint)

    model = Model(inputs=[premise, hypothesis], outputs=pred)
    model.compile(optimizer=OPTIMIZER, loss='categorical_crossentropy', metrics=['accuracy'])

    model.summary()

    print('Training')
    _, tmpfn = tempfile.mkstemp()
    # Save the best model during validation and bail out of training early if we're not improving
    callbacks = [
        # EarlyStopping(patience=PATIENCE),
        ModelCheckpoint(tmpfn, save_best_only=True, save_weights_only=True),
        History(),
    ]
    model.fit(
        [training[0], training[1]],
        training[2],
        batch_size=BATCH_SIZE,
        epochs=MAX_EPOCHS,
        validation_data=([validation[0], validation[1]], validation[2]),
        callbacks=callbacks)

    # Dump the history as a pickle object
    pickle.dump(callbacks[-1].history, open('history_{}.pkl'.format(ACTIVATION), 'wb'))

    # Restore the best found model during validation
    model.load_weights(tmpfn)

    loss, acc = model.evaluate([test[0], test[1]], test[2], batch_size=BATCH_SIZE)
    print('Test loss / test accuracy = {:.4f} / {:.4f}'.format(loss, acc))

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--activation', type=str, default='relu', help='Activation function')
    parser.add_argument('--max_epochs', type=int, default=42, help='Num epochs')
    parser.add_argument('--add_batchnorm', type=bool, default=False, help='Enable/disable Batch Normalization')
    parser.add_argument('--optimizer', type=str, default='adam', help='Optimizer')
    parser.add_argument('--dropout_type', type=str, default='dropout', help='Dropout Type')

    args = parser.parse_args()
    main(args)


#!/Usr/bin/env python
"""
Process the input meds files in chunks

Possible commands are:
    setup - setup jobs
    setup-nbrs - set up jobs for neighbor finding
    collate - combine job outputs into a single file
    verify - verify that all job outputs are present and OK
    clean - clean all outputs from a run
    archive - run after collate to delete intermediate files and tar logs
    link - make symlinks to all final outputs under {run}/output
    collate-link - collate and link the final outputs

You must have the following environment variables set
    $NGMIXER_OUTPUT_DIR - base output dir
    $TMPDIR - temporary directory for work

TODO:
    - work on list of tile ids
"""

from __future__ import print_function
try:
    xrange  # noqa
except Exception:
    xrange = range

import os

import numpy
import meds
import fitsio

import ngmixer
from ngmixer import files

class MegaMixer(dict):
    """
    conf has ngmix config plus
    """
    def __init__(self, ngmix_config, meds_files,
                 psf_map_files=None,
                 walltime='48:00',
                 walltime_e='4:00',
                 ncores=1,
                 missing=False,
                 blind=False,
                 skip_errors=False,
                 clobber=False,
                 seed=None):

        self.ngmix_config = os.path.abspath(
            os.path.expanduser(os.path.expandvars(ngmix_config)))
        if psf_map_files is not None:
            self.psf_map_files = [os.path.abspath(f) for f in psf_map_files]
        else:
            self.psf_map_files = None

        self.missing = missing
        self.walltime = walltime
        self.walltime_e = walltime_e
        self.ncores = ncores

        self.blind = blind
        self.skip_errors = skip_errors
        self.clobber = clobber

        self._load_config()
        self._make_rng(seed=seed)

        self.meds_files = [os.path.abspath(mf) for mf in meds_files]
        meds_list = [meds.MEDS(f) for f in self.meds_files]

        try:
            ngmixer.imageio.medsio.verify_meds(meds_list)
        except AttributeError:
            print(
                "ngmixer.imageio.medsio.verify_meds not found - "
                "probably this version is too old")
            pass

        self.meds_string = ' '.join(self.meds_files)

        self.info = files.get_meds_info(meds_files[0])

        self._set_files()

    def setup(self):
        """
        make directories, write scripts
        """
        self._set_chunk_ranges()
        self._make_scripts()

    def setup_nbrs(self):
        """
        make directories, write scripts
        """
        self._write_nbrs_script()

    def get_collated_file(self):
        return files.get_collated_file_fromfile(
            self.meds_files[0],
            self.conf['run'],
            blind=self.blind
        )

    def collate(self):
        """
        concatenate all the chunk files, adding in
        some calculated information such as mags
        """

        self._set_chunk_ranges()

        cls = ngmixer.megamixer.concat_io.get_concat_class(
            self.conf['collate']['concat_type'],
        )
        flist = []
        for chunk, rng in enumerate(self.chunk_ranges):
            fname = self._get_chunk_file(chunk, rng, ext='.fits')
            flist.append(fname)

        collated_file = files.get_collated_file_fromfile(
            self.meds_files[0],
            self.conf['run'],
            blind=self.blind,
        )

        concatter = cls(
            self.ngmix_config,
            flist,
            collated_file,
            bands=self.conf['jobs']['bands'],
            blind=self.blind,
            clobber=self.clobber,
            skip_errors=self.skip_errors,
        )

        concatter.concat()
        return collated_file

    def _make_scripts(self):
        """
        write the scripts
        """
        for chunk, rng in enumerate(self.chunk_ranges):
            self._write_script(chunk, rng)

    def _write_script(self, chunk, rng):

        self._make_chunk_dir(chunk, rng)

        fmt = self._get_script_template()
        args = {}

        output_file = self._get_chunk_file(chunk, rng, ext='.fits')
        logfile = self._get_chunk_file(chunk, rng, ext='.log')

        args['ngmix_config'] = self.ngmix_config
        if self.psf_map_files is not None:
            psf_maps = ','.join(self.psf_map_files)
            args['psf_map_arg'] = '--psf-map="%s"' % psf_maps
        else:
            args['psf_map_arg'] = ''

        args['meds_files'] = self.meds_string
        args['output_file'] = output_file
        args['logfile'] = logfile
        args['start'] = rng[0]
        args['stop'] = rng[1]

        args['mof_opt'] = ''
        model_nbrs = self.conf.get('model_nbrs', False)
        if model_nbrs:
            fof_file = self._get_fof_file(ext='.fits')
            args['fof_opt'] = '--fof-file=%s' % fof_file

            nbrs_file = self._get_nbrs_file(ext='.fits')
            args['nbrs_opt'] = '--nbrs-file=%s' % nbrs_file
        else:
            args['fof_opt'] = ''
            args['nbrs_opt'] = ''

            if 'correct_meds' in self.conf:
                if self.conf['correct_meds'] is not False:
                    mof_file = self._get_mof_file()
                    args['mof_opt'] = '--mof-file=%s' % mof_file

        if 'forced_photometry' in self.conf:
            models_file = files.get_collated_file_fromfile(
                self.meds_files[0],
                self.conf['forced_photometry']['models_run'],
            )
            args['models_opt'] = '--models="%s"' % models_file

        else:
            args['models_opt'] = ''

        seed = self.rng.randint(low=1, high=1000000000)
        args['seed_opt'] = '--seed=%d' % seed

        scr = fmt.format(**args)

        script_name = self._get_chunk_file(chunk, rng, ext='.sh')

        dowrite = True
        if self.missing:
            if os.path.exists(script_name):
                os.remove(script_name)
            if os.path.exists(output_file):
                dowrite = False

        if dowrite:
            print("writing:", script_name)
            with open(script_name, 'w') as fp:
                fp.write(scr)

            os.system('chmod 755 %s' % script_name)

        return dowrite, script_name

    def _write_nbrs_script(self):

        self._make_nbrs_dir()

        fmt = self._get_nbrs_script_template()
        args = {}

        logfile = self._get_nbrs_file(ext='.log')

        args['ngmix_config'] = self.ngmix_config
        args['meds_file'] = self.meds_files[0]
        args['logfile'] = logfile

        scr = fmt.format(**args)

        script_name = self._get_nbrs_file(ext='.sh')

        print("writing:", script_name)
        with open(script_name, 'w') as fp:
            fp.write(scr)

        os.system('chmod 755 %s' % script_name)
        return script_name

    def _get_mof_file(self):
        """
        mof file used for subtracting neighbor light
        """

        mof_file = files.get_collated_file_fromfile(
            self.meds_files[0],
            self.conf['correct_meds']['mof_run'],
        )
        return mof_file

    def _make_chunk_dir(self, chunk, rng):
        f = files.get_chunk_file_fromfile(
            self.meds_files[0],
            self.conf['run'],
            rng,
        )
        files.makedirs_fromfile(f)

    def _get_chunk_file(self, chunk, rng, ext='.fits'):
        fname = files.get_chunk_file_fromfile(
            self.meds_files[0],
            self.conf['run'],
            rng,
            missing=self.missing,
            ext=ext,
        )
        return expand_path(fname)

    def _get_nbrs_file(self, ext='.fits'):
        fname = files.get_nbrs_file_fromfile(
            self.meds_files[0],
            self.conf['run'],
            ext=ext,
        )
        return expand_path(fname)

    def _get_fof_file(self, ext='.fits'):
        fname = files.get_fof_file_fromfile(
            self.meds_files[0],
            self.conf['run'],
            ext=ext,
        )
        return expand_path(fname)

    def _make_nbrs_dir(self):
        f = self._get_nbrs_file()
        files.makedirs_fromfile(f)

    def _get_script_template(self):
        template = r"""#!/bin/bash

if [[ -n $_CONDOR_SCRATCH_DIR ]]; then
    # the condor system creates a scratch directory for us,
    # and cleans up afterward
    tmpdir=$_CONDOR_SCRATCH_DIR
    export TMPDIR=$tmpdir
else
    # otherwise use the TMPDIR
    tmpdir=$TMPDIR
    mkdir -p $tmpdir
fi

export OMP_NUM_THREADS=1

logfile="{logfile}"

logbase=$(basename $logfile)
tmplog=$tmpdir/$logbase

config="{ngmix_config}"

meds="{meds_files}"

outfile="{output_file}"
start={start}
stop={stop}

ngmixit                                   \
    --fof-range=$start,$stop              \
    --work-dir=$tmpdir                    \
    {psf_map_arg}                         \
    {models_opt}                          \
    {nbrs_opt}                            \
    {fof_opt}                             \
    {mof_opt}                             \
    {seed_opt}                            \
    $config $outfile $meds &> $tmplog

mv -vf $tmplog $logfile
"""

        return template

    def _get_nbrs_script_template(self):
        template = r"""#!/bin/bash

if [[ -n $_CONDOR_SCRATCH_DIR ]]; then
    # the condor system creates a scratch directory for us,
    # and cleans up afterward
    tmpdir=$_CONDOR_SCRATCH_DIR
    export TMPDIR=$tmpdir
else
    # otherwise use the TMPDIR
    tmpdir=$TMPDIR
    mkdir -p $tmpdir
fi

logfile="{logfile}"

logbase=$(basename $logfile)
tmplog=$tmpdir/$logbase

config="{ngmix_config}"
meds="{meds_file}"

ngmixer-meds-make-nbrs-data $config $meds &> $tmplog
mv -vf $tmplog $logfile
"""

        return template

    def _set_files(self):
        files = {
            'ngmix_config': self.ngmix_config,
            'meds_files': self.meds_files,
        }

        self.files = files

    def _get_num(self):

        model_nbrs = self.conf.get('model_nbrs', False)
        if model_nbrs:
            fname = self._get_fof_file()
            fofs = fitsio.read(fname)
            num = numpy.unique(fofs['fofid']).size
        else:
            fname = self.meds_files[0]

            if not os.path.exists(fname):
                raise ngmixer.util.MissingDataError(
                    "missing meds file: %s" % fname)

            with fitsio.FITS(fname) as fits:
                num = fits['object_data'].get_nrows()

        return num

    def _set_chunk_ranges(self):

        if hasattr(self, 'chunk_ranges'):
            return

        # could be number of objects or number of
        # fof groups
        nrows = self._get_num()

        chunksize = self.jobs_conf['chunksize']
        nchunks = nrows//chunksize
        if nchunks*chunksize < nrows:
            nchunks += 1

        chunk_ranges = []
        for chunk in xrange(nchunks):
            sr = chunk*chunksize
            sp = sr + chunksize - 1
            if sp >= nrows:
                sp = nrows-1
            chunk_ranges.append([sr, sp])

        self.chunk_ranges = chunk_ranges

    def _load_config(self):
        self.conf = files.read_config(self.ngmix_config)
        self.jobs_conf = self.conf['jobs']

    def _make_rng(self, seed=None):
        if seed is None:
            if 'global_seed' not in self.conf:
                raise RuntimeError("either set 'global_seed' in the"
                                   "config or send --seed=")
            seed = self.conf['global_seed']

        self.rng = numpy.random.RandomState(seed=seed)

def expand_path(path):
    return os.path.abspath(
        os.path.expandvars(
            os.path.expanduser(
                path,
            )
        )
    )

class ImSimMegaMixer(MegaMixer):
    def _make_scripts(self):
        self.scripts = []
        self.output_files = []
        for chunk, rng in enumerate(self.chunk_ranges):
            _, script_name = self._write_script(chunk, rng)
            self.scripts.append(script_name)
            output_file = self._get_chunk_file(chunk, rng, ext='.fits')
            self.output_files.append(output_file)

    def _write_nbrs_script(self):

        self._make_nbrs_dir()

        fmt = self._get_nbrs_script_template()
        args = {}

        logfile = self._get_nbrs_file(ext='.log')

        args['ngmix_config'] = self.ngmix_config
        args['meds_file'] = self.meds_files[0]
        args['logfile'] = logfile

        scr = fmt.format(**args)

        script_name = self._get_nbrs_file(ext='.sh')

        print("writing:", script_name)
        with open(script_name, 'w') as fp:
            fp.write(scr)

        os.system('chmod 755 %s' % script_name)
        self.nbr_script = script_name


import tensorflow as tf
import tensorflow_addons as tfa
import numpy as np
import pkg_resources

from freedom.neural_nets.transformations import chargenet_trafo, hitnet_trafo, stringnet_trafo, layernet_trafo, domnet_trafo
#from freedom.neural_nets.domnet import combi_activation

class LLH():
    def __init__(self,
                 hitnet_file,
                 chargenet_file=None,
                 stringnet_file=None,
                 layernet_file=None,
                 domnet_file=None,
                 chargenet_batchsize=4096,
                 hitnet_batchsize=4096,
                 stringnet_batchsize=4096,
                 layernet_batchsize=4096,
                 domnet_batchsize=4096,
                 hitnet_trafo=hitnet_trafo,
                 chargenet_trafo=chargenet_trafo
                 ):
        '''
        *net_file : str
            location of * model hdf5 file
        *net_batchsize : int
        '''

        assert (chargenet_file is None) + (stringnet_file is None) + (layernet_file is None) + (domnet_file is None) == 3, 'Choose either chargenet OR stringnet OR layernet OR domnet'

        self.hitnet = tf.keras.models.load_model(hitnet_file, custom_objects={'hitnet_trafo':hitnet_trafo})
        # set to linear output = logit
        self.hitnet.layers[-1].activation = tf.keras.activations.linear
        self.hitnet.compile()

        if chargenet_file is not None:
            self.chargenet = tf.keras.models.load_model(chargenet_file, custom_objects={'chargenet_trafo':chargenet_trafo})
            self.chargenet.layers[-1].activation = tf.keras.activations.linear
            self.chargenet.compile()
            self.stringnet = None
            self.layernet = None
            self.domnet = None

        elif stringnet_file is not None:
            self.stringnet = tf.keras.models.load_model(stringnet_file, custom_objects={'stringnet_trafo':stringnet_trafo})
            self.stringnet.layers[-1].activation = tf.keras.activations.linear
            self.stringnet.compile()
            self.chargenet = None
            self.layernet = None
            self.domnet = None
            if '_reduced_' in stringnet_file:
                self.allowed_strings = np.load(pkg_resources.resource_filename('freedom', 'resources/allowed_strings.npy'))
            else:
                self.allowed_strings = np.arange(86)

        elif layernet_file is not None:
            self.layernet = tf.keras.models.load_model(layernet_file, custom_objects={'layernet_trafo':layernet_trafo})
            self.layernet.layers[-1].activation = tf.keras.activations.linear
            self.layernet.compile()
            self.chargenet = None
            self.stringnet = None
            self.domnet = None
            if '_reduced_' in layernet_file:
                self.allowed_layers = np.load(pkg_resources.resource_filename('freedom', 'resources/allowed_layers.npy'))
            else:
                self.allowed_layers = 'all'

        else:
            self.domnet = tf.keras.models.load_model(domnet_file, custom_objects={'domnet_trafo':domnet_trafo}) #, 'combi_activation':combi_activation
            self.domnet.layers[-1].activation = tf.keras.activations.linear
            self.domnet.compile()
            self.chargenet = None
            self.stringnet = None
            self.layernet = None
            if '_reduced_' in domnet_file:
                self.allowed_DOMs = np.load(pkg_resources.resource_filename('freedom', 'resources/allowed_DOMs.npy'))
            else:
                self.allowed_DOMs = np.arange(5160)

        self.chargenet_batchsize = chargenet_batchsize
        self.hitnet_batchsize = hitnet_batchsize
        self.stringnet_batchsize = stringnet_batchsize
        self.layernet_batchsize = layernet_batchsize
        self.domnet_batchsize = domnet_batchsize

    def __call__(self, event, params):
        """Evaluate LLH for a given event + params

        event : dict containing:
            'total_charge' : array
            'hits' : array shape (n_hits, 5)
                each row is (x, y, z) DOM poitions, time, charge, LC flag, ATWD flag
            'strings' : array shape (86, 5)
                each row is x, y, min(z), charge, nChannels
            'layers' : array shape (n_layers, 4)
                each row is nDOMs, z, charge, nChannels
            'doms' : array shape (5160, 4)
                each row is (x, y, z) DOM poitions, charge
        params : ndarray
            shape (n_likelihood_points, len(labels))

        Returns:
        --------
        total_llh : ndarray
        charge_llh : ndarray
        all_hits_llh : ndarray
        single_hit_llhs : ndarray

        """

        if params.ndim == 1:
            params = np.array([params])
        n_points = params.shape[0]

        # Charge Part
        if self.chargenet is not None:
            inputs = [np.repeat(event['total_charge'][np.newaxis, :], repeats=n_points, axis=0), params]
            charge_llhs = -self.chargenet.predict(inputs, batch_size=self.chargenet_batchsize)[:, 0]
            charge_llh = charge_llhs
        elif self.stringnet is not None:
            strings = event['strings'][self.allowed_strings]
            inputs = []
            inputs.append(np.repeat(strings[:, np.newaxis, :], repeats=n_points, axis=1).reshape(len(strings) * n_points, -1))
            inputs.append(np.repeat(params[np.newaxis, :], repeats=len(strings), axis=0).reshape(len(strings) * n_points, -1))

            charge_llhs = -self.stringnet.predict(inputs, batch_size=self.stringnet_batchsize).reshape(len(strings), n_points)
            charge_llh = np.sum(charge_llhs, axis=0)
        elif self.layernet is not None:
            if self.allowed_layers[0] == 'a':
                layers = event['layers']
            else:
                layers = event['layers'][self.allowed_layers]
            inputs = []
            inputs.append(np.repeat(layers[:, np.newaxis, :], repeats=n_points, axis=1).reshape(len(layers) * n_points, -1))
            inputs.append(np.repeat(params[np.newaxis, :], repeats=len(layers), axis=0).reshape(len(layers) * n_points, -1))

            charge_llhs = -self.layernet.predict(inputs, batch_size=self.layernet_batchsize).reshape(len(layers), n_points)
            charge_llh = np.sum(charge_llhs, axis=0)
        else:
            doms = event['doms'][self.allowed_DOMs]
            inputs = []
            inputs.append(np.repeat(doms[:, np.newaxis, :], repeats=n_points, axis=1).reshape(len(doms) * n_points, -1))
            inputs.append(np.repeat(params[np.newaxis, :], repeats=len(doms), axis=0).reshape(len(doms) * n_points, -1))

            charge_llhs = -self.domnet.predict(inputs, batch_size=self.domnet_batchsize).reshape(len(doms), n_points)
            charge_llh = np.sum(charge_llhs, axis=0)

        # Hit Net
        hits = event['hits']
        n_hits = hits.shape[0]

        if n_hits > 0:
            hit_charges = hits[:, 4]

            inputs = []
            inputs.append(np.repeat(hits[:, np.newaxis, :], repeats=n_points, axis=1).reshape(n_hits * n_points, -1))
            inputs.append(np.repeat(params[np.newaxis, :], repeats=n_hits, axis=0).reshape(n_hits * n_points, -1))

            single_hit_llhs = -self.hitnet.predict(inputs, batch_size=self.hitnet_batchsize).reshape(n_hits, n_points)

            all_hits_llh = hit_charges @ single_hit_llhs

        else:
            single_hit_llhs = np.array([])
            all_hits_llh = 0.

        total_llh = all_hits_llh + charge_llh

        return total_llh, charge_llh, all_hits_llh, charge_llhs, single_hit_llhs#.T

class upgrade_LLH():
    def __init__(self,
                 DOM_hitnet_file=None,
                 mDOM_hitnet_file=None,
                 DEgg_hitnet_file=None,
                 DOM_chargenet_file=None,
                 mDOM_chargenet_file=None,
                 DEgg_chargenet_file=None,
                 all_chargenet_file=None,
                 chargenet_batchsize=4096,
                 hitnet_batchsize=4096,
                 hitnet_trafo=hitnet_trafo,
                 chargenet_trafo=chargenet_trafo
                 ):
        '''
        *net_file : str
            location of * model hdf5 file
        *net_batchsize : int
        '''

        if DOM_hitnet_file is not None:
            self.DOM_hitnet = tf.keras.models.load_model(DOM_hitnet_file, custom_objects={'hitnet_trafo':hitnet_trafo})
            self.DOM_hitnet.layers[-1].activation = tf.keras.activations.linear
            self.DOM_hitnet.compile()
        else:
            self.DOM_hitnet = None

        if mDOM_hitnet_file is not None:
            self.mDOM_hitnet = tf.keras.models.load_model(mDOM_hitnet_file, custom_objects={'hitnet_trafo':hitnet_trafo})
            self.mDOM_hitnet.layers[-1].activation = tf.keras.activations.linear
            self.mDOM_hitnet.compile()
        else:
            self.mDOM_hitnet = None

        if DEgg_hitnet_file is not None:
            self.DEgg_hitnet = tf.keras.models.load_model(DEgg_hitnet_file, custom_objects={'hitnet_trafo':hitnet_trafo})
            self.DEgg_hitnet.layers[-1].activation = tf.keras.activations.linear
            self.DEgg_hitnet.compile()
        else:
            self.DEgg_hitnet = None

        if DOM_chargenet_file is not None:
            self.DOM_chargenet = tf.keras.models.load_model(DOM_chargenet_file, custom_objects={'chargenet_trafo':chargenet_trafo})
            self.DOM_chargenet.layers[-1].activation = tf.keras.activations.linear
            self.DOM_chargenet.compile()
        else:
            self.DOM_chargenet = None

        if mDOM_chargenet_file is not None:
            self.mDOM_chargenet = tf.keras.models.load_model(mDOM_chargenet_file, custom_objects={'chargenet_trafo':chargenet_trafo})
            self.mDOM_chargenet.layers[-1].activation = tf.keras.activations.linear
            self.mDOM_chargenet.compile()
        else:
            self.mDOM_chargenet = None

        if DEgg_chargenet_file is not None:
            self.DEgg_chargenet = tf.keras.models.load_model(DEgg_chargenet_file, custom_objects={'chargenet_trafo':chargenet_trafo})
            self.DEgg_chargenet.layers[-1].activation = tf.keras.activations.linear
            self.DEgg_chargenet.compile()
        else:
            self.DEgg_chargenet = None

        if all_chargenet_file is not None:
            self.all_chargenet = tf.keras.models.load_model(all_chargenet_file, custom_objects={'chargenet_trafo':chargenet_trafo})
            self.all_chargenet.layers[-1].activation = tf.keras.activations.linear
            self.all_chargenet.compile()
        else:
            self.all_chargenet = None

        self.chargenet_batchsize = chargenet_batchsize
        self.hitnet_batchsize = hitnet_batchsize

    def __call__(self, event, params):
        """Evaluate LLH for a given event + params

        event : dict containing:
            'total_charge' : array
            'hits' : array shape (n_hits, 5)
                each row is (x, y, z) DOM poitions, time, charge, LC flag, ATWD flag
        params : ndarray
            shape (n_likelihood_points, len(labels))

        Returns:
        --------
        total_llh : ndarray
        charge_llh : ndarray
        all_hits_llh : ndarray
        single_hit_llhs : ndarray

        """

        if params.ndim == 1:
            params = np.array([params])
        n_points = params.shape[0]

        # Charge Parts
        if self.DOM_chargenet is not None:
            inputs = [np.repeat(event['total_charge_DOM'][np.newaxis, :], repeats=n_points, axis=0), params]
            charge_llh_DOM = -self.DOM_chargenet.predict(inputs, batch_size=self.chargenet_batchsize)[:, 0]
        else:
            charge_llh_DOM = 0.

        if self.mDOM_chargenet is not None:
            inputs = [np.repeat(event['total_charge_mDOM'][np.newaxis, :], repeats=n_points, axis=0), params]
            charge_llh_mDOM = -self.mDOM_chargenet.predict(inputs, batch_size=self.chargenet_batchsize)[:, 0]
        else:
            charge_llh_mDOM = 0.

        if self.DEgg_chargenet is not None:
            inputs = [np.repeat(event['total_charge_DEgg'][np.newaxis, :], repeats=n_points, axis=0), params]
            charge_llh_DEgg = -self.DEgg_chargenet.predict(inputs, batch_size=self.chargenet_batchsize)[:, 0]
        else:
            charge_llh_DEgg = 0.

        if self.all_chargenet is not None:
            inp = np.stack([event['total_charge_DOM'], event['total_charge_mDOM'], event['total_charge_DEgg']]).reshape(6)
            inputs = [np.repeat(inp[np.newaxis, :], repeats=n_points, axis=0), params]
            charge_llh_all = -self.all_chargenet.predict(inputs, batch_size=self.chargenet_batchsize)[:, 0]
        else:
            charge_llh_all = 0.

        # Hit Nets
        hits = event['hits_DOM']
        n_hits = hits.shape[0]
        if n_hits > 0 and self.DOM_hitnet is not None:
            hit_charges = hits[:, 4]

            inputs = []
            inputs.append(np.repeat(hits[:, np.newaxis, :], repeats=n_points, axis=1).reshape(n_hits * n_points, -1))
            inputs.append(np.repeat(params[np.newaxis, :], repeats=n_hits, axis=0).reshape(n_hits * n_points, -1))

            single_hit_llhs_DOM = -self.DOM_hitnet.predict(inputs, batch_size=self.hitnet_batchsize).reshape(n_hits, n_points)

            all_hits_llh_DOM = hit_charges @ single_hit_llhs_DOM

        else:
            single_hit_llhs_DOM = np.array([])
            all_hits_llh_DOM = 0.

        hits = event['hits_mDOM']
        n_hits = hits.shape[0]
        if n_hits > 0 and self.mDOM_hitnet is not None:
            hit_charges = hits[:, 4]

            inputs = []
            inputs.append(np.repeat(hits[:, np.newaxis, :], repeats=n_points, axis=1).reshape(n_hits * n_points, -1))
            inputs.append(np.repeat(params[np.newaxis, :], repeats=n_hits, axis=0).reshape(n_hits * n_points, -1))

            single_hit_llhs_mDOM = -self.mDOM_hitnet.predict(inputs, batch_size=self.hitnet_batchsize).reshape(n_hits, n_points)

            all_hits_llh_mDOM = hit_charges @ single_hit_llhs_mDOM

        else:
            single_hit_llhs_mDOM = np.array([])
            all_hits_llh_mDOM = 0.

        hits = event['hits_DEgg']
        n_hits = hits.shape[0]
        if n_hits > 0 and self.DEgg_hitnet is not None:
            hit_charges = hits[:, 4]

            inputs = []
            inputs.append(np.repeat(hits[:, np.newaxis, :], repeats=n_points, axis=1).reshape(n_hits * n_points, -1))
            inputs.append(np.repeat(params[np.newaxis, :], repeats=n_hits, axis=0).reshape(n_hits * n_points, -1))

            single_hit_llhs_DEgg = -self.DEgg_hitnet.predict(inputs, batch_size=self.hitnet_batchsize).reshape(n_hits, n_points)

            all_hits_llh_DEgg = hit_charges @ single_hit_llhs_DEgg

        else:
            single_hit_llhs_DEgg = np.array([])
            all_hits_llh_DEgg = 0.

        total_llh_DOM = all_hits_llh_DOM + charge_llh_DOM
        total_llh_mDOM = all_hits_llh_mDOM + charge_llh_mDOM
        total_llh_DEgg = all_hits_llh_DEgg + charge_llh_DEgg
        total_llh = total_llh_DOM + total_llh_mDOM + total_llh_DEgg #+ charge_llh_all

        return total_llh, total_llh_DOM, total_llh_mDOM, total_llh_DEgg #, charge_llh_all


import os, sys
base_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')
sys.path.append(base_dir)
sys.path.append(os.path.join(base_dir, 'externals/baselines'))
sys.path.append(os.path.join(base_dir, 'externals/pytorch-a2c-ppo-acktr-gail'))

import numpy as np
from collections import deque
from copy import deepcopy
import time
import torch

import gym
# import a2c_ppo_acktr
from a2c_ppo_acktr import algo, utils
from a2c_ppo_acktr.envs import make_vec_envs, make_env
from a2c_ppo_acktr.model import Policy
from a2c_ppo_acktr.storage import RolloutStorage

from sample import Sample

'''
Evaluate a policy sample.
'''
def evaluation(args, sample):
    eval_env = gym.make(args.env_name)
    objs = np.zeros(args.obj_num)
    ob_rms = sample.env_params['ob_rms']
    policy = sample.actor_critic
    with torch.no_grad():
        for eval_id in range(args.eval_num):
            eval_env.seed(args.seed + eval_id)
            ob = eval_env.reset()
            done = False
            gamma = 1.0
            while not done:
                if args.ob_rms:
                    ob = np.clip((ob - ob_rms.mean) / np.sqrt(ob_rms.var + 1e-8), -10.0, 10.0)
                _, action, _, _ = policy.act(torch.Tensor(ob).unsqueeze(0), None, None, deterministic=True)
                ob, _, done, info = eval_env.step(action)
                objs += gamma * info['obj']
                if not args.raw:
                    gamma *= args.gamma
    eval_env.close()
    objs /= args.eval_num
    return objs

'''
define a MOPG Worker.
Input:
    args: the arguments include necessary ppo parameters.
    task_id: the task_id in all parallel executed tasks.
    device: torch device
    iteration: starting iteration number
    num_updates: number of rl iterations to run.
    start_time: starting time
    results_queue: multi-processing queue to pass the optimized policies back to main process.
    done_event: multi-processing event for process synchronization.
'''
def MOPG_worker(args, task_id, task, device, iteration, num_updates, start_time, results_queue, done_event):
    scalarization = task.scalarization
    env_params, actor_critic, agent = task.sample.env_params, task.sample.actor_critic, task.sample.agent

    weights_str = (args.obj_num * '_{:.3f}').format(*task.scalarization.weights)

    # make envs
    envs = make_vec_envs(env_name=args.env_name, seed=args.seed, num_processes=args.num_processes, \
                        gamma=args.gamma, log_dir=None, device=device, allow_early_resets=False, \
                        obj_rms=args.obj_rms, ob_rms = args.ob_rms)
    if env_params['ob_rms'] is not None:
        envs.venv.ob_rms = deepcopy(env_params['ob_rms'])
    if env_params['ret_rms'] is not None:
        envs.venv.ret_rms = deepcopy(env_params['ret_rms'])
    if env_params['obj_rms'] is not None:
        envs.venv.obj_rms = deepcopy(env_params['obj_rms'])

    # build rollouts data structure
    rollouts = RolloutStorage(num_steps = args.num_steps, num_processes = args.num_processes,
                              obs_shape = envs.observation_space.shape, action_space = envs.action_space,
                              recurrent_hidden_state_size = actor_critic.recurrent_hidden_state_size, obj_num=args.obj_num)
    obs = envs.reset()
    rollouts.obs[0].copy_(obs)
    rollouts.to(device)

    episode_rewards = deque(maxlen=10)
    episode_lens = deque(maxlen=10)
    episode_objs = deque(maxlen=10)   # for each cost component we care
    episode_obj = np.array([None] * args.num_processes)

    total_num_updates = int(args.num_env_steps) // args.num_steps // args.num_processes

    offspring_batch = []

    start_iter, final_iter = iteration, min(iteration + num_updates, total_num_updates)
    for j in range(start_iter, final_iter):
        torch.manual_seed(j)
        if args.use_linear_lr_decay:
            # decrease learning rate linearly
            utils.update_linear_schedule( \
                agent.optimizer, j * args.lr_decay_ratio, \
                total_num_updates, args.lr)

        for step in range(args.num_steps):
            # Sample actions
            with torch.no_grad():
                value, action, action_log_prob, recurrent_hidden_states = actor_critic.act(
                    rollouts.obs[step], rollouts.recurrent_hidden_states[step],
                    rollouts.masks[step])

            obs, _, done, infos = envs.step(action)
            obj_tensor = torch.zeros([args.num_processes, args.obj_num])

            for idx, info in enumerate(infos):
                obj_tensor[idx] = torch.from_numpy(info['obj'])
                episode_obj[idx] = info['obj_raw'] if episode_obj[idx] is None else episode_obj[idx] + info['obj_raw']
                if 'episode' in info.keys():
                    episode_rewards.append(info['episode']['r'])
                    episode_lens.append(info['episode']['l'])
                    if episode_obj[idx] is not None:
                        episode_objs.append(episode_obj[idx])
                        episode_obj[idx] = None

            # If done then clean the history of observations.
            masks = torch.FloatTensor(
                [[0.0] if done_ else [1.0] for done_ in done])
            bad_masks = torch.FloatTensor(
                [[0.0] if 'bad_transition' in info.keys() else [1.0]
                 for info in infos])
            rollouts.insert(obs, recurrent_hidden_states, action,
                            action_log_prob, value, obj_tensor, masks, bad_masks)

        with torch.no_grad():
            next_value = actor_critic.get_value(
                rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
                rollouts.masks[-1]).detach()

        rollouts.compute_returns(next_value, args.use_gae, args.gamma,
                                 args.gae_lambda, args.use_proper_time_limits)

        obj_rms_var = envs.obj_rms.var if envs.obj_rms is not None else None

        value_loss, action_loss, dist_entropy = agent.update(rollouts, scalarization, obj_rms_var)

        rollouts.after_update()

        env_params = {}
        env_params['ob_rms'] = deepcopy(envs.ob_rms) if envs.ob_rms is not None else None
        env_params['ret_rms'] = deepcopy(envs.ret_rms) if envs.ret_rms is not None else None
        env_params['obj_rms'] = deepcopy(envs.obj_rms) if envs.obj_rms is not None else None

        # evaluate new sample
        sample = Sample(env_params, deepcopy(actor_critic), deepcopy(agent))
        objs = evaluation(args, sample)
        sample.objs = objs
        offspring_batch.append(sample)

        if args.rl_log_interval > 0 and (j + 1) % args.rl_log_interval == 0 and len(episode_rewards) > 1:
            if task_id == 0:
                total_num_steps = (j + 1) * args.num_processes * args.num_steps
                end = time.time()
                print(
                    "[RL] Updates {}, num timesteps {}, FPS {}, time {:.2f} seconds"
                    .format(j + 1, total_num_steps,
                            int(total_num_steps / (end - start_time)),
                            end - start_time))

        # put results back every update_iter iterations, to avoid the multi-processing crash
        if (j + 1) % args.update_iter == 0 or j == final_iter - 1:
            offspring_batch = np.array(offspring_batch)
            results = {}
            results['task_id'] = task_id
            results['offspring_batch'] = offspring_batch
            if j == final_iter - 1:
                results['done'] = True
            else:
                results['done'] = False
            results_queue.put(results)
            offspring_batch = []

    envs.close()

    done_event.wait()


#!/usr/bin/env python3

import numpy as np
import cv2

import os
import time

import cv2
import pycuda.autoinit  # For initializing CUDA driver
import pycuda.driver as cuda

from utils.yolo_classes import get_cls_dict
from utils.display import open_window, set_display, show_fps
from utils.visualization import BBoxVisualization
from utils.yolo_with_plugins import TrtYOLO

import rospy
import rospkg
from yolov4_trt_ros.msg import Detector2DArray
from yolov4_trt_ros.msg import Detector2D
from vision_msgs.msg import BoundingBox2D
from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError

class yolov4(object):
    def __init__(self):
        """ Constructor """
        self.init_params()
        self.init_yolo()
        self.cuda_ctx = cuda.Device(0).make_context()
        self.trt_yolo = TrtYOLO(
            (self.model_path + self.model), (self.h, self.w), self.category_num)
        self.tic = time.time()

    def __del__(self):
        """ Destructor """

        self.cuda_ctx.pop()
        del self.trt_yolo
        del self.cuda_ctx

    def clean_up(self):
        """ Backup destructor: Release cuda memory """

        if self.trt_yolo is not None:
            self.cuda_ctx.pop()
            del self.trt_yolo
            del self.cuda_ctx

    def init_params(self):
        """ Initializes ros parameters """

        rospack = rospkg.RosPack()
        package_path = rospack.get_path("yolov4_trt_ros")
        self.model = rospy.get_param("/model", "yolov4Custom")
        self.model_path = rospy.get_param(
            "/model_path", package_path + "/yolo/")
        self.category_num = rospy.get_param("/category_number", 80)
        self.input_shape = rospy.get_param("/input_shape", "416")
        self.conf_th = rospy.get_param("/confidence_threshold", 0.5)
        self.show_img = rospy.get_param("/show_image", True)

    def init_yolo(self):
        """ Initialises yolo parameters required for trt engine """

        if self.model.find('-') == -1:
            self.model = self.model + "-" + self.input_shape

        yolo_dim = self.model.split('-')[-1]

        if 'x' in yolo_dim:
            dim_split = yolo_dim.split('x')
            if len(dim_split) != 2:
                raise SystemExit('ERROR: bad yolo_dim (%s)!' % yolo_dim)
            self.w, self.h = int(dim_split[0]), int(dim_split[1])
        else:
            self.h = self.w = int(yolo_dim)
        if self.h % 32 != 0 or self.w % 32 != 0:
            raise SystemExit('ERROR: bad yolo_dim (%s)!' % yolo_dim)

        cls_dict = get_cls_dict(self.category_num)
        self.vis = BBoxVisualization(cls_dict)

    def detect(self, cv_img):
        boxes, confs, clss = self.trt_yolo.detect(cv_img, self.conf_th)

        cv_img = self.vis.draw_bboxes(cv_img, boxes, confs, clss)
        self.toc = time.time()
        fps = 1.0 / (self.toc - self.tic)
        self.tic = self.toc

        self.publish_detections(boxes, confs, clss)

        if self.show_img:
            cv_img = show_fps(cv_img, fps)
            cv2.imshow("YOLOv4 DETECTION RESULTS", cv_img)
            if cv2.waitKey(0) == ord("q"):
                cv2.destroyAllWindows()

    def publish_detections(self, boxes, confs, clss):
        detection2d = Detector2DArray()
        detection = Detector2D()
        for i in range(len(boxes)):
            # boxes : xmin, ymin, xmax, ymax
            for _ in boxes:
                detection.header.stamp = rospy.Time.now()
                detection.header.frame_id = "camera" # change accordingly
                detection.results.id = clss[i]
                detection.results.score = confs[i]

                detection.bbox.center.x = boxes[i][0] + (boxes[i][2] - boxes[i][0])/2
                detection.bbox.center.y = boxes[i][1] + (boxes[i][3] - boxes[i][1])/2
                detection.bbox.center.theta = 0.0  # change if required

                detection.bbox.size_x = abs(boxes[i][0] - boxes[i][2])
                detection.bbox.size_y = abs(boxes[i][1] - boxes[i][3])

                detection2d.detections.append(detection)

        return detection2d

image = cv2.imread('/home/michelle/Downloads/st_guy.jpeg')

def main():
    yolo = yolov4()
    rospy.init_node('yolov4_toy', anonymous=True)
    image = cv2.imread('/home/michelle/Downloads/st_guy.jpeg')
    detectionList = yolo.detect(image)
    print('hi')
    print(detectionList)
    try:
        rospy.spin()
    except KeyboardInterrupt:
        rospy.on_shutdown(yolo.clean_up())
        print("Shutting down")

if __name__ == '__main__':
    main()