from visual_relationship_dataset import *
import os
import scipy.io as sio
from PIL import Image
import copy
from refine_predictions import refine_equiv

np.set_printoptions(precision=2)
np.set_printoptions(threshold=np.inf)

cwd = os.getcwd()

# swith between GPU and CPU
config = tf.ConfigProto(device_count={'GPU': 1})

img_dir = cwd + '/data/sg_test_images'

# Load training data for prio statistics on the dataset
_, _, _, triples_of_train_data, _, _, _ = get_data("train", True)

# Computing positive and negative examples for predicates and types
idxs_of_positive_examples_of_predicates = {}

for predicate in selected_predicates:
    idxs_of_positive_examples_of_predicates[predicate] = \
        np.where(predicates[triples_of_train_data[:, -1]] == predicate)[0]

prior_stats = np.array([len(idxs_of_positive_examples_of_predicates[pred]) for pred in selected_predicates])
prior_freq = np.true_divide(prior_stats, np.sum(prior_stats))

image_path = sio.loadmat(cwd + '/Visual-Relationship-Detection-master/data/imagePath.mat')
object_detection = sio.loadmat(cwd + '/Visual-Relationship-Detection-master/data/objectDetRCNN.mat')
detection_bboxes = object_detection['detection_bboxes']
detection_labels = object_detection['detection_labels']
detection_confs = object_detection['detection_confs']

for img_id in range(len(detection_bboxes[0])):
    if len(detection_bboxes[0][img_id]) > 0:
        assert np.all(detection_bboxes[0][img_id][:, 0] < detection_bboxes[0][img_id][:, 2])
        assert np.all(detection_bboxes[0][img_id][:, 1] < detection_bboxes[0][img_id][:, 3])

features_detected_bb = []
obj_bboxes_ours = []
sub_bboxes_ours = []
obj_labels_ours = []
sub_labels_ours = []

semantic_feat_vect = np.zeros(len(types))

for pic_idx in range(detection_bboxes.shape[1]):
    detection_bboxes[0, pic_idx] = detection_bboxes[0, pic_idx].astype(np.float)
    features_per_image = np.empty((0, 2 * number_of_features + number_of_extra_features))
    obj_bboxes_ours_per_image = np.array([]).reshape(0, 4)
    sub_bboxes_ours_per_image = np.array([]).reshape(0, 4)
    obj_label_per_image = np.array([])
    sub_label_per_image = np.array([])

    # normalize data
    if len(detection_bboxes[0, pic_idx]) > 0:
        img = Image.open(os.path.join(img_dir, image_path['imagePath'][0, pic_idx][0]).replace('png', 'jpg'))
        width, height = img.size
        normalized_detection_bboxes = copy.deepcopy(detection_bboxes)
        copy.deepcopy
        normalized_detection_bboxes[0, pic_idx][:, -4] /= width
        normalized_detection_bboxes[0, pic_idx][:, -3] /= height
        normalized_detection_bboxes[0, pic_idx][:, -2] /= width
        normalized_detection_bboxes[0, pic_idx][:, -1] /= height

    for bb1_idx in range(len(detection_bboxes[0, pic_idx])):
        for bb2_idx in range(len(detection_bboxes[0, pic_idx])):
            if bb1_idx != bb2_idx:
                bb1 = normalized_detection_bboxes[0, pic_idx][bb1_idx]
                bb2 = normalized_detection_bboxes[0, pic_idx][bb2_idx]
                sub_label_per_image = np.append(sub_label_per_image, detection_labels[0, pic_idx][bb1_idx, 0])
                obj_label_per_image = np.append(obj_label_per_image, detection_labels[0, pic_idx][bb2_idx, 0])

                feat_vect_bb1 = np.hstack((semantic_feat_vect, bb1))
                feat_vect_bb2 = np.hstack((semantic_feat_vect, bb2))
                feat_vect_bb1[detection_labels[0, pic_idx][bb1_idx]] = detection_confs[0, pic_idx][bb1_idx]
                feat_vect_bb2[detection_labels[0, pic_idx][bb2_idx]] = detection_confs[0, pic_idx][bb2_idx]
                feat_vec_pair = np.hstack((feat_vect_bb1, feat_vect_bb2, computing_extended_features(bb1, bb2)))

                features_per_image = np.vstack((features_per_image, feat_vec_pair[np.newaxis, :]))
                sub_bboxes_ours_per_image = np.vstack(
                    (sub_bboxes_ours_per_image, detection_bboxes[0, pic_idx][bb1_idx]))
                obj_bboxes_ours_per_image = np.vstack(
                    (obj_bboxes_ours_per_image, detection_bboxes[0, pic_idx][bb2_idx]))

    features_detected_bb.append(features_per_image)
    obj_bboxes_ours.append(obj_bboxes_ours_per_image)
    sub_bboxes_ours.append(sub_bboxes_ours_per_image)
    obj_labels_ours.append(obj_label_per_image)
    sub_labels_ours.append(sub_label_per_image)

model_list = [
    "models/RWFN_KB_wc_" + str(NUM_ITERATION_TRAIN) + ".ckpt",
    "models/RWFN_KB_nc_" + str(NUM_ITERATION_TRAIN) + ".ckpt"]

for model_type in model_list:

    model = model_type
    obj_bboxes_ours_output = []
    sub_bboxes_ours_output = []
    model_label = model.split("/")[-1][:-5]
    print model.upper()
    predicted_predicates_values_tensor = tf.concat(
        [isInRelation[predicate].tensor() for predicate in selected_predicates], 1)
    saver = tf.train.Saver()
    sess = tf.Session(config=config)
    saver.restore(sess, model)
    rlp_confs_ours = []
    rlp_labels_ours = []

    for pic_idx in range(detection_bboxes.shape[1]):
        if pic_idx % 100 == 0:
            print "Eval img", pic_idx
        values_of_predicates = np.array([], dtype=np.float32).reshape(0, 70)
        values_of_predicates = sess.run(predicted_predicates_values_tensor,
                                        {pairs_of_objects.tensor: features_detected_bb[pic_idx]})

        values_of_predicates = refine_equiv(values_of_predicates, selected_predicates, "max")
        values_of_predicates = np.multiply(values_of_predicates, prior_freq)

        conf_predicates_per_image = values_of_predicates.flatten('F')
        sub_bboxes_ours_output.append(np.tile(sub_bboxes_ours[pic_idx], (len(selected_predicates), 1)))
        obj_bboxes_ours_output.append(np.tile(obj_bboxes_ours[pic_idx], (len(selected_predicates), 1)))
        label_predicates_per_image = np.hstack(
            (np.tile(sub_labels_ours[pic_idx], len(selected_predicates))[:, np.newaxis],
             np.repeat(np.array(range(1, len(selected_predicates) + 1)), len(features_detected_bb[pic_idx]))[:,
             np.newaxis],
             np.tile(obj_labels_ours[pic_idx], len(selected_predicates))[:, np.newaxis]))

        rlp_confs_ours.append(conf_predicates_per_image[:, np.newaxis])
        rlp_labels_ours.append(label_predicates_per_image)

    sess.close()

    sio.savemat(
        cwd + "/Visual-Relationship-Detection-master/results_RWFN/relationship_det_result_" + model_label + ".mat",
        {'sub_bboxes_ours': sub_bboxes_ours_output,
         'obj_bboxes_ours': obj_bboxes_ours_output,
         'rlp_confs_ours': rlp_confs_ours,
         'rlp_labels_ours': rlp_labels_ours})


import argparse
import torch
import torchvision
import torchvision.transforms as transforms
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import skimage.io as io
import numpy as np
import os

import models
import model_deform

parser = argparse.ArgumentParser()
parser.add_argument("--img0Path", type=str, default='./images/f_032242.png', help='input image0 path')
parser.add_argument("--img1Path", type=str, default='./images/f_032246.png', help='input image1 path')
parser.add_argument("--sketPath", type=str, default='./images/f_032244_sket.png', help='input sketch_t path')
parser.add_argument("--saveDir", type=str, default='./images', help='saved image It path')
parser.add_argument("--modelPath", type=str, default='./checkpoints/full_model.ckpt', help='checkpoint path')
parser.add_argument("--frm_num", type=int, default=5, help='number of frame to be interpolated')
args = parser.parse_args()

torch.set_grad_enabled(False)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
torch.cuda.set_device(0)
multiGPUs = [0]

netT   = models.ResNet()
sketExt = models.PWCExtractor()
imagExt = models.PWCExtractor()
flowEst = models.Network()
blenEst = models.blendNet()
flowRef = models.UNet(14, 8)
ImagRef = model_deform.DeformUNet(21, 15)

W = 576
H = 384
flowBackWarp = models.backWarp(W, H)
occlusiCheck = models.occlusionCheck(W, H)

if torch.cuda.device_count() >= 1:
    netT = nn.DataParallel(netT, device_ids=multiGPUs)
    sketExt = nn.DataParallel(sketExt, device_ids=multiGPUs)
    imagExt = nn.DataParallel(imagExt, device_ids=multiGPUs)
    flowEst = nn.DataParallel(flowEst, device_ids=multiGPUs)
    blenEst = nn.DataParallel(blenEst, device_ids=multiGPUs)
    flowRef = nn.DataParallel(flowRef, device_ids=multiGPUs)
    ImagRef = nn.DataParallel(ImagRef, device_ids=multiGPUs)

    flowBackWarp = nn.DataParallel(flowBackWarp, device_ids=multiGPUs)
    occlusiCheck = nn.DataParallel(occlusiCheck, device_ids=multiGPUs)

netT = netT.to(device)
sketExt = sketExt.to(device)
imagExt = imagExt.to(device)
flowEst = flowEst.to(device)
blenEst = blenEst.to(device)
flowRef = flowRef.to(device)
ImagRef = ImagRef.to(device)

flowBackWarp = flowBackWarp.to(device)
occlusiCheck = occlusiCheck.to(device)

model_dict = torch.load(args.modelPath, map_location='cuda:0')
for key, value in model_dict.items():
    print(key)

netT.load_state_dict(model_dict['state_dict_netT'])
sketExt.load_state_dict(model_dict['state_dict_sketExt'])
imagExt.load_state_dict(model_dict['state_dict_imagExt'])
flowEst.load_state_dict(model_dict['state_dict_flowEst'])
blenEst.load_state_dict(model_dict['state_dict_blenEst'])
flowRef.load_state_dict(model_dict['state_dict_flowRef'])
ImagRef.load_state_dict(model_dict['state_dict_ImagRef'])

transform = transforms.Compose([transforms.ToTensor()])
revtransf = transforms.Compose([transforms.ToPILImage()])

L1_lossFn = models.L1Loss()

img0 = io.imread(args.img0Path).astype(np.float32)/255.0
img1 = io.imread(args.img1Path).astype(np.float32)/255.0
sket = io.imread(args.sketPath)[np.newaxis, :, :].astype(np.float32)/255.0

img0 = transform(img0).unsqueeze(0).to(device)
img1 = transform(img1).unsqueeze(0).to(device)
sket = torch.from_numpy(sket).unsqueeze(0).to(device)

# flow estimation
imgt_temp = netT(torch.cat((sket, img0, img1), dim = 1))
featSkt = sketExt(imgt_temp)
featIg0 = imagExt(img0)
featIg1 = imagExt(img1)

f_t0 = flowEst(featSkt, featIg0)
f_t1 = flowEst(featSkt, featIg1)

