import numpy as np

from .name2idx import C, V
from .set_model import param_values, initial_values

def get_search_index():
    """Specify model parameters and/or initial values to optimize
    """
    # parameters
    search_idx_params = [
        C.V1,
        C.Km1,
        C.V5,
        C.Km5,
        C.V10,
        C.Km10,
        C.n10,
        C.p11,
        C.p12,
        C.p13,
        C.V14,
        C.Km14,
        C.V15,
        C.Km15,
        C.KimDUSP,
        C.KexDUSP,
        C.V20,
        C.Km20,
        C.V21,
        C.Km21,
        C.V24,
        C.Km24,
        C.V25,
        C.Km25,
        C.KimRSK,
        C.KexRSK,
        C.V27,
        C.Km27,
        C.V28,
        C.Km28,
        C.V29,
        C.Km29,
        C.V30,
        C.Km30,
        C.V31,
        C.Km31,
        C.n31,
        C.p32,
        C.p33,
        C.p34,
        C.V35,
        C.Km35,
        C.V36,
        C.Km36,
        C.V37,
        C.Km37,
        C.KimFOS,
        C.KexFOS,
        C.V42,
        C.Km42,
        C.V43,
        C.Km43,
        C.V44,
        C.Km44,
        C.p47,
        C.m47,
        C.p48,
        C.p49,
        C.m49,
        C.p50,
        C.p51,
        C.m51,
        C.V57,
        C.Km57,
        C.n57,
        C.p58,
        C.p59,
        C.p60,
        C.p61,
        C.KimF,
        C.KexF,
        C.p63,
        C.KF31,
        C.nF31,
        C.a,
    ]

    # initial values
    search_idx_initials = [
        # V.(variable name)
    ]

    return search_idx_params, search_idx_initials

def get_search_region():
    x = param_values()
    y0 = initial_values()

    search_idx = get_search_index()
    search_param = _init_search_param(search_idx, x, y0)

    search_rgn = np.zeros((2, len(x)+len(y0)))
    # Default: 0.1 ~ 10
    for i, j in enumerate(search_idx[0]):
        search_rgn[0, j] = search_param[i] * 0.1  # lower bound
        search_rgn[1, j] = search_param[i] * 10.  # upper bound

    # Default: 0.5 ~ 2
    for i, j in enumerate(search_idx[1]):
        search_rgn[0, j+len(x)] = \
            search_param[i+len(search_idx[0])] * 0.5  # lower bound
        search_rgn[1, j+len(x)] = \
            search_param[i+len(search_idx[0])] * 2.0  # upper bound

    # search_rgn[:,C.param_name] = [lower_bound,upper_bound]
    # search_rgn[:,V.var_name+len(x)] = [lower_bound,upper_bound]

    # Hill coefficient
    search_rgn[:, C.n10] = [1.00, 4.00]
    search_rgn[:, C.n31] = [1.00, 4.00]
    search_rgn[:, C.n57] = [1.00, 4.00]
    search_rgn[:, C.nF31] = [1.00, 4.00]

    ''' Example ----------------------------------------------------------------
    search_rgn[:, C.V1] = [7.33e-2, 6.60e-01]
    search_rgn[:, C.Km1] = [1.83e+2, 8.50e+2]
    search_rgn[:, C.V5] = [6.48e-3, 7.20e+1]
    search_rgn[:, C.Km5] = [6.00e-1, 1.60e+04]
    search_rgn[:, C.V10] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km10] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.n10] = [1.00, 4.00]
    search_rgn[:, C.p11] = [8.30e-13, 1.44e-2]
    search_rgn[:, C.p12] = [8.00e-8, 5.17e-2]
    search_rgn[:, C.p13] = [1.38e-7, 4.84e-1]
    search_rgn[:, C.V14] = [4.77e-3, 4.77e+1]
    search_rgn[:, C.Km14] = [2.00e+2, 2.00e+6]
    search_rgn[:, C.V15] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km15] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.KimDUSP] = [2.20e-4, 5.50e-1]
    search_rgn[:, C.KexDUSP] = [2.60e-4, 6.50e-1]
    search_rgn[:, C.V20] = [4.77e-3, 4.77e+1]
    search_rgn[:, C.Km20] = [2.00e+2, 2.00e+6]
    search_rgn[:, C.V21] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km21] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.V24] = [4.77e-2, 4.77e+0]
    search_rgn[:, C.Km24] = [2.00e+3, 2.00e+5]
    search_rgn[:, C.V25] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km25] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.KimRSK] = [2.20e-4, 5.50e-1]
    search_rgn[:, C.KexRSK] = [2.60e-4, 6.50e-1]
    search_rgn[:, C.V27] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km27] = [1.00e+2, 1.00e+4]
    search_rgn[:, C.V28] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km28] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.V29] = [4.77e-2, 4.77e+0]
    search_rgn[:, C.Km29] = [2.93e+3, 2.93e+5]
    search_rgn[:, C.V30] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km30] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.V31] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km31] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.n31] = [1.00, 4.00]
    search_rgn[:, C.p32] = [8.30e-13, 1.44e-2]
    search_rgn[:, C.p33] = [8.00e-8, 5.17e-2]
    search_rgn[:, C.p34] = [1.38e-7, 4.84e-1]
    search_rgn[:, C.V35] = [4.77e-3, 4.77e+1]
    search_rgn[:, C.Km35] = [2.00e+2, 2.00e+6]
    search_rgn[:, C.V36] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km36] = [1.00e+2, 1.00e+4]
    search_rgn[:, C.V37] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km37] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.KimFOS] = [2.20e-4, 5.50e-1]
    search_rgn[:, C.KexFOS] = [2.60e-4, 6.50e-1]
    search_rgn[:, C.V42] = [4.77e-3, 4.77e+1]
    search_rgn[:, C.Km42] = [2.00e+2, 2.00e+6]
    search_rgn[:, C.V43] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km43] = [1.00e+2, 1.00e+4]
    search_rgn[:, C.V44] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km44] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.p47] = [1.45e-4, 1.45e+0]
    search_rgn[:, C.m47] = [6.00e-3, 6.00e+1]
    search_rgn[:, C.p48] = [2.70e-3, 2.70e+1]
    search_rgn[:, C.p49] = [5.00e-5, 5.00e-1]
    search_rgn[:, C.m49] = [5.00e-3, 5.00e+1]
    search_rgn[:, C.p50] = [3.00e-3, 3.00e+1]
    search_rgn[:, C.p51] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.m51] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.V57] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.Km57] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.n57] = [1.00, 4.00]
    search_rgn[:, C.p58] = [8.30e-13, 1.44e-2]
    search_rgn[:, C.p59] = [8.00e-8, 5.17e-2]
    search_rgn[:, C.p60] = [1.38e-7, 4.84e-1]
    search_rgn[:, C.p61] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.KimF] = [2.20e-4, 5.50e-1]
    search_rgn[:, C.KexF] = [2.60e-4, 6.50e-1]
    search_rgn[:, C.p63] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.KF31] = [np.exp(-10), np.exp(10)]
    search_rgn[:, C.nF31] = [1.00, 4.00]
    search_rgn[:, C.a] = [1.00e+2, 5.00e+2]

    ----------------------------------------------------------------------------
    '''

    search_rgn = _conv_lin2log(search_rgn, search_idx)
    return search_rgn

def update_param(indiv):
    x = param_values()
    y0 = initial_values()

    search_idx = get_search_index()

    for i,j in enumerate(search_idx[0]):
        x[j] = indiv[i]
    for i,j in enumerate(search_idx[1]):
        y0[j] = indiv[i+len(search_idx[0])]

    # constraints --------------------------------------------------------------
    x[C.V6] = x[C.V5]
    x[C.Km6] = x[C.Km5]
    x[C.KimpDUSP] = x[C.KimDUSP]
    x[C.KexpDUSP] = x[C.KexDUSP]
    x[C.KimpcFOS] = x[C.KimFOS]
    x[C.KexpcFOS] = x[C.KexFOS]
    x[C.p52] = x[C.p47]
    x[C.m52] = x[C.m47]
    x[C.p53] = x[C.p48]
    x[C.p54] = x[C.p49]
    x[C.m54] = x[C.m49]
    x[C.p55] = x[C.p50]
    x[C.p56] = x[C.p51]
    x[C.m56] = x[C.m51]
    # --------------------------------------------------------------------------

    return x, y0

def decode_gene2val(indiv_gene):
    search_rgn = get_search_region()
    indiv = 10**(
        indiv_gene * (
            search_rgn[1, :] - search_rgn[0, :]
        ) + search_rgn[0, :]
    )

    return indiv

def _init_search_param(search_idx, x, y0):
    """Initialize search_param
    """
    if len(search_idx[0]) != len(set(search_idx[0])):
        raise ValueError('Duplicate param name.')
    elif len(search_idx[1]) != len(set(search_idx[1])):
        raise ValueError('Duplicate var name.')
    else:
        pass

    search_param = np.empty(
        len(search_idx[0]) + len(search_idx[1])
    )
    for i, j in enumerate(search_idx[0]):
        search_param[i] = x[j]
    for i, j in enumerate(search_idx[1]):
        search_param[i+len(search_idx[0])] = y0[j]

    if np.any(search_param == 0.):
        message = 'search_param must not contain zero.'
        for _, idx in enumerate(search_idx[0]):
            if x[int(idx)] == 0.:
                raise ValueError(
                    '"C.{}" in search_idx_params: '.format(
                        C.NAMES[int(idx)]
                    ) + message
                )
        for _, idx in enumerate(search_idx[1]):
            if y0[int(idx)] == 0.:
                raise ValueError(
                    '"V.{}" in search_idx_initials: '.format(
                        V.NAMES[int(idx)]
                    ) + message
                )

    return search_param

def _conv_lin2log(search_rgn, search_idx):
    """Convert Linear scale to Logarithmic scale
    """
    for i in range(search_rgn.shape[1]):
        if np.min(search_rgn[:, i]) < 0.0:
            message = 'search_rgn[lb,ub] must be positive.'
            if i <= C.NUM:
                raise ValueError(
                    '"C.{}": '.format(C.NAMES[i]) + message
                )
            else:
                raise ValueError(
                    '"V.{}": '.format(V.NAMES[i-C.NUM]) + message
                )
        elif np.min(search_rgn[:, i]) == 0 and np.max(search_rgn[:, i]) != 0:
            message = 'lower_bound must be larger than 0.'
            if i <= C.NUM:
                raise ValueError(
                    '"C.{}" '.format(C.NAMES[i]) + message
                )
            else:
                raise ValueError(
                    '"V.{}" '.format(V.NAMES[i-C.NUM]) + message
                )
        elif search_rgn[1, i] - search_rgn[0, i] < 0.0:
            message = 'lower_bound < upper_bound'
            if i <= C.NUM:
                raise ValueError(
                    '"C.{}" : '.format(C.NAMES[i]) + message
                )
            else:
                raise ValueError(
                    '"V.{}" : '.format(V.NAMES[i-C.NUM]) + message
                )
    difference = list(
        set(
            np.where(
                np.any(search_rgn != 0., axis=0)
            )[0]
        ) ^ set(
            np.append(
                search_idx[0],
                [C.NUM + idx for (_, idx) in enumerate(search_idx[1])]
            )
        )
    )
    if len(difference) > 0:
        message = 'in both search_idx and search_rgn'
        for _, idx in enumerate(difference):
            if idx <= C.NUM:
                raise ValueError(
                    'Set "C.{}" '.format(C.NAMES[int(idx)]) + message
                )
            else:
                raise ValueError(
                    'Set "V.{}" '.format(V.NAMES[int(idx-C.NUM)]) + message
                )
    search_rgn = search_rgn[:, np.any(search_rgn != 0., axis=0)]

    return np.log10(search_rgn)


# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# 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.

import os
import datetime
import numpy as np
import nnabla as nn
import nnabla.functions as F
import nnabla.solvers as S
from nnabla.ext_utils import get_extension_context
from nnabla.monitor import Monitor
from tecogan_model import get_tecogan_model, get_frvsr_model, get_common_monitors, get_tecogan_monitors
from utils import CommunicatorWrapper, save_checkpoint, load_checkpoint
from utils.utils import ExponentialMovingAverage
from args import get_config

