'''
.. module:: controller
   :synopsis: Control algorithms for lateral control of vehicle dynamics
.. moduleauthor:: Yongkyun Shin <github.com/yongkyuns>

This module is a collection of algorithms for control of lateral vehicle dynamics.
Currently implemented algorithms are PID controller, Model-Predictive controller,
and Deep-Deterministic Policy Gradient (DDPG) controller.
'''

import numpy as np
from scipy import sparse
import cvxpy as cp
import torch
from .learning.ddpg_agent import Agent
from .learning.model import Actor
import sys, os
# import acado

# DDPG Params
STATE_SIZE = 4
ACTION_SIZE = 1
RANDOM_SEED = 1
CHECK_POINT = 'checkpoint_actor.pth'

POS = 0
ANG = 1
NOW = 0
FUTURE = 4

class Controller:
    '''
    Generic controller parent class with some common variables (e.g. sample time, number of states, etc.)
    '''

    def __init__(self, NU=1, NY=1, T=1, dt=0.01):
        self.reference = np.zeros((NY, T))
        self.measured_output = np.zeros((NY, T))
        self.control_input = np.zeros(NU)
        self.dt = dt  # Controller sample time
        self.T = T   # Number of control points along the path

    def control(self):
        pass

class PathController(Controller):
    '''
    PathController has projection horizon in extension to generic Controller class.
    Projection horizon defines the preview, or lookahead distances w.r.t. vehicle
    '''
    def __init__(self, NU=1, NY=1, T=1, dt=0.01, proj_horizon=0.5):
        super().__init__(NU=NU, NY=NY, T=T, dt=dt)
        # [sec] Use T control points along desired path for proj_horizon time
        self.proj_horizon = proj_horizon

    @property
    def proj_step(self):
        return self.proj_horizon / self.T

class DDPG(PathController):
    '''
    Experimental controller which uses reinforcement learning algorithm.
    '''
    def __init__(self, check_point=CHECK_POINT, dt=0.01, T=1, NY=1, proj_horizon=0.5):
        super().__init__(dt=dt, T=T, proj_horizon=proj_horizon)
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        if torch.cuda.is_available():
            trained_model = torch.load(CHECK_POINT)
        else:
            trained_model = torch.load(
                CHECK_POINT, map_location={'cuda:0': 'cpu'})

        self.agent = Agent(state_size=STATE_SIZE,
                           action_size=ACTION_SIZE, random_seed=RANDOM_SEED)
        self.agent.actor_local = Actor(
            STATE_SIZE, ACTION_SIZE, RANDOM_SEED).to(device)
        self.agent.actor_local.load_state_dict(trained_model)

    def get_state(self, yref):
        # current/future pos reference, current/future heading reference
        states = np.array([yref[POS, NOW], yref[POS, FUTURE],
                           yref[ANG, NOW], yref[ANG, FUTURE]])
        states = np.reshape(states, [ACTION_SIZE, STATE_SIZE])
        return states

    def control(self, yref):
        states = self.get_state(yref)
        action = self.agent.act(states, add_noise=False)
        return float(action)

class PID(PathController):
    '''
    PID controller with variable number of control points along desired path.
    Control parameters are lateral position and heading. For each control point & parameter,
    there is a corresponding normalized weight which adds up to 1. Therefore, a single PID
    controller is used to control all of the controlled states and points along desired path.

    ================  ==================================================
    **Arguments:**
    NU                (int) Number of controlled states
    NY                (int) Number of outputs
    T                 (int) Number of reference points along trajectory
    dt                (float) Sample time [sec] of the controller
    P_Gain            (float) Proportional gain of the controller
    I_Gain            (float) Integral gain of the controller
    D_Gain            (float) Derivative gain of the controller
    weight_split_T    (list) Determine how to split weight on reference points (e.g. [0.5, 0.5] for equal split on all reference points)
    weight_split_Y    (list) Determine how to split weight on measured output (e.g. [0.5, 0.5] for equal split on position and heading control)

    **Variables:**
    weight_table      (numpy array) NY-by-T array of normalized weight for each controlled point & parameter
    ================  ==================================================
    '''

    def __init__(self, NU=1,
                       NY=1,
                       T=1,
                       dt=0.01,
                       P_Gain=1,
                       I_Gain=0.1,
                       D_Gain=0.01,
                       weight_split_T=None,
                       weight_split_Y=None,
                       proj_horizon=0.5):
        super().__init__(T=T, dt=dt, proj_horizon=proj_horizon)
        self.P_Gain = P_Gain
        self.I_Gain = I_Gain
        self.D_Gain = D_Gain

        self.weight_table = self.__calc_weight_table(
            T, NY, weight_split_T, weight_split_Y)

        self._integral = np.zeros(T)
        self._error_old = np.zeros(T)

    def __calc_weight_table(self, T, NY, T_split, Y_split):
        '''
        Update normalized weight table for PID control.
        '''
        if T_split == None and Y_split == None:
            table = np.ones((NY, T)) / (NY*T)
        else:
            table = np.zeros((NY, T))
            vec_T = np.linspace(T_split[0], T_split[1], num=T)
            vec_Y = Y_split
            for i in range(NY):
                for j in range(T):
                    table[i, j] = vec_Y[i] * vec_T[j]
            if np.sum(table) > 0:
                table /= np.sum(table)
        return table

    def control(self, error_table):
        '''
        Calculate control input
        '''
        u = np.zeros(self.T)
        for i in range(self.T):
            e = 0
            for j in range(self.weight_table.shape[0]):
                e += self.weight_table[j, i] * error_table[j, i]
            self._integral[i] += e * self.dt
            u[i] = (self.P_Gain * e) + (self.I_Gain * self._integral[i]
                                        ) + (self.D_Gain * (e-self._error_old[i])/self.dt)
            self._error_old[i] = e
        return np.sum(u)