f_0t = flowEst(featIg0, featSkt)
f_1t = flowEst(featIg1, featSkt)

I_t0 = flowBackWarp(img0, f_t0)
I_t1 = flowBackWarp(img1, f_t1)

O_t0 = occlusiCheck(f_t0, f_0t)
O_t1 = occlusiCheck(f_t1, f_1t)

W_0 = blenEst(torch.cat((I_t0, I_t1, O_t0, O_t1, sket), dim = 1))

W_1 = 1 - W_0
I_t = W_0 * I_t0 + W_1 * I_t1

imgall = []
frm_list = [x / (args.frm_num + 1) for x in range(args.frm_num + 2)]

for frm_index, i in enumerate(frm_list):
    if i == 0:
        imgall.append(img0)
    if i == 1:
        imgall.append(img1)
    if i == 0.5:
        imgall.append(I_t)

    if i != 0 and i < 0.5:
        t = i * 2
        co_eff = [-t * (1 - t), t * t, (1 - t) * (1 - t), -t * (1 - t)]
        f_lr = f_0t
        f_rl = f_t0

        f_k0 = co_eff[0] * f_lr + co_eff[1] * f_rl
        f_kt = co_eff[2] * f_lr + co_eff[3] * f_rl
        f_0k = t * f_0t
        f_tk = (1 - t) * f_t0

        I_k0 = flowBackWarp(img0, f_k0)
        I_kt = flowBackWarp(I_t, f_kt)

        refflows = flowRef(torch.cat((I_k0, I_kt, f_k0, f_0k, f_kt, f_tk), dim = 1))

        f_k0 = refflows[:, 0:2, :, :] + f_k0
        f_0k = refflows[:, 2:4, :, :] + f_0k
        f_kt = refflows[:, 4:6, :, :] + f_kt
        f_tk = refflows[:, 6:8, :, :] + f_tk

        I_k0 = flowBackWarp(img0, f_k0)
        I_kt = flowBackWarp(I_t, f_kt)

        O_k0 = occlusiCheck(f_k0, f_0k)
        O_kt = occlusiCheck(f_kt, f_tk)

        W_0 = blenEst(torch.cat((I_k0, I_kt, O_k0, O_kt, flowBackWarp(sket, f_kt)), dim = 1))
        W_1 = 1 - W_0
        I_k = W_0 * I_k0 + W_1 * I_kt

        imgall.append(I_k)

    if i != 1 and i > 0.5:
        t = (i - 0.5) * 2
        co_eff = [-t * (1 - t), t * t, (1 - t) * (1 - t), -t * (1 - t)]
        f_lr = f_t1
        f_rl = f_1t

        f_kt = co_eff[0] * f_lr + co_eff[1] * f_rl
        f_k1 = co_eff[2] * f_lr + co_eff[3] * f_rl
        f_tk = t * f_t1
        f_1k = (1 - t) * f_1t

        I_kt = flowBackWarp(I_t, f_kt)
        I_k1 = flowBackWarp(img1, f_k1)

        refflows = flowRef(torch.cat((I_k1, I_kt, f_k1, f_1k, f_kt, f_tk), dim = 1))

        f_k1 = refflows[:, 0:2, :, :] + f_k1
        f_1k = refflows[:, 2:4, :, :] + f_1k
        f_kt = refflows[:, 4:6, :, :] + f_kt
        f_tk = refflows[:, 6:8, :, :] + f_tk

        I_kt = flowBackWarp(I_t, f_kt)
        I_k1 = flowBackWarp(img1, f_k1)
        O_kt = occlusiCheck(f_kt, f_tk)
        O_k1 = occlusiCheck(f_k1, f_1k)

        W_0 = blenEst(torch.cat((I_kt, I_k1, O_kt, O_k1, flowBackWarp(sket, f_kt)), dim = 1))
        W_1 = 1 - W_0
        I_k = W_0 * I_kt + W_1 * I_k1

        imgall.append(I_k)

imgref = torch.clamp(ImagRef(torch.cat(imgall, dim = 1)) + torch.cat(imgall[1:-1], dim = 1), 0, 1)

for i in range(len(imgall) - 2):
    io.imsave('%s%s%s%s' % (args.saveDir, '/f_output_', str(i + 1), '.png'), np.asarray(revtransf(imgref[:,(i*3):(i+1)*3,:,:].cpu()[0])))


from __future__  import print_function
import argparse
parser = argparse.ArgumentParser(description='Fast Neural Style - generator')
parser.add_argument('--model', default='./model/pre_trained.model', type=str)
parser.add_argument('--input', type=str, help='input image')
parser.add_argument('--output', type=str, help='output image')
args = parser.parse_args()

import numpy as np
from PIL import Image
from datetime import datetime

import torch
from torch.autograd import Variable

import net

use_cuda = torch.cuda.is_available()

style_model = net.StylePart()
style_model.load_state_dict(torch.load(args.model))

# Origin Image
if use_cuda:
	style_model.cuda()

print('=> Load Origin Image')
start_time = datetime.now().strftime('%H:%M:%S')
img = Image.open(args.input)
img = np.array(img)  # PIL->numpy
img = np.array(img[..., ::-1])  # RGB->BGR
img = img.transpose(2, 0, 1)  # HWC->CHW
img = img.reshape((1, ) + img.shape)  # CHW->BCHW
img = torch.from_numpy(img).float()
img = Variable(img, volatile=True)

if use_cuda:
	img = img.cuda()

print('=> Generating')
style_model.eval()
output = style_model(img)
end_time = datetime.now().strftime('%H:%M:%S')
print('=> Save ')
output = output.data.cpu().clamp(0, 255).byte().numpy()
output = output[0].transpose((1, 2, 0))
output = output[..., ::-1]
output = Image.fromarray(output)
output.save(args.output)
print('Start Time: %s\nEnd Time: %s'%(start_time,end_time))