def main():
    conf = get_config()
    extension_module = conf.nnabla_context.context
    ctx = get_extension_context(
        extension_module, device_id=conf.nnabla_context.device_id)
    comm = CommunicatorWrapper(ctx)
    nn.set_default_context(comm.ctx)
    print("comm rank", comm.rank)

    # data iterators for train and val data
    from data_loader import data_iterator_sr, get_sample_name_grid, nn_data_gauss_down_quad

    sample_names = get_sample_name_grid(conf)
    num_samples = len(sample_names[0])
    print("No of training samples :", num_samples)

    tar_size = conf.train.crop_size
    tar_size = (conf.train.crop_size * 4) + int(1.5 * 3.0) * \
        2  # crop_size * 4, and Gaussian blur margin

    data_iterator_train = data_iterator_sr(
        conf, num_samples, sample_names, tar_size, shuffle=True)

    if comm.n_procs > 1:
        data_iterator_train = data_iterator_train.slice(
            rng=None, num_of_slices=comm.n_procs, slice_pos=comm.rank)

    train_hr = nn.Variable(
        (conf.train.batch_size, conf.train.rnn_n, conf.train.crop_size*4, conf.train.crop_size*4, 3))
    data_hr = nn.Variable(
        (conf.train.batch_size, conf.train.rnn_n, tar_size, tar_size, 3))
    train_lr = nn_data_gauss_down_quad(data_hr.reshape(
        (conf.train.batch_size * conf.train.rnn_n, tar_size, tar_size, 3)))
    train_lr = F.reshape(
        train_lr, (conf.train.batch_size, conf.train.rnn_n, conf.train.crop_size, conf.train.crop_size, 3))

    # setting up monitors for logging
    monitor_path = './nnmonitor' + \
        str(datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
    monitor = Monitor(monitor_path)
    common_monitor = get_common_monitors(monitor)

    # Change max_iter and learning_rate when batch size or no. of gpu devices change.
    div_factor = conf.train.batch_size * comm.n_procs
    max_iter = (conf.train.max_iter * 4) // div_factor
    learning_rate = conf.train.learning_rate * \
        (conf.train.batch_size / 4) * comm.n_procs

    if comm.rank == 0:
        print("maximum iterations", max_iter)

    scope_name = 'frvsr/'
    if conf.train.tecogan:
        scope_name = 'tecogan/'
        if not conf.train.checkpoint:
            print('loading pretrained FRVSR model',
                  conf.train.pre_trained_frvsr_weights)
            with nn.parameter_scope(scope_name):
                nn.load_parameters(conf.train.pre_trained_frvsr_weights)
                params_from_pre_trained_model = []
                for key, val in nn.get_parameters().items():
                    params_from_pre_trained_model.append(scope_name + key)

            network = get_tecogan_model(conf, train_lr, train_hr, scope_name)
            params_from_graph = nn.get_parameters()

            # Set the Generator parameters which are not in FRVSR to zero,
            # as done in orig implementation.
            for key, val in params_from_graph.items():
                if key in params_from_pre_trained_model or key.startswith('vgg') or key.startswith('disc'):
                    continue
                print(key)
                val.data.zero()  # fill with zero

        else:
            network = get_tecogan_model(conf, train_lr, train_hr, scope_name)

        # Define discriminator optimizer/solver
        solver_disc = S.Adam(alpha=learning_rate,
                             beta1=conf.train.beta, eps=conf.train.adameps)
        # Set discriminator Parameters
        with nn.parameter_scope("discriminator"):
            solver_disc.set_parameters(nn.get_parameters())

        # setting up monitors for TecoGAN
        tecogan_monitor = get_tecogan_monitors(monitor)

    else:
        network = get_frvsr_model(conf, train_lr, train_hr, scope_name)

    # Define generator and fnet optimizer/solver
    solver_gen = S.Adam(alpha=learning_rate,
                        beta1=conf.train.beta, eps=conf.train.adameps)
    solver_fnet = S.Adam(alpha=learning_rate,
                         beta1=conf.train.beta, eps=conf.train.adameps)

    # Set generator and fnet Parameters
    with nn.parameter_scope(scope_name + "generator"):
        solver_gen.set_parameters(nn.get_parameters())
    with nn.parameter_scope(scope_name + "fnet"):
        solver_fnet.set_parameters(nn.get_parameters())

    if conf.train.tecogan:
        solver_dict = {"gen": solver_gen,
                       "fnet": solver_fnet, "disc": solver_disc}
    else:
        solver_dict = {"gen": solver_gen, "fnet": solver_fnet}

    start_point = 0
    if conf.train.checkpoint:
        # Load optimizer/solver information and model weights from checkpoint
        start_point = load_checkpoint(conf.train.checkpoint, solver_dict)

    # Exponential Moving Average Calculation for tb
    ema = ExponentialMovingAverage(conf.train.decay)
    tb = 0

    # Create output directory if it doesn't exist
    if not os.path.exists(conf.data.output_dir):
        os.makedirs(conf.data.output_dir)

    # Training loop.
    for i in range(start_point, max_iter):
        # Get Training Data
        data_hr.d, train_hr.d = data_iterator_train.next()

        if conf.train.tecogan:
            network.t_discrim_loss.forward(clear_no_need_grad=True)
            if np.less(tb, 0.4):  # train gen with d
                # Compute grads for discriminator and update
                solver_disc.zero_grad()
                # Stop back-propagation from t_discrim_loss to generator
                network.t_gen_output.need_grad = False
                if comm.n_procs > 1:
                    all_reduce_callback = comm.get_all_reduce_callback()
                    network.t_discrim_loss.backward(clear_buffer=True,
                                                    communicator_callbacks=all_reduce_callback)
                else:
                    network.t_discrim_loss.backward(clear_buffer=True)
                solver_disc.update()  # Update grads
                # Enable back propagation from fnet_loss to Generator
                network.t_gen_output.need_grad = True

        # Compute grads for fnet and generator together using fnet_loss
        solver_fnet.zero_grad()
        solver_gen.zero_grad()
        # Apply forward and backward propagation on fnet_loss
        network.fnet_loss.forward(clear_no_need_grad=True)
        if comm.n_procs > 1:
            all_reduce_callback = comm.get_all_reduce_callback()
            network.fnet_loss.backward(clear_buffer=True,
                                       communicator_callbacks=all_reduce_callback)
        else:
            network.fnet_loss.backward(clear_buffer=True)
        # Update grads for fnet and generator
        solver_gen.update()
        solver_fnet.update()

        if conf.train.tecogan:
            if comm.n_procs > 1:
                comm.all_reduce([network.t_discrim_real_loss.data,
                                 network.t_adversarial_loss.data], division=True, inplace=True)
            t_balance = F.mean(network.t_discrim_real_loss.data) + \
                network.t_adversarial_loss.data
            if i == 0:
                ema.register(t_balance)
            else:
                tb = ema(t_balance)
            if comm.rank == 0:
                tecogan_monitor.monitor_pp_loss.add(
                    i, network.pp_loss.d.copy())
                tecogan_monitor.monitor_vgg_loss.add(
                    i, network.vgg_loss.d.copy())
                tecogan_monitor.monitor_sum_layer_loss.add(
                    i, network.sum_layer_loss.d.copy())
                tecogan_monitor.monitor_adv_loss.add(
                    i, network.t_adversarial_loss.d.copy())
                tecogan_monitor.monitor_disc_loss.add(
                    i, network.t_discrim_loss.d.copy())
                tecogan_monitor.monitor_tb.add(i, tb)

        if comm.rank == 0:
            common_monitor.monitor_content_loss.add(
                i, network.content_loss.d.copy())
            common_monitor.monitor_gen_loss.add(i, network.gen_loss.d.copy())
            common_monitor.monitor_warp_loss.add(i, network.warp_loss.d.copy())
            common_monitor.monitor_lr.add(i, learning_rate)
            common_monitor.monitor_time.add(i)
            if (i % conf.train.save_freq) == 0:
                # Save intermediate model parameters
                with nn.parameter_scope(scope_name):
                    nn.save_parameters(os.path.join(
                        conf.data.output_dir, "model_param_%08d.h5" % i))

                # Save intermediate check_points
                save_checkpoint(conf.data.output_dir, i, solver_dict)

    # save final Generator and Fnet network parameters
    if comm.rank == 0:
        with nn.parameter_scope(scope_name):
            nn.save_parameters(os.path.join(
                conf.data.output_dir, "model_param_%08d.h5" % i))

if __name__ == "__main__":
    main()


#!/usr/bin/env python
"""Sample script of recurrent neural network language model.

This code is ported from following implementation written in Torch.
https://github.com/tomsercu/lstm

"""
from __future__ import print_function
import argparse
import math
import sys
import time

import numpy as np
import six

import six.moves.cPickle as pickle

import chainer
from chainer import cuda
import chainer.functions as F
import chainer.links as L
from chainer import optimizers
from chainer import serializers

import net
import dataset
import os

parser = argparse.ArgumentParser()
parser.add_argument('--initmodel', '-m', default='',
                    help='Initialize the model from given file')
parser.add_argument('--resume', '-r', default='',
                    help='Resume the optimization from snapshot')
parser.add_argument('--gpu', '-g', default=-1, type=int,
                    help='GPU ID (negative value indicates CPU)')
parser.add_argument('--epoch', '-e', default=10, type=int,
                    help='number of epochs to learn')
parser.add_argument('--unit', '-u', default=650, type=int,
                    help='number of units')
parser.add_argument('--batchsize', '-b', type=int, default=32,
                    help='learning minibatch size')
parser.add_argument('--seqlen', '-l', type=int, default=200,
                    help='sequence length')
parser.add_argument('--gradclip', '-c', type=int, default=5,
                    help='gradient norm threshold to clip')
parser.add_argument('--test', dest='test', action='store_true')
parser.add_argument('--check', dest='check', action='store_true')
parser.add_argument('--input', '-i', default='data/reviews.txt',
                    help='input text file')
parser.add_argument('--label', default='data/labels.txt.',
                    help='label file')
parser.set_defaults(test=False)

args = parser.parse_args()
xp = cuda.cupy if args.gpu >= 0 else np

n_epoch = args.epoch   # number of epochs
n_units = args.unit  # number of units per layer
batchsize = args.batchsize   # minibatch size
seq_len = args.seqlen   # sequence length
grad_clip = args.gradclip    # gradient norm threshold to clip

ds = dataset.Dataset(args.input, label=args.label, seq_len=seq_len)

# Load data
train, test = ds.get_inputs_and_labels()
print('#data size =', train[0].shape, train[1].shape)

# Get vocabulary dictionary
input_vocab, label_vocab = ds.get_vocab()
print('#vocab size =', len(input_vocab), len(label_vocab))

if args.test:
    print('TEST mode!')
    train = (train[0][0:100], train[1][0:100])
    test = (test[0][0:100], test[1][0:100:])
    print(train[0].shape)

with open(os.path.join(os.path.dirname(args.input), 'input_vocab.bin'), 'wb') as f:
    pickle.dump(input_vocab, f)
with open(os.path.join(os.path.dirname(args.input), 'label_vocab.bin'), 'wb') as f:
    pickle.dump(label_vocab, f)
print('save vocabulary dictionary')

# Prepare RNNLM model, defined in net.py
lm = net.RNNLM(len(input_vocab), len(label_vocab), n_units)
model = L.Classifier(lm)
if args.gpu >= 0:
    cuda.get_device(args.gpu).use()
    model.to_gpu()

# Setup optimizer
#optimizer = optimizers.SGD(lr=1.)
optimizer = optimizers.Adam()
optimizer.setup(model)
optimizer.add_hook(chainer.optimizer.GradientClipping(grad_clip))

# Init/Resume
if args.initmodel:
    print('Load model from', args.initmodel)
    serializers.load_npz(args.initmodel, model)
if args.resume:
    print('Load optimizer state from', args.resume)
    serializers.load_npz(args.resume, optimizer)

def evaluate(dataset):
    # Evaluation routine
    evaluator = model.copy()  # to use different state
    evaluator.predictor.reset_state()  # initialize state
    evaluator.predictor.train = False  # dropout does nothing

    sum_log_perp = 0
    sum_accuracy = 0
    jump = len(dataset[0]) // batchsize

    for i in six.moves.range(0, jump * batchsize, batchsize):
        evaluator.predictor.reset_state()  # initialize state
        x_batch = dataset[0][i:i + batchsize]
        t_batch = dataset[1][i:i + batchsize]
        for j in range(seq_len):
            x = chainer.Variable(xp.asarray([x_batch[i][j] for i in range(batchsize)]), volatile='on')
            t = chainer.Variable(xp.asarray([t_batch[i] for i in range(batchsize)]), volatile='on')
            loss = evaluator(x, t)
            if j == seq_len - 1:
                evaluator.accuracy.to_cpu()
                sum_log_perp += loss.data
                sum_accuracy += evaluator.accuracy.data
    return math.exp(float(sum_log_perp) / jump), sum_accuracy / jump

def predict(dataset):
    evaluator = model
    evaluator.predictor.reset_state()  # initialize state
    evaluator.predictor.train = False  # dropout does nothing

    count = 0
    match = 0
    jump = len(dataset[0])

    for i in six.moves.range(jump):
        evaluator.predictor.reset_state()  # initialize state
        x_data = dataset[0][i]
        for j in range(seq_len):
            x = chainer.Variable(xp.asarray([x_data[j]]), volatile='on')
            pred = F.softmax(model.predictor(x))
            if j == seq_len - 1:
                if args.gpu >= 0:
                    pred_data = cuda.to_cpu(pred.data)
                else:
                    pred_data = pred.data
                indice = pred_data[0].argsort()[-1:][::-1]
                y = dataset[1][i]
                count += 1
                if indice[0] == y:
                    match += 1
    return match / count

# Learning loop
whole_len = train[0].shape[0]
print('whole_len=', whole_len)
seq_len = train[0].shape[1]
print('seq_len=', seq_len)
jump = whole_len // batchsize
cur_log_perp = xp.zeros(())
epoch = 0
start_at = time.time()
cur_at = start_at
accum_loss = 0
print('going to train {} iterations'.format(jump * n_epoch))

iteration = 0
for epoch in six.moves.range(1, n_epoch + 1):
    # print('epoch', epoch)
    for i in six.moves.range(0, jump * batchsize, batchsize):
        iteration += 1
        x_batch = train[0][i:i + batchsize]
        t_batch = train[1][i:i + batchsize]
        for j in range(seq_len):
            x = chainer.Variable(xp.asarray([x_batch[i][j] for i in range(batchsize)]))
            t = chainer.Variable(xp.asarray([t_batch[i] for i in range(batchsize)]))
            loss_i = model(x, t)

            if j == seq_len - 1:  # Run truncated BPTT
                # Last sequence output matters
                accum_loss += loss_i
                cur_log_perp += loss_i.data
                if iteration % 5 == 0:
                    print('epoch=', epoch, 'iteration=', iteration, 'loss=', loss_i.data)

                model.zerograds()
                accum_loss.backward()
                accum_loss.unchain_backward()  # truncate
                accum_loss = 0
                optimizer.update()
                model.predictor.reset_state()  # initialize state
            sys.stdout.flush()

        if iteration % 1000 == 0:
            now = time.time()
            throuput = 10000. / (now - cur_at)
            perp = math.exp(float(cur_log_perp) / 1000)
            print('iter {} training perplexity: {:.2f} ({:.2f} iters/sec)'.format(
                iteration, perp, throuput))
            cur_at = now
            cur_log_perp.fill(0)

    print('evaluate...')
    now = time.time()
    perp, acc = evaluate(test)
    print('epoch {} validation perplexity: {:.2f} validation accuracy: {:.2f}'.format(epoch, perp, float(acc)))  # acc = <class 'cupy.core.core.ndarray'>
    cur_at += time.time() - now  # skip time of evaluation

    # if epoch >= 6:
    #     optimizer.lr /= 1.2
    #     print('learning rate =', optimizer.lr)

# Evaluate on test dataset
print('test...')
test_perp, test_acc = evaluate(test)
print('test perplexity:', test_perp)
print('test accuracy:', test_acc)

if args.check:
    test_acc = predict(test)
    print('test accuracy(check):', test_acc)

# Save the model and the optimizer
print('save the model')
serializers.save_npz(os.path.join(os.path.dirname(args.input), 'rnnlm.model'), model)


import os
import sys

sys.path.insert(0, "./")

from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

import json
from tqdm import tqdm

from comet_ml import Experiment
import numpy as np
from sklearn.metrics import pairwise

import tensorflow as tf

import logging
from logging import handlers

from constants import CONSTANTS

from fewshot.experiments.oc_fewshot import evaluate, get_stats
from fewshot.experiments.utils import get_config
from fewshot.experiments.build_model import build_net, build_pretrain_net

from fewshot.experiments.oc_fewshot import (
    ExperimentConfig,
    EpisodeConfig,
    EnvironmentConfig,
    get_data_fs,
    get_dataiter_continual,
    get_dataiter_sim,
    # build_net,
)
from fewshot.experiments.get_data_iter import get_dataiter
from fewshot.models.nets.net import Net, ContainerModule
from fewshot.models.modules.nnlib import Linear
from my_code.finetune import get_loaders, build_datasets, init_writers, evaluate_many

from dotenv import load_dotenv

def build_model(config, data_config, add_last_relu=None, num_classes=None):
    if add_last_relu is not None:
        if config.backbone_class == "resnet_12_backbone":
            config.resnet_config.add_last_relu = add_last_relu
        else:
            config.c4_config.add_last_relu = add_last_relu

    # config.memory_net_config.radius_init = 10
    # config.memory_net_config.radius_init_write = 10

    if num_classes is not None:
        config.num_classes = num_classes

    config.num_steps = data_config.maxlen
    config.memory_net_config.max_classes = data_config.nway
    config.memory_net_config.max_stages = data_config.nstage
    config.memory_net_config.max_items = data_config.maxlen
    config.fix_unknown = data_config.fix_unknown  # Assign fix unknown ID.
    model = build_pretrain_net(config)

    return model

def main(args):
    logger, comet, writer, logdir = init_writers(
        args.logdir,
        args.exp_name,
        args.logcomet,
        # args.comet_api_key,
        args.comet_project_name,
        args.seed,
        args.exist_ok,
        args,
    )

    if args.data_config_fp is None:
        args.data_config_fp = CONSTANTS.data_configs[args.dataset]
    if args.env_config_fp is None:
        args.env_config_fp = CONSTANTS.env_configs[args.dataset]
    config = get_config(args.model_config_fp, ExperimentConfig)
    data_config = get_config(args.data_config_fp, EpisodeConfig)
    env_config = get_config(args.env_config_fp, EnvironmentConfig)

    env_config.results = "./results/oc-fewshot/"
    if args.dataset == "omniglot":
        env_config.data_folder = "./data/roaming-omniglot/"

    if config.backbone_class == "resnet_12_backbone":
        bb_config = config.resnet_config
    else:
        bb_config = config.c4_config

    datasets, data = build_datasets(
        # args.batch_size,
        1,
        args.dataset,
        data_config,
        env_config,
        bb_config.data_format == "NCHW",
        args.seed,
    )

    if args.dataset == "tim":
        num_classes = 351
    elif args.dataset == "rooms":
        num_classes = None
    elif args.dataset == "omniglot":
        num_classes = len(datasets["train_fs"].cls_dict)

    model = build_model(config, data_config, args.add_last_relu, num_classes=num_classes)
    # mem_model = build_net(
    #     config,
    #     backbone=model.backbone,
    #     # learn_temp=args.learn_temp,
    # )

    dataset_name = args.dataset
    if args.use_dataset_stats:
        mean=CONSTANTS.datasets_stats[dataset_name].mean
        std=CONSTANTS.datasets_stats[dataset_name].std
    else:
        mean = std = None
    # mean = std = None

    data = get_dataiter(
    #   dataset,
        {k.replace("_fs", ""): v for k, v in datasets.items()},
        # config.optimizer_config.batch_size,
        args.batch_size,
        nchw=model.backbone.config.data_format == 'NCHW',
        data_aug=True,
        random_box=data_config.random_box,
        # seed=args.seed,
        mean=mean,
        std=std,
        da_prep2=args.da_prep2,
        min_object_covered=args.min_object_covered,
    )

    learning_rate_fn = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
        args.lr_decay_steps, args.lr_list,
    )
    optimizer = tf.optimizers.SGD(
        learning_rate=learning_rate_fn,
        momentum=0.9,
    )

    loss_fn = lambda y, outputs: tf.reduce_mean(
        tf.nn.sparse_softmax_cross_entropy_with_logits(logits=outputs, labels=y)
    )

    @tf.function
    def train_step(x, y):
        with tf.GradientTape(persistent=True) as tape:
            outputs = model(x)
            _loss = loss_fn(y, outputs)
            _wd = args.weight_decay * tf.add_n(
                [tf.reduce_sum(w ** 2) * (1 / 2) for w in model.regularized_weights()]
            )
            loss = _loss + _wd

        gradients = tape.gradient(loss, model.var_to_optimize())
        optimizer.apply_gradients(zip(gradients, model.var_to_optimize()))

        # if args.loss == "sup":
        acc = tf.reduce_mean(
            tf.cast(tf.equal(tf.argmax(outputs, -1), y), tf.float32)
        )
        # else:
        #     acc = 0.0

        return acc, _loss, _wd

    acc = None
    loss = None
    alpha = 0.99

    num_steps = args.num_steps
    pbar = tqdm(range(num_steps), total=num_steps, ncols=0, mininterval=3)
    i = -1
    train_loader = data["train"].get_dataset()
    while i < (num_steps - 1):
        for b in train_loader:
            x = b["x"]
            y= b["y"]
            i += 1
            if i >= num_steps:
                break
            _acc, _loss, wd = train_step(x, tf.cast(y, tf.int64))

            if acc is None:
                acc = _acc
                loss = _loss
            else:
                acc = alpha * acc + (1 - alpha) * _acc
                loss = alpha * loss + (1 - alpha) * _loss

            if (i % 50) == 0:
                _acc = acc.numpy()
                _loss = loss.numpy()
                _wd = wd.numpy()

                pbar.set_postfix(
                    acc=_acc,
                    loss=_loss,
                    wd=_wd,
                )

                tf.summary.scalar("train/acc_ema", _acc, step=i + 1)
                tf.summary.scalar("train/loss_ema", _loss, step=i + 1)
                tf.summary.scalar("train/wd_ema", _wd, step=i + 1)

            if (i == 0) or (((i + 1) % 5000) == 0) or ((i + 1) == num_steps):
                logger.info("")
                logger.info(f"Step {i}")
                logger.info(f"Loss {loss.numpy()}")
                logger.info(f"Acc {acc.numpy()}")
                logger.info(f"Wd {wd.numpy()}")

            if (i + 1) in args.lr_decay_steps:
                model.save(f"{logdir}/model-{i}.pkl")
                # save_optimizer(f"{logdir}/optimizer-{i}.pkl", optimizer)
                # if args.num_layers_mlp > 0:
                #     save_mlp(f"{logdir}/mlp-{i}.pkl", mlp)

            if (i % 5000) == 0:
                model.save(f"{logdir}/model-curr.pkl")
                # save_optimizer(f"{logdir}/optimizer-curr.pkl", optimizer)
                # if args.num_layers_mlp > 0:
                #     save_mlp(f"{logdir}/mlp-curr.pkl", mlp)

            pbar.update()

    model.save(f"{logdir}/model-last.pkl")
    # save_optimizer(f"{logdir}/optimizer-last.pkl", optimizer)
    # if args.num_layers_mlp > 0:
    #     save_mlp(f"{logdir}/mlp-last.pkl", mlp)