class MPC(PathController):
    '''
    Generic model predictive controller class. Uses linearized state-space model.
    Convex optimization is solved through cvxpy module.

    ================  ==================================================
    **Arguments:**
    getModel          (function) Function which should return linearized state-space dynamics model as numpy array
    NU                (int) Number of controlled states
    NY                (int) Number of outputs
    T                 (int) Number of reference points along trajectory
    dt                (float) Sample time [sec] of the controller
    output_weight     (float) Weight for control of the output
    input_weight      (float) Weight for input constraint control
    input_lim         (float) Control input limit
    input_rate_lim    (float) Control input rate limit
    ================  ==================================================
    '''

    def __init__(self, getModel,
                 T=5,
                 dt=0.01,
                 output_weight=None,
                 input_weight=None,
                 input_lim=None,
                 input_rate_lim=None,
                 proj_horizon=0.5):

        super().__init__(NU=1, NY=1, T=T, dt=dt, proj_horizon=proj_horizon)
        self.getModel = getModel
        self.output_weight = output_weight
        self.input_weight = input_weight
        self.input_lim = input_lim
        self.input_rate_lim = input_rate_lim

        self.setup_problem()

    def setup_problem(self):
        A, B, C, _ = self.getModel(sample_time=self.dt)
        NU = B.shape[1]  # Number of inputs to state-space model
        NX = A.shape[1]  # Number of states
        NY = C.shape[0]  # Number of outputs
        T = self.T  # Prediction horizon
        Q_DEFAULT = 5
        R_DEFAULT = 1

        # Constraints
        if self.input_lim is not None:
            umax = np.array(self.input_lim)
            umin = -1 * umax
        if self.input_rate_lim is not None:
            urmax = np.array([[1], [1]]) * 1
            urmin = -1 * urmax

        # Objective function
        if self.output_weight is None:
            Q = sparse.eye(NY) * Q_DEFAULT
        else:
            Q = sparse.diags(self.output_weight)

        if self.input_weight is None:
            R = sparse.eye(NU) * R_DEFAULT
        else:
            R = sparse.diags(self.input_weight)

        # Define problem
        self.u = cp.Variable((NU, T+1)) # T has one extra buffer element at the end
        x = cp.Variable((NX, T+1))
        y = cp.Variable((NY, T+1))
        self.y_ref = cp.Parameter((NY, T))
        # self.y_ref.value = np.zeros((NY, T))
        self.x_init = cp.Parameter(NX)
        # self.x_init.value = np.zeros((NX,1))
        objective = 0
        constraints = [x[:, 0] == self.x_init]

        for k in range(T):
            objective += cp.quad_form(y[:, k] -
                                         self.y_ref[:, k], Q) + cp.quad_form(self.u[:, k], R)

            constraints += [x[:, k+1] == A*x[:, k] + B*self.u[:, k]]
            constraints += [y[:, k] == C*x[:, k]]

            if self.input_rate_lim is not None:
                constraints += [urmin <= self.u[:, k+1] -
                                self.u[:, k], self.u[:, k+1] - self.u[:, k] <= urmax]
            if self.input_lim is not None:
                constraints += [umin <= self.u[:, k], self.u[:, k] <= umax]

        self.prob = cp.Problem(cp.Minimize(objective), constraints)

    def update_values(self, x_init, y_ref):
        self.y_ref.value = y_ref
        self.x_init.value = x_init

    def control(self,y_desired, init_state):
        self.update_values(init_state, y_desired)
        # self.prob.solve(solver=cp.OSQP, max_iter=100,
                #    verbose=False, warm_start=True)
        self.prob.solve(cp.ECOS)
        return self.u[:, 0].value[0]

    def control_with_acado(self, y_desired, init_state):
        A, B, C, _ = self.getModel(sample_time=self.dt)
        NU = B.shape[1]    # Number of inputs to state-space model
        NX = A.shape[1]    # Number of states
        NYN = NX#C.shape[0]   # Number of outputs
        NY = NX+NU        # Control problem
        T = self.T         # Prediction horizon

        x0=np.zeros((1,NX))
        X=np.zeros((T+1,NX))
        X[0:T,1] = y_desired[0,:]
        X[0:T,3] = y_desired[1,:]
        U=np.zeros((T,NU))
        Y=np.concatenate((X[0:T,:],np.zeros((T,NU))),axis=1)
        # Y = np.zeros((T,NY))
        yN=np.zeros((1,NYN))
        x0[0,:]=init_state  # initial state
        yN[0,:]=Y[-1,:NYN]         # reference terminal state

        # Objective function
        Q_DEFAULT = 5
        R_DEFAULT = 1
        # Q = sparse.eye(NY) * Q_DEFAULT
        # Qf = sparse.eye(NYN) * Q_DEFAULT
        # R = sparse.eye(NU) * R_DEFAULT
        Q = np.diag([Q_DEFAULT,0,Q_DEFAULT,0,R_DEFAULT])
        Qf = np.diag([Q_DEFAULT,0,Q_DEFAULT,0])

        X, U = acado.mpc(0, 1, x0, X,U,Y,yN, np.transpose(np.tile(Q,T)), Qf, 0)

        if U[0] != 0 or U[4] != 0:
            print('NOT ZERO!!!!')
        return 0

    # def control2(self, y_desired, init_state):
    #     A, B, C, _ = self.getModel(sample_time=self.dt)

    #     NU = B.shape[1]  # Number of inputs to state-space model
    #     NX = A.shape[1]  # Number of states
    #     NY = C.shape[0]  # Number of outputs
    #     T = self.T  # Prediction horizon
    #     Q_DEFAULT = 5
    #     R_DEFAULT = 1

    #     # Constraints
    #     if self.input_lim is not None:
    #         umax = np.array(self.input_lim)
    #         umin = -1 * umax
    #     if self.input_rate_lim is not None:
    #         urmax = np.array([[1], [1]]) * 1
    #         urmin = -1 * urmax

    #     # Objective function
    #     if self.output_weight is None:
    #         Q = sparse.eye(NY) * Q_DEFAULT
    #     else:
    #         Q = sparse.diags(self.output_weight)

    #     if self.input_weight is None:
    #         R = sparse.eye(NU) * R_DEFAULT
    #     else:
    #         R = sparse.diags(self.input_weight)

    #     # Define problem
    #     u = cp.Variable((NU, T+1)) # T has one extra buffer element at the end
    #     x = cp.Variable((NX, T+1))
    #     y = cp.Variable((NY, T+1))
    #     yref = cp.Parameter((NY, T))
    #     x_init = cp.Parameter(NX)
    #     objective = 0
    #     constraints = [x[:, 0] == x_init]

    #     yref.value = y_desired
    #     x_init.value = init_state

    #     for k in range(T):
    #         objective += cp.quad_form(y[:, k] -
    #                                      yref[:, k], Q) + cp.quad_form(u[:, k], R)

    #         constraints += [x[:, k+1] == A*x[:, k] + B*u[:, k]]
    #         constraints += [y[:, k] == C*x[:, k]]

    #         if self.input_rate_lim is not None:
    #             constraints += [urmin <= u[:, k+1] -
    #                             u[:, k], u[:, k+1] - u[:, k] <= urmax]
    #         if self.input_lim is not None:
    #             constraints += [umin <= u[:, k], u[:, k] <= umax]

    #     prob = cp.Problem(cp.Minimize(objective), constraints)
    #     prob.solve(solver=cp.OSQP, max_iter=100,
    #                verbose=False, warm_start=True)

    #     return u[:, 0].value[0]

def main():
    pass

if __name__ == '__main__':
    main()


"""
A selection of functions for extracting vectors
Encoder + vocab expansion
"""
import theano
import theano.tensor as tensor
from theano.sandbox.rng_mrg import MRG_RandomStreams as RandomStreams

import pickle as pkl
import numpy
import nltk

from collections import OrderedDict, defaultdict
from nltk.tokenize import word_tokenize
from scipy.linalg import norm
from gensim.models import Word2Vec as word2vec
from sklearn.linear_model import LinearRegression

from .utils import load_params, init_tparams
from .model import init_params, build_encoder, build_encoder_w2v

#-----------------------------------------------------------------------------#
# Specify model and dictionary locations here
#-----------------------------------------------------------------------------#
path_to_model = '/Users/chrisn/mad-science/pythia/data/book_corpus/model/corpus.npz'
path_to_dictionary = '/Users/chrisn/mad-science/pythia/data/book_corpus/model/vocab.pickle'
path_to_word2vec =  '/Users/chrisn/mad-science/pythia/data/stackexchange/model/word2vecAnime.bin'
#-----------------------------------------------------------------------------#

def load_model(embed_map=None):
    """
    Load all model components + apply vocab expansion
    """
    # Load the worddict
    print('Loading dictionary...')
    with open(path_to_dictionary, 'rb') as f:
        worddict = pkl.load(f)

    # Create inverted dictionary
    print('Creating inverted dictionary...')
    word_idict = dict()
    for kk, vv in worddict.items():
        word_idict[vv] = kk
    word_idict[0] = '<eos>'
    word_idict[1] = 'UNK'

    # Load model options
    print('Loading model options...')
    with open('%s.pkl'%path_to_model, 'rb') as f:
        options = pkl.load(f)

    # Load parameters
    print('Loading model parameters...')
    params = init_params(options)
    params = load_params(path_to_model, params)
    tparams = init_tparams(params)

    # Extractor functions
    print('Compiling encoder...')
    trng = RandomStreams(1234)
    trng, x, x_mask, ctx, emb = build_encoder(tparams, options)
    f_enc = theano.function([x, x_mask], ctx, name='f_enc')
    f_emb = theano.function([x], emb, name='f_emb')
    trng, embedding, x_mask, ctxw2v = build_encoder_w2v(tparams, options)
    f_w2v = theano.function([embedding, x_mask], ctxw2v, name='f_w2v')

    # Load word2vec, if applicable
    if embed_map == None:
        print('Loading word2vec embeddings...')
        embed_map = load_googlenews_vectors(path_to_word2vec)

    # Lookup table using vocab expansion trick
    print('Creating word lookup tables...')
    table = lookup_table(options, embed_map, worddict, word_idict, f_emb)

    # Store everything we need in a dictionary
    print('Packing up...')
    model = {}
    model['options'] = options
    model['table'] = table
    model['f_w2v'] = f_w2v

    return model