# Copyright 2014-2016 Nervana Systems Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Definition of the GPU layers
These layers are mainly used for old benchmarking code,
but they also cache all the computed params for complex layers.
TODO: clean up merge with CPU layers
TODO: remove any non-param caching code, neon layers should replace benchmark code.
"""
from __future__ import division
from builtins import object, str
import logging
from math import ceil
import numpy as np
from operator import mul
import pycuda.driver as drv
from pycuda.tools import context_dependent_memoize
import sys
from neon import logger as neon_logger
from neon.backends import kernel_specs
from neon.backends import convolution

logger = logging.getLogger(__name__)

if sys.version_info >= (3, 0):
    from functools import reduce

class Layer(object):

    """
    GPU Layer base class
    """

    def __init__(self, lib, dtype, N, dtypeU=None):

        if hasattr(dtype, 'type'):
            self.dtype = dtype
        else:
            self.dtype = np.dtype(dtype)

        self.N      = N
        self.dtypeU = dtype if dtypeU is None else dtypeU
        self.lib    = lib
        self.flops  = 0
        self.sizeI  = 0
        self.sizeO  = 0
        self.sizeF  = 0
        self.weights   = None
        self.fprop_in  = None
        self.fprop_out = None
        self.bprop_in  = None
        self.bprop_out = None

        self.learning_rate = 0.0

    def init_activations(self, fprop_out=None):

        if fprop_out is not None:
            self.fprop_out = fprop_out
        else:
            self.fprop_out = self.lib.empty(self.dimO, dtype=self.dtype)

        self.act_stats = self.lib.empty((self.dimO2[0], 1), dtype=np.float32)

    def init_deltas(self, shared=None):

        if shared is None:
            self.bprop_out = self.lib.empty(self.dimI, dtype=self.dtype)
        else:
            self.bprop_out = shared[0].share(self.dimI)
            shared.reverse()

        self.delta_stats = self.lib.empty((self.dimI2[0], 1), dtype=np.float32)

    def init_weights(self, loc=0.0, scale=0.1, shared=None, zeros=False):

        if self.sizeF > 0:
            if zeros:
                self.weights  = self.lib.zeros(self.dimF, dtype=self.dtype)
            else:
                weights       = np.random.normal(loc, scale, self.dimF)
                self.weights  = self.lib.array(weights, dtype=self.dtype)

            if shared is None:
                self.updat_out = self.lib.empty(self.dimF, dtype=self.dtypeU)
            else:
                self.updat_out = shared.share(self.dimF, dtype=self.dtypeU)

            self.weight_stats = self.lib.empty((self.dimF2[0], 1), dtype=np.float32)

    def scale_weights(self, scale):

        mean = self.get_activation_mean()
        self.weights[:] *= scale / mean

    def fprop(self, fprop_in, scale_weights=0):
        if self.fprop_in is None and fprop_in:
            self.fprop_in = fprop_in.reshape(self.dimI)
        return self.fprop_in

    def bprop(self, bprop_in, beta=0):
        return bprop_in

    # fprop relu happens inside of the conv and gemm kernels
    def bprop_relu(self, bprop_in):

        bprop_in[:] = bprop_in * (self.fprop_out > 0)
        return bprop_in

    def grad_descent(self):
        self.weights[:] += self.updat_out * self.learning_rate

    def get_activation_mean(self):
        return self._get_mean(self.fprop_out, self.act_stats, self.dimO2)

    def get_activation_max(self):
        return self._get_max(self.fprop_out, self.act_stats, self.dimO2)

    def get_delta_mean(self):
        return self._get_mean(self.bprop_out, self.delta_stats, self.dimI2)

    def get_delta_max(self):
        return self._get_max(self.bprop_out, self.delta_stats, self.dimI2)

    def get_update_mean(self):
        if self.sizeF > 0:
            return self._get_mean(self.updat_out, self.weight_stats, self.dimF2)
        return 0

    def get_update_max(self):
        if self.sizeF > 0:
            return self._get_max(self.updat_out, self.weight_stats, self.dimF2)
        return 0

    def get_weight_mean(self):
        if self.sizeF > 0:
            return self._get_mean(self.weights, self.weight_stats, self.dimF2)
        return 0

    def get_weight_max(self):
        if self.sizeF > 0:
            return self._get_max(self.weights, self.weight_stats, self.dimF2)
        return 0

    def _get_mean(self, ary, buf, shape):

        buf1    = buf[0:1, 0:1]
        buf[:]  = self.lib.sum(abs(ary.reshape(shape)), axis=1)
        buf1[:] = self.lib.sum(buf, axis=0) * (1.0 / ary.size)
        return float(buf1.get()[0, 0])

    def _get_max(self, ary, buf, shape):

        buf1    = buf[0:1, 0:1]
        buf[:]  = self.lib.max(abs(ary.reshape(shape)), axis=1)
        buf1[:] = self.lib.max(buf, axis=0)
        return float(buf1.get()[0, 0])

    def fprop_stats(self):
        neon_logger.display("fprop:%10.5f mean %11.5f max %s"
              % (self.get_activation_mean(), self.get_activation_max(), self))

    def bprop_stats(self):
        if self.bprop_out is not None:
            neon_logger.display("bprop:%10.5f mean %11.5f max %s"
                  % (self.get_delta_mean(), self.get_delta_max(), self))

        if self.weights is not None:
            up_mean, up_max = (self.get_update_mean(), self.get_update_max())
            wt_mean, wt_max = (self.get_weight_mean(), self.get_weight_max())
            rt_mean, rt_max = (0.0001 * up_mean / wt_mean, 0.0001 * up_max / wt_max)
            neon_logger.display("updat:%10.5f mean %11.5f max %s" % (up_mean, up_max, self))
            neon_logger.display("weigh:%10.5f mean %11.5f max" % (wt_mean, wt_max))
            neon_logger.display("ratio:%10.5f mean %11.5f max" % (rt_mean, rt_max))

    @staticmethod
    def create(lib, conf, prev_layer, dtype):

        config     = dict(conf)
        layer_type = config.pop("layer")

        # merge dtype specific settings
        config["dtype"] = dtype

        # merge shared params
        config.update(config.pop("common", {}))

        # Propagate the fixed and calculated dimensions
        if prev_layer is not None:
            config["N"] = prev_layer.N

            if layer_type is FullLayer:
                config["nIn"] = prev_layer.nOut
            elif layer_type is PoolLayer and type(prev_layer) is FullLayer:
                config["C"] = prev_layer.nOut
            elif layer_type is BatchNorm and type(prev_layer) is FullLayer:
                config["nIn"] = prev_layer.nOut
            else:
                config["C"] = prev_layer.K
                config["D"] = prev_layer.M
                config["H"] = prev_layer.P
                config["W"] = prev_layer.Q

                if layer_type is Inception:
                    partitions  = config.pop("partitions")
                    config["K"] = 0

                    config["partitions"] = []
                    for part in partitions:
                        layer_sequence = []
                        part_prev_layer = prev_layer
                        for layer_conf in part:
                            part_prev_layer = Layer.create(lib, layer_conf, part_prev_layer, dtype)
                            layer_sequence.append(part_prev_layer)

                        last = layer_sequence[-1]
                        config["partitions"].append(layer_sequence)
                        config["K"] += last.K
                        if "P" in config:
                            assert config["P"] == last.P and config["Q"] == last.Q
                        else:
                            config["M"] = last.M
                            config["P"] = last.P
                            config["Q"] = last.Q

        # Instantiate the layer
        return layer_type(lib, **config)

class DataLayer(Layer):

    """
    Input data layer.
    """

    def __init__(self, lib, dtype, N, C, D=1, H=1, W=1):

        super(DataLayer, self).__init__(lib, dtype, N)

        self.C = C
        self.K = C
        self.M = D
        self.P = H
        self.Q = W
        self.DHW = (D, H, W)
        self.dimI   = (C, D, H, W, N)
        self.dimO   = (C, D, H, W, N)
        self.dimI2  = (C * D * H * W, N)
        self.dimO2  = (C * D * H * W, N)
        self.sizeO  = reduce(mul, self.dimO, 1)
        self.sizeI  = self.sizeO

    def init_data(self, ary=None):
        if ary is None:
            self.fprop_out.fill(0)
        else:
            self.fprop_out.set(ary)

    def init_deltas(self, shared=None):
        pass

    def init_weights(self, loc=0.0, scale=0.1, shared=None, zeros=False):
        pass

    def fprop(self, fprop_in, scale_weights=0):
        return self.fprop_out

    def __str__(self):
        return "DataLayer: NCK: (%d, %d, %d) DHW:%s" % (self.N, self.C, self.K, self.DHW)

class FullLayer(Layer):

    """
    Fully connnected layer.
    """

    def __init__(self, lib, dtype, N, nIn, nOut, relu=False):

        super(FullLayer, self).__init__(lib, dtype, N)

        self.nIn    = nIn
        self.nOut   = nOut
        self.flops  = N * nIn * nOut * 2.0
        self.dimI   = (nIn, N)
        self.dimI2  = (nIn, N)
        self.dimO   = (nOut, N)
        self.dimO2  = (nOut, N)
        self.dimF   = (nOut, nIn)
        self.dimF2  = (nOut, nIn)
        self.sizeI  = nIn  * N
        self.sizeO  = nOut * N
        self.sizeF  = nIn  * nOut
        self.relu   = relu

    def fprop(self, fprop_in, scale_weights=0):

        fprop_in = super(FullLayer, self).fprop(fprop_in)
        self.lib.compound_dot(self.weights, fprop_in, self.fprop_out, self.relu)

        if scale_weights:
            self.scale_weights(scale_weights)
            self.fprop(fprop_in)

        return self.fprop_out

    def bprop(self, bprop_in, beta=0):

        if self.relu:
            self.bprop_relu(bprop_in)
        self.lib.compound_dot(self.weights.T, bprop_in, self.bprop_out)

        self.lib.compound_dot(bprop_in, self.fprop_in.T, self.updat_out)
        self.grad_descent()

        return self.bprop_out

    def __str__(self):
        return "FullLayer: N, nIn, nOut: (%d, %d, %d)" % (self.N, self.nIn, self.nOut)

class ConvLayer(Layer):

    """
    ConvLayer parameter object.
    This then is passed as an argument to all the convolution operations.

    N: Number of images in mini-batch
    C: Number of input feature maps
    K: Number of output feature maps

    D: Depth  of input image
    H: Height of input image
    W: Width  of input image

    T: Depth  of filter kernel
    R: Height of filter kernel
    S: Width  of filter kernel

    padding: amount of zero-padding around the given edge
    strides: factor to step the filters by in a given direction
    """

    def __init__(self, lib, dtype,
                 N, C, K,
                 D=1, H=1, W=1,
                 T=1, R=1, S=1,
                 pad_d=0, pad_h=0, pad_w=0,
                 str_d=1, str_h=1, str_w=1):

        super(ConvLayer, self).__init__(lib, dtype, N, np.float32)

        # Compute the output spatial dimensions
        M = lib.output_dim(D, T, pad_d, str_d)
        P = lib.output_dim(H, R, pad_h, str_h)
        Q = lib.output_dim(W, S, pad_w, str_w)

        self.C = C
        self.K = K
        self.M = M
        self.P = P
        self.Q = Q
        self.NCK = (N, C, K)
        self.TRS = (T, R, S)
        self.DHW = (D, H, W)
        self.MPQ = (M, P, Q)
        self.padding = (pad_d, pad_h, pad_w)
        self.strides = (str_d, str_h, str_w)

        self.all_params = (N, C, K, D, H, W, T, R, S, pad_d, pad_h, pad_w, str_d, str_h, str_w)

        self.dimI   = (C, D, H, W, N)
        self.dimF   = (C, T, R, S, K)
        self.dimFb  = (K, T, R, S, C)
        self.dimO   = (K, M, P, Q, N)
        self.dimI2  = (C * D * H * W, N)
        self.dimF2  = (C * T * R * S, K)
        self.dimF2t = (K, C * T * R * S)
        self.dimO2  = (K * M * P * Q, N)
        self.dimS   = (K, 1)
        self.sizeI  = reduce(mul, self.dimI, 1)
        self.sizeF  = reduce(mul, self.dimF, 1)
        self.sizeO  = reduce(mul, self.dimO, 1)
        self.nOut   = reduce(mul, self.MPQ, 1) * K

        # flop count for benchmarking
        self.flops = P * Q * M * K * N * C * R * S * T * 2.0

        args = (lib, self.dtype, N, C, K, D, H, W, T, R, S, M, P, Q,
                pad_d, pad_h, pad_w, str_d, str_h, str_w)

        #lib.enable_winograd = 0

        ####### Cuda C ###########
        if lib.use_cudac_kernels:

            #3D conv not supported yet
            if T > 1 or D > 1:
                raise ValueError("3D Convolution not supported by CUDA C kernels and pre-Maxwell GPUs")

            # TODO small C bprop?
            self.fprop_kernels = convolution.FpropCuda(*args)
            self.bprop_kernels = convolution.BpropCuda(*args)
            self.updat_kernels = convolution.UpdateCuda(*args)

        ####### Winograd ###########
        elif lib.enable_winograd and R == 3 and S == 3 and all(x == 1 for x in (D,M,T,str_w,str_h,str_d)):
            from .winograd_conv import (FpropWinograd_2x2_3x3, BpropWinograd_2x2_3x3, UpdateWinograd_3x3_2x2,
                                        FpropWinograd_4x4_3x3, BpropWinograd_4x4_3x3, UpdateWinograd_3x3_4x4)

            # Temp for now till we can autotune
            # 2 is safer for fp16 without batchnorm
            if dtype == np.float32 and lib.enable_winograd == 4:
                winograd = 4
            else:
                winograd = 2

            if C < 8:
                self.fprop_kernels = convolution.FpropDirect(*args)
            elif winograd == 4 and H * W < 112 * 112:
                self.fprop_kernels = FpropWinograd_4x4_3x3(*args)
            else:
                self.fprop_kernels = FpropWinograd_2x2_3x3(*args)

            if winograd == 4 and H * W < 112 * 112:
                self.bprop_kernels = BpropWinograd_4x4_3x3(*args)
            else:
                self.bprop_kernels = BpropWinograd_2x2_3x3(*args)

            if N >= 4 and (C < 8 or H * W > 112 * 112):
                self.updat_kernels = convolution.UpdateDirect(*args)
            elif winograd == 4:
                self.updat_kernels = UpdateWinograd_3x3_4x4(*args)
            else:
                self.updat_kernels = UpdateWinograd_3x3_2x2(*args)

#        elif lib.enable_winograd and not lib.deterministic and N > 1 and \
#            R == 5 and S == 5 and all(x == 1 for x in (D,M,T,str_w,str_h,str_d)):
#
#            from .winograd_conv import (FpropWinograd_2x2_5x5, BpropWinograd_2x2_5x5)
#
#            self.fprop_kernels = FpropWinograd_2x2_5x5(*args)
#            self.bprop_kernels = BpropWinograd_2x2_5x5(*args)
#            if N >= 4:
#                self.updat_kernels = convolution.UpdateDirect(*args)

        ####### Direct ###########
        else:

            self.fprop_kernels = convolution.FpropDirect(*args)
            self.bprop_kernels = convolution.BpropDirect(*args)
            if N >= 4:
                self.updat_kernels = convolution.UpdateDirect(*args)

        #logger.debug("%s: %s, %s, %s", str(self), str(self.fprop_kernels), str(self.bprop_kernels), str(self.updat_kernels))

    def init_activations(self, fprop_out=None):

        super(ConvLayer, self).init_activations(fprop_out)

        if self.bsum:
            self.batch_sum = self.lib.empty(self.dimS, dtype=np.float32)
        else:
            self.batch_sum = None

    def fprop(self, fprop_in, scale_weights=0):
        """
        Conv Layer forward propagation.

        Arguments:
            fprop_in (Tensor): Inputs
            scale_weights (float): Scale weights by scale/mean if nonzero

        Returns:
            fprop_out (Tensor): Output activations
        or
            (self.fprop_out, self.batch_sum) (tuple): Tuple with batch_sum
                added as the second entry.
        """
        fprop_in = super(ConvLayer, self).fprop(fprop_in)
        self.lib.fprop_conv(self, fprop_in, self.weights, self.fprop_out, bsum=self.batch_sum)

        if scale_weights:
            self.scale_weights(scale_weights)
            self.fprop(fprop_in)

        if self.bsum:
            return (self.fprop_out, self.batch_sum)
        return self.fprop_out

    def bprop(self, bprop_in, beta=0):

        if self.relu:
            self.bprop_relu(bprop_in)
        if self.bprop_out is not None:
            self.lib.bprop_conv(self, self.weights, bprop_in, self.bprop_out, beta=beta)

        self.lib.update_conv(self, self.fprop_in, bprop_in, self.updat_out)
        self.grad_descent()

        return self.bprop_out

    def __str__(self):
        return ("ConvLayer: NCK: (%3d, %3d, %3d) HW:%s" %
                (self.N, self.C, self.K, self.DHW[1:3]))

# Add Deconv class
class DeconvLayer(ConvLayer):

    """
    DeconvLayer parameter object.
    This then is passed as an argument to all the convolution operations.

    N: Number of images in mini-batch
    C: Number of output feature maps
    K: Number of input feature maps

    P: Height of input
    Q: Width of input

    D: Depth  of output image
    H: Height of output image
    W: Width  of output image

    T: Depth  of filter kernel
    R: Height of filter kernel
    S: Width  of filter kernel

    padding: amount of zero-padding around the given edge
    strides: factor to step the filters by in a given direction
    """

    def __init__(self, lib, dtype,
                 N, C, K,
                 P, Q,
                 R=1, S=1,
                 pad_d=0, pad_h=0, pad_w=0,
                 str_d=1, str_h=1, str_w=1):

        # Set T and D to be consts.
        D = T = 1

        # Cannot get exact, e.g. because not unique
        H = (P - 1) * str_h - 2 * pad_h + R
        W = (Q - 1) * str_w - 2 * pad_w + S

        super(DeconvLayer, self).__init__(
            lib, dtype,
            N, C, K,
            D, H, W,
            T, R, S,
            pad_d, pad_h, pad_w,
            str_d, str_h, str_w)

        self.nOut = reduce(mul, self.DHW, 1) * C
        self.H = H
        self.W = W

    def __str__(self):
        return ("DeconvLayer: NCK: (%d, %d, %d) DHW:%s TRS:%s MPQ:%s" %
                (self.N, self.C, self.K, self.DHW, self.TRS, self.MPQ))

class PoolLayer(Layer):

    """
    PoolLayer parameter object.
    This then is passed as an argument to all pooling kernels.

    op: max, avg, l2 pooling
    N: Number of images in mini-batch

    C: Number of input feature maps
    D: Depth  of input image
    H: Height of input image
    W: Width  of input image

    J: Size of feature map pooling window (maxout n_pieces)
    T: Depth  of pooling window
    R: Height of pooling window
    S: Width  of pooling window

    padding: amount of zero-padding around the given image or feature map edge
    strides: factor to step the window by in a given direction (overlap allowed)

    Leave spatial dimensions at 1 to allow feature map pooling in the fc layers.

    """

    def __init__(self, lib, dtype,
                 op, N, C,
                 D=1, H=1, W=1,
                 J=1, T=1, R=1, S=1,
                 pad_c=0, pad_d=0, pad_h=0, pad_w=0,
                 str_c=None, str_d=None, str_h=None, str_w=None):

        super(PoolLayer, self).__init__(lib, dtype, N)

        if self.dtype.type is np.float16:
            clss = "hpool"
        elif self.dtype.type is np.float32:
            clss = "spool"
        else:
            raise TypeError("Type not supported.")

        # default to non-overlapping
        if str_c is None:
            str_c = J
        if str_d is None:
            str_d = T
        if str_h is None:
            str_h = R
        if str_w is None:
            str_w = S

#        if str_c < J or str_d < T or str_h < R or str_w < S:
#            self.overlap = (ceil(float(J) / str_c) *
#                            ceil(float(T) / str_d) *
#                            ceil(float(R) / str_h) *
#                            ceil(float(S) / str_w))
#        else:
#            self.overlap = 0.0

        self.overlap = 1.0

        # TODO: detect other forms of gaps
        if str_c > J or str_d > T or str_h > R or str_w > S:
            self.gaps = 1
        else:
            self.gaps = 0

        bprop_zero = self.overlap or self.gaps

        # Compute the output dimensions
        K = lib.output_dim(C, J, pad_c, str_c, pooling=True)
        M = lib.output_dim(D, T, pad_d, str_d, pooling=True)
        P = lib.output_dim(H, R, pad_h, str_h, pooling=True)
        Q = lib.output_dim(W, S, pad_w, str_w, pooling=True)

        self.op   = op
        self.C    = C
        self.K    = K
        self.M    = M
        self.P    = P
        self.Q    = Q
        self.JTRS = (J, T, R, S)
        self.DHW  = (D, H, W)
        self.MPQ  = (M, P, Q)
        self.padding = (pad_c, pad_d, pad_h, pad_w)
        self.strides = (str_c, str_d, str_h, str_w)

        self.dimI   = (C, D, H, W, N)
        self.dimO   = (K, M, P, Q, N)
        self.dimF2  = None
        self.dimI2  = (C * D * H * W, N)
        self.dimO2  = (K * M * P * Q, N)
        self.sizeI  = reduce(mul, self.dimI, 1)
        self.sizeO  = reduce(mul, self.dimO, 1)
        self.nOut   = reduce(mul, self.MPQ, 1) * K

        # precompute some multiplications for fast constant memory access
        WN   = W * N
        HWN  = H * WN
        DHWN = D * HWN
        DH   = D * H
        RS   = R * S
        RST  = T * RS
        JRST = J * RST
        QN   = Q * N
        PQN  = P * QN
        MPQN = M * PQN

        assert JRST + 32 < 2**16, "Integer division is faster with 16bit numerators"

        sb_large = {
            #SB  shlP maskP shrP shlQ maskQ shrQ maskN shrN
            1  : (0,   0x00, 0,   0,   0x00, 0,   0xfff, 32), # 1x1  nnnnn
            2  : (0,   0x00, 0,   1,   0x10, 4,   0x00f,  4), # 1x2  xnnnn
            4  : (0,   0x00, 0,   2,   0x18, 3,   0x007,  3), # 1x4  xxnnn
            8  : (0,   0x00, 0,   3,   0x1c, 2,   0x003,  2), # 1x8  xxxnn
            16 : (0,   0x00, 0,   4,   0x1e, 1,   0x001,  1), # 1x16 xxxxn
            32 : (0,   0x00, 0,   5,   0x1f, 0,   0x000,  0), # 1x32 xxxxx
        }
        sb_medium = {
            #SB  shlP maskP shrP shlQ maskQ shrQ maskN shrN
            8  : (1,   0x10, 4,   2,   0x0c, 2,   0x003,  2), # 2x4  yxxnn
            16 : (1,   0x10, 4,   3,   0x0e, 1,   0x001,  1), # 2x8  yxxxn
            32 : (1,   0x10, 4,   4,   0x0f, 0,   0x000,  0), # 2x16 yxxxx
        }
        sb_small = {
            #SB  shlP maskP shrP shlQ maskQ shrQ maskN shrN
            16 : (2,   0x18, 3,   2,   0x06, 1,   0x001,  1), # 4x4  yyxxn
            32 : (2,   0x18, 3,   3,   0x07, 0,   0x000,  0), # 4x8  yyxxx
        }

        if N == 1:
            super_block = 0
        elif N < 32:
            super_block = len(bin(N - 1)) - 2
        else:
            super_block = 5
        super_block = 1 << (5 - super_block)

        # try to minimize the zero overlap in the superblock
        # but maximize the x dim of the superblock for more contiguous memory access
        if super_block < 8 or Q > 64:
            sb_params = sb_large.get(super_block)
        elif super_block < 16 or Q > 32:
            sb_params = sb_medium.get(super_block)
        else:
            sb_params = sb_small.get(super_block)

        supP = _ceil_div(P, 1 << sb_params[0])
        supQ = _ceil_div(Q, 1 << sb_params[3])

        # precompute the magic numbers and shift amounts for integer division
        magic_RST = _magic32(JRST + 32, RST)
        magic_RS  = _magic32(RST + 32, RS)
        magic_S   = _magic32(RS + 32, S)
        magic_P   = _magic32(M * supP, supP)

        fprop_name = "fprop_" + op
        bprop_name = "bprop_" + op

        threads = 32 if super_block > 1 else N

        self.fprop_kernel = [fprop_name, (supQ, supP * M, K), (threads, 1, 1), _flatten([
            N, W, H, D, C, WN, HWN, DHWN,
            P, Q, magic_P, QN, PQN, MPQN,
            pad_c, pad_d, pad_h, pad_w,
            str_c, str_d, str_h, str_w,
            S, RS, RST, JRST, magic_S, magic_RS, magic_RST,
            supP, supQ, sb_params ])]

        lut_size = JRST
        if lut_size % 4 != 0:
            lut_size += 4 - lut_size % 4

        self.bprop_lut_size = self.fprop_lut_size = super_block * lut_size * 4

        if self.overlap > 0:

            # we have a special kernel to handle the overlapping avg pooling
            bprop_name += "_overlap"

            magic_str_w = _magic32(W + S, str_w)
            magic_str_h = _magic32(H + R, str_h)
            magic_str_d = _magic32(D + T, str_d)
            magic_str_c = _magic32(C + J, str_c)

            if super_block > 1:

                bprop_name += "_smallN"

                if super_block < 8 or W > 64:
                    sb_params = sb_large.get(super_block)
                elif super_block < 16 or W > 32:
                    sb_params = sb_medium.get(super_block)
                else:
                    sb_params = sb_small.get(super_block)

                supH = _ceil_div(H, 1 << sb_params[0])
                supW = _ceil_div(W, 1 << sb_params[3])

                magic_H = _magic32(D * supH, supH)

                maxLutSize = \
                    _ceil_div(S, str_w) * \
                    _ceil_div(R, str_h) * \
                    _ceil_div(T, str_d) * \
                    _ceil_div(J, str_c)

                #neon_logger.display((supW, D*supH, C), sb_params, maxLutSize)

                self.bprop_kernel = [bprop_name, (supW, D * supH, C), (threads, 1, 1), _flatten([
                    N, W, H, D, C, WN, HWN, DHWN, magic_H,
                    pad_w, pad_h, pad_d, pad_c,
                    str_w, str_h, str_d, str_c,
                    magic_str_w, magic_str_h, magic_str_d, magic_str_c,
                    S, R, T, J, RS, RST, JRST, magic_S, magic_RS, magic_RST,
                    Q, P, M, K, QN, PQN, MPQN,
                    supH, supW, sb_params, maxLutSize])]

                lut_size = maxLutSize
                if lut_size % 4 != 0:
                    lut_size += 4 - lut_size % 4

                self.bprop_lut_size = super_block * lut_size * 4 * 2

            else:

                # The overlap kernel can be much more efficient if we aren't doing superblocking
                magic_H = _magic32(DH, H)

                self.bprop_kernel = [bprop_name, (W, DH, C), (threads, 1, 1), _flatten([
                    N, W, H, D, C, WN, HWN, DHWN, magic_H,
                    pad_w, pad_h, pad_d, pad_c,
                    str_w, str_h, str_d, str_c,
                    magic_str_w, magic_str_h, magic_str_d, magic_str_c,
                    S, R, T, J, RS, RST, JRST, magic_S, magic_RS, magic_RST,
                    Q, P, M, K, QN, PQN, MPQN])]

                self.bprop_lut_size = lut_size * 4 * 2
        else:
            self.bprop_kernel = [bprop_name, (supQ, supP * M, K), (threads, 1, 1), _flatten([
                N, W, H, D, C, WN, HWN, DHWN,
                P, Q, magic_P, QN, PQN, MPQN,
                pad_c, pad_d, pad_h, pad_w,
                str_c, str_d, str_h, str_w,
                S, RS, RST, JRST, magic_S, magic_RS, magic_RST,
                supP, supQ, sb_params])]

    def init_activations(self, fprop_out=None):

        super(PoolLayer, self).init_activations(fprop_out)

        self.argmax = self.lib.empty(self.dimO, dtype=np.uint8)

    def fprop(self, fprop_in, scale_weights=0):
        """ Used for benchmarking only"""
        fprop_in = super(PoolLayer, self).fprop(fprop_in)
        self.lib.fprop_pool(self, fprop_in, self.fprop_out, argmax=self.argmax)
        return self.fprop_out

    def bprop(self, bprop_in, beta=0):
        self.lib.bprop_pool(self, bprop_in, self.bprop_out, argmax=self.argmax)
        return self.bprop_out

    def __str__(self):
        return ("PoolLayer: NCK: (%d, %d, %d) DHW:%s JTRS:%s MPQ:%s op: %s " %
                (self.N, self.C, self.K, self.DHW, self.JTRS, self.MPQ, self.op))

class Inception(Layer):

    """
    GoogLeNet inception assembly layer.

    """

    def __init__(self, lib, dtype, partitions,
                 N, C, K,
                 D=1, H=1, W=1,
                 M=1, P=1, Q=1):

        super(Inception, self).__init__(lib, dtype, N)

        self.partitions = partitions

        self.C = C
        self.K = K
        self.M = M
        self.P = P
        self.Q = Q
        self.NCK = (N, C, K)
        self.DHW = (D, H, W)
        self.MPQ = (M, P, Q)

        self.dimI  = (C, D, H, W, N)
        self.dimO  = (K, M, P, Q, N)
        self.dimI2 = (C * D * H * W, N)
        self.dimO2 = (K * M * P * Q, N)
        self.sizeI = reduce(mul, self.dimI, 1)
        self.sizeO = reduce(mul, self.dimO, 1)
        self.nOut  = reduce(mul, self.MPQ, 1) * K

        self.sizeF = 0
        self.flops = 0
        for part in partitions:
            for layer in part:
                self.flops += layer.flops
                self.sizeF  = max(self.sizeF, layer.sizeF)
                if self.sizeF == layer.sizeF:
                    self.dimF = layer.dimF

    def __str__(self):
        out = "Inception: NCK: (%d, %d, %d) DHW:%s MPQ:%s\n" % (self.N, self.C, self.K,
                                                                self.DHW, self.MPQ)
        for i, part in enumerate(self.partitions):
            out += "  Part%d:\n" % (i + 1)
            for layer in part:
                out += "    %s\n" % layer
        return out.rstrip()

    def init_activations(self, fprop_out=None):

        super(Inception, self).init_activations(fprop_out)
        K = 0
        for part in self.partitions:
            for layer in part:
                if layer is part[-1]:
                    layer.init_activations(self.fprop_out[K:K + layer.K, ...])
                    K += layer.K
                else:
                    layer.init_activations()

    def init_deltas(self, shared=None):

        super(Inception, self).init_deltas(shared)

        shared_deltas = shared[1:3] if shared else None

        for part in self.partitions:
            for layer in part:
                if layer is part[0]:
                    layer.bprop_out   = self.bprop_out
                    layer.delta_stats = self.delta_stats
                else:
                    layer.init_deltas(shared=shared_deltas)

    def init_weights(self, loc=0.0, scale=0.1, shared=None, zeros=False):

        for part in self.partitions:
            for layer in part:
                layer.init_weights(loc, scale, shared, zeros)

    def fprop(self, fprop_in, scale_weights=0):

        fprop_in = super(Inception, self).fprop(fprop_in)

        for part in self.partitions:
            part_fprop_in = fprop_in
            for layer in part:
                part_fprop_in = layer.fprop(part_fprop_in, scale_weights)

        return self.fprop_out

    def bprop(self, bprop_in):

        K = self.K
        for part in self.partitions[::-1]:
            part_bprop_in = bprop_in[K - part[-1].K:K, ...]
            K -= part[-1].K
            for layer in part[::-1]:
                if part is not self.partitions[-1] and layer is part[0]:
                    beta = 1.0
                else:
                    beta = 0.0

                part_bprop_in = layer.bprop(part_bprop_in, beta)

        return self.bprop_out

    def fprop_stats(self):
        for part in self.partitions:
            for layer in part:
                layer.fprop_stats()

    def bprop_stats(self):
        for part in self.partitions[::-1]:
            for layer in part[::-1]:
                layer.bprop_stats()

class BatchNorm(Layer):

    """
    Batch Normalization Layer
    """

    def __init__(self, lib, dtype, N, C=None, D=1, H=1, W=1, nIn=None, rho=0.99, eps=1e-6,
                 relu=False, bsum=False):
        """
        Batch Normalization layer

        Arguments:
            lib (Class): NervanaGPU instance
            dtype (Dtype): Data type
            N (int): batch size
            C (int): Number of input feature maps
            D (int): Depth  of input feature maps
            H (int): Height of input feature maps
            W (int): Width  of input feature maps
            nIn (int): Number on inputs for fully connected layer
            rho (float): Exponential window averaging factor
            eps (float): Constant added for numerical stability
            relu (bool): Flag for rectified linear activation function
            bsum (bool): PQN sum precomputed in conv kernel
        """
        super(BatchNorm, self).__init__(lib, dtype, N)

        self.rho  = rho
        self.eps  = eps
        self.relu = relu
        self.bsum = bsum

        if C is not None:

            self.C     = C
            self.K     = C
            self.M     = D
            self.P     = H
            self.Q     = W
            self.dimI  = (C, D, H, W, N)
            self.dimO  = (C, D, H, W, N)
            self.dimO2 = (C * D * H * W, N)
            self.dim2  = (C, D * H * W * N)
            self.nOut  = C * D * H * W

        elif nIn is not None:

            self.nOut  = nIn
            self.K     = nIn
            self.dimI  = (nIn, N)
            self.dimO  = (nIn, N)
            self.dimO2 = (nIn, N)
            self.dim2  = (nIn, N)

        else:
            raise ValueError("missing C or nIn")

        self.rcp_depth = 1.0 / self.dim2[1]

    def __str__(self):
        return ("BatchNorm: (%d, %d)" % self.dim2)

    def init_activations(self, fprop_out=None):

        if fprop_out is not None:
            self.fprop_out = fprop_out.reshape(self.dim2)
        else:
            self.fprop_out = self.lib.empty(self.dim2, dtype=self.dtype)

        self.xvar = self.lib.empty((self.K, 1), dtype=self.dtype)
        if not self.bsum:
            self.xsum = self.lib.empty((self.K, 1), dtype=np.float32)

    def init_deltas(self, shared=None):
        pass

    def init_weights(self, loc=0.0, scale=0.1, shared=None, zeros=False):

        lib = self.lib

        self.beta  = lib.zeros((self.K, 1), dtype=self.dtype)
        self.gamma = lib.ones((self.K, 1),  dtype=self.dtype)

        self.gmean = lib.zeros((self.K, 1), dtype=self.dtype)
        self.gvar  = lib.zeros((self.K, 1), dtype=self.dtype)

        self.grad_beta  = lib.zeros((self.K, 1), dtype=self.dtype)
        self.grad_gamma = lib.zeros((self.K, 1), dtype=self.dtype)

    def fprop(self, fprop_in, scale_weights=0):
        """
        Batch normalization forward pass. Uses a compound kernel call.
        """
        if type(fprop_in) is tuple:
            fprop_in, bsum = fprop_in
        else:
            bsum = None

        if self.fprop_in is None:
            self.fprop_in = fprop_in.reshape(self.dim2)

        if bsum is None:
            self.xsum[:] = self.lib.sum(self.fprop_in, axis=1)
        else:
            self.xsum = bsum

        self.lib.compound_fprop_bn(
            self.fprop_in, self.xsum, self.xvar, self.gmean, self.gvar,
            self.gamma, self.beta, self.fprop_out, self.eps, self.rho, self.relu)

        return self.fprop_out

    def bprop(self, bprop_in, beta=0):

        if self.bprop_in is None:
            self.bprop_in = bprop_in.reshape(self.dim2)

        if self.relu:
            self.bprop_relu(self.bprop_in)

        self.lib.compound_bprop_bn(
            self.bprop_in, self.grad_gamma, self.grad_beta, self.fprop_in,
            self.xsum, self.xvar, self.gamma, self.eps)

        return self.bprop_in

    def fprop_stats(self):
        pass

    def bprop_stats(self):
        pass

# Magic numbers and shift amounts for integer division
# Suitable for when nmax*magic fits in 32 bits
# Shamelessly pulled directly from:
# http://www.hackersdelight.org/hdcodetxt/magicgu.py.txt
def _magic32(nmax, d):
    nc = ((nmax + 1) // d) * d - 1
    nbits = len(bin(nmax)) - 2
    for p in range(0, 2 * nbits + 1):
        if 2 ** p > nc * (d - 1 - (2 ** p - 1) % d):
            m = (2 ** p + d - 1 - (2 ** p - 1) % d) // d
            return (m, p)
    raise ValueError("Can't find magic number for division")

# Magic numbers and shift amounts for integer division
# Suitable for when nmax*magic fits in 64 bits and the shift
# lops off the lower 32 bits
def _magic64(d):
    # 3 is a special case that only ends up in the high bits
    # if the nmax is 0xffffffff
    # we can't use 0xffffffff for all cases as some return a 33 bit
    # magic number
    nmax = 0xffffffff if d == 3 else 0x7fffffff
    magic, shift = _magic32(nmax, d)
    if magic != 1:
        shift -= 32
    return (magic, shift)

# flatten a nested list of lists or values
def _flatten(lst):
    return sum(([x] if not isinstance(x, (list, tuple))
                else _flatten(x) for x in lst), [])

def _ceil_div(x, y):
    return -(-x // y)

@context_dependent_memoize
def _get_sm_count():
    attributes = drv.Context.get_device().get_attributes()
    return attributes[drv.device_attribute.MULTIPROCESSOR_COUNT]


# -*- coding: utf-8 -*-
from __future__ import division
from __future__ import print_function
import random
import pprint
import sys
import time
import numpy as np
from optparse import OptionParser
import pickle

import cv2
import os
from keras import backend as K
from keras.optimizers import Adam, SGD, RMSprop
from keras.layers import Input
from keras.models import Model
from keras_frcnn import config, data_generators
from keras_frcnn import losses as losses
import keras_frcnn.roi_helpers as roi_helpers
from keras.utils import generic_utils
from keras_frcnn import march_ori as march

from keras_frcnn.pascal_voc_parser import get_data
def get_output_length(input_length):
    # zero_pad
    input_length += 6
    # apply 4 strided convolutions
    filter_sizes = [7, 3, 1, 1]
    stride = 2
    for filter_size in filter_sizes:
        input_length = (input_length - filter_size + stride) // stride
    return input_length
part_class_mapping={'head':0,'wings':1,'legs':2,'back':3,'belly':4,'breast':5,'tail':6}
bird_map_path = '/home/e813/dataset/CUBbird/CUB_200_2011/CUB_200_2011/c.pkl'
with open(bird_map_path,'r') as f:
    bird_class_mapping=pickle.load(f)
#pprint.pprint(bird_class_mapping)
def read_prepare_input(img_path):
    img = cv2.imread(img_path)
    return img
os.environ['CUDA_VISIBLE_DEVICES']= '1'

cfg = config.Config()
sys.setrecursionlimit(40000)

if cfg.network == 'vgg':
    # cfg.network = 'vgg'
    from keras_frcnn import vgg as nn
    print('use vgg')
elif cfg.network == 'resnet50':
    from keras_frcnn import resnet as nn
    print('use resnet50')
else:
    print('Not a valid model')
    raise ValueError

#cfg.base_net_weights = cfg.ori_res50_withtop

all_imgs, classes_count, bird_class_count = get_data(cfg.train_path,part_class_mapping)
data_lei = march.get_voc_label(all_imgs, classes_count, part_class_mapping, bird_class_count, bird_class_mapping,config= cfg,trainable=True)
#pprint.pprint(classes_count)
#pprint.pprint(part_class_mapping)
# 这里的类在match里边定义
if 'bg' not in classes_count:
    classes_count['bg'] = 0
    part_class_mapping['bg'] = len(part_class_mapping)

cfg.class_mapping = part_class_mapping

print('Training images per class:')
pprint.pprint(classes_count)
pprint.pprint(part_class_mapping)
print('Num classes (including bg) = {}'.format(len(classes_count)))
print('Training bird per class:')
pprint.pprint(bird_class_count)
print('total birds class is {}'.format(len(bird_class_count)))
print('bird_class_mapping')
pprint.pprint(bird_class_mapping)

config_output_filename = cfg.config_filepath
with open(config_output_filename, 'wb') as config_f:
    pickle.dump(cfg, config_f)
    print('Config has been written to {}, and can be loaded when testing to ensure correct results'.format(
        config_output_filename))
random.shuffle(all_imgs)

num_imgs = len(all_imgs)

train_imgs = [s for s in all_imgs if s['imageset'] == 'trainval']
val_imgs = [s for s in all_imgs if s['imageset'] == 'test']
print('Num train samples {}'.format(len(train_imgs)))
print('Num val samples {}'.format(len(val_imgs)))

#data_gen_train = data_generators.get_anchor_gt(train_imgs, classes_count, cfg, nn.get_img_output_length,K.image_dim_ordering(), mode='train')
#data_gen_val = data_generators.get_anchor_gt(val_imgs, classes_count, cfg, nn.get_img_output_length,K.image_dim_ordering(), mode='val')
input_shape_img = (None, None, 3)

img_input = Input(shape=input_shape_img)
#roi_input = Input(shape=(None, 4))  # roiinput是什么,要去看看清楚
part_roi_input = Input(shape=[None,4])

# define the base network (resnet here, can be VGG, Inception, etc)
shared_layers = nn.nn_base(img_input, trainable=True)  # 共享网络层的输出.要明确输出的size
#bird_classifier_output = nn.fg_classifier(shared_layers,bird_rois_input0,bird_rois_input1,bird_rois_input2,bird_rois_input3,bird_rois_input4,bird_rois_input5,bird_rois_input6,nb_classes=200, trainable=True)
holyclass_out = nn.fine_layer(shared_layers, part_roi_input,nb_classes=200)

class_holyimg_out = nn.fine_layer_hole(shared_layers, part_roi_input,num_rois=1,nb_classes=200)

model_holyclassifier = Model([img_input,part_roi_input],holyclass_out)
#model_classifier_holyimg = Model([img_input,part_roi_input],class_holyimg_out)
cfg.base_net_weights = '/media/e813/D/weights/kerash5/frcnn/TST_holy_img/model_part{}.hdf5'.format(18)
#cfg.base_net_weights = cfg.holy_img_weight_path+'part_headonly'+str(8)+'.hdf5'
try:
    print('loading weights from {}'.format(cfg.base_net_weights))
    #model_rpn.load_weights(cfg.base_net_weights, by_name=True)
    #model_classifier.load_weights(cfg.base_net_weights, by_name=True)
    #model_birdclassifier.load_weights(cfg.base_net_weights, by_name=True)
    model_holyclassifier.load_weights(cfg.base_net_weights, by_name=True)
    #model_classifier_holyimg.load_weights(cfg.base_net_weights,by_name=True)
except:
    print('Could not load pretrained model weights. Weights can be found in the keras application folder \
		https://github.com/fchollet/keras/tree/master/keras/applications')
    raise ValueError('load wrong')
optimizer = Adam(lr=1e-5)
lossfn_list =[]
for i in range(7):
    lossfn_list.append(losses.holy_loss())
model_holyclassifier.compile(optimizer=optimizer,loss=lossfn_list)
#model_classifier_holyimg.compile(optimizer=optimizer,loss=lossfn_list)

max_epoch=32
step= 0
now_epoch = 0
tru_s=np.zeros([7],dtype=np.int16)
fal_s=np.zeros([7],dtype=np.int16)
holy_pre_result = []
while 1:
    step+=1
    img_np,boxnp, labellist,img_path = data_lei.next_batch(1)
    #print(img_np.shape)
    #print(boxnp.shape)
    #input_img = read_prepare_input(img_path)
    #print(boxnp.shape)
    #print(img_path)
    #print(index)
    #exit(4)
    #print(labellist)
    #exit(4)
    #holynet_loss = model_holyclassifier.train_on_batch([img_np,boxnp],labellist)
    netoutput = model_holyclassifier.predict_on_batch([img_np,boxnp])
    one_sample = {'netout':netoutput,'label':labellist}
    holy_pre_result.append(one_sample)
    for i in range(7):
        net_predict = netoutput[i][0,:]
        label = labellist[i][0,1:]
        pre_idx = np.argmax(net_predict)
        lab_idx = np.argmax(label)
        #one_sample = {'pre':net_predict,'label':labellist}
        if labellist[i][0, 0] == 1:
            if pre_idx ==lab_idx:
                tru_s[i]+=1
            else:
                fal_s[i]+=1
    '''for j in range(boxnp.shape[1]):
        x1_out = boxnp[0, j, 0]
        y1_out = boxnp[0, j, 1]
        w_out = boxnp[0, j, 2]
        h_out = boxnp[0, j, 3]
        w = w_out * (size_w_ori / size_w_map)
        h = h_out * (size_h_ori / size_h_map)
        x1 = x1_out * (size_w_ori / size_w_map)
        y1 = y1_out * (size_h_ori / size_h_map)
        if labellist[j][0, 0] == 1:
            for key, value in part_class_mapping.items():
                if value == j:
                    pname = key
            #print('{}  {}  {}'.format(key, pre_idx, lab_idx))
            cv2.rectangle(img_ori, (int(x1), int(y1)), (int(x1 + w), int(y1 + w)), (0, 255, 0), 2)
            #cv2.rectangle(img_ori, (int(x1), int(y1) - 20),
            #            (int(x1 + w), int(y1)), (125, 125, 125), -1)
            cv2.putText(img_ori, '{} p:{} l:{} C :{}'.format(pname, pre_idx, lab_idx, net_predict[pre_idx]), (int(x1) + 5, int(y1) - 7), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                        (0, 0, 0), 1, cv2.LINE_AA)
            print('part:{}  predict idx:{}  lab_idx:{} C :{}'.format(pname, pre_idx, lab_idx, net_predict[pre_idx]))'''

    # print()
    # continue
    # annot = img_annot
    # img_path = annot['filepath']
    #img_np = cv2.imread(img_path)
    # print(annot)
    # cv2.imshow('ori_label', img_np)
    #predict = model_holyclassifier.predict([img_np,boxnp])
    #predict =np.mean(predict,axis=1)
    #print(predict[0])
    #for i in range(7):
    #    print(np.argmax(predict[i],axis=1))
    #result = np.max(predict,axis=1)
    #print(predict.shape)
    #print(result)
    #exit()
    #holynet_loss = model_holyclassifier.train_on_batch([img_np, boxnp], labellist)
    #holynet_loss = model_holyclassifier.train_on_batch([img_np, boxnp], labellist)
    #holynet_loss = model_holyclassifier.train_on_batch([img_np, boxnp], labellist)
    #A = model_holyclassifier.predict([img_np,boxnp])
    #print(holynet_loss)
    #exit()
    #print(boxnp)
    #print('step is {} batch_index is {} and epoch is {}'.format(step,data_lei.batch_index,data_lei.epoch))
    #print(holynet_loss)

    #break
    if data_lei.epoch!= now_epoch:
        #break
        arc = np.zeros([7],dtype=np.float32)
        for j in range(7):
            arc[j]= float(tru_s[j])/float(tru_s[j]+fal_s[j])
        print(arc)
        break
result_class_path = '/media/e813/D/weights/kerash5/frcnn/TST_holy_img/part_result_train{}.pkl'.format(18)
with open(result_class_path,'wb') as f:
    pickle.dump(holy_pre_result,f)

def get_output_length(input_length):
    # zero_pad
    input_length += 6
    # apply 4 strided convolutions
    filter_sizes = [7, 3, 1, 1]
    stride = 2
    for filter_size in filter_sizes:
        input_length = (input_length - filter_size + stride) // stride
    return input_length


#!/usr/bin/env python

import argparse
import pandas as pd
from Bio import SeqIO
import os
import numpy as np
from collections import OrderedDict
import gzip
from tqdm import tqdm
from abnumber import Chain
import re
import requests
import time
import sys

T20_REGEX = re.compile('T20 Score:	([0-9.]+)	')

def get_t20_online(seq):
    chain = Chain(seq, scheme='imgt')
    chain_type = 'vh' if chain.chain_type == 'H' else ('vl' if chain.chain_type == 'L' else 'vk')
    html = None
    for retry in range(5):
        url = f'https://dm.lakepharma.com/cgi-bin/blast.py?chain={chain_type}&region=1&output=3&seqs={seq}'
        try:
            request = requests.get(url)
            if request.ok:
                html = request.text
                break
        except Exception as e:
            print(e)
        time.sleep(0.5 + retry * 5)
        print('Retry', retry+1)
    if not html:
        sys.exit(1)
    matches = T20_REGEX.findall(html)
    time.sleep(1)
    if not matches:
        print(html)
        raise ValueError(f'Error calling url {url}')
    return float(matches[0])

def get_multiple_t20_online(queries):
    results = []
    for query in tqdm(queries):
        t20 = get_t20_online(query.seq)
        results.append(OrderedDict(
            id=query.id,
            description=query.description,
            t20=t20
        ))
    return pd.DataFrame(results)

def get_multiple_t20_custom_db(queries, targets, num=20):
    from skbio.alignment import StripedSmithWaterman
    from skbio.alignment._pairwise import blosum50
    results = []
    for query in tqdm(queries):
        aligner = StripedSmithWaterman(
            str(query.seq),
            protein=True,
            substitution_matrix=blosum50
        )
        best_identities = []
        best_aligned_query = None
        best_aligned_target = None

        for target in targets:
            alignment = aligner(target)
            query_length = len(alignment.query_sequence)
            identity = sum([a == b for a, b in zip(alignment.aligned_query_sequence, alignment.aligned_target_sequence)]) / query_length
            if len(best_identities) < num:
                best_identities.append(identity)
                best_identities = sorted(best_identities)
            elif identity > best_identities[0]:
                best_identities = best_identities[1:] + [identity]
                best_identities = sorted(best_identities)

            if identity == best_identities[-1]:
                best_aligned_query = alignment.aligned_query_sequence
                best_aligned_target = alignment.aligned_target_sequence
        result = OrderedDict(
            query_seq=best_aligned_query,
            id=query.id,
            description=query.description,
            best_seq=best_aligned_target,
            best_identity=best_identities[-1]
        )
        result[f'T{num}'] = np.mean(best_identities)
        results.append(result)
    return pd.DataFrame(results)

if __name__ == "__main__":
    # Parse command line
    parser = argparse.ArgumentParser()
    parser.add_argument("query", help="Query FASTA sequence file.")
    parser.add_argument("output", help="Output TSV file path.")
    parser.add_argument("--db", help="Custom human gzipped TXT sequence database file.")
    parser.add_argument("-n", "--num", default=20, type=int, help="Number of closest sequences to evaluate.")

    options = parser.parse_args()

    queries = SeqIO.parse(options.query, 'fasta')

    if options.db:
        with gzip.open(options.target, 'rt') as f:
            targets = [l.strip() for l in f.readlines()]

        df = get_multiple_t20_custom_db(queries, targets, num=options.num)
    else:
        assert options.num == 20, 'You need to provide a custom --db to run with n != 20'
        print('Note: The sequences will be processed through lakepharma T20 service! Sleeping for 10s, press Ctrl+C to cancel...')
        time.sleep(10)
        print('Processing...')
        df = get_multiple_t20_online(queries)

    df.to_csv(options.output, sep='\t', index=False)
    print(f'Saved to: {options.output}')


"""
An example of Adversarial Inverse Reinforcement Learning.
"""

import argparse
import copy
import json
import os
from pprint import pprint
import pickle

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import gym

import machina as mc
from machina.pols import GaussianPol, CategoricalPol, MultiCategoricalPol
from machina.algos import airl, behavior_clone
from machina.vfuncs import DeterministicSVfunc, DeterministicSAVfunc
from machina.envs import GymEnv, C2DEnv
from machina.traj import Traj
from machina.traj import epi_functional as ef
from machina.samplers import EpiSampler
from machina import logger
from machina.utils import measure, set_device

from simple_net import PolNet, PolNetLSTM, VNet, DiscrimNet, VNetLSTM

parser = argparse.ArgumentParser()
parser.add_argument('--log', type=str, default='garbage',
                    help='Directory name of log.')
parser.add_argument('--env_name', type=str,
                    default='Pendulum-v0', help='Name of environment.')
parser.add_argument('--c2d', action='store_true',
                    default=False, help='If True, action is discretized.')
parser.add_argument('--record', action='store_true',
                    default=False, help='If True, movie is saved.')
parser.add_argument('--seed', type=int, default=256)
parser.add_argument('--max_epis', type=int,
                    default=100000000, help='Number of episodes to run.')
parser.add_argument('--num_parallel', type=int, default=4,
                    help='Number of processes to sample.')
parser.add_argument('--cuda', type=int, default=-1, help='cuda device number.')

parser.add_argument('--expert_dir', type=str, default='../data/expert_epis',
                    help='Directory path storing file of expert trajectory.')
parser.add_argument('--expert_fname', type=str,
                    default='Pendulum-v0_100epis.pkl', help='Name of pkl file of expert trajectory')

parser.add_argument('--max_steps_per_iter', type=int, default=50000,
                    help='Number of steps to use in an iteration.')
parser.add_argument('--batch_size', type=int, default=50000)
parser.add_argument('--discrim_batch_size', type=int, default=32)
parser.add_argument('--pol_lr', type=float, default=1e-4,
                    help='Policy learning rate.')
parser.add_argument('--vf_lr', type=float, default=1e-3,
                    help='Value function learning rate.')
parser.add_argument('--discrim_lr', type=float, default=3e-4,
                    help='Discriminator learning rate.')

parser.add_argument('--epoch_per_iter', type=int, default=50,
                    help='Number of epoch in an iteration')
parser.add_argument('--discrim_step', type=int, default=10)

parser.add_argument('--gamma', type=float, default=0.995,
                    help='Discount factor.')
parser.add_argument('--lam', type=float, default=0.97,
                    help='Tradeoff value of bias variance.')
parser.add_argument('--pol_ent_beta', type=float, default=0,
                    help='Entropy coefficient for policy.')

parser.add_argument('--max_grad_norm', type=float, default=10,
                    help='Value of maximum gradient norm.')

parser.add_argument('--pol_h1', type=int, default=100,
                    help='Hidden size of layer1 of policy.')
parser.add_argument('--pol_h2', type=int, default=100,
                    help='Hidden size of layer2 of policy.')
parser.add_argument('--vf_h1', type=int, default=32,
                    help='Hidden size of layer1 of value function.')
parser.add_argument('--vf_h2', type=int, default=32,
                    help='Hidden size of layer2 of value function.')
parser.add_argument('--discrim_h1', type=int, default=100,
                    help='Hidden size of layer1 of discriminator.')
parser.add_argument('--discrim_h2', type=int, default=100,
                    help='Hidden size of layer2 of discriminator.')

parser.add_argument('--rl_type', type=str,
                    choices=['trpo', 'ppo_clip', 'ppo_kl'], default='trpo', help='Choice for Reinforcement Learning algorithms.')

parser.add_argument('--clip_param', type=float, default=0.2,
                    help='Value of clipping liklihood ratio.')
parser.add_argument('--kl_targ', type=float, default=0.01,
                    help='Target value of kl divergence.')
parser.add_argument('--init_kl_beta', type=float,
                    default=1, help='Initial kl coefficient.')

parser.add_argument('--pretrain', action='store_true', default=False,
                    help='If True, policy is pretrained by behavioral cloning.')
parser.add_argument('--bc_batch_size', type=int, default=256)
parser.add_argument('--bc_epoch', type=int, default=1000)

parser.add_argument('--rew_type', type=str,
                    choices=['adv', 'rew'], default='rew', help='Choice for reward type.')
args = parser.parse_args()

if not os.path.exists(args.log):
    os.makedirs(args.log)

with open(os.path.join(args.log, 'args.json'), 'w') as f:
    json.dump(vars(args), f)
pprint(vars(args))

if not os.path.exists(os.path.join(args.log, 'models')):
    os.makedirs(os.path.join(args.log, 'models'))

np.random.seed(args.seed)
torch.manual_seed(args.seed)

device_name = 'cpu' if args.cuda < 0 or args.rl_type == 'trpo' else "cuda:{}".format(
    args.cuda)
device = torch.device(device_name)
set_device(device)

score_file = os.path.join(args.log, 'progress.csv')
logger.add_tabular_output(score_file)
logger.add_tensorboard_output(args.log)

env = GymEnv(args.env_name, log_dir=os.path.join(
    args.log, 'movie'), record_video=args.record)
env.env.seed(args.seed)
if args.c2d:
    env = C2DEnv(env)

observation_space = env.observation_space
action_space = env.action_space

pol_net = PolNet(observation_space, action_space)
if isinstance(action_space, gym.spaces.Box):
    pol = GaussianPol(observation_space, action_space, pol_net)
elif isinstance(action_space, gym.spaces.Discrete):
    pol = CategoricalPol(observation_space, action_space, pol_net)
elif isinstance(action_space, gym.spaces.MultiDiscrete):
    pol = MultiCategoricalPol(observation_space, action_space, pol_net)
else:
    raise ValueError('Only Box, Discrete, and MultiDiscrete are supported')

vf_net = VNet(observation_space)
vf = DeterministicSVfunc(observation_space, vf_net)

if args.rew_type == 'rew':
    rewf_net = VNet(observation_space, h1=args.discrim_h1, h2=args.discrim_h2)
    rewf = DeterministicSVfunc(observation_space, rewf_net)
    shaping_vf_net = VNet(
        observation_space, h1=args.discrim_h1, h2=args.discrim_h2)
    shaping_vf = DeterministicSVfunc(observation_space, shaping_vf_net)
    optim_discrim = torch.optim.Adam(
        list(rewf_net.parameters()) + list(shaping_vf_net.parameters()), args.discrim_lr)
    advf = None
elif args.rew_type == 'adv':
    advf_net = DiscrimNet(observation_space, action_space,
                          h1=args.discrim_h1, h2=args.discrim_h2)
    advf = DeterministicSAVfunc(
        observation_space, action_space, advf_net)
    optim_discrim = torch.optim.Adam(advf_net.parameters(), args.discrim_lr)
    rewf = None
    shaping_vf = None
else:
    raise ValueError('Only rew and adv are supported')

sampler = EpiSampler(env, pol, num_parallel=args.num_parallel, seed=args.seed)

optim_pol = torch.optim.Adam(pol_net.parameters(), args.pol_lr)
optim_vf = torch.optim.Adam(vf_net.parameters(), args.vf_lr)

with open(os.path.join(args.expert_dir, args.expert_fname), 'rb') as f:
    expert_epis = pickle.load(f)
expert_traj = Traj()
expert_traj.add_epis(expert_epis)
expert_traj = ef.add_next_obs(expert_traj)
expert_traj.register_epis()
expert_rewards = [np.sum(epi['rews']) for epi in expert_epis]
expert_mean_rew = np.mean(expert_rewards)
logger.log('expert_score={}'.format(expert_mean_rew))
logger.log('expert_num_epi={}'.format(expert_traj.num_epi))

total_epi = 0
total_step = 0
max_rew = -1e6

if args.rl_type == 'ppo_kl':
    kl_beta = args.init_kl_beta

if args.pretrain:
    with measure('bc pretrain'):
        for _ in range(args.bc_epoch):
            _ = behavior_clone.train(
                expert_traj, pol, optim_pol, args.bc_batch_size
            )

while args.max_epis > total_epi:
    with measure('sample'):
        epis = sampler.sample(pol, max_steps=args.max_steps_per_iter)
    with measure('train'):
        agent_traj = Traj()
        agent_traj.add_epis(epis)
        agent_traj = ef.add_next_obs(agent_traj)
        agent_traj = ef.compute_pseudo_rews(agent_traj, rew_giver=rewf if args.rew_type ==
                                            'rew' else advf, state_only=True if args.rew_type == 'rew' else False)
        agent_traj = ef.compute_vs(agent_traj, vf)
        agent_traj = ef.compute_rets(agent_traj, args.gamma)
        agent_traj = ef.compute_advs(agent_traj, args.gamma, args.lam)
        agent_traj = ef.centerize_advs(agent_traj)
        agent_traj = ef.compute_h_masks(agent_traj)
        agent_traj.register_epis()

        if args.rl_type == 'trpo':
            result_dict = airl.train(agent_traj, expert_traj, pol, vf, optim_vf, optim_discrim,
                                     rewf=rewf, shaping_vf=shaping_vf, advf=advf,
                                     rl_type=args.rl_type,
                                     epoch=args.epoch_per_iter,
                                     batch_size=args.batch_size, discrim_batch_size=args.discrim_batch_size,
                                     discrim_step=args.discrim_step,
                                     pol_ent_beta=args.pol_ent_beta, gamma=args.gamma)
        elif args.rl_type == 'ppo_clip':
            result_dict = airl.train(agent_traj, expert_traj, pol, vf, optim_vf, optim_discrim,
                                     rewf=rewf, shaping_vf=shaping_vf, advf=advf,
                                     rl_type=args.rl_type,
                                     epoch=args.epoch_per_iter,
                                     batch_size=args.batch_size,
                                     discrim_batch_size=args.discrim_batch_size,
                                     discrim_step=args.discrim_step,
                                     pol_ent_beta=args.pol_ent_beta,
                                     optim_pol=optim_pol,
                                     clip_param=args.clip_param, max_grad_norm=args.max_grad_norm, gamma=args.gamma)

        else:
            result_dict = airl.train(agent_traj, expert_traj, pol, vf, optim_vf, optim_discrim,
                                     rewf=rewf, shaping_vf=shaping_vf, advf=advf,
                                     rl_type=args.rl_type,
                                     pol_ent_beta=args.pol_ent_beta,
                                     epoch=args.epoch_per_iter,
                                     batch_size=args.batch_size,
                                     discrim_batch_size=args.discrim_batch_size,
                                     discrim_step=args.discrim_step,
                                     optim_pol=optim_pol,
                                     kl_beta=kl_beta, kl_targ=args.kl_targ, max_grad_norm=args.max_grad_norm, gamma=args.gamma)
            kl_beta = result_dict['new_kl_beta']

    total_epi += agent_traj.num_epi
    step = agent_traj.num_step
    total_step += step
    rewards = [np.sum(epi['real_rews']) for epi in epis]
    mean_rew = np.mean(rewards)
    pseudo_rewards = [np.sum(epi['rews']) for epi in epis]
    result_dict['PseudoReward'] = pseudo_rewards
    logger.record_results(args.log, result_dict, score_file,
                          total_epi, step, total_step,
                          rewards,
                          plot_title=args.env_name)

    mean_rew = np.mean([np.sum(path['real_rews']) for path in epis])
    if mean_rew > max_rew:
        torch.save(pol.state_dict(), os.path.join(
            args.log, 'models', 'pol_max.pkl'))
        torch.save(vf.state_dict(), os.path.join(
            args.log, 'models', 'vf_max.pkl'))
        if args.rew_type == 'rew':
            torch.save(rewf.state_dict(), os.path.join(
                args.log, 'models', 'rewf_max.pkl'))
            torch.save(shaping_vf.state_dict(), os.path.join(
                args.log, 'models', 'shaping_vf_max.pkl'))
        else:
            torch.save(advf.state_dict(), os.path.join(
                args.log, 'models', 'advf_max.pkl'))
        torch.save(optim_pol.state_dict(), os.path.join(
            args.log, 'models', 'optim_pol_max.pkl'))
        torch.save(optim_vf.state_dict(), os.path.join(
            args.log, 'models', 'optim_vf_max.pkl'))
        torch.save(optim_discrim.state_dict(), os.path.join(
            args.log, 'models', 'optim_discrim_max.pkl'))
        max_rew = mean_rew

    torch.save(pol.state_dict(), os.path.join(
        args.log, 'models', 'pol_last.pkl'))
    torch.save(vf.state_dict(), os.path.join(
        args.log, 'models', 'vf_last.pkl'))
    if args.rew_type == 'rew':
        torch.save(rewf.state_dict(), os.path.join(
            args.log, 'models', 'rewf_last.pkl'))
        torch.save(shaping_vf.state_dict(), os.path.join(
            args.log, 'models', 'shaping_vf_last.pkl'))
    else:
        torch.save(advf.state_dict(), os.path.join(
            args.log, 'models', 'advf_last.pkl'))
    torch.save(optim_pol.state_dict(), os.path.join(
        args.log, 'models', 'optim_pol_last.pkl'))
    torch.save(optim_vf.state_dict(), os.path.join(
        args.log, 'models', 'optim_vf_last.pkl'))
    torch.save(optim_discrim.state_dict(), os.path.join(
        args.log, 'models', 'optim_discrim_last.pkl'))
    del agent_traj
del sampler


#! /usr/bin/env python
# coding=utf-8

import os
import time
import shutil
import numpy as np
import tensorflow as tf
import core.utils as utils
from tqdm import tqdm
from core.dataset import Dataset
from core.yolov3 import YOLOV3
from core.config import cfg

class YoloTrain(object):
    def __init__(self):  # 从config文件获取到一些变量
        self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE
        self.classes = utils.read_class_names(cfg.YOLO.CLASSES)
        self.num_classes = len(self.classes)
        self.learn_rate_init = cfg.TRAIN.LEARN_RATE_INIT
        self.learn_rate_end = cfg.TRAIN.LEARN_RATE_END
        self.first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
        self.second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
        self.warmup_periods = cfg.TRAIN.WARMUP_EPOCHS
        self.initial_weight = cfg.TRAIN.INITIAL_WEIGHT
        self.time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
        self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY
        self.max_bbox_per_scale = 150
        self.train_logdir = "./data/log/train"  # 日志保存地址
        self.trainset = Dataset('train')
        self.testset = Dataset('test')
        self.steps_per_period = len(self.trainset)
        self.sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))

        with tf.name_scope('define_input'): # 定义输入层
            self.input_data = tf.placeholder(dtype=tf.float32, name='input_data')
            self.label_sbbox = tf.placeholder(dtype=tf.float32, name='label_sbbox')
            self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox')
            self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox')
            self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes')
            self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes')
            self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes')
            self.trainable = tf.placeholder(dtype=tf.bool, name='training')

        with tf.name_scope("define_loss"): # 定义损失函数
            self.model = YOLOV3(self.input_data, self.trainable)
            self.net_var = tf.global_variables()
            self.giou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss(
                self.label_sbbox, self.label_mbbox, self.label_lbbox,
                self.true_sbboxes, self.true_mbboxes, self.true_lbboxes)
            self.loss = self.giou_loss + self.conf_loss + self.prob_loss

        with tf.name_scope('learn_rate'): # 定义学习率
            self.global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step')
            warmup_steps = tf.constant(self.warmup_periods * self.steps_per_period,
                                       dtype=tf.float64, name='warmup_steps')
            train_steps = tf.constant((self.first_stage_epochs + self.second_stage_epochs) * self.steps_per_period,
                                      dtype=tf.float64, name='train_steps')
            self.learn_rate = tf.cond(
                pred=self.global_step < warmup_steps,
                true_fn=lambda: self.global_step / warmup_steps * self.learn_rate_init,
                false_fn=lambda: self.learn_rate_end + 0.5 * (self.learn_rate_init - self.learn_rate_end) *
                                 (1 + tf.cos(
                                     (self.global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))
            )
            global_step_update = tf.assign_add(self.global_step, 1.0)

        '''
        warmup_steps作用：
        神经网络在刚开始训练的过程中容易出现loss=NaN的情况，为了尽量避免这个情况，因此初始的学习率设置得很低
        但是这又使得训练速度变慢了。因此，采用逐渐增大的学习率，从而达到既可以尽量避免出现nan，又可以等训练过程稳定了再增大训练速度的目的。
        '''

        with tf.name_scope("define_weight_decay"):  # 指数平滑，可以让算法在最后不那么震荡，结果更有鲁棒性
            moving_ave = tf.train.ExponentialMovingAverage(self.moving_ave_decay).apply(tf.trainable_variables())

        # 指定需要恢复的参数。层等信息, 位置提前，减少模型体积。
        with tf.name_scope('loader_and_saver'):
            variables_to_restore = [v for v in self.net_var if
                                    v.name.split('/')[0] not in ['conv_sbbox', 'conv_mbbox', 'conv_lbbox']]
            self.loader = tf.train.Saver(variables_to_restore)
            self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=5)

        with tf.name_scope("define_first_stage_train"): # 第一阶段训练，只训练指定层
            self.first_stage_trainable_var_list = []
            for var in tf.trainable_variables():
                var_name = var.op.name
                var_name_mess = str(var_name).split('/')
                if var_name_mess[0] in ['conv_sbbox', 'conv_mbbox', 'conv_lbbox']:
                    self.first_stage_trainable_var_list.append(var)

            first_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate).minimize(self.loss,
                                                                                var_list=self.first_stage_trainable_var_list)
            with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
                with tf.control_dependencies([first_stage_optimizer, global_step_update]):
                    with tf.control_dependencies([moving_ave]):
                        self.train_op_with_frozen_variables = tf.no_op()

        with tf.name_scope("define_second_stage_train"):    # 第二阶段训练，释放所有层
            second_stage_trainable_var_list = tf.trainable_variables()
            second_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate).minimize(self.loss,
                                                                                var_list=second_stage_trainable_var_list)

            with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
                with tf.control_dependencies([second_stage_optimizer, global_step_update]):
                    with tf.control_dependencies([moving_ave]):
                        self.train_op_with_all_variables = tf.no_op()

        with tf.name_scope('summary'):
            tf.summary.scalar("learn_rate", self.learn_rate)
            tf.summary.scalar("giou_loss", self.giou_loss)
            tf.summary.scalar("conf_loss", self.conf_loss)
            tf.summary.scalar("prob_loss", self.prob_loss)
            tf.summary.scalar("total_loss", self.loss)

            logdir = "./data/log/"  #  日志保存地址
            if os.path.exists(logdir): shutil.rmtree(logdir)
            os.mkdir(logdir)
            self.write_op = tf.summary.merge_all()
            self.summary_writer = tf.summary.FileWriter(logdir, graph=self.sess.graph)

    def train(self):
        self.sess.run(tf.global_variables_initializer())
        try:
            print('=> Restoring weights from: %s ... ' % self.initial_weight)
            self.loader.restore(self.sess, self.initial_weight)
        except:
            print('=> %s does not exist !!!' % self.initial_weight)
            print('=> Now it starts to train YOLOV3 from scratch ...')
            self.first_stage_epochs = 0

        # 阶段学习率
        for epoch in range(1, 1 + self.first_stage_epochs + self.second_stage_epochs):
            if epoch <= self.first_stage_epochs:
                train_op = self.train_op_with_frozen_variables
            else:
                train_op = self.train_op_with_all_variables

            # tqdm is a visualization tool that displays an Iterable object in a progree bar
            pbar = tqdm(self.trainset)
            train_epoch_loss, test_epoch_loss = [], []

            for train_data in pbar:
                _, summary, train_step_loss, global_step_val = self.sess.run(
                    [train_op, self.write_op, self.loss, self.global_step], feed_dict={
                        self.input_data: train_data[0],
                        self.label_sbbox: train_data[1],
                        self.label_mbbox: train_data[2],
                        self.label_lbbox: train_data[3],
                        self.true_sbboxes: train_data[4],
                        self.true_mbboxes: train_data[5],
                        self.true_lbboxes: train_data[6],
                        self.trainable: True,
                    })

                train_epoch_loss.append(train_step_loss)
                self.summary_writer.add_summary(summary, global_step_val)
                pbar.set_description("train loss: %.2f" % train_step_loss)

            for test_data in self.testset:
                test_step_loss = self.sess.run(self.loss, feed_dict={
                    self.input_data: test_data[0],
                    self.label_sbbox: test_data[1],
                    self.label_mbbox: test_data[2],
                    self.label_lbbox: test_data[3],
                    self.true_sbboxes: test_data[4],
                    self.true_mbboxes: test_data[5],
                    self.true_lbboxes: test_data[6],
                    self.trainable: False,
                })

                test_epoch_loss.append(test_step_loss)

            train_epoch_loss, test_epoch_loss = np.mean(train_epoch_loss), np.mean(test_epoch_loss)
            ckpt_file = "./checkpoint/yolov3_train_loss=%.4f.ckpt" % train_epoch_loss
            log_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
            print("=> Epoch: %2d Time: %s Train loss: %.2f Test loss: %.2f Saving %s ..."
                  % (epoch, log_time, train_epoch_loss, test_epoch_loss, ckpt_file))
            self.saver.save(self.sess, ckpt_file, global_step=epoch)

if __name__ == '__main__': YoloTrain().train()