def parse_args():
    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)

    parser.add_argument("--seed", default=0, type=int)
    # parser.add_argument("--num_test_steps", type=int, default=-1)
    # parser.add_argument(
    #     "--split", choices=["trainval_fs", "val_fs", "test_fs"], default="val_fs"
    # )
    # parser.add_argument("--nway", type=int)
    # parser.add_argument("--maxlen", type=int)

    parser.add_argument("--use_dataset_stats", default=False, action="store_true")
    parser.add_argument("--da_prep2", default=False, action="store_true")
    parser.add_argument("--min_object_covered", default=0.2, type=float)
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--lr_list", type=float, nargs="+", default=[5e-2, 5e-3, 5e-4])
    parser.add_argument("--lr_decay_steps", type=int, nargs="+", default=[100000, 150000])
    parser.add_argument("--weight_decay", type=float, default=5e-4)
    parser.add_argument("--num_steps", type=int, default=200000)
    parser.add_argument("--num_train_test_steps", type=int, default=500)
    parser.add_argument("--num_test_test_steps", type=int, default=2000)
    parser.add_argument("--test_interval", type=int, default=100000)

    parser.add_argument(
        "--dataset", choices=["omniglot", "rooms", "tim"], required=True
    )
    parser.add_argument("--data_config_fp")
    parser.add_argument("--env_config_fp")
    parser.add_argument("--model_config_fp")

    parser.add_argument("--model_dir")
    # parser.add_argument("--model_fname")

    parser.add_argument("--exist_ok", action="store_true", default=False)
    parser.add_argument("--logdir")
    parser.add_argument("--exp_name")

    parser.add_argument("--learn_temp", default=False, action="store_true")
    parser.add_argument("--add_last_relu", type=int, choices=[0, 1])
    parser.add_argument("--hidden_dim", type=int)
    # parser.add_argument("--adapt", default=False, action="store_true")

    parser.add_argument("--logcomet", action="store_true", default=False)
    parser.add_argument("--comet_project_name", default="wandering")
    # parser.add_argument("--comet_api_key", default="wandering")

    return parser.parse_args()