def encode(model, X, use_norm=True, verbose=True, batch_size=128, use_eos=False):
    """
    Encode sentences in the list X. Each entry will return a vector
    """
    # first, do preprocessing
    X = preprocess(X)

    # word dictionary and init
    d = defaultdict(lambda : 0)
    for w in list(model['table'].keys()):
        d[w] = 1
    features = numpy.zeros((len(X), model['options']['dim']), dtype='float32')

    # length dictionary
    ds = defaultdict(list)
    captions = [s.split() for s in X]
    for i,s in enumerate(captions):
        ds[len(s)].append(i)

    # Get features. This encodes by length, in order to avoid wasting computation
    for k in list(ds.keys()):
        if verbose:
            print(k)
        numbatches = int(len(ds[k]) / batch_size) + 1
        for minibatch in range(numbatches):
            caps = ds[k][minibatch::numbatches]

            if use_eos:
                embedding = numpy.zeros((k+1, len(caps), model['options']['dim_word']), dtype='float32')
            else:
                embedding = numpy.zeros((k, len(caps), model['options']['dim_word']), dtype='float32')
            for ind, c in enumerate(caps):
                caption = captions[c]
                for j in range(len(caption)):
                    if d[caption[j]] > 0:
                        embedding[j,ind] = model['table'][caption[j]]
                    else:
                        embedding[j,ind] = model['table']['UNK']
                if use_eos:
                    embedding[-1,ind] = model['table']['<eos>']
            if use_eos:
                ff = model['f_w2v'](embedding, numpy.ones((len(caption)+1,len(caps)), dtype='float32'))
            else:
                ff = model['f_w2v'](embedding, numpy.ones((len(caption),len(caps)), dtype='float32'))
            if use_norm:
                for j in range(len(ff)):
                    ff[j] /= norm(ff[j])
            for ind, c in enumerate(caps):
                features[c] = ff[ind]

    return features

def preprocess(text):
    """
    Preprocess text for encoder
    """
    X = []
    sent_detector = nltk.data.load('tokenizers/punkt/english.pickle')
    for t in text:
        sents = sent_detector.tokenize(t)
        result = ''
        for s in sents:
            tokens = word_tokenize(s)
            result += ' ' + ' '.join(tokens)
        X.append(result)
    return X

def load_googlenews_vectors(path_to_word2vec):
    """
    load the word2vec GoogleNews vectors
    """
    embed_map = word2vec.load_word2vec_format(path_to_word2vec, binary=True)
    return embed_map

def lookup_table(options, embed_map, worddict, word_idict, f_emb, use_norm=False):
    """
    Create a lookup table from linear mapping of word2vec into RNN word space
    """
    wordvecs = get_embeddings(options, word_idict, f_emb)
    clf = train_regressor(options, embed_map, wordvecs, worddict)
    table = apply_regressor(clf, embed_map, use_norm=use_norm)

    for i in range(options['n_words']):
        w = word_idict[i]
        table[w] = wordvecs[w]
        if use_norm:
            table[w] /= norm(table[w])
    return table

def get_embeddings(options, word_idict, f_emb, use_norm=False):
    """
    Extract the RNN embeddings from the model
    """
    d = OrderedDict()
    for i in range(options['n_words']):
        caption = [i]
        ff = f_emb(numpy.array(caption).reshape(1,1)).flatten()
        if use_norm:
            ff /= norm(ff)
        d[word_idict[i]] = ff
    return d

def train_regressor(options, embed_map, wordvecs, worddict):
    """
    Return regressor to map word2vec to RNN word space
    """
    # Gather all words from word2vec that appear in wordvecs
    d = defaultdict(lambda : 0)
    for w in list(embed_map.vocab.keys()):
        d[w] = 1
    shared = OrderedDict()
    count = 0
    for w in list(worddict.keys())[:options['n_words']-2]:
        if d[w] > 0:
            shared[w] = count
            count += 1

    # Get the vectors for all words in 'shared'
    wordvec_size = embed_map.vector_size
    w2v = numpy.zeros((len(shared), wordvec_size), dtype='float32')
    sg = numpy.zeros((len(shared), options['dim_word']), dtype='float32')
    for w in list(shared.keys()):
        w2v[shared[w]] = embed_map[w]
        sg[shared[w]] = wordvecs[w]

    clf = LinearRegression()
    clf.fit(w2v, sg)
    return clf

def apply_regressor(clf, embed_map, use_norm=False):
    """
    Map words from word2vec into RNN word space
    """
    wordvecs = OrderedDict()
    for i, w in enumerate(embed_map.vocab.keys()):
        if '_' not in w:
            wordvecs[w] = clf.predict(embed_map[w]).astype('float32')
            if use_norm:
                wordvecs[w] /= norm(wordvecs[w])
    return wordvecs

import numpy,from scipy,import theano,from theano


# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for searcher_dataloader."""

import numpy as np
import tensorflow as tf
from tensorflow_examples.lite.model_maker.core.data_util import searcher_dataloader
from tensorflow_lite_support.python.task.core.proto import base_options_pb2
from tensorflow_lite_support.python.task.vision import image_embedder
from tensorflow_examples.lite.model_maker.core import test_util

_BaseOptions = base_options_pb2.BaseOptions
_ImageEmbedder = image_embedder.ImageEmbedder
_ImageEmbedderOptions = image_embedder.ImageEmbedderOptions
_SearcherDataLoader = searcher_dataloader.DataLoader

class SearcherDataloaderTest(tf.test.TestCase):

  def setUp(self):
    super().setUp()
    self.tflite_path = test_util.get_test_data_path(
        "mobilenet_v2_035_96_embedder_with_metadata.tflite")

  def test_concat_dataset(self):
    dataset1 = np.random.rand(2, 1280)
    dataset2 = np.random.rand(2, 1280)
    dataset3 = np.random.rand(1, 1280)
    metadata = ["0", "1", b"\x11\x22", b"\x33\x44", "4"]
    data = _SearcherDataLoader(embedder_path=self.tflite_path)
    data._dataset = dataset1
    data._cache_dataset_list = [dataset2, dataset3]
    data._metadata = metadata

    self.assertTrue((data.dataset == np.vstack([dataset1, dataset2,
                                                dataset3])).all())
    self.assertEqual(data.metadata, metadata)

  def test_append(self):
    dataset1 = np.random.rand(4, 1280)
    metadata1 = ["0", "1", b"\x11\x22", "3"]
    data_loader1 = _SearcherDataLoader(embedder_path=self.tflite_path)
    data_loader1._dataset = dataset1
    data_loader1._metadata = metadata1

    dataset2 = np.random.rand(2, 1280)
    metadata2 = [b"\x33\x44", "5"]
    data_loader2 = _SearcherDataLoader(embedder_path=self.tflite_path)
    data_loader2._dataset = dataset2
    data_loader2._metadata = metadata2

    data_loader1.append(data_loader2)
    self.assertEqual(data_loader1.dataset.shape, (6, 1280))
    self.assertEqual(data_loader1.metadata,
                     ["0", "1", b"\x11\x22", "3", b"\x33\x44", "5"])

  def test_init(self):
    # Initializes from embedder first and then generates the embedding dataset
    # from raw input leverage embedder later.
    data = _SearcherDataLoader(embedder_path=self.tflite_path)
    self.assertEqual(data.dataset.shape, (0,))
    self.assertEqual(data.metadata, [])

    # Initializes directly from the dataset and metadata.
    dataset = np.random.rand(2, 512)
    metadata = ["0", b"\66"]
    data = _SearcherDataLoader(dataset=dataset, metadata=metadata)
    self.assertTrue((data.dataset == dataset).all())
    self.assertEqual(data.metadata, metadata)