if __name__ == "__main__":
    load_dotenv()
    main(parse_args())


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

from __future__ import print_function, division
import argparse
import caffe
import cv2
import json
import numba
import numpy as np
from os.path import dirname, exists, join, splitext
import sys

import util

__author__ = 'Fisher Yu'
__copyright__ = 'Copyright (c) 2016, Fisher Yu'
__email__ = 'i@yf.io'
__license__ = 'MIT'

class Dataset(object):
    def __init__(self, dataset_name):
        self.work_dir = dirname(__file__)
        info_path = join(self.work_dir, 'datasets', dataset_name + '.json')
        if not exists(info_path):
            raise IOError("Do not have information for dataset {}"
                          .format(dataset_name))
        with open(info_path, 'r') as fp:
            info = json.load(fp)
        self.palette = np.array(info['palette'], dtype=np.uint8)
        self.mean_pixel = np.array(info['mean'], dtype=np.float32)
        self.dilation = info['dilation']
        self.zoom = info['zoom']
        self.name = dataset_name
        self.model_name = 'dilation{}_{}'.format(self.dilation, self.name)
        self.model_path = join(self.work_dir, 'models',
                               self.model_name + '_deploy.prototxt')

    @property
    def pretrained_path(self):
        p = join(dirname(__file__), 'pretrained',
                 self.model_name + '.caffemodel')
        if not exists(p):
            download_path = join(self.work_dir, 'pretrained',
                                 'download_{}.sh'.format(self.name))
            raise IOError('Pleaes run {} to download the pretrained network '
                          'weights first'.format(download_path))
        return p

def predict(dataset_name, input_path, output_path):
    dataset = Dataset(dataset_name)
    net = caffe.Net(dataset.model_path, dataset.pretrained_path, caffe.TEST)
    label_margin = 186

    input_dims = net.blobs['data'].shape
    batch_size, num_channels, input_height, input_width = input_dims
    caffe_in = np.zeros(input_dims, dtype=np.float32)
    image = cv2.imread(input_path, 1).astype(np.float32) - dataset.mean_pixel

    image_size = image.shape
    output_height = input_height - 2 * label_margin
    output_width = input_width - 2 * label_margin
    image = cv2.copyMakeBorder(image, label_margin, label_margin,
                               label_margin, label_margin,
                               cv2.BORDER_REFLECT_101)

    num_tiles_h = image_size[0] // output_height + \
                  (1 if image_size[0] % output_height else 0)
    num_tiles_w = image_size[1] // output_width + \
                  (1 if image_size[1] % output_width else 0)

    prediction = []
    for h in range(num_tiles_h):
        col_prediction = []
        for w in range(num_tiles_w):
            offset = [output_height * h,
                      output_width * w]
            tile = image[offset[0]:offset[0] + input_height,
                         offset[1]:offset[1] + input_width, :]
            margin = [0, input_height - tile.shape[0],
                      0, input_width - tile.shape[1]]
            tile = cv2.copyMakeBorder(tile, margin[0], margin[1],
                                      margin[2], margin[3],
                                      cv2.BORDER_REFLECT_101)
            caffe_in[0] = tile.transpose([2, 0, 1])
            out = net.forward_all(**{net.inputs[0]: caffe_in})
            prob = out['prob'][0]
            col_prediction.append(prob)
        # print('concat row')
        col_prediction = np.concatenate(col_prediction, axis=2)
        prediction.append(col_prediction)
    prob = np.concatenate(prediction, axis=1)
    if dataset.zoom > 1:
        prob = util.interp_map(prob, dataset.zoom, image_size[1], image_size[0])
    prediction = np.argmax(prob.transpose([1, 2, 0]), axis=2)
    color_image = dataset.palette[prediction.ravel()].reshape(image_size)
    color_image = cv2.cvtColor(color_image, cv2.COLOR_RGB2BGR)
    print('Writing', output_path)
    cv2.imwrite(output_path, color_image)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('dataset', nargs='?',
                        choices=['pascal_voc', 'camvid', 'kitti', 'cityscapes'])
    parser.add_argument('input_path', nargs='?', default='',
                        help='Required path to input image')
    parser.add_argument('-o', '--output_path', default=None)
    parser.add_argument('--gpu', type=int, default=-1,
                        help='GPU ID to run CAFFE. '
                             'If -1 (default), CPU is used')
    args = parser.parse_args()
    if args.input_path == '':
        raise IOError('Error: No path to input image')
    if not exists(args.input_path):
        raise IOError("Error: Can't find input image " + args.input_path)
    if args.gpu >= 0:
        caffe.set_mode_gpu()
        caffe.set_device(args.gpu)
        print('Using GPU ', args.gpu)
    else:
        caffe.set_mode_cpu()
        print('Using CPU')
    if args.output_path is None:
        args.output_path = '{}_{}.png'.format(
                splitext(args.input_path)[0], args.dataset)
    predict(args.dataset, args.input_path, args.output_path)

if __name__ == '__main__':
    main()


import pyrealsense2 as rs
import numpy as np
import cv2
import torch
from show_example import showbbox
from train import get_model_instance_segmentation, get_transform
import transforms as T
import torchvision

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

def eval_show(img):
    transform1 = T.Compose([
        T.ToTensor(),  # range [0, 255] -> [0.0,1.0]
    ]
    )
    num_class = 2
    model = get_model_instance_segmentation(num_class)
    model.load_state_dict(torch.load("test.pth"))
    model.eval()
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    model.to(device)
    xx, _ = transform1(img, 0)
    xx = showbbox(model, xx)
    return xx

# Configure depth and color streams
pipeline = rs.pipeline()
config = rs.config()

# Get device product line for setting a supporting resolution
pipeline_wrapper = rs.pipeline_wrapper(pipeline)
pipeline_profile = config.resolve(pipeline_wrapper)
device1 = pipeline_profile.get_device()
device_product_line = str(device1.get_info(rs.camera_info.product_line))

config.enable_stream(rs.stream.depth, 1024, 768, rs.format.z16, 30)

if device_product_line == 'L500':
    config.enable_stream(rs.stream.color, 1920, 1080, rs.format.bgr8, 30)
else:
    config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)

# Start streaming
pipeline.start(config)

try:
    while True:

        # Wait for a coherent pair of frames: depth and color
        frames = pipeline.wait_for_frames()
        depth_frame = frames.get_depth_frame()
        color_frame = frames.get_color_frame()
        if not depth_frame or not color_frame:
            continue

        # Convert images to numpy arrays
        depth_image = np.asanyarray(depth_frame.get_data())
        color_image = np.asanyarray(color_frame.get_data())

        # Apply colormap on depth image (image must be converted to 8-bit per pixel first)
        depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(depth_image, alpha=0.03), cv2.COLORMAP_JET)

        image = eval_show(color_image)

        depth_colormap_dim = depth_colormap.shape
        color_colormap_dim = color_image.shape

        # If depth and color resolutions are different, resize color image to match depth image for display
        # if depth_colormap_dim != color_colormap_dim:
        #     resized_color_image = cv2.resize(color_image, dsize=(depth_colormap_dim[1], depth_colormap_dim[0]), interpolation=cv2.INTER_AREA)
        #     images = np.hstack((resized_color_image, depth_colormap))
        # else:
        #     images = np.hstack((color_image, depth_colormap))

        # Show images
        cv2.namedWindow('RealSense', cv2.WINDOW_NORMAL)
        cv2.imshow('RealSense', image)
        cv2.waitKey(1)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

finally:

    # Stop streaming
    pipeline.stop()


import torch
import torch.nn as nn
import torch.nn.functional as F
from pointnet2_lib.pointnet2.pointnet2_modules import PointnetSAModule
from lib.utils.bbox_transform import decode_center_target, decode_bbox_target_stage_2, center_box2box, box2center_box, refine_box
import pointnet2_lib.pointnet2.pytorch_utils as pt_utils
import lib.utils.loss_utils as loss_utils
from lib.config import cfg
from lib.net.transformer import Transformer
import numpy as np
import matplotlib.pyplot as plt
import lib.utils.kitti_utils as kitti_utils
import lib.utils.roipool3d.roipool3d_utils as roipool3d_utils

class RCNNNet(nn.Module):
    def __init__(self, num_classes, num_point=512, input_channels=0, use_xyz=True):
        super().__init__()

        self.SA_modules = nn.ModuleList()
        self.ATT_modules = nn.ModuleList()
        channel_in = input_channels
        self.MEAN_SIZE = torch.from_numpy(cfg.CLS_MEAN_SIZE[0]).cuda()

        #todo use statics feature num
        #self.rcnn_input_channel = 3 + int(cfg.RCNN.USE_INTENSITY) + int(cfg.RCNN.USE_MASK) + int(cfg.RCNN.USE_DEPTH)
        self.rcnn_input_channel = 5
        self.input_tansformer = Transformer(num_point, 3)
        self.xyz_up_layer = pt_utils.SharedMLP([3] + cfg.RCNN.XYZ_UP_LAYER,
                                               bn=cfg.RCNN.USE_BN)

        #self.feature_tansformer = Transformer(num_point, cfg.RCNN.XYZ_UP_LAYER[-1])

        self.feature_up_layer = pt_utils.SharedMLP([self.rcnn_input_channel-3] + cfg.RCNN.XYZ_UP_LAYER,
                                               bn=cfg.RCNN.USE_BN)
        c_out = cfg.RCNN.XYZ_UP_LAYER[-1]
        self.merge_down_layer = pt_utils.SharedMLP([c_out * 2, c_out], bn=cfg.RCNN.USE_BN)

        for k in range(cfg.RCNN.SA_CONFIG.NPOINTS.__len__()):

            if cfg.ATTENTION:
                self.ATT_modules.append(pt_utils.SharedMLP([channel_in], bn=cfg.RCNN.USE_BN, activation=nn.ReLU(inplace=True)))

            mlps = [channel_in] + cfg.RCNN.SA_CONFIG.MLPS[k]

            npoint = cfg.RCNN.SA_CONFIG.NPOINTS[k] if cfg.RCNN.SA_CONFIG.NPOINTS[k] != -1 else None
            self.SA_modules.append(
                PointnetSAModule(
                    npoint=npoint,
                    radius=cfg.RCNN.SA_CONFIG.RADIUS[k],
                    nsample=cfg.RCNN.SA_CONFIG.NSAMPLE[k],
                    mlp=mlps,
                    use_xyz=use_xyz,
                    bn=cfg.RCNN.USE_BN
                )
            )
            channel_in = mlps[-1]

            # class SharedMLP(nn.Sequential):
            #
            #     def __init__(
            #             self,
            #             args: List[int],
            #             *,
            #             bn: bool = False,
            #             activation=nn.ReLU(inplace=True),
            #             preact: bool = False,
            #             first: bool = False,
            #             name: str = "",
            #             instance_norm: bool = False,
            #     ):

        # classification layer
        cls_channel = 1 if num_classes == 2 else num_classes
        cls_layers = []
        pre_channel = channel_in
        for k in range(0, cfg.RCNN.CLS_FC.__len__()):
            cls_layers.append(pt_utils.Conv1d(pre_channel, cfg.RCNN.CLS_FC[k], bn=cfg.RCNN.USE_BN))
            pre_channel = cfg.RCNN.CLS_FC[k]
        cls_layers.append(pt_utils.Conv1d(pre_channel, cls_channel, activation=None))
        if cfg.RCNN.DP_RATIO >= 0:
            cls_layers.insert(1, nn.Dropout(cfg.RCNN.DP_RATIO))
        self.cls_layer = nn.Sequential(*cls_layers)

        if cfg.RCNN.LOSS_CLS == 'SigmoidFocalLoss':
            self.cls_loss_func = loss_utils.SigmoidFocalClassificationLoss(alpha=cfg.RCNN.FOCAL_ALPHA[0],
                                                                           gamma=cfg.RCNN.FOCAL_GAMMA)
        elif cfg.RCNN.LOSS_CLS == 'BinaryCrossEntropy':
            self.cls_loss_func = F.binary_cross_entropy
        elif cfg.RCNN.LOSS_CLS == 'CrossEntropy':
            cls_weight = torch.from_numpy(cfg.RCNN.CLS_WEIGHT).float()
            self.cls_loss_func = nn.CrossEntropyLoss(ignore_index=-1, reduce=False, weight=cls_weight)
        else:
            raise NotImplementedError

        # regression layer
        per_loc_bin_num = int(cfg.RCNN.LOC_SCOPE / cfg.RCNN.LOC_BIN_SIZE) * 2
        loc_y_bin_num = int(cfg.RCNN.LOC_Y_SCOPE / cfg.RCNN.LOC_Y_BIN_SIZE) * 2
        reg_channel = per_loc_bin_num * 4 + cfg.RCNN.NUM_HEAD_BIN * 2 + 3
        reg_channel += (1 if not cfg.RCNN.LOC_Y_BY_BIN else loc_y_bin_num * 2)

        reg_layers = []
        pre_channel = channel_in
        for k in range(0, cfg.RCNN.REG_FC.__len__()):
            reg_layers.append(pt_utils.Conv1d(pre_channel, cfg.RCNN.REG_FC[k], bn=cfg.RCNN.USE_BN))
            pre_channel = cfg.RCNN.REG_FC[k]
        reg_layers.append(pt_utils.Conv1d(pre_channel, reg_channel, activation=None))
        if cfg.RCNN.DP_RATIO >= 0:
            reg_layers.insert(1, nn.Dropout(cfg.RCNN.DP_RATIO))
        self.reg_layer = nn.Sequential(*reg_layers)

        # IOU estimation
        # IOU layer
        if cfg.IOUN.ENABLED:
            self.cascade = cfg.CASCADE
            self.can_xyz_up_layer = nn.ModuleList()
            self.can_feature_up_layer = nn.ModuleList()
            self.can_merge_down_layer = nn.ModuleList()
            self.SA_score_modules = nn.ModuleList()
            self.ATT_score_modules = nn.ModuleList()
            self.IOU_layer = nn.ModuleList()
            self.ICL_layer = nn.ModuleList()
            self.ref_layer = nn.ModuleList()
            for i in range(self.cascade):
                for p in self.parameters():
                    p.requires_grad = False

                self.can_xyz_up_layer.append(pt_utils.SharedMLP([3] + cfg.RCNN.XYZ_UP_LAYER,
                                                              bn=cfg.RCNN.USE_BN).cuda())
                self.can_feature_up_layer.append(pt_utils.SharedMLP([2] + cfg.RCNN.XYZ_UP_LAYER,
                                                                bn=cfg.RCNN.USE_BN).cuda())
                c_out = cfg.RCNN.XYZ_UP_LAYER[-1]
                self.can_merge_down_layer.append(pt_utils.SharedMLP([c_out * 2, c_out], bn=cfg.RCNN.USE_BN))

                iou_channel_in = input_channels
                for k in range(cfg.IOUN.SA_CONFIG.NPOINTS.__len__()):

                    mlps = [iou_channel_in] + cfg.IOUN.SA_CONFIG.MLPS[k]

                    if cfg.ATTENTION:
                        self.ATT_score_modules.append(pt_utils.SharedMLP([iou_channel_in], bn=cfg.RCNN.USE_BN, activation=nn.ELU(inplace=True)))

                    npoint = cfg.IOUN.SA_CONFIG.NPOINTS[k] if cfg.IOUN.SA_CONFIG.NPOINTS[k] != -1 else None
                    self.SA_score_modules.append(
                        PointnetSAModule(
                            npoint=npoint,
                            radius=cfg.IOUN.SA_CONFIG.RADIUS[k],
                            nsample=cfg.IOUN.SA_CONFIG.NSAMPLE[k],
                            mlp=mlps,
                            use_xyz=use_xyz,
                            bn=cfg.IOUN.USE_BN
                        ).cuda()
                    )
                    iou_channel_in = mlps[-1]

                IOU_channel = 1
                IOU_layers = []
                pre_channel = iou_channel_in
                for k in range(0, cfg.IOUN.CLS_FC.__len__()):
                    IOU_layers.append(pt_utils.Conv1d(pre_channel, cfg.IOUN.CLS_FC[k], bn=cfg.IOUN.USE_BN))
                    pre_channel = cfg.IOUN.CLS_FC[k]
                IOU_layers.append(pt_utils.Conv1d(pre_channel, IOU_channel, activation=None))
                if cfg.IOUN.DP_RATIO >= 0:
                    IOU_layers.insert(1, nn.Dropout(cfg.IOUN.DP_RATIO))
                self.IOU_layer.append(nn.Sequential(*IOU_layers).cuda())

                ICL_channel = 1
                ICL_layers = []
                pre_channel = iou_channel_in
                for k in range(0, cfg.IOUN.CLS_FC.__len__()):
                    ICL_layers.append(pt_utils.Conv1d(pre_channel, cfg.IOUN.CLS_FC[k], bn=cfg.IOUN.USE_BN))
                    pre_channel = cfg.IOUN.CLS_FC[k]
                ICL_layers.append(pt_utils.Conv1d(pre_channel, ICL_channel, activation=None))
                if cfg.IOUN.DP_RATIO >= 0:
                    ICL_layers.insert(1, nn.Dropout(cfg.IOUN.DP_RATIO))
                self.ICL_layer.append(nn.Sequential(*ICL_layers).cuda())

                per_loc_bin_num = int(cfg.IOUN.LOC_SCOPE / cfg.IOUN.LOC_BIN_SIZE) * 2
                loc_y_bin_num = int(cfg.IOUN.LOC_Y_SCOPE / cfg.IOUN.LOC_Y_BIN_SIZE) * 2
                ref_channel = 7

                ref_layers = []
                pre_channel = iou_channel_in
                for k in range(0, cfg.IOUN.REG_FC.__len__()):
                    ref_layers.append(pt_utils.Conv1d(pre_channel, cfg.IOUN.REG_FC[k], bn=cfg.IOUN.USE_BN))
                    pre_channel = cfg.IOUN.REG_FC[k]
                ref_layers.append(pt_utils.Conv1d(pre_channel, ref_channel, activation=None))
                if cfg.IOUN.DP_RATIO >= 0:
                    ref_layers.insert(1, nn.Dropout(cfg.IOUN.DP_RATIO))
                self.ref_layer.append(nn.Sequential(*ref_layers).cuda())

        self.init_weights(weight_init='xavier')

    def init_weights(self, weight_init='xavier'):
        if weight_init == 'kaiming':
            init_func = nn.init.kaiming_normal_
        elif weight_init == 'xavier':
            init_func = nn.init.xavier_normal_
        elif weight_init == 'normal':
            init_func = nn.init.normal_
        else:
            raise NotImplementedError

        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv1d):
                if weight_init == 'normal':
                    init_func(m.weight, mean=0, std=0.001)
                else:
                    init_func(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
        nn.init.normal_(self.reg_layer[-1].conv.weight, mean=0, std=0.001)

    def _break_up_pc(self, pc):
        xyz = pc[..., 0:3].contiguous()
        features = (
            pc[..., 3:].transpose(1, 2).contiguous()
            if pc.size(-1) > 3 else None
        )

        return xyz, features

    def get_rotation_matrix(self, batch_size, ry):
        Rot_y = torch.zeros((batch_size, 3, 3)).cuda()
        Rot_y[:, 0, 0] = torch.cos(ry[:, 0])
        Rot_y[:, 0, 2] = torch.sin(ry[:, 0])
        Rot_y[:, 1, 1] = 1
        Rot_y[:, 2, 0] = -torch.sin(ry[:, 0])
        Rot_y[:, 2, 2] = torch.cos(ry[:, 0])
        return Rot_y

    def forward(self, input_data):
        """
        :param input_data: input dict
        :return:
        """
        if cfg.RCNN.ROI_SAMPLE_JIT:
            xyz, raw_features = input_data['cur_box_point'], torch.cat((input_data['cur_box_reflect'], input_data['train_mask']), dim=-1)
            features = torch.cat((xyz, raw_features), dim=-1)

        if cfg.RCNN.USE_RPN_FEATURES:

            # xyz_input = pts_input[..., 0:self.rcnn_input_channel].transpose(1, 2).unsqueeze(dim=3)
            xyz_input = xyz.transpose(1, 2).unsqueeze(dim=3)
            raw_features_input = raw_features.transpose(1, 2).unsqueeze(dim=3)

            uper_xyz = self.xyz_up_layer(xyz_input)
            uper_feature = self.feature_up_layer(raw_features_input)

            #use rpn feature
            if 'cur_pts_feature' in input_data.keys():
                uper_feature = input_data['cur_pts_feature'].transpose(1, 2).unsqueeze(dim=3)

            merged_feature = torch.cat((uper_xyz, uper_feature), dim=1)
            merged_feature = self.merge_down_layer(merged_feature)

            l_xyz, l_features = [xyz], [merged_feature.squeeze(dim=3)]
        else:
            l_xyz, l_features = [xyz], [features]

        for i in range(len(self.SA_modules)):

            if cfg.ATTENTION:
                mean_channels = np.sqrt(float(l_features[i].shape[1]))
                context = self.ATT_modules[i](l_features[i].unsqueeze(dim=3)).squeeze(dim=3)
                attention = F.softmax(torch.bmm(l_features[i].transpose(1,2),l_features[i])/mean_channels,dim=1)
                l_features[i] = torch.bmm(
                    context,
                    attention
                    )+l_features[i]

            li_xyz, li_features = self.SA_modules[i](l_xyz[i], l_features[i])
            l_xyz.append(li_xyz)
            l_features.append(li_features)

        rcnn_cls = self.cls_layer(l_features[-1]).transpose(1, 2).contiguous().squeeze(dim=1)  # (B, 1 or 2)
        rcnn_reg = self.reg_layer(l_features[-1]).transpose(1, 2).contiguous().squeeze(dim=1)  # (B, C)

        #predict box
        roi_boxes3d = torch.zeros((rcnn_reg.shape[0], 3)).cuda()
        pred_boxes3d_ce = decode_bbox_target_stage_2(roi_boxes3d.view(-1, 3), rcnn_reg.view(-1, rcnn_reg.shape[-1]),
                                                     anchor_size=self.MEAN_SIZE,
                                                     loc_scope=cfg.RCNN.LOC_SCOPE,
                                                     loc_bin_size=cfg.RCNN.LOC_BIN_SIZE,
                                                     num_head_bin=cfg.RCNN.NUM_HEAD_BIN,
                                                     get_xz_fine=False,
                                                     loc_y_scope=cfg.RCNN.LOC_Y_SCOPE,
                                                     loc_y_bin_size=cfg.RCNN.LOC_Y_BIN_SIZE,
                                                     get_ry_fine=False).view(-1, 1, 7).detach()

        #center regression
        #pred_boxes3d = center_box2box(pred_boxes3d_ce)
        #ground regression
        pred_boxes3d = pred_boxes3d_ce.clone()
        pred_boxes3d_ce = box2center_box(pred_boxes3d_ce)

        ret_dict = {'rcnn_cls': rcnn_cls,
                    'rcnn_reg': rcnn_reg,
                    'pred_boxes3d':pred_boxes3d} #, 'T1':T1, 'T2':T2}

        if cfg.IOUN.ENABLED:

            SA_SCORE_layer=0
            for c in range(self.cascade):
                xyz = input_data['cur_box_point'].clone()
                if c !=0:
                    rcnn_ref = rcnn_ref.view(rcnn_ref.shape[0], 1, rcnn_ref.shape[-1])
                    pred_boxes3d_ce = refine_box(pred_boxes3d_ce.view(-1,7), rcnn_ref.view(rcnn_ref.shape[0], rcnn_ref.shape[-1]))
                    pred_boxes3d_ce = pred_boxes3d_ce.view(-1,1,7)

                #noise the prediction
                if 'iou_trans' in input_data.keys():
                    iou_trans_noise = input_data['iou_trans'][...,c]
                    iou_scale_noise = input_data['iou_scale'][...,c]
                    iou_ry_noise = input_data['iou_ry'][...,c]

                    # location
                    pred_boxes3d_ce[:, :, 0:3] += iou_trans_noise
                    # size
                    pred_boxes3d_ce[:, :, 3:6] *= iou_scale_noise
                    # angle
                    pred_boxes3d_ce[:, :, 6] += iou_ry_noise[:,:,0]

                #canional the points by box
                xyz[:, :, 0] = xyz[:, :, 0] - pred_boxes3d_ce[:, :, 0]
                xyz[:, :, 1] = xyz[:, :, 1] - pred_boxes3d_ce[:, :, 1]
                xyz[:, :, 2] = xyz[:, :, 2] - pred_boxes3d_ce[:, :, 2]

                Rot_y = self.get_rotation_matrix(xyz.shape[0], -pred_boxes3d_ce[:, :, 6])
                canional_xyz = torch.einsum('ijk,ikl->ijl', xyz, Rot_y.permute(0, 2, 1))

                # norm transfer
                extend_factor = 1.2
                canional_xyz[:, :, 0] = canional_xyz[:, :, 0] / (pred_boxes3d_ce[:, :, 5] / 2)
                canional_xyz[:, :, 1] = canional_xyz[:, :, 1] / (pred_boxes3d_ce[:, :, 3] / 2)
                canional_xyz[:, :, 2] = canional_xyz[:, :, 2] / (pred_boxes3d_ce[:, :, 4] / 2)
                canional_mask = torch.max(torch.abs(canional_xyz), dim=-1)[0] > extend_factor
                canional_xyz[canional_mask] = canional_xyz[canional_mask] = 0.0

                #network
                can_xyz_input = canional_xyz.transpose(1, 2).unsqueeze(dim=3)
                uper_can_xyz = self.can_xyz_up_layer[c](can_xyz_input)
                uper_can_feature = self.can_feature_up_layer[c](raw_features_input)
                can_merged_feature = torch.cat((uper_can_xyz, uper_can_feature), dim=1)
                can_merged_feature = self.can_merge_down_layer[c](can_merged_feature)

                #no merge
                # can_merged_feature = uper_can_xyz

                l_xyz, l_features = [canional_xyz], [can_merged_feature.squeeze(dim=3)]

                for i in range(int(len(self.SA_score_modules)/cfg.CASCADE)):

                    if cfg.ATTENTION:
                        mean_channels = np.sqrt(float(l_features[i].shape[1]))
                        context = self.ATT_score_modules[SA_SCORE_layer+i](l_features[i].unsqueeze(dim=3)).squeeze(dim=3)
                        attention = F.softmax(torch.bmm(l_features[i].transpose(1,2),l_features[i])/mean_channels,dim=1)
                        l_features[i] = torch.bmm(
                            context,
                            attention
                            )+l_features[i]

                    li_xyz, li_features = self.SA_score_modules[SA_SCORE_layer+i](l_xyz[i], l_features[i])
                    l_xyz.append(li_xyz)
                    l_features.append(li_features)
                SA_SCORE_layer +=i+1

                rcnn_iou = self.IOU_layer[c](l_features[-1]).transpose(1, 2).contiguous().squeeze(dim=1)  # (B, 1 or 2)
                rcnn_ref = self.ref_layer[c](l_features[-1]).transpose(1, 2).contiguous().squeeze(dim=1)  # (B, C)
                ioun_cls = self.ICL_layer[c](l_features[-1]).transpose(1, 2).contiguous().squeeze(dim=1)

                pred_boxes3d = center_box2box(pred_boxes3d_ce).view(-1,1,7)

                refined_box = refine_box(pred_boxes3d.view(-1, 7),
                                             rcnn_ref.view(rcnn_ref.shape[0], rcnn_ref.shape[-1])).view(-1,1,7)

                ret_dict.update({'rcnn_iou': rcnn_iou,
                                 'rcnn_ref': rcnn_ref,
                                 'ioun_cls': ioun_cls,
                                 'pred_boxes3d':pred_boxes3d,
                                 'refined_box':refined_box})

        ret_dict.update(input_data)
        return ret_dict


# -*- coding: utf-8 -*-
# Copyright © 2017 Apple Inc. All rights reserved.
#
# Use of this source code is governed by a BSD-3-clause license that can
# be found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause
from __future__ import print_function as _
from __future__ import division as _
from __future__ import absolute_import as _

from ._pre_trained_models import _get_cache_dir
import turicreate.toolkits._tf_utils as _utils

def _create_feature_extractor(model_name):
    from platform import system
    from ._internal_utils import _mac_ver
    from ._pre_trained_models import IMAGE_MODELS
    from turicreate import extensions

    # If we don't have Core ML, use a TensorFlow model.
    if system() != "Darwin" or _mac_ver() < (10, 13):
        ptModel = IMAGE_MODELS[model_name]()
        return TensorFlowFeatureExtractor(ptModel)

    download_path = _get_cache_dir()

    result = extensions.__dict__["image_deep_feature_extractor"]()
    result.init_options({"model_name": model_name, "download_path": download_path})
    return result

class ImageFeatureExtractor(object):
    def __init__(self, ptModel):
        """
        Parameters
        ----------
        ptModel: ImageClassifierPreTrainedModel
            An instance of a pre-trained model.
        """
        pass

    def extract_features(self, dataset, feature):
        """
        Parameters
        ----------
        dataset: SFrame
            SFrame with data to extract features from
        feature: str
            Name of the column in `dataset` containing the features
        """
        pass

    def get_coreml_model(self):
        """
        Returns
        -------
        model:
            Return the underlying model in Core ML format
        """
        pass

class TensorFlowFeatureExtractor(ImageFeatureExtractor):
    def __init__(self, ptModel):
        """
        Parameters
        ----------
        ptModel: ImageClassifierPreTrainedModel
            An instance of a pre-trained model.
        """

        # Suppresses verbosity to only errors
        _utils.suppress_tensorflow_warnings()

        from tensorflow import keras

        self.gpu_policy = _utils.TensorFlowGPUPolicy()
        self.gpu_policy.start()

        self.ptModel = ptModel

        self.input_shape = ptModel.input_image_shape
        self.coreml_data_layer = ptModel.coreml_data_layer
        self.coreml_feature_layer = ptModel.coreml_feature_layer

        model_path = ptModel.get_model_path("tensorflow")
        self.model = keras.models.load_model(model_path)

    def __del__(self):
        self.gpu_policy.stop()

    def extract_features(self, dataset, feature, batch_size=64, verbose=False):
        from array import array
        import turicreate as tc
        import numpy as np

        # Only expose the feature column to the SFrame-to-NumPy loader.
        image_sf = tc.SFrame({"image": dataset[feature]})

        # Encapsulate state in a dict to sidestep variable scoping issues.
        state = {}
        state["num_started"] = 0  # Images read from the SFrame
        state["num_processed"] = 0  # Images processed by TensorFlow
        state["total"] = len(dataset)  # Images in the dataset

        # We should be using SArrayBuilder, but it doesn't accept ndarray yet.
        # TODO: https://github.com/apple/turicreate/issues/2672
        # out = _tc.SArrayBuilder(dtype = array.array)
        state["out"] = tc.SArray(dtype=array)

        if verbose:
            print("Performing feature extraction on resized images...")

        # Provide an iterator-like interface around the SFrame.
        def has_next_batch():
            return state["num_started"] < state["total"]

        # Yield a numpy array representing one batch of images.
        def next_batch():
            # Compute the range of the SFrame to yield.
            start_index = state["num_started"]
            end_index = min(start_index + batch_size, state["total"])
            state["num_started"] = end_index

            # Allocate a numpy array with the desired shape.
            # TODO: Recycle the ndarray instances we're allocating below with
            # _np.zeros, instead of creating new ones every time.
            num_images = end_index - start_index
            shape = (num_images,) + self.ptModel.input_image_shape
            batch = np.zeros(shape, dtype=np.float32)

            # Resize and load the images.
            tc.extensions.sframe_load_to_numpy(
                image_sf,
                batch.ctypes.data,
                batch.strides,
                batch.shape,
                start_index,
                end_index,
            )

            # TODO: Converge to NCHW everywhere.
            batch = batch.transpose(0, 2, 3, 1)  # NCHW -> NHWC

            if self.ptModel.input_is_BGR:
                batch = batch[:, :, :, ::-1]  # RGB -> BGR

            return batch

        # Encapsulates the work done by TensorFlow for a single batch.
        def handle_request(batch):
            y = self.model.predict(batch)
            tf_out = [i.flatten() for i in y]
            return tf_out

        # Copies the output from TensorFlow into the SArrayBuilder and emits
        # progress.
        def consume_response(tf_out):
            sa = tc.SArray(tf_out, dtype=array)
            state["out"] = state["out"].append(sa)

            state["num_processed"] += len(tf_out)
            if verbose:
                print(
                    "Completed {num_processed:{width}d}/{total:{width}d}".format(
                        width=len(str(state["total"])), **state
                    )
                )

        # Just iterate through the image batches, converting them into batches
        # of feature vectors. Note: the helper functions defined above are
        # designed to support a multi-threaded approach, allowing one Python
        # thread to drive TensorFlow computation and another to drive SFrame
        # traversal and image resizing. But test failures reveal an interaction
        # between the multi-threaded approach we used for MXNet and our usage of
        # TensorFlow, which we need to resolve before switching to the pipelined
        # implementation below.
        while has_next_batch():
            images_in_numpy = next_batch()
            predictions_from_tf = handle_request(images_in_numpy)
            consume_response(predictions_from_tf)

        # # Create a dedicated thread for performing TensorFlow work, using two
        # # FIFO queues for communication back and forth with this thread, with
        # # the goal of keeping TensorFlow busy throughout.
        # request_queue = _Queue()
        # response_queue = _Queue()
        # def tf_worker():
        #     from tensorflow import keras
        #     model_path = self.ptModel.get_model_path('tensorflow')
        #     self.model = keras.models.load_model(model_path)
        #
        #     while True:
        #         batch = request_queue.get()  # Consume request
        #         if batch is None:
        #             # No more work remains. Allow the thread to finish.
        #             return
        #         response_queue.put(handle_request(batch))  # Produce response
        # tf_worker_thread = _Thread(target=tf_worker)
        # tf_worker_thread.start()

        # try:
        #     # Attempt to have two requests in progress at any one time (double
        #     # buffering), so that the iterator is creating one batch while
        #     # TensorFlow performs inference on the other.
        #     if has_next_batch():
        #         request_queue.put(next_batch())  # Produce request
        #         while has_next_batch():
        #             request_queue.put(next_batch())  # Produce request
        #             consume_response(response_queue.get())
        #         consume_response(response_queue.get())
        # finally:
        #     # Tell the worker thread to shut down.
        #     request_queue.put(None)

        return state["out"]

    def get_coreml_model(self):
        import coremltools

        model_path = self.ptModel.get_model_path("coreml")
        return coremltools.models.MLModel(model_path)


#!/usr/bin/env python
# coding: utf-8
# Two separate models for artifact classification and HFO with spike classification
###
import torch
import torch.nn as nn
from random import random
import numpy as np
import matplotlib.pyplot as plt
import time
import os

from sklearn.metrics import f1_score
import sys
import random

from src.dataloader import create_patient_loader
from src.dataloader_spike import create_patient_loader_90
from src.utilities import *
from src.config import inference_parse
from src.meter import TrainingMeter, Meter
from src.model import NeuralCNN
from patient_info import seizure_free_patient_names

def validate(val_loader, model, criterion, computing_device, long_eeg_flag, fn=None):
    start = time.time()
    meter_s = Meter("spike")
    model_s = model["spike"]
    model_s.eval()
    for _ , (image, waveform, intensity, label, info, start_end) in enumerate(val_loader, 0):

        with torch.no_grad():
            _ , s_, train_s = create_sa_labels(image, waveform, intensity, label, info, start_end, computing_device)
            if not train_s:
                print(label.sum())
                continue
            outputs_s = model_s(s_["inputs"]).squeeze()
            if not long_eeg_flag:
                s_["label"] = s_["label"].squeeze()[:,0] ## change it to binary #90 inference 10 comment it !!! 90 inference 90 hold
            if outputs_s.dim() == 0:
                outputs_s = outputs_s.unsqueeze(0)
                s_["label"] = s_["label"].unsqueeze(0)
            #print(outputs_s.shape, s_["label"].shape)
            loss_s = criterion(outputs_s, s_["label"])
            if not fn:
                meter_s.update_loss(loss_s.detach().cpu().numpy())
                meter_s.update_outputs(outputs_s.detach().cpu().numpy(), s_[
                                       "label"].cpu().numpy())
            else:
                meter_s.add(s_["spectrum"],
                            s_["label"].detach().cpu().numpy(),
                            s_["channel_name"],
                            s_["start_end"],
                            s_["intensity"],
                            s_["waveform"],
                            outputs_s.detach().cpu().numpy())
    acc_s = meter_s.accuracy()
    if fn is not None:
        loss_s = 0
        meter_s.dump_prediction(os.path.join(fn, f"spikes.pkl"))
        #meter_s.dump_pickle(os.path.join(fn, f"spikes.pkl"))
    else:
        loss_s = meter_s.loss()
    f1_s = meter_s.f1()
    print('Inference: Time %.3f, loss_s: %.3f, accuracy_s: %.3f , f1_s: %0.3f' %
          (time.time()-start, loss_s, acc_s, f1_s))

    return loss_s, acc_s, f1_s

def create_sa_labels(spectrum, waveform, intensity, label, channel_name, start_end, computing_device):
    channel_name = np.array(channel_name)
    label = label.squeeze().float()
    s_ ={}
    #s_spectrum_norm = torch.log(spectrum)
    inputs_s = torch.stack([spectrum,waveform,intensity], dim=1, out=None).to(computing_device).float()
    label_s= label.to(computing_device)
    #print(inputs_s.shape)
    s_ = {
        "inputs": expand_dim(inputs_s, 4),
        "spectrum":spectrum,
        "label": expand_dim(label_s, 1),
        "intensity":intensity,
        "waveform":waveform,
        "channel_name":channel_name,
        "start_end":start_end
    }

    return None , s_ , True

def pipeline(args, patient_name):
    data_dir = args.data_dir
    res_dir = os.path.join(args.work_dir, args.res_dir)
    trained_model_path = ""
    model_dir = os.path.join(args.work_dir, args.model_dir)
    for fn in os.listdir(model_dir):
        if fn.endswith("model_s.pth"):
            trained_model_path= os.path.join(model_dir, fn)

    """
    Modified from train_reverse.py:
    Load the trained model and test on the training patient
    """

    model_spike = NeuralCNN(num_classes=2)

    batch_size = 512         # Number of samples in each minibatch
    use_cuda = torch.cuda.is_available()

    # Setup GPU optimization if CUDA is supported
    if use_cuda:
        print("cuda is supported")
        computing_device = torch.device(args.device)
        extras = {"num_workers": 1, "pin_memory": True}
    else:  # Otherwise, train on the CPU
        computing_device = torch.device("cpu")
        extras = False

    model_spike = model_spike.to(computing_device)
    model = {"spike": model_spike}

    start_time = time.time()
    if not args.long_eeg:
        test_loader = create_patient_loader(data_dir, patient_name,batch_size,shuffle=False,extras=extras)
    else:
        test_loader = create_patient_loader_90(data_dir, patient_name,batch_size,shuffle=False,extras=extras)
    end_time = time.time()
    epoch_mins, epoch_secs = epoch_time(start_time, end_time)
    print(f'Prepare dataset | Time: {epoch_mins}m {epoch_secs}s')

    criterion = nn.BCELoss(reduction="none").to(computing_device)

    print("-----------------testing ----------------")

    if trained_model_path is not None:
        model['spike'].load_state_dict(torch.load(trained_model_path)['state_dict'])
        model['spike'].eval()

    print("patient_names is", patient_name)
    stats_folder = os.path.join(res_dir, patient_name)
    clean_folder(stats_folder)
    loss_s, acc_s, _ = validate(test_loader, model, criterion,
                                 computing_device, args.long_eeg ,fn=stats_folder)
    return loss_s, acc_s

def patient_fold(args):
    torch.manual_seed(0)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(0)
    random.seed(0)
    clean_folder(args.res_dir)

    p_names = os.listdir(args.data_dir)
    if args.long_eeg:
        p_names = list(set(p_names) - set(seizure_free_patient_names))
    p_names = os.listdir(args.data_dir)
    for p_name in p_names:
        loss_s, acc_s = pipeline(args, p_name)

if __name__ == "__main__":
    args = inference_parse(sys.argv[1:])
    print(args)
    patient_fold(args)


import numpy as np
import torch
import math
import torch.utils.data as Data
from torch.autograd import Variable

from EvalBox.Evaluation.evaluation import Evaluation
from EvalBox.Evaluation.evaluation import MIN_COMPENSATION

class NBC(Evaluation):
    def __init__(self, outputs_origin,outputs_adv, device, model, **kwargs):
        '''
        @description:
            该方法计算的神经元边界覆盖率。
            在参数设置时应注意
                1. Data_path中前两个，需要放置训练集而非测试集，用于生成神经元阈值范围
                2. IS_WHITE为True，白盒攻击
                3. IS_PYTHORCH_WHITE为True
                4. IS_COMPARE_MODEL为False
        @param {
            model:
            device:
            kwargs:
        }
        @return: None
        '''
        super(NBC, self).__init__(outputs_origin,outputs_adv, device)
        self.model = model
        self._parsing_parameters(**kwargs)

    def _parsing_parameters(self, **kwargs):
        '''
        @description:
        @param {
            batch_size:
        }
        @return:
        '''
        self.batch_size = kwargs.get('batch_size', 64)
        self.module_name = []
        self.clean_feature_map = []
        self.upper_feature_map = []
        self.lower_feature_map = []
        self.nbc_feature_map = []
        self.features_out_hook = []

    def for_hook(self, module, fea_in, fea_out):
        self.features_out_hook.append(fea_out.data.cpu().numpy())

    def output_hook(self):
        print("*"*5+"hook record features"+"*"*5)
        for i in range(len(self.module_name)):
            print(self.module_name[i])
            print(self.features_out_hook[i].size())
        print("*"*5+"hook record features"+"*"*5)

    def evaluate(self,adv_xs=None, cln_xs=None, cln_ys=None,adv_ys=None,target_preds=None, target_flag=False):
        '''
        @description:
        @param {
            adv_xs: 攻击样本
            cln_xs：原始样本
            cln_ys: 原始类别，非目标攻击下原始样本的类型
            adv_ys: 攻击样本的预测类别
            target_preds： 目标攻击下希望原始样本攻击的目标类别
            target_flag：是否是目标攻击
        }
        @return: nbc {}
        '''
        #print("target_flag",target_flag)
        total = len(adv_xs)
        print("total",total)
        device = self.device
        self.model = self.model.eval().to(device)
        assert len(adv_xs) == len(adv_ys), 'examples and labels do not match.'
        assert len(cln_xs) == len(cln_ys), 'examples and labels do not match.'

        # 监听所有中间层特征
        hook_handle_list = []
        for name, module in self.model._modules.items():
            self.module_name.append(name)
            handle = module.register_forward_hook(self.for_hook)
            hook_handle_list.append(handle)

        cln_dataset = Data.TensorDataset(cln_xs, cln_ys)
        adv_dataset = Data.TensorDataset(adv_xs, adv_ys)
        cln_loader = Data.DataLoader(cln_dataset, batch_size=self.batch_size, num_workers=3)
        adv_loader = Data.DataLoader(adv_dataset, batch_size=self.batch_size, num_workers=3)

        for x, y in cln_loader:
            x, y = Variable(x.to(device)), Variable(y.to(device))
            with torch.no_grad():
                pred = self.model(x)
        self.clean_feature_map = self.features_out_hook.copy()
        self.features_out_hook.clear()

        for x, y in adv_loader:
            x, y = Variable(x.to(device)), Variable(y.to(device))
            with torch.no_grad():
                pred = self.model(x)

        for handle in hook_handle_list:
            handle.remove()

        # 输出feature map =  feature_num * N * neuron_num * output
        # self.output_feature_map()

        self.lower_feature_map = [[] for i in range(len(self.module_name))]
        self.upper_feature_map = [[] for i in range(len(self.module_name))]
        self.nbc_feature_map = [[] for i in range(len(self.module_name))]
        for key in range(len(self.module_name)):
            self.lower_feature_map[key] = [0 for i in range(len(self.clean_feature_map[key][0]))]
            self.upper_feature_map[key] = [0 for i in range(len(self.clean_feature_map[key][0]))]
            self.nbc_feature_map[key] = [0 for i in range(len(self.clean_feature_map[key][0]))]

        # 获取训练集产生的[lower, upper]
        # key表示特征名
        for key in range(len(self.module_name)):
            # k 表示第k个测试样本
            for k in range(len(self.clean_feature_map[key])):
                # c表示第c个神经元
                for c in range(len(self.clean_feature_map[key][k])):
                    # x表示key层的神经元c在第k个测试时的输出结果
                    if k == 0:
                        self.lower_feature_map[key][c] = self.clean_feature_map[key][k][c].min()
                        self.upper_feature_map[key][c] = self.clean_feature_map[key][k][c].max()
                    lower_bound = self.clean_feature_map[key][k][c].min()
                    upper_bound = self.clean_feature_map[key][k][c].max()
                    if self.lower_feature_map[key][c] > lower_bound:
                        self.lower_feature_map[key][c] = lower_bound
                    if self.upper_feature_map[key][c] < upper_bound:
                        self.upper_feature_map[key][c] = upper_bound

        # 输出上下界结果
        # for key in range(len(self.module_name)):
        #     print("======", self.module_name[key], "======")
        #     for c in range(len(self.clean_feature_map[key][0])):
        #         print(self.lower_feature_map[key][c], self.upper_feature_map[key][c])

        total = 0
        covered = 0
        # 统计测试集测试结果
        for key in range(len(self.module_name)):
            # k 表示第k个测试样本
            for k in range(len(self.features_out_hook[key])):
                # c表示第c个神经元
                for c in range(len(self.features_out_hook[key][k])):
                    if self.nbc_feature_map[key][c] == 1:
                        continue
                    # x表示key层的神经元c在第k个测试时的输出结果
                    # 对应上下界分别为 self.lower_feature_map[key][c]  self.upper_feature_map[key][c]
                    if self.features_out_hook[key][k][c].min() < self.lower_feature_map[key][c] or \
                      self.features_out_hook[key][k][c].max() > self.upper_feature_map[key][c]:
                        self.nbc_feature_map[key][c] = 1
        total = 0
        covered = 0
        for key in range(len(self.module_name)):
            # c表示第c个神经元
            for c in range(len(self.features_out_hook[key][0])):
                if self.nbc_feature_map[key][c] == 1:
                    covered += 1
                total += 2
        return covered / total


# Borrow a lot from openai baselines:
# https://github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.py

import cv2
import gym
import numpy as np
from collections import deque
from copy import deepcopy
from torch import float32
import matplotlib.pyplot as plt

from ding.envs import RamWrapper, NoopResetWrapper, MaxAndSkipWrapper, EpisodicLifeWrapper, FireResetWrapper, WarpFrameWrapper, ClipRewardWrapper, FrameStackWrapper

class ScaledFloatFrameWrapper(gym.ObservationWrapper):
    """Normalize observations to -1~1.

    :param gym.Env env: the environment to wrap.
    """

    def __init__(self, env):
        super().__init__(env)
        low = np.min(env.observation_space.low)
        high = np.max(env.observation_space.high)
        self.bias = low
        self.scale = high - low
        self.observation_space = gym.spaces.Box(low=-1., high=1., shape=env.observation_space.shape, dtype=np.float32)

    def observation(self, observation):
        # use fixed scale and bias temporarily
        return (observation - 128) / 128
        # return (observation - self.bias) / self.scale

class FrameStackWrapperRam(gym.Wrapper):
    """Stack n_frames last frames.
    :param gym.Env env: the environment to wrap.
    :param int n_frames: the number of frames to stack.
    """

    def __init__(
        self,
        env,
        n_frames,
        pomdp={
            "noise_scale": 0.01,
            "zero_p": 0.2,
            "duplicate_p": 0.2,
            "reward_noise": 0.01
        },
        render=False
    ):
        super().__init__(env)
        self.n_frames = n_frames
        self.n_dims = env.observation_space.shape[0]
        self._pomdp = pomdp
        self._render = render
        self.frames = deque([], maxlen=n_frames)
        self._images = deque([], maxlen=n_frames)
        self.viewer = None

        shape = (n_frames * self.n_dims, )
        self.observation_space = gym.spaces.Box(
            low=np.min(env.observation_space.low),
            high=np.max(env.observation_space.high),
            shape=shape,
            dtype=env.observation_space.dtype
        )

    def reset(self):
        obs = self.env.reset()
        for _ in range(self.n_frames):
            self.frames.append(obs)
        return self._get_ob()

    def step(self, action):
        obs, reward, done, info = self.env.step(action)
        self.frames.append(obs)
        reward = reward + self._pomdp["reward_noise"] * np.random.randn()

        if self._render:
            _img = self.env.unwrapped._get_image()
            _img = _img.mean(axis=-1, keepdims=True).astype(np.uint8)
            self._images.append(_img)
            self.render()

        return self._get_ob(), reward, done, info

    def render(self):
        from gym.envs.classic_control import rendering
        state = np.stack(self._images, axis=0)
        obs = self._pomdp_preprocess(state, img=True).astype(np.uint8)
        obs = np.tile(obs[-1], (1, 1, 3))
        if self.viewer is None:
            self.viewer = rendering.SimpleImageViewer()
        self.viewer.imshow(obs)
        return self.viewer.isopen

    def _get_ob(self):
        # the original wrapper use `LazyFrames` but since we use np buffer,
        # it has no effect
        state = np.stack(self.frames, axis=0)
        obs = self._pomdp_preprocess(state)

        return obs.flatten()

    def _pomdp_preprocess(self, state, img=False):
        obs = deepcopy(state)
        # POMDP process
        if np.random.random() > (1 - self._pomdp["duplicate_p"]):
            update_end_point = np.random.randint(
                1, self.n_frames
            )  # choose a point from that point we can't get new observation
            _s = (self.n_frames - update_end_point, 1, 1, 1)
            obs[update_end_point:, ] = np.tile(obs[update_end_point, ], _s)

        if img:
            pomdp_noise_mask = self._pomdp["noise_scale"] * np.random.randn(*obs.shape) * 128
        else:
            pomdp_noise_mask = self._pomdp["noise_scale"] * np.random.randn(*obs.shape)

        # Flickering Atari game
        obs = obs * int(np.random.random() > self._pomdp["zero_p"]) + pomdp_noise_mask
        return obs.astype(np.float32)

def wrap_deepmind(
    env_id,
    episode_life=True,
    clip_rewards=True,
    pomdp={},
    frame_stack=4,
    scale=True,
    warp_frame=True,
    use_ram=False,
    render=False,
    only_info=False
):
    """Configure environment for DeepMind-style Atari. The observation is
    channel-first: (c, h, w) instead of (h, w, c).

    :param str env_id: the atari environment id.
    :param bool episode_life: wrap the episode life wrapper.
    :param bool clip_rewards: wrap the reward clipping wrapper.
    :param int frame_stack: wrap the frame stacking wrapper.
    :param bool scale: wrap the scaling observation wrapper.
    :param bool warp_frame: wrap the grayscale + resize observation wrapper.
    :param float pomdp: parameter to control POMDP prepropress,
    :return: the wrapped atari environment.
    """
    assert 'NoFrameskip' in env_id
    if not only_info:
        env = gym.make(env_id)
        env = RamWrapper(env)
        env = NoopResetWrapper(env, noop_max=30)
        env = MaxAndSkipWrapper(env, skip=4)
        if episode_life:
            env = EpisodicLifeWrapper(env)
        if 'FIRE' in env.unwrapped.get_action_meanings():
            env = FireResetWrapper(env)
        if warp_frame:
            env = WarpFrameWrapper(env)
        if scale:
            env = ScaledFloatFrameWrapper(env)
        if clip_rewards:
            env = ClipRewardWrapper(env)

        if frame_stack:
            if use_ram:
                env = FrameStackWrapperRam(env, frame_stack, pomdp, render)
            else:
                env = FrameStackWrapper(env, frame_stack)

        return env
    else:
        wrapper_info = RamWrapper.__name__ + '\n'
        wrapper_info += NoopResetWrapper.__name__ + '\n'
        wrapper_info += MaxAndSkipWrapper.__name__ + '\n'
        if episode_life:
            wrapper_info = EpisodicLifeWrapper.__name__ + '\n'
        if 'Pong' in env_id or 'Qbert' in env_id or 'SpaceInvader' in env_id or 'Montezuma' in env_id:
            wrapper_info = FireResetWrapper.__name__ + '\n'
        if warp_frame:
            wrapper_info = WarpFrameWrapper.__name__ + '\n'
        if scale:
            wrapper_info = ScaledFloatFrameWrapper.__name__ + '\n'
        if clip_rewards:
            wrapper_info = ClipRewardWrapper.__name__ + '\n'

        if frame_stack:
            if use_ram:
                wrapper_info = FrameStackWrapperRam.__name__ + '\n'
            else:
                wrapper_info = FrameStackWrapper.__name__ + '\n'

        return wrapper_info


#%%
import os

import torch
from torch.nn import BCELoss
import pickle
import json
import argparse

from data_loader.dataset import DataSet
from modules.model import GGNNSum
import numpy as np
from sklearn.metrics import f1_score, precision_score, recall_score, accuracy_score

def run_one(output_dir, model, loss_function, name, num_batches, get_fn, idx, logger_fn=print):
    all_predictions = []
    all_targets = []
    all_loss = []
    final = []
    with torch.no_grad():
        for l in range(num_batches):
            graph, targets = get_fn()
            graph.cuda(device='cuda:0')
            graph.graph = graph.graph.to('cuda:0')

            targets = targets.cuda()
            predictions = model(graph, cuda=True)
            batch_loss = loss_function(predictions, targets)
            all_predictions.extend(
                predictions.ge(torch.ones_like(predictions) / 2)
                .detach().cpu().int().numpy().tolist())
            all_targets.extend(targets.detach().cpu().int().numpy().tolist())
            all_loss.append(batch_loss.detach().cpu().item())

            output = model.get_graph_embeddings(graph)
            final += list(zip(output.tolist(), [int(i) for i in targets.tolist()]))

    out = []
    for f in final:
        out.append({'graph_feature':f[0], 'target':f[1]})

    dst_filepath = os.path.join(output_dir, name + f'_GGNNoutput_graph-{idx}.pkl')
    logger_fn(f'Saving to {dst_filepath}')
    with open(dst_filepath, 'wb') as of:
        pickle.dump(out, of)

    out = {
        "loss": np.mean(all_loss),
        "accuracy": accuracy_score(all_targets, all_predictions),
        "precision": precision_score(all_targets, all_predictions),
        "recall": recall_score(all_targets, all_predictions),
        "f1": f1_score(all_targets, all_predictions),
    }
    json_filepath = os.path.join(output_dir, name + f'_result-{idx}.json')
    logger_fn(f'Saving intermediate results to {json_filepath}')
    with open(json_filepath, 'w') as of:
        json.dump(out, of, indent=2)

    logger_fn(f'DONE: {name} BATCHES')

def main(raw_args=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_dir', help='path to GGNNSum model to be used', required=True)
    parser.add_argument('--dataset', help='path to data to be processed (i.e. processed.bin)', required=True)
    parser.add_argument('--output_dir', help='location to place data after ggnn processing', default='data/after_ggnn/chrome_debian/')
    parser.add_argument('--batch_size', help='batch size', type=int, default=128)
    args = parser.parse_args(raw_args)

    os.makedirs(args.output_dir, exist_ok=True)

    #%% data and model imports
    # (data, augmented_data) = pickle.load(open(args.dataset, 'rb'))
    saved_dataset = pickle.load(open(args.dataset, 'rb'))

    with open(os.path.join(args.model_dir, f'splits-ggnn.json')) as f:
        splits = json.load(f)
    print('Data & Models Loaded')
    print('='*83)

    for split in splits:
        i = split["idx"]
        print(f'Fold: {i}')
        model = GGNNSum(input_dim=169, output_dim=200, num_steps=6, max_edge_types=5)
        loss_function = BCELoss(reduction='sum')
        state_dict = torch.load(os.path.join(args.model_dir, f'ggnn-model-{i}.pth'))
        model.load_state_dict(state_dict, strict=False)
        model = model.to('cuda:0')
        model.eval()

        dataset = saved_dataset.to_dataset(split["roll"], split["splits"], args.batch_size)
        # current_data = np.roll(data, split["roll"])
        # train_data, valid_data, test_data = np.split(current_data, split["splits"])
        # dataset = DataSet(train_data=train_data, valid_data=valid_data, test_data=test_data,
        #                   augmented=augmented_data,
        #                   batch_size=args.batch_size, n_ident=args.node_tag, g_ident=args.graph_tag,
        #                   l_ident=args.label_tag, inf=False)

        run_one(args.output_dir, model, loss_function, 'test', dataset.initialize_test_batches(), dataset.get_next_test_batch, i)
        run_one(args.output_dir, model, loss_function, 'valid', dataset.initialize_valid_batches(), dataset.get_next_valid_batch, i)
        run_one(args.output_dir, model, loss_function, 'train', dataset.initialize_train_batches(), dataset.get_next_train_batch, i)

if __name__ == "__main__":
    main()