if __name__ == "__main__":
  tf.test.main()


# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import argparse
import os
import time
import numpy as np
import timm
from loguru import logger
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.multiprocessing as mp
from basecls.utils import registers, set_nccl_env, set_num_threads
import benchmark_timm as benchmark

class X: pass
DEVICE = dict(cpu="cpu", gpu="cuda")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-a",
        "--arch",
        default="resnet50",
        help="model architecture (default: resnet50)",
    )
    parser.add_argument(
        "-n",
        "--ngpus",
        default=1,
        type=int,
        help="number of GPUs per node (default: None, use all available GPUs)",
    )
    parser.add_argument(
        "-b",
        "--batch-size",
        metavar="SIZE",
        default=64,
        type=int,
        help="batch size for single GPU (default: 64)",
    )
    parser.add_argument(
        "--trace",
        action='store_true',
        default=False,
        help="whether use trace or not (default: False)",
    )
    parser.add_argument(
        "-s",
        "--steps",
        default=200,
        type=int,
        help="number of train steps (default: 200)",
    )
    parser.add_argument(
        "--lr",
        metavar="LR",
        default=0.025,
        help="learning rate for single GPU (default: 0.025)",
    )
    parser.add_argument("--momentum", default=0.9, help="momentum (default: 0.9)")
    parser.add_argument(
        "--weight-decay", default=1e-4, help="weight decay (default: 0.9)"
    )
    parser.add_argument(
        "-p",
        "--print-freq",
        default=1,
        type=int,
        metavar="N",
        help="print frequency (default: 1)",
    )
    parser.add_argument(
        "-m",
        "--mode",
        default="normal",
        type=str,
        choices=["normal", "mp"],
        help="Quantization Mode\n"
        "normal: no quantization, using float32\n"
        "mp: input type of conv and bn is fp16",
    )

    args = parser.parse_args()
    new_args = X()
    new_args.model = args.arch
    new_args.mode = "train"
    new_args.device = "gpu"
    new_args.amp = 0 if args.mode == "normal" else 2
    new_args.fastrun = False
    new_args.trace = args.trace
    new_args.batch_size = args.batch_size
    new_args.channel = 3
    new_args.height = 224
    new_args.width = 224
    new_args.world_size = args.ngpus
    new_args.warm_iters = 50
    new_args.total_iters = args.steps
    new_args.log_seconds = 2

    set_nccl_env()
    set_num_threads()

    if new_args.world_size == 1:
        worker(0, new_args)
    else:
        mp.spawn(worker, nprocs=new_args.world_size, args=(new_args,))

def worker(local_rank, args):
    if args.world_size > 1:
        dist.init_process_group(
            backend="nccl" if args.device == "gpu" else "gloo",
            init_method="tcp://127.0.0.1:6789",
            world_size=args.world_size,
            rank=local_rank,
        )

    if local_rank != 0:
        logger.remove()
    logger.info(f"args: {args}")

    if args.fastrun:
        logger.info("Using fastrun mode...")
        torch.backends.cudnn.benchmark = True

    device = torch.device(DEVICE[args.device], local_rank)
    if args.device == "gpu":
        torch.cuda.set_device(local_rank)

    model = timm.create_model(args.model, pretrained=False).to(device)
    if args.world_size > 1:
        model = nn.parallel.DistributedDataParallel(model)

    dataloader = benchmark.FakeDataLoader(
        args.batch_size,
        (args.height, args.width),
        args.channel,
        length=args.warm_iters + args.total_iters,
        num_classes=1000,
    )

    if args.mode == "train":
        BenchCls = benchmark.TrainBench
    elif args.mode == "eval":
        BenchCls = benchmark.EvalBench
    else:
        raise NotImplementedError(f"Benchmark mode '{args.mode}' not supported")

    bench = BenchCls(model, dataloader, device, args.trace, args.amp)
    avg_time = bench.benchmark(args.warm_iters, args.log_seconds)

    if local_rank == 0:
        print("="*20, "summary", "="*20)
        print(" benchmark: timm")
        print("      arch: {}".format(args.model))
        print("train_mode: {}".format("normal" if args.amp == 0 else "mp"))
        print(" batchsize: {}".format(args.batch_size))
        print("      #GPU: {}".format(args.world_size))
        print("  avg time: {:.3f} seconds".format(avg_time/1000))

if __name__ == "__main__":
    main()


# -*- coding: utf-8 -*-
"""
UDP server implemented as the low-level python code here.

This script is used in connection with the LabVIEW code for sending / receiving commands and, moreover,
images. This script is used for testing of capabilities of entire image transfer.
"""
# %% Imports
import socket
import time
import numpy as np

# %% Parameters (global)
host = 'localhost'
port = 5005  # Port for this script
port_listener = 5100
st_n_bytes = 1024

# %% Decoding pixel values - redundant explicit method

# %% Encoding list of numbers into string for sending it - RETAINED for needs of back communication
def encodeNumbers(numbersInList: list) -> str:
    """Encode list of numbers to a string for sending it to a client."""
    sending_string = str(numbersInList)
    # Make a pour string with numbers - below removing brackets and white spaces
    sending_string = sending_string.replace('[', '')
    sending_string = sending_string.replace(']', '')
    sending_string = sending_string.replace(',', '')
    sending_string = sending_string.encode()  # Default encdoing - utf-8
    return sending_string

# %% UDP communication with LabVIEW code for transferring an image
with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
    s.bind((host, port))
    # Below - from TCP script
    # s.listen()
    # (connection, address) = s.accept()
    flag = True
    print("Python UDP Server launched")
    with s:
        while flag:
            (data, address) = s.recvfrom(st_n_bytes)  # Receive 1024 bytes of data from LV client (text commands)
            command = str(data, encoding='utf-8')  # Conversion to a normal string
            if "Ping" in command:
                print(command, '- received command')
                sendingString = "Echo"
                sendingString = sendingString.encode()  # to utf-8
                s.sendto(sendingString, address)

            if "Image" in command:
                print(command, '- received command')
                # sendingString = "Receiving an Image"
                # sendingString = sendingString.encode()  # to utf-8
                # s.sendto(sendingString, address)  # Confirmation to a client - not neccessary
                (width, address) = s.recvfrom(st_n_bytes)  # Height and width could be flipped!
                (height, address) = s.recvfrom(st_n_bytes)   # Height and width could be flipped!
                width = int(str(width, encoding='utf-8'))
                height = int(str(height, encoding='utf-8'))
                print("Sizes of an image [pixels]:", width, "x", height)  # debug
                img = np.zeros((height, width), dtype="uint16")  # image initialization (container)

                # (img_max_size, address) = s.recvfrom(st_n_bytes)  # estimation of image size in bytes
                # Really, after tests it was defined that maximum packet size for transfer through UDP in Windows
                # is about 8000 kb, so this number will be used as a constant to reduce transferring numbers

                img_max_size = 110*100*8  # ... bytes to be transferred ( ~ 80000 should be transferred w/t errors)
                # 11000 may seem as magical number but I haven't found the way to access it through some call
                # 11000 = 110 * 100 pixels that could be transferred w/t error through single UDP port (Win)
                n_rows_tr = 11000 // width
                n_chunks = height // n_rows_tr
                n_last_transfer = height - (n_chunks * n_rows_tr)
                if n_last_transfer > 0:
                    n_chunks += 1
                print("# of chunks:", n_chunks)
                print("# for last transfer:", n_last_transfer)
                print("# of rows for transfer", n_rows_tr)
                # s.settimeout(2)
                # Below - receiving an image in chunks
                for i in range(n_chunks):
                    (raw_chunk, address) = s.recvfrom(img_max_size)
                    raw_chunk = (str(raw_chunk, encoding='utf-8'))  # Converting byte string to normal one
                    string_chunk = raw_chunk.split()  # split to numbers using standard whitespace as a separator
                    # print("string chunk:", string_chunk)
                    if i < (n_chunks - 1):
                        for j in range(n_rows_tr):
                            # Transfer by a row to cumulative image
                            img[i*n_rows_tr + j, :] = np.uint16(string_chunk[width*j:width*(j+1)])
                    if (i == (n_chunks - 1)) and (n_last_transfer > 0):
                        for j in range(n_last_transfer):
                            img[i*n_rows_tr + j, :] = np.uint16(string_chunk[width*j:width*(j+1)])
                    elif (i == (n_chunks - 1)):
                        for j in range(n_rows_tr):
                            img[i*n_rows_tr + j, :] = np.uint16(string_chunk[width*j:width*(j+1)])
                print("Image receiving completed")
                sendingString = "An Image received"
                sendingString = sendingString.encode()  # to utf-8
                s.sendto(sendingString, address)

            elif "QUIT" in command:
                print(command, '- received command')
                # Do the quiting action with a delay
                s.sendto(sendingString, address)
                time.sleep(0.2)  # A delay for preventing of closing connection automatically by Python (causing errors)
                flag = False
                break

        # %% For correctly representation of an image - it should be performed after quiting of image transfer
        # For testing if the transfer of an image perfromed correctly
        try:
            if np.size(img, axis=0) > 0:
                from skimage import io
                from skimage.util import img_as_ubyte
                import matplotlib.pyplot as plt
                if np.max(img) < 256:
                    img = np.uint8(img)
                img_show = img_as_ubyte(img)
                plt.figure("Transferred image")
                io.imshow(img_show)
        except NameError:
            print("No image transferred")


###########################################################################
# Created by: CASIA IVA
# Email: jliu@nlpr.ia.ac.cn
# Copyright (c) 2018
###########################################################################
from __future__ import division
import os
import numpy as np
import torch
import torch.nn as nn
from torch.nn import Module, Sequential, Conv2d, Parameter, Softmax
from torch.nn.functional import upsample, normalize
from ...nn import PAM_Module
from ...nn import CAM_Module
from .base import BaseNet
from copy import deepcopy

__all__ = ['DANet_Dep', 'get_danet_dep']

class DANet_Dep(BaseNet):
    r"""Fully Convolutional Networks for Semantic Segmentation

    Parameters
    ----------
    nclass : int
        Number of categories for the training dataset.
    backbone : string
        Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50',
        'resnet101' or 'resnet152').
    norm_layer : object
        Normalization layer used in backbone network (default: :class:`mxnet.gluon.nn.BatchNorm`;

    Reference:

        Long, Jonathan, Evan Shelhamer, and Trevor Darrell. "Fully convolutional networks
        for semantic segmentation." *CVPR*, 2015

    """

    def __init__(self, nclass, backbone, dep_encoder = 'cnn', aux=False, se_loss=False, norm_layer=nn.BatchNorm2d, **kwargs):
        super(DANet_Dep, self).__init__(nclass, backbone, aux, se_loss, norm_layer=norm_layer, **kwargs)
        self.head = DANetHead(2048, nclass, norm_layer)

        # depth encoder
        self.dep_encoder = deepcopy(self.pretrained)

        self.dep_encoder = dep_encoder
        if dep_encoder == 'cnn':
            self.dep_layer0 = nn.Sequential(
                          nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False),
                          norm_layer(64),
                          nn.ReLU(inplace=True),
                          nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
            )

    def forward(self, x, dep):
        imsize = x.size()[2:]
        _, _, c3, c4 = self.base_forward(x)

        if self.dep_encoder == 'cnn':
            d0 = self.dep_layer0(dep)
            d1 = self.dep_encoder.layer1(d0)
            d2 = self.dep_encoder.layer2(d1)
            d3 = self.dep_encoder.layer3(d2)
            d = self.dep_encoder.layer4(d3)
        else:
            d = torch.nn.functional.interpolate(dep, imsize, mode='bilinear', align_corners=False)

        x = self.head(c4, d)
        x = list(x)
        x[0] = upsample(x[0], imsize, **self._up_kwargs)
        x[1] = upsample(x[1], imsize, **self._up_kwargs)
        x[2] = upsample(x[2], imsize, **self._up_kwargs)

        outputs = [x[0]]
        outputs.append(x[1])
        outputs.append(x[2])
        return tuple(outputs)

class DANetHead(nn.Module):
    def __init__(self, in_channels, out_channels, norm_layer):
        super(DANetHead, self).__init__()
        inter_channels = in_channels // 4

        # spatial attention
        self.conv5a = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
                                    norm_layer(inter_channels),
                                    nn.ReLU())
        self.sa = PAM_Module_Dep(inter_channels)

        self.conv51 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
                                    norm_layer(inter_channels),
                                    nn.ReLU())
        # preprocess for depth
        self.conv5a_d = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
                                      norm_layer(inter_channels),
                                      nn.ReLU())
        # spatial attention output
        self.conv6 = nn.Sequential(nn.Dropout2d(0.1, False), nn.Conv2d(inter_channels, out_channels, 1))

        # channel attention
        self.conv5c = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
                                    norm_layer(inter_channels),
                                    nn.ReLU())

        self.sc = CAM_Module(inter_channels)

        self.conv52 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
                                    norm_layer(inter_channels),
                                    nn.ReLU())

        # channel attention output
        self.conv7 = nn.Sequential(nn.Dropout2d(0.1, False), nn.Conv2d(inter_channels, out_channels, 1))

        # overall output
        self.conv8 = nn.Sequential(nn.Dropout2d(0.1, False), nn.Conv2d(inter_channels, out_channels, 1))

    def forward(self, x, dep):
        # spatial attention
        feat1 = self.conv5a(x)
        dep = self.conv5a_d(dep)

        sa_feat = self.sa(feat1, dep)
        sa_conv = self.conv51(sa_feat)
        sa_output = self.conv6(sa_conv)

        # channel attention
        feat2 = self.conv5c(x)
        sc_feat = self.sc(feat2)
        sc_conv = self.conv52(sc_feat)
        sc_output = self.conv7(sc_conv)

        feat_sum = sa_conv + sc_conv
        sasc_output = self.conv8(feat_sum)

        output = [sasc_output, sa_output, sc_output]
        return tuple(output)

def get_danet_dep(dataset='pascal_voc', backbone='resnet50', pretrained=False,
                  root='../../encoding/models/pretrain', **kwargs):
    r"""DANet model from the paper `"Dual Attention Network for Scene Segmentation"
    <https://arxiv.org/abs/1809.02983.pdf>`
    """
    acronyms = {
        'pascal_voc': 'voc',
        'pascal_aug': 'voc',
        'pcontext': 'pcontext',
        'ade20k': 'ade',
        'cityscapes': 'cityscapes',
    }
    # infer number of classes
    from ...datasets import datasets, VOCSegmentation, VOCAugSegmentation, ADE20KSegmentation
    model = DANet_Dep(datasets[dataset.lower()].NUM_CLASS, backbone=backbone, root=root, **kwargs)
    if pretrained:
        from ..model_store import get_model_file
        model.load_state_dict(torch.load(
            get_model_file('fcn_%s_%s' % (backbone, acronyms[dataset]), root=root)),
            strict=False)
    return model

class PAM_Module_Dep(Module):
    """ Position attention module"""
    # Ref from SAGAN
    def __init__(self, in_dim):
        super(PAM_Module_Dep, self).__init__()
        self.chanel_in = in_dim

        self.query_conv = Conv2d(in_channels=in_dim, out_channels=in_dim//8, kernel_size=1)
        self.key_conv = Conv2d(in_channels=in_dim, out_channels=in_dim//8, kernel_size=1)

        self.query_d_conv = Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
        self.key_d_conv = Conv2d(in_channels=in_dim, out_channels=in_dim//8, kernel_size=1)

        self.value_conv = Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
        self.gamma = Parameter(torch.zeros(1))

        self.softmax = Softmax(dim=-1)

    def forward(self, x, dep):
        """
            inputs :
                x : input feature maps( B X C X H X W)
            returns :
                out : attention value + input feature
                attention: B X (HxW) X (HxW)
        """
        m_batchsize, C, height, width = x.size()
        query = torch.cat((self.query_conv(x), self.query_d_conv(dep)), 1)            # [B, 128, h, w]
        proj_query = query.view(m_batchsize, -1, width*height).permute(0, 2, 1)       # [B, hw, 128]
        key = torch.cat((self.key_conv(x), self.key_d_conv(dep)), 1)                  # [B, 128, h, w]
        proj_key = key.view(m_batchsize, -1, width*height)                            # [B, 128, hw]

        energy = torch.bmm(proj_query, proj_key)                                              # [B, hw, hw]
        attention = self.softmax(energy)                                                      # [B, hw, hw]

        proj_value = self.value_conv(x).view(m_batchsize, -1, width*height)                   # [B, 512, hw]
        out = torch.bmm(proj_value, attention.permute(0, 2, 1))                               # [B, 512, hw]
        out = out.view(m_batchsize, C, height, width)                                         # [B, 512, h, w]

        out = self.gamma*out + x
        return out


from torch import nn
import torch
import numpy as np
import torch.optim as optim
from transformers import BertModel,BertConfig
import src.config.ModelConfig as ModelConfig
import torch.nn.functional as F
from src.utils.common import flatten_lists
class MRCBert(nn.Module):
    def __init__(self,num_labels,vocab_size=0,emb_size=0,hidden_size=0,use_pretrained=True):
        super(MRCBert,self).__init__()
        if use_pretrained:
            self.bertconfig=BertConfig.from_pretrained(ModelConfig.BERT_BASE_CHINESE,author='lingze',num_labels=num_labels)
            self.bert=BertModel.from_pretrained(ModelConfig.BERT_BASE_CHINESE,config=self.bertconfig)
            self.vocab_size=self.bertconfig.vocab_size
            self.emb_size=self.bertconfig.hidden_size
            self.hidden_size=self.bertconfig.hidden_size
        else:
            self.vocab_size=vocab_size
            self.emb_size=emb_size
            self.hidden_size=hidden_size
            self.bertconfig=BertConfig(vocab_size=self.vocab_size,hidden_size=self.hidden_size,\
                num_labels=num_labels,author='lingze')
            self.bert=BertModel(config=self.bertconfig)
        self.dropout=nn.Dropout(self.bertconfig.hidden_dropout_prob)
        self.ll=nn.Linear(self.hidden_size,num_labels)

    def forward(self,input_ids,attention_mask,token_type_ids,berttokenizer):
        '''
        input_ids:输入数据 [Batch_size,max_length]
        token_type_ids: 标定word属于哪个句子，只能取1,0.
        attention_mask:Mask to avoid performing attention on padding token indices. Mask Value selected in [0,1] [Batch_size,max_length]
        '''
        SEP=berttokenizer.convert_tokens_to_ids('[SEP]')
        CLS=berttokenizer.convert_tokens_to_ids('[CLS]')

        hidden_score=self.bert(input_ids=input_ids,\
            attention_mask=attention_mask,token_type_ids=token_type_ids)[0]
        #[batch_size,sent_len,hidden_size]
        # import pdb; pdb.set_trace()
        #这里我们需要利用token_type_ids 丢掉Query Representation,同时可以去掉Padding部分
        mask=(token_type_ids==0) & (attention_mask==1) & (input_ids !=SEP) & (input_ids !=CLS)

        mask=mask.unsqueeze(2).expand(-1,-1,self.hidden_size)
        hidden_score=hidden_score.masked_select(mask).reshape(-1,self.hidden_size)
        hidden_score=self.dropout(hidden_score)
        #加入Linear层
        return self.ll(hidden_score)

    #labels [batch_size,seq_len] 不加padding
    def cal_loss(self,score,labels,weight):
        flatten_labels=flatten_lists(labels)
        return F.cross_entropy(score,labels,weight=weight)

    def predict(self,input_ids,attention_mask,token_type_ids,berttokenizer):

        score=self.forward(input_ids=input_ids,attention_mask=attention_mask,\
            token_type_ids=token_type_ids,berttokenizer=berttokenizer)
        #score [flatten_size, num_labels] 可以通过attention_mask,token_type_ids将其还原
        labels=torch.argmax(score,dim=1)
        #labels [flatten_size]
        return labels


#
# Copyright (c) 2021, NVIDIA 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 pytest

import merlin.models.tf as mm
from merlin.io.dataset import Dataset
from merlin.schema import Tags

def test_topk_index(ecommerce_data: Dataset):
    import tensorflow as tf

    from merlin.models.tf.metrics.evaluation import ItemCoverageAt, PopularityBiasAt

    model: mm.RetrievalModel = mm.TwoTowerModel(
        ecommerce_data.schema, query_tower=mm.MLPBlock([64, 128])
    )
    model.compile(run_eagerly=False, optimizer="adam")
    model.fit(ecommerce_data, epochs=1, batch_size=50)

    item_features = ecommerce_data.schema.select_by_tag(Tags.ITEM).column_names
    item_dataset = ecommerce_data.to_ddf()[item_features].drop_duplicates().compute()
    item_dataset = Dataset(item_dataset)
    recommender = model.to_top_k_recommender(item_dataset, k=20)
    NUM_ITEMS = 1001
    item_frequency = tf.sort(
        tf.random.uniform((NUM_ITEMS,), minval=0, maxval=NUM_ITEMS, dtype=tf.int32)
    )
    eval_metrics = [
        PopularityBiasAt(item_freq_probs=item_frequency, is_prob_distribution=False, k=10),
        ItemCoverageAt(num_unique_items=NUM_ITEMS, k=10),
    ]
    batch = mm.sample_batch(ecommerce_data, batch_size=10, include_targets=False)
    _, top_indices = recommender(batch)
    assert top_indices.shape[-1] == 20
    _, top_indices = recommender(batch, k=10)
    assert top_indices.shape[-1] == 10

    for metric in eval_metrics:
        metric.update_state(predicted_ids=top_indices)
    for metric in eval_metrics:
        results = metric.result()
        metric.reset_state()
        assert results >= 0

def test_topk_index_duplicate_indices(ecommerce_data: Dataset):
    model: mm.RetrievalModel = mm.TwoTowerModel(
        ecommerce_data.schema, query_tower=mm.MLPBlock([64, 128])
    )
    model.compile(run_eagerly=True, optimizer="adam")
    model.fit(ecommerce_data, epochs=1, batch_size=50)
    item_features = ecommerce_data.schema.select_by_tag(Tags.ITEM).column_names
    item_dataset = ecommerce_data.to_ddf()[item_features].compute()
    item_dataset = Dataset(item_dataset)

    with pytest.raises(ValueError) as excinfo:
        _ = model.to_top_k_recommender(item_dataset, k=20)
    assert "Please make sure that `data` contains unique indices" in str(excinfo.value)

def test_topk_recommender_outputs(ecommerce_data: Dataset, batch_size=100):
    import numpy as np
    import tensorflow as tf

    import merlin.models.tf.dataset as tf_dataloader
    from merlin.models.tf.blocks.core.index import IndexBlock
    from merlin.models.utils.dataset import unique_rows_by_features

    def numpy_recall(labels, top_item_ids, k):
        return np.equal(np.expand_dims(labels, -1), top_item_ids[:, :k]).max(axis=-1).mean()

    model = mm.TwoTowerModel(
        ecommerce_data.schema,
        query_tower=mm.MLPBlock([64]),
        samplers=[mm.InBatchSampler()],
    )

    model.compile("adam", metrics=[mm.RecallAt(10)])
    model.fit(ecommerce_data, batch_size=batch_size, epochs=3)
    eval_metrics = model.evaluate(
        ecommerce_data, item_corpus=ecommerce_data, batch_size=batch_size, return_dict=True
    )

    # Manually compute top-k ids for a given batch
    batch = mm.sample_batch(ecommerce_data, batch_size=batch_size, include_targets=False)
    item_dataset = unique_rows_by_features(ecommerce_data, Tags.ITEM, Tags.ITEM_ID)
    candidates_dataset_df = IndexBlock.get_candidates_dataset(
        block=model.retrieval_block.item_block(), data=item_dataset, id_column="item_id"
    )
    item_tower_ids, item_tower_embeddings = IndexBlock.extract_ids_embeddings(
        candidates_dataset_df, check_unique_ids=True
    )
    batch_query_tower_embeddings = model.retrieval_block.query_block()(batch)
    batch_user_scores_all_items = tf.matmul(
        batch_query_tower_embeddings, item_tower_embeddings, transpose_b=True
    )
    top_scores, top_indices = tf.math.top_k(batch_user_scores_all_items, k=10)
    top_ids = tf.gather(item_tower_ids, top_indices)

    # Get top-k ids from the topk_recommender_model
    topk_recommender_model = model.to_top_k_recommender(ecommerce_data, k=10)
    topk_predictions, topk_items = topk_recommender_model(batch)

    # Assert top-k items from top-k recommender are the same as the manually computed top-k items
    tf.debugging.assert_equal(top_ids, topk_items)

    # Compute recall using top-k recommender
    data_dl = tf_dataloader.BatchedDataset(
        ecommerce_data,
        batch_size=batch_size,
        shuffle=False,
    )
    topk_output = topk_recommender_model.predict(data_dl)
    topk_predictions, topk_items = topk_output
    test_df = ecommerce_data.to_ddf()
    positive_item_ids = np.array(test_df["item_id"].compute().values.tolist())
    recall_at_10 = numpy_recall(positive_item_ids, topk_items, k=10)

    np.isclose(recall_at_10, eval_metrics["recall_at_10"], rtol=1e-6)


# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# This is a application of Votenet to the BenchBot Semantic SLAM problem

import os
import sys
import numpy as np
import argparse
import importlib
import time
import torch
import torch.nn as nn
import torch.optim as optim

from scipy.spatial.transform import Rotation

# Import votenet helpers explicitly (as the package does not offer any pip
# install tools)
ROOT_DIR = '/benchbot_submission/votenet'
sys.path.append(os.path.join(ROOT_DIR, 'utils'))
sys.path.append(os.path.join(ROOT_DIR, 'models'))
from pc_util import random_sampling, read_ply
from ap_helper import parse_predictions
from nms import nms_2d_faster, nms_3d_faster, nms_3d_faster_samecls

# Default settings
DATASET = 'sunrgbd'  # supported values: sunrgbd, scannet
NUM_POINTS = 40000  # number of points used in sampling the point cloud

def preprocess_point_cloud(point_cloud):
    ''' Prepare the numpy point cloud (N,3) for forward pass '''
    point_cloud = point_cloud[:, 0:3]  # do not use color for now
    floor_height = np.percentile(point_cloud[:, 2], 0.99)
    height = point_cloud[:, 2] - floor_height
    point_cloud = np.concatenate(
        [point_cloud, np.expand_dims(height, 1)], 1)  # (N,4) or (N,7)
    point_cloud = random_sampling(point_cloud, NUM_POINTS)
    pc = np.expand_dims(point_cloud.astype(np.float32), 0)  # (1,40000,4)
    return pc

def votenet_build():
    # Set file paths and dataset config
    demo_dir = os.path.join(ROOT_DIR, 'demo_files')
    if DATASET == 'sunrgbd':
        sys.path.append(os.path.join(ROOT_DIR, 'sunrgbd'))
        from sunrgbd_detection_dataset import DC  # dataset config

        checkpoint_path = os.path.join(demo_dir,
                                       'pretrained_votenet_on_sunrgbd.tar')
        pc_path = os.path.join(demo_dir, 'input_pc_sunrgbd.ply')
    elif DATASET == 'scannet':
        sys.path.append(os.path.join(ROOT_DIR, 'scannet'))
        from scannet_detection_dataset import DC  # dataset config

        checkpoint_path = os.path.join(demo_dir,
                                       'pretrained_votenet_on_scannet.tar')
        pc_path = os.path.join(demo_dir, 'input_pc_scannet.ply')
    else:
        print('Unkown dataset %s. Exiting.' % (DATASET))
        exit(-1)

    eval_config_dict = {
        'remove_empty_box': True,
        'use_3d_nms': True,
        'nms_iou': 0.25,
        'use_old_type_nms': False,
        'cls_nms': False,
        'per_class_proposal': False,
        'conf_thresh': 0.5,
        'dataset_config': DC
    }

    # Init the model and optimzier
    MODEL = importlib.import_module('votenet')  # import network module
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    net = MODEL.VoteNet(num_proposal=256,
                        input_feature_dim=1,
                        vote_factor=1,
                        sampling='seed_fps',
                        num_class=DC.num_class,
                        num_heading_bin=DC.num_heading_bin,
                        num_size_cluster=DC.num_size_cluster,
                        mean_size_arr=DC.mean_size_arr).to(device)
    print('Constructed model.')

    # Load checkpoint
    optimizer = optim.Adam(net.parameters(), lr=0.001)
    checkpoint = torch.load(checkpoint_path)
    net.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    epoch = checkpoint['epoch']
    print("Loaded checkpoint %s (epoch: %d)" % (checkpoint_path, epoch))

    net.device = device
    net.eval_config_dict = eval_config_dict

    type2class = {
        'bed': 0,
        'table': 1,
        'sofa': 2,
        'chair': 3,
        'toilet': 4,
        'desk': 5,
        'dresser': 6,
        'night_stand': 7,
        'bookshelf': 8,
        'bathtub': 9
    }
    class2type = {type2class[t]: t for t in type2class}
    # merge 'table:1' and 'desk:5' into 'table'
    class2type[5] = 'table'
    net.type2class = type2class
    net.class2type = class2type

    return net

def votenet_forward(net, point_cloud):
    # Load and preprocess input point cloud
    net.eval()  # set model to eval mode (for bn and dp)
    # point_cloud = read_ply(pc_path)
    pc = preprocess_point_cloud(point_cloud)
    # print('Loaded point cloud data: %s' % (pc_path))

    # Model inference
    inputs = {'point_clouds': torch.from_numpy(pc).to(net.device)}
    tic = time.time()
    with torch.no_grad():
        end_points = net(inputs)
    toc = time.time()
    print('Inference time: %f' % (toc - tic))
    end_points['point_clouds'] = inputs['point_clouds']
    try:
        pred_map_cls = parse_predictions(end_points, net.eval_config_dict)
        print('Finished detection. %d object detected.' %
              (len(pred_map_cls[0])))
        return pred_map_cls
    except:
        print('Finished detection. 0 objects detected.')
        return [None]

def compute_world_T_camera(pose_rotation, pose_translation):
    """
    This function computes the transformation matrix T that transforms one point
    in the camera coordinate system to the world coordinate system. This
    function needs the scipy package - scipy.spatial.transform.Rotation
    P_w = P_r * T, where P_r = [x, y, z] is a row vector in the camera
    coordinate (x - right, y - down, z - forward) and P_w the location of the
    point in the world coordinate system (x - right, y - forward, z - up)

    :param pose_rotation: [quat_x, quat_y, quat_z, quat_w], a quaternion that
                          represents a rotation
    :param pose_translation: [tr_x, tr_y, tr_z], the
                             location of the robot origin in the world coordinate
                             system
    :return: transformation matrix
    """
    # compute the rotation matrix
    # note that the scipy...Rotation accepts the quaternion in scalar-last format
    pose_quat = pose_rotation
    rot = Rotation.from_quat(pose_quat)
    rotation_matrix = np.eye(4)
    rotation_matrix[0:3][:, 0:3] = rot.as_dcm().T

    # build the tranlation matrix
    translation_matrix = np.eye(4)
    translation_matrix[3][0:3] = np.array(pose_translation)

    # compute the tranformation matrix
    transformation_matrix = rotation_matrix @ translation_matrix

    # # convert robot's right-handed coordinate system to unreal world's left-handed coordinate system
    # transformation_matrix[:, 1] = -1.0 * transformation_matrix[:, 1]

    return transformation_matrix

def convert_depth_to_pointcloud(observations):
    # convert depth image into point cloud
    depth_image = observations['image_depth']
    image_size = depth_image.shape
    nx = np.arange(0, image_size[1])
    ny = np.arange(0, image_size[0])
    x, y = np.meshgrid(nx, ny)
    f = np.array([480], dtype=np.float32)
    cx = image_size[1] / 2.0
    cy = image_size[0] / 2.0
    x3 = (x - cx) * depth_image * 1 / f
    y3 = (y - cy) * depth_image * 1 / f
    # pack the point cloud according to the votenet camera coordinate system
    # which is actually the same as benchbot's
    point3d = np.stack((x3.flatten(), depth_image.flatten(), -y3.flatten()),
                       axis=1)

    return point3d

def votenet_detection(net, observations):
    # build the transformation matrix w_T_c that transforms
    # the point in the camera coordinate system into the world system
    # P_w = P_c * w_T_c, P_c = [x, y, z, 1] a row vector denoting a point location in the camera coordinate
    camera_pose = observations['poses']['camera']
    w_T_c = compute_world_T_camera(camera_pose['rotation_xyzw'],
                                   camera_pose['translation_xyz'])
    # print(np.array([0, 0, 1, 1], dtype=np.float) @ w_T_c)

    point3d = convert_depth_to_pointcloud(observations)

    # forward the point cloud into the votenet trained on SUNRGBD dataset
    detections = votenet_forward(net, point3d)

    # parse the detection results for current frame
    frame_results = []
    if detections[0]:
        detections = detections[0]
        for detection in detections:
            # ignore those classes that are not in ACRV classes
            if detection[0] > net.type2class['desk']:
                continue
            result = {
                "class": net.class2type[detection[0]],
                "class_ID": net.type2class[net.class2type[detection[0]]],
                "confidence": np.float64(detection[2]),
                "box_corners_camera": detection[1]
            }

            # convert the 3D box corners in the camera coordinate system to
            # the world coordinate system
            box_corners_camera = detection[1]
            box_corners_camera = np.concatenate(
                (box_corners_camera,
                 np.ones((box_corners_camera.shape[0], 1),
                         dtype=box_corners_camera.dtype)),
                axis=1)
            box_corners_world = box_corners_camera @ w_T_c
            box_corners_world = box_corners_world[:, 0:3]

            result["box_corners_world"] = box_corners_world
            # compute the 3D box's center roughly
            box_center_world = np.sum(box_corners_world, axis=0) / 8.0
            result["centroid"] = box_center_world
            # compute the 3D box's extent roughly
            box_extent_world = np.sum(
                np.abs(box_corners_world - box_center_world), axis=0) / 4.0
            result["extent"] = box_extent_world

            result["box_corners_camera"] = detection[1]
            box_center_camera = np.sum(detection[1], axis=0) / 8.0
            result["centroid_camera"] = box_center_camera
            box_extent_camera = np.sum(
                np.abs(detection[1] - box_center_camera), axis=0) / 4.0
            result["extent_camera"] = box_extent_camera

            frame_results.append(result)

    return frame_results

def votenet_nms(all_results, net, class_list):
    nms_iou_cls = 0.1  # this threshold for removing duplicates that are of the same class
    boxes_world = []
    results_list = []
    for results in all_results:
        if not results:
            continue
        for result in results:
            box_world = np.zeros((1, 8), dtype=np.float32)
            box_corners_world = result["box_corners_world"]
            box_world[0, 0] = np.min(box_corners_world[:, 0])
            box_world[0, 1] = np.min(box_corners_world[:, 1])
            box_world[0, 2] = np.min(box_corners_world[:, 2])
            box_world[0, 3] = np.max(box_corners_world[:, 0])
            box_world[0, 4] = np.max(box_corners_world[:, 1])
            box_world[0, 5] = np.max(box_corners_world[:, 2])
            box_world[0, 6] = result["confidence"]
            box_world[0, 7] = result["class_ID"]
            boxes_world.append(box_world)
            results_list.append(result)
    boxes_world = np.concatenate(boxes_world)

    pick = nms_3d_faster_samecls(boxes_world, nms_iou_cls,
                                 net.eval_config_dict['use_old_type_nms'])
    assert (len(pick) > 0)
    results_intermediate = [results_list[p] for p in pick]

    nms_iou = 0.25  # this threshold is used to remove duplicates that are not of the same class
    boxes_world = []
    results_list = []
    for result in results_intermediate:
        if not result:
            continue
        box_world = np.zeros((1, 7), dtype=np.float32)
        box_corners_world = result["box_corners_world"]
        box_world[0, 0] = np.min(box_corners_world[:, 0])
        box_world[0, 1] = np.min(box_corners_world[:, 1])
        box_world[0, 2] = np.min(box_corners_world[:, 2])
        box_world[0, 3] = np.max(box_corners_world[:, 0])
        box_world[0, 4] = np.max(box_corners_world[:, 1])
        box_world[0, 5] = np.max(box_corners_world[:, 2])
        box_world[0, 6] = result["confidence"]
        # box_world[0, 7] = result["class_ID"]
        boxes_world.append(box_world)
        results_list.append(result)
    boxes_world = np.concatenate(boxes_world)

    pick = nms_3d_faster(boxes_world, nms_iou,
                         net.eval_config_dict['use_old_type_nms'])
    assert (len(pick) > 0)

    results_final = [results_list[p] for p in pick]
    results_final = jsonify(results_final)

    for r in results_final:
        r['label_probs'] = [0] * len(class_list)
        if r['class'] in class_list:
            r['label_probs'][class_list.index(r['class'])] = r['confidence']

    return results_final

def jsonify(data_list):
    json_list = []
    for data in data_list:
        json_data = dict()
        for key, value in data.items():
            if isinstance(value, list):  # for lists
                value = [
                    jsonify(item) if isinstance(item, dict) else item
                    for item in value
                ]
            if isinstance(value, dict):  # for nested lists
                value = jsonify(value)
            if isinstance(key, int):  # if key is integer: > to string
                key = str(key)
            if type(
                    value
            ).__module__ == 'numpy':  # if value is numpy.*: > to python list
                value = value.tolist()
            json_data[key] = value
        json_list.append(json_data)
    return json_list