import mxnet as mx
import numpy as np
import sys
from io_func.feat_io import DataReadStream

# The interface of a data iter that works for bucketing
#
# DataIter
#   - default_bucket_key: the bucket key for the default symbol.
#
# DataBatch
#   - provide_data: same as DataIter, but specific to this batch
#   - provide_label: same as DataIter, but specific to this batch
#   - bucket_key: the key for the bucket that should be used for this batch

def read_content(path):
    with open(path) as input:
        content = input.read()
        content = content.replace('\n', ' <eos> ').replace('. ', ' <eos> ')
        return content

class SimpleBatch(object):
    def __init__(self, data_names, data, label_names, label, bucket_key,
                 utt_id=None, utt_len=0, effective_sample_count=None):
        self.data = data
        self.label = label
        self.data_names = data_names
        self.label_names = label_names
        self.bucket_key = bucket_key
        self.utt_id = utt_id
        self.utt_len = utt_len
        self.effective_sample_count = effective_sample_count

        self.pad = 0
        self.index = None  # TODO: what is index?

    @property
    def provide_data(self):
        return [(n, x.shape) for n, x in zip(self.data_names, self.data)]

    @property
    def provide_label(self):
        return [(n, x.shape) for n, x in zip(self.label_names, self.label)]

class SimpleIter(mx.io.DataIter):
    """DataIter used in Calculate Statistics (in progress).

    Parameters
    ----------
    pad_zeros : bool
        Default `False`. Control the behavior of padding when we run
        out of the whole dataset. When true, we will pad with all-zeros.
        When false, will pad with a random sentence in the dataset.
        Usually, for training we would like to use `False`, but
        for testing use `True` so that the evaluation metric can
        choose to ignore the padding by detecting the zero-labels.
    """
    def __init__(self, train_sets, batch_size,
            init_states, delay=5, feat_dim=40, label_dim=1955, label_mean_sets=None,
            data_name='data', label_name='softmax_label', has_label=True, load_label_mean=True):

        self.train_sets=train_sets
        self.label_mean_sets=label_mean_sets
        self.train_sets.initialize_read()

        self.data_name = data_name
        if has_label:
            self.label_name = label_name

        features = []
        labels = []
        utt_lens = []
        utt_ids = []
        buckets = []
        self.has_label = has_label

        if label_mean_sets is not None:
            self.label_mean_sets.initialize_read()
            (feats, tgts, utt_id) = self.label_mean_sets.load_next_seq()

            self.label_mean = feats/np.sum(feats)
            for i,v in enumerate(feats):
                if v <= 1.0:
                    self.label_mean[i] = 1

        sys.stderr.write("Loading data...\n")
        buckets_map = {}
        n = 0
        while True:
            (feats, tgts, utt_id) = self.train_sets.load_next_seq()
            if utt_id is None:
                break
            if tgts is None and self.has_label:
                continue
            if feats.shape[0] == 0:
                continue
            features.append(feats)
            utt_lens.append(feats.shape[0])
            utt_ids.append(utt_id)
            if self.has_label:
                labels.append(tgts+1)
            if feats.shape[0] not in buckets:
                buckets_map[feats.shape[0]] = feats.shape[0]

        for k, v in buckets_map.iteritems():
            buckets.append(k)

        buckets.sort()
        i_max_bucket = len(buckets)-1
        max_bucket = buckets[i_max_bucket]
        self.buckets = buckets
        self.data = [[] for k in buckets]
        self.utt_id = [[] for k in buckets]
        self.utt_lens = [[] for k in buckets]
        self.feat_dim = feat_dim
        self.default_bucket_key = max(buckets)

        for i, feats in enumerate(features):
            if has_label:
                tgts = labels[i]
            utt_len = utt_lens[i]
            utt_id = utt_ids[i]

            for i, bkt in enumerate(buckets):
                if bkt >= utt_len:
                    i_bucket = i
                    break

            if self.has_label:
                self.data[i_bucket].append((feats,tgts))
            else:
                self.data[i_bucket].append(feats)
            self.utt_id[i_bucket].append(utt_id)
            self.utt_lens[i_bucket].append(utt_len)

        # Get the size of each bucket, so that we could sample
        # uniformly from the bucket
        bucket_sizes = [len(x) for x in self.data]

        self.batch_size = batch_size
        # convert data into ndarrays for better speed during training

        data = [np.zeros((len(x), buckets[i], self.feat_dim))
                if len(x) % self.batch_size == 0
                else np.zeros(((len(x)/self.batch_size + 1) * self.batch_size, buckets[i], self.feat_dim))
                for i, x in enumerate(self.data)]

        label = [np.zeros((len(x), buckets[i]))
                 if len(x) % self.batch_size == 0
                 else np.zeros(((len(x)/self.batch_size + 1) * self.batch_size, buckets[i]))
                 for i, x in enumerate(self.data)]

        utt_id = [[] for k in buckets]
        for i, x in enumerate(data):
            utt_id[i] = ["GAP_UTT"] * len(x)
        utt_lens = [[] for k in buckets]
        for i, x in enumerate(data):
            utt_lens[i] = [0] * len(x)

        for i_bucket in range(len(self.buckets)):
            for j in range(len(self.data[i_bucket])):
                sentence = self.data[i_bucket][j]
                if self.has_label:
                    sentence[1][delay:] = sentence[1][:-delay]
                    sentence[1][:delay] = sentence[1][0] # broadcast assignment
                    data[i_bucket][j, :len(sentence[0])] = sentence[0]
                    label[i_bucket][j, :len(sentence[1])] = sentence[1]
                else:
                    data[i_bucket][j, :len(sentence)] = sentence
                    # borrow this place to pass in sentence length. TODO: use a less hacky way.
                    label[i_bucket][j, :len(sentence)] += len(sentence)

                utt_id[i_bucket][j] = self.utt_id[i_bucket][j]
                utt_lens[i_bucket][j] = self.utt_lens[i_bucket][j]

        self.data = data
        self.label = label
        self.utt_id = utt_id
        self.utt_lens = utt_lens

        # Get the size of each bucket, so that we could sample
        # uniformly from the bucket
        bucket_sizes = [len(x) for x in self.data]

        sys.stderr.write("Summary of dataset ==================\n")
        for bkt, sz in zip(buckets, bucket_sizes):
            sys.stderr.write("bucket of len %3d : %d samples\n" % (bkt, sz))

        bucket_size_tot = float(sum(bucket_sizes))

        self.bucket_sizes = bucket_sizes
        self.make_data_iter_plan()

        self.init_states = init_states
        self.init_state_arrays = [mx.nd.zeros(x[1]) for x in init_states]

        self.provide_data = [(data_name, (batch_size, self.default_bucket_key, self.feat_dim))] + init_states
        self.provide_label = []
        if has_label:
            self.provide_label = [(label_name, (self.batch_size, self.default_bucket_key))]

    def make_data_iter_plan(self):
        "make a random data iteration plan"
        # truncate each bucket into multiple of batch-size
        bucket_n_batches = []
        for i in range(len(self.data)):
            bucket_n_batches.append(len(self.data[i]) / self.batch_size)
            self.data[i] = self.data[i][:int(bucket_n_batches[i]*self.batch_size),:]
            self.label[i] = self.label[i][:int(bucket_n_batches[i]*self.batch_size)]

        bucket_plan = np.hstack([np.zeros(n, int)+i for i, n in enumerate(bucket_n_batches)])
        np.random.shuffle(bucket_plan)

        bucket_idx_all = [np.random.permutation(len(x)) for x in self.data]

        self.bucket_plan = bucket_plan
        self.bucket_idx_all = bucket_idx_all
        self.bucket_curr_idx = [0 for x in self.data]

        self.data_buffer = []
        self.label_buffer = []
        for i_bucket in range(len(self.data)):
            data = mx.nd.zeros((self.batch_size, self.buckets[i_bucket], self.feat_dim))
            label = mx.nd.zeros((self.batch_size, self.buckets[i_bucket]))
            self.data_buffer.append(data)
            self.label_buffer.append(label)

    def __iter__(self):
        init_state_names = [x[0] for x in self.init_states]
        data_names = [self.data_name] + init_state_names
        label_names = []
        if self.has_label:
            label_names = [self.label_name]

        for i_bucket in self.bucket_plan:
            data = self.data_buffer[i_bucket]
            label = self.label_buffer[i_bucket]

            i_idx = self.bucket_curr_idx[i_bucket]
            idx = self.bucket_idx_all[i_bucket][i_idx:i_idx+self.batch_size]
            self.bucket_curr_idx[i_bucket] += self.batch_size
            data[:] = self.data[i_bucket][idx]
            label[:] = self.label[i_bucket][idx]
            data_all = [data] + self.init_state_arrays
            label_all = [label]
            utt_id = np.array(self.utt_id[i_bucket])[idx]
            utt_len = np.array(self.utt_lens[i_bucket])[idx]
            effective_sample_count = mx.nd.sum(label)
            data_batch = SimpleBatch(data_names, data_all, label_names, label_all,
                                     self.buckets[i_bucket], utt_id, utt_len,
                                     effective_sample_count=effective_sample_count)
            yield data_batch

    def reset(self):
        self.bucket_curr_idx = [0 for x in self.data]

class TruncatedSentenceIter(mx.io.DataIter):
    """DataIter used in Truncated-BPTT.

    Each sentence is split into chunks of fixed lengths. The states are
    forwarded during forward, but the backward is only computed within
    chunks. This mechanism does not require bucketing, and it sometimes
    avoid gradient exploding problems in very long sequences.

    Parameters
    ----------
    pad_zeros : bool
        Default `False`. Control the behavior of padding when we run
        out of the whole dataset. When true, we will pad with all-zeros.
        When false, will pad with a random sentence in the dataset.
        Usually, for training we would like to use `False`, but
        for testing use `True` so that the evaluation metric can
        choose to ignore the padding by detecting the zero-labels.
    """
    def __init__(self, train_sets, batch_size, init_states, truncate_len=20, delay=5,
                 feat_dim=40, data_name='data', label_name='softmax_label',
                 has_label=True, do_shuffling=True, pad_zeros=False):

        self.train_sets = train_sets
        self.train_sets.initialize_read()

        self.data_name = data_name
        self.label_name = label_name

        self.feat_dim = feat_dim
        self.has_label = has_label
        self.batch_size = batch_size
        self.truncate_len = truncate_len
        self.delay = delay

        self.do_shuffling = do_shuffling
        self.pad_zeros = pad_zeros

        self.data = [mx.nd.zeros((batch_size, truncate_len, feat_dim))]
        self.label = None
        if has_label:
            self.label = [mx.nd.zeros((batch_size, truncate_len))]

        self.init_state_names = [x[0] for x in init_states]
        self.init_state_arrays = [mx.nd.zeros(x[1]) for x in init_states]

        self.provide_data = [(data_name, self.data[0].shape)] + init_states
        self.provide_label = []
        if has_label:
            self.provide_label = [(label_name, self.label[0].shape)]

        self._load_data()
        self._make_data_plan()

    def _load_data(self):
        sys.stderr.write('Loading data into memory...\n')
        self.features = []
        self.labels = []
        self.utt_ids = []

        seq_len_tot = 0.0
        while True:
            (feats, tgs, utt_id) = self.train_sets.load_next_seq()
            if utt_id is None:
                break
            if tgs is None and self.has_label:
                continue
            if feats.shape[0] == 0:
                continue

            if self.has_label and self.delay > 0:
                # delay the labels
                tgs[self.delay:] = tgs[:-self.delay]
                tgs[:self.delay] = tgs[0]  # boradcast assign
            self.features.append(feats)
            if self.has_label:
                self.labels.append(tgs+1)
            self.utt_ids.append(utt_id)
            seq_len_tot += feats.shape[0]

        sys.stderr.write('    %d utterances loaded...\n' % len(self.utt_ids))
        sys.stderr.write('    avg-sequence-len = %.0f\n' % (seq_len_tot/len(self.utt_ids)))

    def _make_data_plan(self):
        if self.do_shuffling:
            # TODO: should we group utterances of similar length together?
            self._data_plan = np.random.permutation(len(self.features))
        else:
            # we might not want to do shuffling for testing for example
            self._data_plan = np.arange(len(self.features))

    def __iter__(self):
        assert len(self._data_plan) >= self.batch_size, \
            "Total number of sentences smaller than batch size, consider using smaller batch size"
        utt_idx = self._data_plan[:self.batch_size]
        utt_inside_idx = [0] * self.batch_size

        next_utt_idx = self.batch_size
        is_pad = [False] * self.batch_size
        pad = 0

        np_data_buffer = np.zeros((self.batch_size, self.truncate_len, self.feat_dim))
        np_label_buffer = np.zeros((self.batch_size, self.truncate_len))
        utt_id_buffer = [None] * self.batch_size

        data_names = [self.data_name] + self.init_state_names
        label_names = [self.label_name]

        # reset states
        for state in self.init_state_arrays:
            state[:] = 0.1

        while True:
            effective_sample_count = self.batch_size * self.truncate_len
            for i, idx in enumerate(utt_idx):
                fea_utt = self.features[idx]
                if utt_inside_idx[i] >= fea_utt.shape[0]:
                    # we have consumed this sentence

                    # reset the states
                    for state in self.init_state_arrays:
                        state[i:i+1] = 0.1
                    # load new sentence
                    if is_pad[i]:
                        # I am already a padded sentence, just rewind to the
                        # beginning of the sentece
                        utt_inside_idx[i] = 0
                    elif next_utt_idx >= len(self.features):
                        # we consumed the whole dataset, simply repeat this sentence
                        # and set pad
                        pad += 1
                        is_pad[i] = True
                        utt_inside_idx[i] = 0
                    else:
                        # move to the next sentence
                        utt_idx[i] = self._data_plan[next_utt_idx]
                        idx = utt_idx[i]
                        fea_utt = self.features[idx]
                        utt_inside_idx[i] = 0
                        next_utt_idx += 1

                if is_pad[i] and self.pad_zeros:
                    np_data_buffer[i] = 0
                    np_label_buffer[i] = 0
                    effective_sample_count -= self.truncate_len
                else:
                    idx_take = slice(utt_inside_idx[i],
                                     min(utt_inside_idx[i]+self.truncate_len,
                                         fea_utt.shape[0]))
                    n_take = idx_take.stop - idx_take.start
                    np_data_buffer[i][:n_take] = fea_utt[idx_take]
                    np_label_buffer[i][:n_take] = self.labels[idx][idx_take]
                    if n_take < self.truncate_len:
                        np_data_buffer[i][n_take:] = 0
                        np_label_buffer[i][n_take:] = 0
                        effective_sample_count -= self.truncate_len - n_take

                    utt_inside_idx[i] += n_take

                utt_id_buffer[i] = self.utt_ids[idx]

            if pad == self.batch_size:
                # finished all the senteces
                break

            self.data[0][:] = np_data_buffer
            self.label[0][:] = np_label_buffer
            data_batch = SimpleBatch(data_names, self.data + self.init_state_arrays,
                                     label_names, self.label, bucket_key=None,
                                     utt_id=utt_id_buffer,
                                     effective_sample_count=effective_sample_count)

            # Instead of using the 'pad' property, we use an array 'is_pad'. Because
            # our padded sentence could be in the middle of a batch. A sample is pad
            # if we are running out of the data set and they are just some previously
            # seen data to be filled for a whole batch. In prediction, those data
            # should be ignored
            data_batch.is_pad = is_pad

            yield data_batch

    def reset(self):
        self._make_data_plan()

class BucketSentenceIter(mx.io.DataIter):
    def __init__(self, train_sets, buckets, batch_size,
                 init_states, delay=5, feat_dim=40,
                 data_name='data', label_name='softmax_label', has_label=True):

        self.train_sets = train_sets
        self.train_sets.initialize_read()

        self.data_name = data_name
        self.label_name = label_name

        buckets.sort()
        i_max_bucket = len(buckets)-1
        max_bucket = buckets[i_max_bucket]

        if has_label != True:
            buckets = [ i for i in range(1,max_bucket) ]
            i_max_bucket = len(buckets)-1
            max_bucket = buckets[i_max_bucket]

        self.buckets = buckets
        self.data = [[] for k in buckets]
        self.utt_id = [[] for k in buckets]
        self.feat_dim = feat_dim
        self.default_bucket_key = max(buckets)
        self.has_label = has_label

        sys.stderr.write("Loading data...\n")
        T_OVERLAP = buckets[0]/2
        n = 0
        while True:
            (feats, tgts, utt_id) = self.train_sets.load_next_seq()
            if utt_id is None:
                break
            if tgts is None and self.has_label:
                continue
            if feats.shape[0] == 0:
                continue

            # we split sentence into overlapping segments if it is
            # longer than the largest bucket
            t_start = 0
            t_end = feats.shape[0]
            while t_start < t_end:
                if t_end - t_start > max_bucket:
                    t_take = max_bucket
                    i_bucket = i_max_bucket
                else:
                    for i, bkt in enumerate(buckets):
                        if bkt >= t_end-t_start:
                            t_take = t_end-t_start
                            i_bucket = i
                            break

                n += 1
                if self.has_label:
                    self.data[i_bucket].append((feats[t_start:t_start+t_take],
                                                tgts[t_start:t_start+t_take]+1))
                else:
                    self.data[i_bucket].append(feats[t_start:t_start+t_take])

                self.utt_id[i_bucket].append(utt_id)
                t_start += t_take
                if t_start >= t_end:
                    # this sentence is consumed
                    break
                t_start -= T_OVERLAP

        # Get the size of each bucket, so that we could sample
        # uniformly from the bucket
        bucket_sizes = [len(x) for x in self.data]

        self.batch_size = batch_size
        # convert data into ndarrays for better speed during training

        data = [np.zeros((len(x), buckets[i], self.feat_dim))
                if len(x) % self.batch_size == 0
                else np.zeros(((len(x)/self.batch_size + 1) * self.batch_size, buckets[i],
                               self.feat_dim))
                for i, x in enumerate(self.data)]

        label = [np.zeros((len(x), buckets[i]))
                 if len(x) % self.batch_size == 0
                 else np.zeros(((len(x)/self.batch_size + 1) * self.batch_size, buckets[i]))
                 for i, x in enumerate(self.data)]

        utt_id = [[] for k in buckets]
        for i, x in enumerate(data):
            utt_id[i] = ["GAP_UTT"] * len(x)

        for i_bucket in range(len(self.buckets)):
            for j in range(len(self.data[i_bucket])):
                sentence = self.data[i_bucket][j]
                if self.has_label:
                    sentence[1][delay:] = sentence[1][:-delay]
                    sentence[1][:delay] = sentence[1][0]  # broadcast assignment
                    data[i_bucket][j, :len(sentence[0])] = sentence[0]
                    label[i_bucket][j, :len(sentence[1])] = sentence[1]
                else:
                    data[i_bucket][j, :len(sentence)] = sentence
                    # borrow this place to pass in sentence length. TODO: use a less hacky way.
                    label[i_bucket][j, :len(sentence)] += len(sentence)

                utt_id[i_bucket][j] = self.utt_id[i_bucket][j]

        self.data = data
        self.label = label
        self.utt_id = utt_id

        # Get the size of each bucket, so that we could sample
        # uniformly from the bucket
        bucket_sizes = [len(x) for x in self.data]

        sys.stderr.write("Summary of dataset ==================\n")
        for bkt, sz in zip(buckets, bucket_sizes):
            sys.stderr.write("bucket of len %3d : %d samples\n" % (bkt, sz))

        self.bucket_sizes = bucket_sizes
        self.make_data_iter_plan()

        self.init_states = init_states
        self.init_state_arrays = [mx.nd.zeros(x[1]) for x in init_states]

        self.provide_data = [(data_name, (batch_size, self.default_bucket_key, self.feat_dim))] + \
            init_states
        self.provide_label = [(label_name, (self.batch_size, self.default_bucket_key))]

    def make_data_iter_plan(self):
        "make a random data iteration plan"
        # truncate each bucket into multiple of batch-size
        bucket_n_batches = []
        for i in range(len(self.data)):
            bucket_n_batches.append(len(self.data[i]) / self.batch_size)
            self.data[i] = self.data[i][:int(bucket_n_batches[i]*self.batch_size), :]
            self.label[i] = self.label[i][:int(bucket_n_batches[i]*self.batch_size)]

        bucket_plan = np.hstack([np.zeros(n, int)+i for i, n in enumerate(bucket_n_batches)])
        np.random.shuffle(bucket_plan)

        bucket_idx_all = [np.random.permutation(len(x)) for x in self.data]

        self.bucket_plan = bucket_plan
        self.bucket_idx_all = bucket_idx_all
        self.bucket_curr_idx = [0 for x in self.data]

        self.data_buffer = []
        self.label_buffer = []
        for i_bucket in range(len(self.data)):
            data = mx.nd.zeros((self.batch_size, self.buckets[i_bucket], self.feat_dim))
            label = mx.nd.zeros((self.batch_size, self.buckets[i_bucket]))
            self.data_buffer.append(data)
            self.label_buffer.append(label)

    def __iter__(self):
        init_state_names = [x[0] for x in self.init_states]
        data_names = [self.data_name] + init_state_names
        label_names = [self.label_name]

        for i_bucket in self.bucket_plan:
            data = self.data_buffer[i_bucket]
            label = self.label_buffer[i_bucket]

            i_idx = self.bucket_curr_idx[i_bucket]
            idx = self.bucket_idx_all[i_bucket][i_idx:i_idx+self.batch_size]
            self.bucket_curr_idx[i_bucket] += self.batch_size
            data[:] = self.data[i_bucket][idx]
            label[:] = self.label[i_bucket][idx]
            data_all = [data] + self.init_state_arrays
            label_all = [label]
            utt_id = np.array(self.utt_id[i_bucket])[idx]
            effective_sample_count = mx.nd.sum(label)
            data_batch = SimpleBatch(data_names, data_all, label_names, label_all,
                                     self.buckets[i_bucket], utt_id,
                                     effective_sample_count=effective_sample_count)
            yield data_batch

    def reset(self):
        self.bucket_curr_idx = [0 for x in self.data]


from xml.etree import ElementTree

import os
import numpy as np

import mujoco_py as mjp

class MujocoConfig():
    """ A wrapper on the Mujoco simulator to generate all the kinematics and
    dynamics calculations necessary for controllers.
    """

    def __init__(self, xml_file):
        """ Loads the Mujoco model from the specified xml file

        Parameters
        ----------
        xml_file: string
            the name of the arm model to load
            The string passed in is parsed such that everything up to the first
            underscore is used for the arm directory, and the full string is
            used to load the xml within that folder.

            EX: 'myArm' and 'myArm_with_gripper' will both look in the
            'myArm' directory, however they will load myArm.xml and
            myArm_with_gripper.xml, respectively
        """

        current_dir = os.path.dirname(__file__)
        self.xml_file = os.path.join(
            current_dir, xml_file.split('_')[0], '%s.xml' % xml_file)
        self.model = mjp.load_model_from_path(self.xml_file)

        # get access to some of our custom arm parameters from the xml definition
        tree = ElementTree.parse(self.xml_file)
        root = tree.getroot()
        for custom in root.findall('custom/numeric'):
            name = custom.get('name')
            if name == 'N_JOINTS':
                self.N_JOINTS = int(custom.get('data'))
            elif name == 'START_ANGLES':
                START_ANGLES = custom.get('data').split(' ')
                self.START_ANGLES = np.array(
                    [float(angle) for angle in START_ANGLES])

    def _connect(self, sim):
        """ Called by the interface once the Mujoco simulation is created,
        this connects the config to the simulator so it can access the
        kinematics and dynamics information calculated by Mujoco.

        Parameters
        ----------
        sim: MjSim
            The Mujoco Simulator object created by the Mujoco Interface class
        """
        # get access to the Mujoco simulation
        self.sim = sim

        self.JOINT_NAMES = self.sim.model.joint_names
        self.N_JOINTS = len(self.JOINT_NAMES)

        # NOTE: We assume that every body defined in the xml _after_
        # 'base_link' is part of the robot
        self.base_link_index = self.sim.model.body_name2id('base_link')
        self.N_LINKS = len(self.sim.model.body_parentid) - self.base_link_index

        # a place to store data returned from Mujoco
        self._g = np.zeros(self.N_JOINTS)
        self._J3N = np.zeros(3 * self.N_JOINTS)
        self._J6N = np.zeros((6, self.N_JOINTS))
        self._MNN_vector = np.zeros(self.N_JOINTS**2)
        self._MNN = np.zeros((self.N_JOINTS, self.N_JOINTS))
        self._R9 = np.zeros(9)
        self._R = np.zeros((3, 3))
        self._x = np.ones(4)

    def g(self, q=None):
        """ Returns qfrc_bias variable, which stores the effects of Coriolis,
        centrifugal, and gravitational forces

        Parameters
        ----------
        q: float numpy.array, optional (Default: None)
            The joint angles of the robot. If None the current state is
            retrieved from the Mujoco simulator
        """
        return -self.sim.data.qfrc_bias

    def dJ(self, name, q=None, dq=None, x=None):
        """ Returns the derivative of the Jacobian wrt to time

        Parameters
        ----------
        name: string
            The name of the Mujoco body to retrieve the Jacobian for
        q: float numpy.array, optional (Default: None)
            The joint angles of the robot. If None the current state is
            retrieved from the Mujoco simulator
        dq: float numpy.array, optional (Default: None)
            The joint velocities of the robot. If None the current state is
            retrieved from the Mujoco simulator
        x: float numpy.array, optional (Default: None)
        """
        # TODO if ever required
        # Note from Emo in Mujoco forums:
        # 'You would have to use a finate-difference approximation in the
        # general case, check differences.cpp'
        raise NotImplementedError

    def J(self, name, q=None, x=None):
        """ Returns the Jacobian for the specified Mujoco body

        Parameters
        ----------
        name: string
            The name of the Mujoco body to retrieve the Jacobian for
        q: float numpy.array, optional (Default: None)
            The joint angles of the robot. If None the current state is
            retrieved from the Mujoco simulator
        x: float numpy.array, optional (Default: None)
        """
        if x is not None and not np.allclose(x, 0):
            raise Exception('x offset currently not supported: ', x)

        # get the position Jacobian hstacked (1 x N_JOINTS*3)
        self._J3N[:] = self.sim.data.get_body_jacp(name)
        self._J6N[:3] = self._J3N.reshape((3, self.N_JOINTS))
        # get the rotation Jacobian hstacked (1 x N_JOINTS*3)
        self._J3N[:] = self.sim.data.get_body_jacr(name)
        self._J6N[3:] = self._J3N.reshape((3, self.N_JOINTS))

        return self._J6N

    def M(self, q=None):
        """ Returns the inertia matrix in task space

        Parameters
        ----------
        q: float numpy.array, optional (Default: None)
            The joint angles of the robot. If None the current state is
            retrieved from the Mujoco simulator
        """
        # stored in mjData.qM, stored in custom sparse format,
        # convert qM to a dense matrix with mj_fullM
        mjp.cymj._mj_fullM(self.model, self._MNN_vector, self.sim.data.qM)
        # TODO: there's a shape like _MNN function or some such, right?
        self._MNN = self._MNN_vector.reshape((self.N_JOINTS, self.N_JOINTS))
        return self._MNN

    def R(self, name, q=None):
        """ Returns the rotation matrix of the specified body
        """
        mjp.cymj._mju_quat2Mat(
            self._R9, self.sim.data.get_body_xquat(name))
        self._R = self._R9.reshape((3, 3))
        return self._R

    def quaternion(self, name, q=None):
        """ Returns the quaternion of the specified body
        """
        return self.sim.data.get_body_xquat(name)

    def C(self, q=None):
        """ NOTE: The Coriolis and centrifugal effects (and gravity) are
        already accounted for by Mujoco in the qfrc_bias variable. There's
        no easy way to separate these, so all are returned by the g function.
        To prevent accounting for these effects twice, this function will
        return an error instead of qfrc_bias again.
        """
        raise Exception('Coriolis and centrifugal effects already accounted '
                        + 'for in the term return by the gravity function.')

    def T(self, name, q=None, x=None):
        """ Get the transform matrix of the specified body.

        Parameters
        ----------
        name: string
            The name of the Mujoco body to retrieve the Jacobian for
        q: float numpy.array, optional (Default: None)
            The joint angles of the robot. If None the current state is
            retrieved from the Mujoco simulator
        x: float numpy.array, optional (Default: None)
        """
        # TODO if ever required
        raise NotImplementedError

    def Tx(self, name, q=None, x=None):
        """ Returns the Cartesian coordinates of the specified Mujoco body

        Parameters
        ----------
        name: string
            The name of the Mujoco body to retrieve the Jacobian for
        q: float numpy.array, optional (Default: None)
            The joint angles of the robot. If None the current state is
            retrieved from the Mujoco simulator
        x: float numpy.array, optional (Default: None)
        """
        if x is not None and not np.allclose(x, 0):
            raise Exception('x offset currently not supported: ', x)
        return self.sim.data.get_body_xpos(name)

    def T_inv(self, name, q=None, x=None):
        """  Get the inverse transform matrix of the specified body.

        Parameters
        ----------
        name: string
            The name of the Mujoco body to retrieve the Jacobian for
        q: float numpy.array, optional (Default: None)
            The joint angles of the robot. If None the current state is
            retrieved from the Mujoco simulator
        x: float numpy.array, optional (Default: None)
        """
        # TODO if ever required
        raise NotImplementedError


""" running training and evaluation code for document-level sentiment analysis project

    Created by Chenyang Lyu
"""

import torch
import torch.nn as nn
from torch.nn import DataParallel
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, TensorDataset
from torch.nn import CrossEntropyLoss, MSELoss

import argparse
import sklearn.metrics as metric
import glob
import logging
import os
import random
import numpy as np
import json

from os import listdir
from os.path import isfile, join

from tqdm import tqdm, trange

from transformers import (
    WEIGHTS_NAME,
    AdamW,
    BertModel,
    BertTokenizer,
    BertConfig,
    BertForSequenceClassification,
    RobertaConfig,
    RobertaModel,
    RobertaTokenizer,
    RobertaForSequenceClassification,
    get_linear_schedule_with_warmup,
)

from modeling import SecondPretrainedBert, IncrementalContextBert, FocalLoss, IncrementalContextRoberta
from utils import eval_to_file, load_data
from pargs import Arguments

try:
    from torch.utils.tensorboard import SummaryWriter
except ImportError:
    from tensorboardX import SummaryWriter

logger = logging.getLogger(__name__)
label_list = ["1", "2", "3", "4", "5"]
label_list_imdb = [str(i) for i in range(1, 11)]

MODEL_CLASSES = {
    "bert": (BertConfig, BertTokenizer, BertForSequenceClassification),
    "roberta": (RobertaConfig, RobertaTokenizer, RobertaForSequenceClassification),
}

args = Arguments('document-level-sa')

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

def second_train(args, train_dataset, model):
    train_sampler = RandomSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)

    for batch in tqdm(train_dataloader, desc="Second Train Iterating"):
        model.eval()
        batch = tuple(t.to(args.device) for t in batch)

        with torch.no_grad():
            inputs = {"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}
            if args.model_type != "roberta":
                inputs["token_type_ids"] = (
                    batch[2] if args.model_type in ["bert", "xlnet"] else None
                )

            user, product = batch[4].view(-1, 1), batch[5].view(-1, 1)
            user_product = torch.cat([user, product], dim=1)
            _ = model(inputs, user_product)

def train(args, train_dataset, model, tokenizer, freeze=True, dev_set=None, eval_set=None, global_s=0):
    """ Training the model """
    tb_writer = SummaryWriter()

    num_labels = len(args.label_list)

    train_sampler = RandomSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)

    t_total = len(train_dataloader) * args.num_train_epochs

    if args.is_second_training and freeze:
        model.embedding.weight.require_grad = False

    # Prepare optimizer for training
    if args.is_incremental and args.model_type == "roberta":
        large_lr = ["embedding.weight"]
        optimizer_group_parameters = [
            {
                "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in large_lr)],
                "lr": args.learning_rate,
             },
            {
                "params": [p for n, p in model.named_parameters() if any(nd in n for nd in large_lr)],
                "lr": args.amply * args.learning_rate,
            }
        ]
    else:
        no_decay = ["bias", "LayerNorm.weight"]
        optimizer_group_parameters = [
            {
                "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
                "weight_decay": args.weight_decay
             },
            {
                "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
                "weight_decay": 0.0
            }
        ]

    optimizer = AdamW(optimizer_group_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=int(args.warmup_steps*t_total), num_training_steps=t_total)

    # Train
    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
    logger.info("  Instantaneous batch size per GPU = %d", args.train_batch_size)
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)

    global_step = global_s
    epochs_trained = 0
    steps_trained_in_current_epoch = 0

    # Check if continuing training from a checkpoint
    # if os.path.exists(args.model_name_or_path):
    #     global_step = int(args.model_name_or_path.split("-")[-1].split("/")[0])
    #     epochs_trained = global_step // len(train_dataloader // args.gradient_accumulation_steps)
    #     steps_trained_in_current_epoch = global_step % (len(train_dataloader) // args.gradient_accumulation_steps)
    #
    #     logger.info("  Continuing training from epoch %d", epochs_trained)
    #     logger.info("  Continuing training from global step %d", global_step)
    #     logger.info("  Will skip the first %d steps in the first epoch", steps_trained_in_current_epoch)

    tr_loss, logging_loss = 0.0, 0.0
    model.zero_grad()
    train_iterator = trange(int(args.num_train_epochs), desc="Epoch")
    set_seed(args)  # Added here for reproductibility (even between python 2 and 3)
    up_indices, new_embeddings = None, None

    loss_fct = CrossEntropyLoss()
    if args.is_focal_loss:
        loss_fct = FocalLoss(gamma=3, alpha=args.alpha)

    for _ in train_iterator:
        epoch_iterator = tqdm(train_dataloader, desc="Iteration")
        for step, batch in enumerate(epoch_iterator):

            # Skip past any already trained steps
            if steps_trained_in_current_epoch > 0:
                steps_trained_in_current_epoch -= 1

            model.train()
            batch = tuple(t.to(args.device) for t in batch)

            user, product = batch[4].view(-1, 1), batch[5].view(-1, 1)
            user_product = torch.cat([user, product], dim=1)

            inputs = {'input_ids':      batch[0],
                      'attention_mask': batch[1]}
            if args.model_type != 'distilbert':
                inputs['token_type_ids'] = batch[2] if args.model_type in ['bert', 'xlnet'] else None  # XLM, DistilBERT and RoBERTa don't use segment_ids

            if args.is_incremental:
                logits, up_indices, new_embeddings = model(inputs, user_product, up_indices, new_embeddings)
            else:
                inputs['labels'] = batch[3]
                outputs = model(**inputs)
                logits = outputs[1]  # model outputs are always tuple in transformers (see doc)

            loss = loss_fct(logits.view(-1, num_labels), batch[3].view(-1))

            if args.n_gpu > 1:
                loss = loss.mean() # mean() to average on multi-gpu parallel training
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            loss.backward()

            tr_loss += loss.item()
            if (step + 1) % args.gradient_accumulation_steps == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

                optimizer.step()
                scheduler.step()  # Update learning rate schedule
                model.zero_grad()
                global_step += 1

                if args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    logs = {}
                    if global_step % (5*args.logging_steps == 0) and dev_set is not None:
                        results = evaluate(args, model, dev_set, tokenizer)
                        for key, value in results.items():
                            eval_key = 'eval_{}'.format(key)
                            logs[eval_key] = value

                    loss_scalar = (tr_loss - logging_loss) / args.logging_steps
                    learning_rate_scalar = scheduler.get_lr()[0]
                    logs['learning_rate'] = learning_rate_scalar
                    logs['loss'] = loss_scalar
                    logging_loss = tr_loss

                    for key, value in logs.items():
                        tb_writer.add_scalar(key, value, global_step)
                    print(json.dumps({**logs, **{'step': global_step}}))

                # if args.save_steps > 0 and global_step % args.save_steps == 0:
                #     # Save model checkpoint
                #     output_dir = os.path.join(args.output_dir, 'checkpoint-{}'.format(global_step))
                #     if not os.path.exists(output_dir):
                #         os.makedirs(output_dir)
                #     model_to_save = model.module if hasattr(model, 'module') else model  # Take care of distributed/parallel training
                #     model_to_save.save_pretrained(output_dir)
                #     torch.save(args, os.path.join(output_dir, 'training_args.bin'))
                #     logger.info("Saving model checkpoint to %s", output_dir)

    tb_writer.close()
    global_step = 1 if global_step ==0 else global_step

    return global_step, tr_loss / global_step

def evaluate(args, model, eval_dataset, tokenizer, prefix=""):
    # Loop to handle MNLI double evaluation (matched, mis-matched)
    eval_task_names = (args.task_name,)
    eval_outputs_dirs = (args.output_dir,)

    results = {}
    for eval_task, eval_output_dir in zip(eval_task_names, eval_outputs_dirs):

        if not os.path.exists(eval_output_dir):
            os.makedirs(eval_output_dir)

        args.eval_batch_size = args.train_batch_size
        # Note that DistributedSampler samples randomly
        eval_sampler = SequentialSampler(eval_dataset)
        eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)

        # Eval!
        logger.info("***** Running evaluation {} *****".format(prefix))
        logger.info("  Num examples = %d", len(eval_dataset))
        logger.info("  Batch size = %d", args.eval_batch_size)
        eval_loss = 0.0
        nb_eval_steps = 0
        preds = None
        out_label_ids = None
        up_indices, new_embeddings = None, None
        for batch in tqdm(eval_dataloader, desc="Evaluating"):
            model.eval()
            batch = tuple(t.to(args.device) for t in batch)

            user, product = batch[4].view(-1, 1), batch[5].view(-1, 1)
            user_product = torch.cat([user, product], dim=1)

            with torch.no_grad():
                inputs = {'input_ids': batch[0], 'attention_mask': batch[1]}
                if args.model_type != 'roberta':
                    inputs['token_type_ids'] = batch[2] if args.model_type in ['bert',
                                                                               'xlnet'] else None  # XLM, DistilBERT and RoBERTa don't use segment_ids

                if args.is_incremental:
                    logits, up_indices, new_embeddings = model(inputs, user_product)
                else:
                    inputs['labels'] = batch[3]
                    outputs = model(**inputs)
                    logits = outputs[1]  # model outputs are always tuple in transformers (see doc)
            nb_eval_steps += 1
            if preds is None:
                preds = logits.detach().cpu().numpy()
                out_label_ids = batch[3].detach().cpu().numpy()
            else:
                preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
                out_label_ids = np.append(out_label_ids, batch[3].detach().cpu().numpy(), axis=0)

        if args.output_mode == "classification":
            preds = np.argmax(preds, axis=1)
        elif args.output_mode == "regression":
            preds = np.squeeze(preds)
        result = {'acc': (preds == out_label_ids).mean()}
        results.update(result)

        eval_to_file(args.eval_out_file, preds, out_label_ids)
        print("accuracy: ", metric.accuracy_score(out_label_ids, preds))
        print("precision: ", metric.precision_score(out_label_ids, preds, average='macro'))
        print("recall: ", metric.recall_score(out_label_ids, preds, average='macro'))
        print("F1: ", metric.f1_score(out_label_ids, preds, average='macro'))
        print("Mean Squared Error: ", metric.mean_squared_error(out_label_ids, preds))
        print("Root Mean Squared Error: ", metric.mean_squared_error(out_label_ids, preds, squared=False))

    return results

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

    parser.add_argument("--task_name", type=str, default="yelp-2013",
                        help="the name of the training task (the dataset name)")
    parser.add_argument("--model_size", type=str, default="base",
                        help="the size of pre-trained model")
    parser.add_argument("--model_type", type=str, default="bert",
                        help="the type of pre-trained model")
    parser.add_argument("--epochs", type=int, default=1,
                        help="the numebr of training epochs")
    parser.add_argument("--incremental", action="store_true",
                        help="use incremental mode")
    parser.add_argument("--second_train", action="store_true",
                        help="use second train mode")
    parser.add_argument("--do_train", action="store_true",
                        help="whether to train the model or not")
    parser.add_argument("--do_eval", action="store_true",
                        help="whether to evaluate the model or not")
    parser.add_argument("--weight_decay", type=float, default=0.0,
                        help="the weight decay rate")
    parser.add_argument("--learning_rate", type=float, default=3e-5,
                        help="the learning rate used to train the model")
    parser.add_argument("--warmup_steps", type=float, default=0.05,
                        help="the warm_up step rate")
    parser.add_argument("--max_seq_length", type=int, default=512,
                        help="the maximum sequence length used to load dataset")
    parser.add_argument("--seed", type=int, default=1,
                        help="the random seed used in model initialization and dataloader")
    parser.add_argument("--batch_size", type=int, default=8,
                        help="the batch size used in training and evaluation")
    parser.add_argument("--device", type=int, default=0,
                        help="the device id used for training and evaluation")
    parser.add_argument("--do_shrink", action="store_true",
                        help="whether to shrink the embedding size in order to reduce the amount of parameters")
    parser.add_argument("--inner_size", type=int, default=128,
                        help="the inner embedding size when do_shrink is true")
    parser.add_argument("--attention_heads", type=int, default=8,
                        help="the attention heads used in multi head attention function")
    parser.add_argument("--is_focal_loss", action="store_true",
                        help="whether to use focal loss function or not")

    arguments, _ = parser.parse_known_args()

    args.task_name = arguments.task_name
    args.model_size = arguments.model_size
    args.num_train_epochs = arguments.epochs
    args.is_incremental = arguments.incremental
    args.is_second_training = arguments.second_train
    args.do_train = arguments.do_train
    args.do_eval = arguments.do_eval
    args.weight_decay = arguments.weight_decay
    args.learning_rate = arguments.learning_rate
    args.warmup_steps = arguments.warmup_steps
    args.max_seq_length = arguments.max_seq_length
    args.seed = arguments.seed
    args.train_batch_size = arguments.batch_size
    args.device = torch.device("cuda:" + str(arguments.device))
    args.model_type = arguments.model_type
    args.is_focal_loss = arguments.is_focal_loss
    args.do_shrink = arguments.do_shrink

    if args.task_name == 'imdb':
        args.label_list = label_list_imdb
    else:
        args.label_list = label_list

    if args.is_incremental and args.is_second_training:
        raise ValueError("Incremental and second training modes can't be applied at the same time!")

    if args.model_type == 'bert':
        args.model_name_or_path = "bert-" + args.model_size + "-uncased"
    if args.model_type == 'roberta':
        args.model_name_or_path = "roberta-" + args.model_size

    if args.is_incremental:
        model_type = "incremental"
    else:
        model_type = 'vanilla'

    if args.is_second_training:
        model_type = 'second_train'

    if args.do_shrink:
        model_type += "_shrink"

    output_dir = "trained_model/" + args.model_name_or_path + "_" + args.task_name + '_' + model_type + "_epochs_" + str(args.num_train_epochs) + "_lr_" + \
        str(args.learning_rate) + "_weight-decay_" + str(args.weight_decay) + "_warmup_" + str(args.warmup_steps) + "_mql_" + str(args.max_seq_length) + '_shrink_' + str(arguments.inner_size) + \
        "_seed_" + str(args.seed) + "/"
    eval_dir = "eval_results/" + args.model_name_or_path + "_" + args.task_name + '_' + model_type + "_epochs_" + str(args.num_train_epochs) + "_lr_" + \
        str(args.learning_rate) + "_weight-decay_" + str(args.weight_decay) + "_warmup_" + str(args.warmup_steps) + "_mql_" + str(args.max_seq_length) + '_shrink_' + str(arguments.inner_size) + \
        "_seed_" + str(args.seed) + ".log"

    args.output_dir = output_dir
    args.eval_out_file = eval_dir

    config_class, tokenizer_class, model_class = MODEL_CLASSES[args.model_type]

    tokenizer = tokenizer_class.from_pretrained(
        args.model_name_or_path,
        do_lower_case=True,
    )

    num_labels = len(args.label_list)

    data_dirs = ["data/document-level-sa-dataset/" + args.task_name + "-seg-20-20.train.ss",
                 "data/document-level-sa-dataset/" + args.task_name + "-seg-20-20.dev.ss",
                 "data/document-level-sa-dataset/" + args.task_name + "-seg-20-20.test.ss",
                 ]
    train_dataset, dev_dataset, test_dataset, up_vocab = load_data(args, data_dirs, tokenizer)

    # long_dev_dataset, long_test_dataset = load_dev_and_eval_data(args, data_dirs[1:], tokenizer, up_vocab)

    set_seed(args)
    args.model_type = args.model_type.lower()
    config = config_class.from_pretrained(
        args.model_name_or_path,
        num_labels=num_labels,
    )
    config.output_attention = True
    config.output_hidden_states = True
    config.inner_size = arguments.inner_size
    config.do_shrink = arguments.do_shrink
    config.attention_heads = arguments.attention_heads
    if args.is_incremental:
        model_class = IncrementalContextBert if args.model_type == "bert" else IncrementalContextRoberta
        model = model_class.from_pretrained(
            args.model_name_or_path,
            num_embeddings=up_vocab.word_count,
            up_vocab=up_vocab,
            config=config,
        )
    else:
        model = model_class.from_pretrained(args.model_name_or_path, config=config)

    model.to(args.device)

    logger.info("Training/evaluation parameters %s", args)

    # Training
    if args.do_train:
        global_step, tr_loss = train(args, train_dataset, model, tokenizer, freeze=False, dev_set=dev_dataset, eval_set=test_dataset)
        logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)

    if args.is_second_training:
        st_model = SecondPretrainingBert.from_pretrained(
            args.model_name_or_path,
            num_embeddings=up_vocab.word_count,
            config=config
        )
        st_model.bert = model
        st_model.to(args.device)
        second_train(args, train_dataset, st_model)

        embedding_matrix = st_model.embedding_matrix

        new_model = IncrementalContextBert.from_pretrained(
            args.model_name_or_path,
            num_embeddings=up_vocab.word_count,
            up_vocab=up_vocab,
            config=config,
        )

        new_model.embedding = embedding_matrix

        args.is_incremental = True
        model = new_model
        model.to(args.device)
        global_step, tr_loss = train(args, train_dataset, model, tokenizer)
        logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)

    # Saving best-practices: if you use defaults names for the model, you can reload it using from_pretrained()
    if args.do_train:
        # Create output directory if needed
        if not os.path.exists(args.output_dir):
            os.makedirs(args.output_dir)

        logger.info("Saving model checkpoint to %s", args.output_dir)
        # Save a trained model, configuration and tokenizer using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        model_to_save = model.module if hasattr(model,
                                                'module') else model  # Take care of distributed/parallel training
        model_to_save.save_pretrained(args.output_dir)
        tokenizer.save_pretrained(args.output_dir)

        # Good practice: save your training arguments together with the trained model
        torch.save(args, os.path.join(args.output_dir, 'training_args.bin'))

        # Load a trained model and vocabulary that you have fine-tuned
        # model = model_class.from_pretrained(args.output_dir)
        # tokenizer = tokenizer_class.from_pretrained(args.output_dir)
        # model.to(args.device)

    # Evaluation
    results = {}
    if args.do_eval:
        tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case)
        checkpoints = [args.output_dir]
        if args.eval_all_checkpoints:
            checkpoints = list(
                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME, recursive=True)))
            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce logging
        logger.info("Evaluate the following checkpoints: %s", checkpoints)
        for checkpoint in checkpoints:
            global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
            prefix = checkpoint.split('/')[-1] if checkpoint.find('checkpoint') != -1 else ""

            if args.is_incremental or args.is_second_training:
                model = model_class.from_pretrained(
                    checkpoint,
                    num_embeddings=up_vocab.word_count,
                    up_vocab=up_vocab,
                    config=config,
                )
            else:
                model = model_class.from_pretrained(checkpoint)
            model.to(args.device)
            result = evaluate(args, model, dev_dataset, tokenizer, prefix=prefix)
            result = evaluate(args, model, test_dataset, tokenizer, prefix=prefix)
            result = dict((k + '_{}'.format(global_step), v) for k, v in result.items())
            results.update(result)

    return results

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
"""
Script to train Vgg_16 with one single attention layer inserted at different positions.

Specify class_name from [ave, canidae, cloth, felidae, kitchen, land_trans] when run the script.

Make sure you have the class_name+'_models' directory before training.

"""

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import matplotlib
matplotlib.use("Agg")

import keras
from keras.applications.vgg16 import VGG16, preprocess_input, decode_predictions
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
from keras.layers import Dense, Flatten, Reshape, Concatenate, Lambda
from keras import optimizers
from keras.models import Sequential, Model
from keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard
from keras import backend as K
from keras.layers import Layer

import numpy as np
import pandas as pd
import h5py

import random
import math

from custom_layer_constraints import CustomConstraint, SinglyConnected

#Forked from Ken's
import keras_custom_objects as KO
from custom_generator import create_good_generator

import argparse
parser = argparse.ArgumentParser(
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--class_name', type=str, default='kitchen')
args = parser.parse_args()
class_name = args.class_name
print (class_name)
assert (class_name in ['ave', 'canidae', 'cloth', 'felidae', 'kitchen', 'land_trans']), "only support class from ['ave', 'canidae', 'cloth', 'felidae', 'kitchen', 'land_trans']"

#Magic numbers
num_epochs = 100
bs = 64
img_rows = 224
img_cols = 224
layer_indice = ['21', '03', '06', '10', '14', '18', '20']  #positions to insert attention layer,use dtr here for writing file name properly(layer03 rather than layer3)
patience = 2

# Custom generator from Ken's codes
imagenet_train = '/mnt/fast-data16/datasets/ILSVRC/2012/clsloc/train/'
ImageGen = ImageDataGenerator(fill_mode='nearest',
                              horizontal_flip=True,
                              rescale=None,
                              preprocessing_function=preprocess_input,
                              data_format="channels_last",
                              validation_split=0.1
                              )

class_csv_path = 'groupings-csv/' + class_name + '_Imagenet.csv'
df_classes = pd.read_csv(class_csv_path, usecols=['wnid'])
classes = sorted([i for i in df_classes['wnid']])

good_train_generator, steps = create_good_generator(ImageGen,
                                                    imagenet_train,
                                                    batch_size=bs,
                                                    target_size = (img_rows, img_cols),
                                                    class_mode='sparse',
                                                    subset= 'training',
                                                    classes=classes)

good_validation_generator, steps_val = create_good_generator(ImageGen,
                                                    imagenet_train,
                                                    batch_size=bs,
                                                    target_size = (img_rows, img_cols),
                                                    class_mode='sparse',
                                                    subset= 'validation',
                                                    classes=classes)

#Build the model

def get_ATT_model(str_layer_index, class_name):
  '''
  The function is to insert an attention layer
  between frozen_model.layers[layer_index] and frozen_model.layers[layer_index+1].
  frozen_model is set to be VGG16.

  Input str_layer_index is a string. This is to have proper model_file_name, e.g. layer03 rather than layer3
  '''
  layer_index = int(str_layer_index)
  frozen_model = VGG16(weights = 'imagenet', include_top=True, input_shape = (img_rows, img_cols, 3))
  for layer in frozen_model.layers:
    layer.trainable = False

  last = frozen_model.layers[layer_index].output

  x = SinglyConnected(kernel_constraint= CustomConstraint())(last)

  for layer in frozen_model.layers[layer_index+1:]:
    x = layer(x)

  model = KO.CustomModel(frozen_model.input, x)

  model.summary()

  model_file_name = class_name+ '_models/layer' + str_layer_index + '-' + class_name + '-model.h5'

  return model, model_file_name

for str_layer_index in layer_indice:

  model, file_name = get_ATT_model(str_layer_index, class_name)
  model.compile(loss='sparse_categorical_crossentropy',
                optimizer='adam',
                metrics=['accuracy'])

  print('This model is going to be saved in: ' + file_name)

  #Callbacks
  custom_earlystopping = KO.RelativeEarlyStopping(monitor='val_loss',
                                                  min_perc_delta=0.001,  # perc means percentage
                                                  patience=patience,
                                                  verbose=1,
                                                  mode='min'
                                                  )

  log_path = './'+class_name+'_logs'
  tensorboard = TensorBoard(log_dir = log_path, update_freq = bs*100)
  callbacks = [tensorboard,
               custom_earlystopping,
               ModelCheckpoint(filepath=file_name, monitor='val_loss', save_best_only=True)]

  print('Now training model with layer index: '+ str_layer_index)

  model.fit_generator_custom(
          good_train_generator,
          steps_per_epoch=steps,
          epochs=num_epochs,
          verbose = 1,
          callbacks = callbacks,
          validation_data=good_validation_generator,
          validation_steps=steps_val,
          #max_queue_size=20,
          #workers = 6,
          use_multiprocessing=False
          )


#===================================
#        Chengwen Liu              #
#      liuchw2010@gmail.com        #
#   University of Texas at Austin  #
#===================================

import os
import sys
import time
import numpy as np

def prepare(xyz):
  cmdstr = f"python {converterdir}/lToTXYZ.py -i {xyz} -t xyz"
  os.system(cmdstr)
  with open("tinker.key", 'w') as f:
    f.write("parameters ttt.prm\n")
    f.write("bondterm only\n")
    f.write("openmp-threads 8\n")
  return

def genprm():
  elements, types = np.loadtxt(txyz, usecols=(1,5), dtype='str', unpack=True, skiprows=1)
  headers = open(f"{amoebaheaddir}/amoebabio18_header.prm").readlines()
  tmp = []
  with open("ttt.prm",'w') as f:
    for line in headers:
      f.write(line)
    for e,t in zip(elements, types):
      if t not in tmp:
        f.write(f'atom    {t:>6s} {t:>6s} {e:>6s} "{e}" {ptable[e][0]:>5d} {ptable[e][1]:10.3f} {ptable[e][2]:10d}\n')
        tmp.append(t)
  cmdstr = f"{tinkerdir}/analyze.x {txyz} -k tinker.key > ana_err.dat"
  os.system(cmdstr)
  lines = open("ana_err.dat").readlines()
  bondlist = []
  for line in lines:
    if "Bond   " in line:
      d = line.split()
      a1,a2 = int(d[-2]), int(d[-1])
      if a1 > a2:
        a1, a2 = a2, a1
      if [a1,a2] not in bondlist:
        bondlist.append([a1,a2])
  with open("ttt.prm", 'a') as f:
    f.write("# ASSIGN\n")
    for b in bondlist:
      f.write(f"bond {b[0]:>6d} {b[1]:>6d} 0.0 0.0 \n")
  cmdstr = f"python {assignprmdir}/lAssignAMOEBAplusPRM.py -xyz {txyz} -key ttt.prm -potent bonded -konly no"
  os.system(cmdstr)
  newkeyname = f"{txyz.replace('.txyz', '_new.key')}"
  bonds = []
  lines = open("ttt.prm").readlines()
  for l in lines:
    if "bond " not in l:
      bonds.append(l)
  lines = open(newkeyname)
  for l in lines:
    if "bond " in l:
      bonds.append(l)
  with open("ttt.prm", 'w') as f:
    for b in bonds:
      f.write(b)
  return

def optimize():
  cmdstr = f"{tinkerdir}/minimize.x {txyz} -k tinker.key 0.5"
  os.system(cmdstr)
  return

def main():
  global txyz
  global ptable
  global tinkerdir
  global converterdir
  global amoebaheaddir
  global assignprmdir

  tinkerdir="/home/liuchw/Softwares/tinkers/Tinker-latest/source-C8"
  converterdir="/home/liuchw/bin"
  amoebaheaddir="/home/liuchw/shared/forcefield"
  assignprmdir="/home/liuchw/bin/AssignPRM"

  ptable = {'H':[1,   1.008, 1], 'C':[6,  12.011, 4], 'N':[7, 14.007, 3],  'O':[8,  15.999, 2], \
            'S':[16, 32.066, 2], 'P':[15, 30.974, 4], 'F':[9, 18.998, 1], 'Cl':[17, 35.453, 1]}
  xyz = sys.argv[1]
  txyz = xyz.replace('xyz', 'txyz')
  if os.path.isfile(txyz):
    print(f"Warning: using the existing txyz file {txyz} to assign parameters")
  else:
    prepare(xyz)
  genprm()
  optimize()

if __name__ == "__main__":
  main()


# -*- coding: utf-8 -*-
"""SAMR-DNLP.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/10tebUo2wA5Mj3NEKSCZ-iQjMQiGZ31Q-

#**Sentiment Analysis based on Movie Review**
"""

import numpy as np
import pandas as pd
from sklearn.neural_network import MLPClassifier
import matplotlib.pyplot as plt
import zipfile

zip = zipfile.ZipFile('train.tsv.zip')
zip.extractall()

zip = zipfile.ZipFile('test.tsv.zip')
zip.extractall()

df_train=pd.read_csv('train.tsv',sep='\t')
df_train=df_train.dropna(how='any')
df_train=df_train.drop(columns=['PhraseId','SentenceId'])
df_train

df_train['Sentiment'].unique()

import re
import nltk
nltk.download('stopwords')
from nltk.corpus import stopwords
from nltk.stem.porter import PorterStemmer
corpus = []
corpus1=[]
corpus2=[]
for i in range(0, 156060):
  review = re.sub('[^a-zA-Z]', ' ', df_train['Phrase'][i])
  review = review.lower()
  review = review.split()
  ps = PorterStemmer()
  all_stopwords = stopwords.words('english')
  all_stopwords.remove('not')
  review = [ps.stem(word) for word in review if not word in set(all_stopwords)]
  review = ' '.join(review)
  corpus.append(review)

from sklearn.feature_extraction.text import CountVectorizer
cv = CountVectorizer(max_features = 1500)
X = cv.fit_transform(corpus).toarray()
y=df_train.iloc[:,-1].values

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.20, random_state = 0)

model=MLPClassifier(alpha=0.01, batch_size=256, epsilon=1e-08, learning_rate='adaptive', max_iter=100)
model.fit(X_train,y_train)

res = model.predict(X_test)
res=np.round(res)
np.set_printoptions(precision=2)
print(res)

from sklearn.metrics import accuracy_score
print("Accuracy Score for the algorithm=>{}%".format(round(accuracy_score(y_test,res)*100),2))

"""#**Testing the Model**"""

df_test=pd.read_csv('test.tsv',sep='\t')
df_test=df_test.dropna(how='any')
df_test=df_test.drop(columns=['PhraseId','SentenceId'])
df_test

import re
import nltk
nltk.download('stopwords')
from nltk.corpus import stopwords
from nltk.stem.porter import PorterStemmer
corpus = []
corpus1=[]
corpus2=[]
for i in range(0, 66292):
  review = re.sub('[^a-zA-Z]', ' ', df_train['Phrase'][i])
  review = review.lower()
  review = review.split()
  ps = PorterStemmer()
  all_stopwords = stopwords.words('english')
  all_stopwords.remove('not')
  review = [ps.stem(word) for word in review if not word in set(all_stopwords)]
  review = ' '.join(review)
  corpus.append(review)

from sklearn.feature_extraction.text import CountVectorizer
cv = CountVectorizer(max_features = 1500)
X = cv.fit_transform(corpus).toarray()

res = model.predict(X_test)
res=np.round(res)
np.set_printoptions(precision=2)
print(res)

for i in range(len(df_test)):
  for j in range(len(res)):
    if i==j:
      print("The catgory for this phrase{0} is{1}" .format(df_test['Phrase'][i],res[i]))


#----import lib----
import cv2
import numpy as np
import matplotlib.pyplot as plt
import time

class YoloObjectDetection(object):
    """docstring for YoloObject detection"""
    def __init__(self,model,weights,lable_file):
        self.model=model
        self.weights=weights
        self.net = cv2.dnn.readNetFromDarknet(self.model,self.weights)
        self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
        self.lable_file=lable_file

    """read labels from file"""
    def read_label(self):
        # label name
        lines = []
        with open(lable_file) as f:
            for line in f:
                lines.append(line.rstrip())
        return lines

    """feed the input towards network"""
    def feed_network(self,img):

        blob = cv2.dnn.blobFromImage(img, 1/255.0, (416, 416),(0,0,0),swapRB=True, crop=False)
        # determine the output layer
        output_layers=self.net.getUnconnectedOutLayers()
        layers = self.net.getLayerNames()
        layers = [layers[i - 1] for i in output_layers]
        self.net.setInput(blob)
        outputs = self.net.forward(layers)

        return outputs

    """detection and draw bounding boxes"""
    def detection(self,img,outputs):
        width=img.shape[1]
        height=img.shape[0]
        boxes=[]
        box_confidences=[]
        class_ids=[]
        label_name=self.read_label()
        colors = np.random.randint(0, 255, size=(len(label_name), 3), dtype='uint8')
        for output in outputs:
            for out in output:
                scores=out[5:]
                class_id=np.argmax(scores)
                confidence=scores[class_id]

                if confidence>0.5:
                    width_ = int(out[2] * width)
                    height_ = int(out[4] * height)
                    x = int((out[0]*width)-width_/2)
                    y = int((out[1]*height)-height_/2)
                    boxes.append([x,y,width_,height_])
                    box_confidences.append(float(confidence))
                    class_ids.append(class_id)

        indices = cv2.dnn.NMSBoxes(boxes, box_confidences, 0.5, 0.4)

        if len(indices) > 0:

            for i in indices.flatten():
                (x, y) = (boxes[i][0], boxes[i][1])
                (w, h) = (boxes[i][2], boxes[i][3])
                color = [int(c) for c in colors[class_ids[i]]]

                cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
                label = "{}: {:.4f}".format(label_name[class_ids[i]], box_confidences[i])
                cv2.putText(img, label, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

        return img

if __name__ == "__main__":

    # Loading YOLO v4 model confi and weights
    model="dnn/yolov4/yolov4.cfg"
    weights="dnn/yolov4/yolov4.weights"
    lable_file='dnn/label_yolo.txt'

    opencv_dnn= YoloObjectDetection(model,weights,lable_file)
    cap=cv2.VideoCapture(0)
    frame_rate = 10
    prev = 0

    while (cap.isOpened):
        time_elapsed = time.time() - prev
        ret, img=cap.read()
        if time_elapsed > 1.0/frame_rate:
            prev = time.time()
            boxes=opencv_dnn.feed_network(img)
            img=opencv_dnn.detection(img,boxes)
            cv2.imshow("Image", img)
            if cv2.waitKey(27) & 0xFF == ord('q'):
                 break

    cap.release()
    cv2.destroyAllWindows()


from __future__ import absolute_import
import numpy as np
import sklearn.preprocessing
import ctypes
import faiss
import os
import time
import gc
import resource
import threading
import json

from multiprocessing.pool import ThreadPool

from benchmark.algorithms.base import BaseANN
from benchmark.datasets import DATASETS, download_accelerated

def unwind_index_ivf(index):
    if isinstance(index, faiss.IndexPreTransform):
        assert index.chain.size() == 1
        vt = index.chain.at(0)
        index_ivf, vt2 = unwind_index_ivf(faiss.downcast_index(index.index))
        assert vt2 is None
        return index_ivf, vt
    if hasattr(faiss, "IndexRefine") and isinstance(index, faiss.IndexRefine):
        return unwind_index_ivf(faiss.downcast_index(index.base_index))
    if isinstance(index, faiss.IndexIVF):
        return index, None
    else:
        return None, None

def rate_limited_iter(l):
    'a thread pre-processes the next element'
    pool = ThreadPool(1)
    res = None

    def next_or_None():
        try:
            return next(l)
        except StopIteration:
            return None

    while True:
        res_next = pool.apply_async(next_or_None)
        if res is not None:
            res = res.get()
            if res is None:
                return
            yield res
        res = res_next

def build_index(buildthreads, by_residual, maxtrain, clustering_niter,
     indexkey, indexfile, add_bs, add_splits, ds, train_on_gpu=True, quantizer_on_gpu_add=True):

    nq, d = ds.nq, ds.d
    nb, d = ds.nq, ds.d

    if buildthreads == -1:
        print("Build-time number of threads:", faiss.omp_get_max_threads())
    else:
        print("Set build-time number of threads:", buildthreads)
        faiss.omp_set_num_threads(buildthreads)

    metric_type = (
            faiss.METRIC_L2 if ds.distance() == "euclidean" else
            faiss.METRIC_INNER_PRODUCT if ds.distance() in ("ip", "angular") else
            1/0
    )
    index = faiss.index_factory(d, indexkey, metric_type)

    index_ivf, vec_transform = unwind_index_ivf(index)
    if vec_transform is None:
        vec_transform = lambda x: x
    else:
        vec_transform = faiss.downcast_VectorTransform(vec_transform)

    if by_residual != -1:
        by_residual = by_residual == 1
        index_ivf.by_residual   # check if field exists
        index_ivf.by_residual = by_residual

    if index_ivf:
        print("Update add-time parameters")
        # adjust default parameters used at add time for quantizers
        # because otherwise the assignment is inaccurate
        quantizer = faiss.downcast_index(index_ivf.quantizer)
        if isinstance(quantizer, faiss.IndexRefine):
            print("   update quantizer k_factor=", quantizer.k_factor, end=" -> ")
            quantizer.k_factor = 32 if index_ivf.nlist < 1e6 else 64
            print(quantizer.k_factor)
            base_index = faiss.downcast_index(quantizer.base_index)
            if isinstance(base_index, faiss.IndexIVF):
                print("   update quantizer nprobe=", base_index.nprobe, end=" -> ")
                base_index.nprobe = (
                    16 if base_index.nlist < 1e5 else
                    32 if base_index.nlist < 4e6 else
                    64)
                print(base_index.nprobe)

    index.verbose = True
    if index_ivf:
        index_ivf.verbose = True
        index_ivf.quantizer.verbose = True
        index_ivf.cp.verbose = True

    if maxtrain == 0:
        if 'IMI' in indexkey:
            maxtrain = int(256 * 2 ** (np.log2(index_ivf.nlist) / 2))
        elif index_ivf:
            maxtrain = 50 * index_ivf.nlist
        else:
            # just guess...
            maxtrain = 256 * 100
        maxtrain = max(maxtrain, 256 * 100)
        print("setting maxtrain to %d" % maxtrain)

    # train on dataset
    print(f"getting first {maxtrain} dataset vectors for training")

    xt2 = next(ds.get_dataset_iterator(bs=maxtrain))

    print("train, size", xt2.shape)
    assert np.all(np.isfinite(xt2))

    t0 = time.time()

    if (isinstance(vec_transform, faiss.OPQMatrix) and
        isinstance(index_ivf, faiss.IndexIVFPQFastScan)):
        print("  Forcing OPQ training PQ to PQ4")
        ref_pq = index_ivf.pq
        training_pq = faiss.ProductQuantizer(
            ref_pq.d, ref_pq.M, ref_pq.nbits
        )
        vec_transform.pq
        vec_transform.pq = training_pq

    if clustering_niter >= 0:
        print(("setting nb of clustering iterations to %d" %
                clustering_niter))
        index_ivf.cp.niter = clustering_niter

    train_index = None
    if train_on_gpu:
        print("add a training index on GPU")
        train_index = faiss.index_cpu_to_all_gpus(
                faiss.IndexFlatL2(index_ivf.d))
        index_ivf.clustering_index = train_index

    index.train(xt2)
    print("  Total train time %.3f s" % (time.time() - t0))

    if train_index is not None:
        del train_index
        index_ivf.clustering_index = None
        gc.collect()

    print("adding")

    t0 = time.time()

    if not quantizer_on_gpu_add:
        i0 = 0
        for xblock in ds.get_dataset_iterator(bs=add_bs):
            i1 = i0 + len(xblock)
            print("  adding %d:%d / %d [%.3f s, RSS %d kiB] " % (
                i0, i1, ds.nb, time.time() - t0,
                faiss.get_mem_usage_kb()))
            index.add(xblock)
            i0 = i1
    elif True:
        quantizer_gpu = faiss.index_cpu_to_all_gpus(index_ivf.quantizer)

        nsplit = add_splits

        def produce_batches(sno):
            for xblock in ds.get_dataset_iterator(bs=add_bs, split=(nsplit, sno)):
                _, assign = quantizer_gpu.search(xblock, 1)
                yield xblock, assign.ravel()

        i0 = 0
        for sno in range(nsplit):
            print(f"============== SPLIT {sno}/{nsplit}")

            stage2 = rate_limited_iter(produce_batches(sno))
            for xblock, assign in stage2:
                i1 = i0 + len(xblock)
                print("  adding %d:%d / %d [%.3f s, RSS %d kiB] " % (
                    i0, i1, ds.nb, time.time() - t0,
                    faiss.get_mem_usage_kb()))
                index.add_core(
                    len(xblock),
                    faiss.swig_ptr(xblock),
                    None,
                    faiss.swig_ptr(assign)
                )
                i0 = i1
        del quantizer_gpu
        gc.collect()

    print("  add in %.3f s" % (time.time() - t0))
    if indexfile:
        print("storing", indexfile)
        faiss.write_index(index, indexfile)

    return index

class IndexQuantizerOnGPU:
    """ run query quantization on GPU """

    def __init__(self, index, search_bs):
        self.search_bs = search_bs
        index_ivf, vec_transform = unwind_index_ivf(index)
        self.index_ivf = index_ivf
        self.vec_transform = vec_transform.apply
        self.quantizer_gpu = faiss.index_cpu_to_all_gpus(self.index_ivf.quantizer)

    def produce_batches(self, x, bs):
        n = len(x)
        nprobe = self.index_ivf.nprobe
        ivf_stats = faiss.cvar.indexIVF_stats
        for i0 in range(0, n, bs):
            xblock = x[i0:i0 + bs]
            t0 = time.time()
            D, I = self.quantizer_gpu.search(xblock, nprobe)
            ivf_stats.quantization_time += 1000 * (time.time() - t0)
            yield i0, xblock, D, I

    def search(self, x, k):
        bs = self.search_bs
        if self.vec_transform:
            x = self.vec_transform(x)
        nprobe = self.index_ivf.nprobe
        n, d = x.shape
        assert self.index_ivf.d == d
        D = np.empty((n, k), dtype=np.float32)
        I = np.empty((n, k), dtype=np.int64)

        sp = faiss.swig_ptr
        stage2 = rate_limited_iter(self.produce_batches(x, bs))
        t0 = time.time()
        for i0, xblock, Dc, Ic in stage2:
            ni = len(xblock)
            self.index_ivf.search_preassigned(
                ni, faiss.swig_ptr(xblock),
                k, sp(Ic), sp(Dc),
                sp(D[i0:]), sp(I[i0:]),
                False
            )

        return D, I

    def range_search(self, x, radius):
        bs = self.search_bs
        if self.vec_transform:
            x = self.vec_transform(x)
        nprobe = self.index_ivf.nprobe
        n, d = x.shape
        assert self.index_ivf.d == d

        sp = faiss.swig_ptr
        rsp = faiss.rev_swig_ptr
        stage2 = rate_limited_iter(self.produce_batches(x, bs))
        t0 = time.time()
        all_res = []
        nres = 0
        for i0, xblock, Dc, Ic in stage2:
            ni = len(xblock)
            res = faiss.RangeSearchResult(ni)

            self.index_ivf.range_search_preassigned(
                ni, faiss.swig_ptr(xblock),
                radius, sp(Ic), sp(Dc),
                res
            )
            all_res.append((ni, res))
            lims = rsp(res.lims, ni + 1)
            nres += lims[-1]
        nres = int(nres)
        lims = np.zeros(n + 1, int)
        I = np.empty(nres, int)
        D = np.empty(nres, 'float32')

        n0 = 0
        for ni, res in all_res:
            lims_i = rsp(res.lims, ni + 1)
            nd = int(lims_i[-1])
            Di = rsp(res.distances, nd)
            Ii = rsp(res.labels, nd)
            i0 = int(lims[n0])
            lims[n0: n0 + ni + 1] = lims_i + i0
            I[i0:i0 + nd] = Ii
            D[i0:i0 + nd] = Di
            n0 += ni

        return lims, D, I

class FaissT3(BaseANN):
    def __init__(self, metric, index_params):
        self._index_params = index_params
        self._metric = metric

    def index_name(self, name):
        return f"data/{name}.{self._index_params['indexkey']}.faissindex"

    def fit(self, dataset):
        index_params = self._index_params

        ds = DATASETS[dataset]()
        d = ds.d

        # get build parameters
        buildthreads = index_params.get("buildthreads", -1)
        by_residual = index_params.get("by_residual", -1)
        maxtrain = index_params.get("maxtrain", 0)
        clustering_niter = index_params.get("clustering_niter", -1)
        indexkey = index_params.get("indexkey", "IVF1048576,SQ8")
        add_bs = index_params.get("add_bs", 100000)
        add_splits = index_params.get("add_splits", 1)
        indexfile = self.index_name(dataset)

        # determine how we use the GPU
        #search_type = ds.search_type()
        #if search_type == "knn":
        #    train_on_gpu = True
        #    quantizer_on_gpu_add = True
        #else: #range
        #    train_on_gpu = False
        #    quantizer_on_gpu_add = False

        index = build_index(buildthreads, by_residual, maxtrain, clustering_niter, indexkey,
                        indexfile, add_bs, add_splits, ds)

        index_ivf, vec_transform = unwind_index_ivf(index)
        if vec_transform is None:
            vec_transform = lambda x: x
        if index_ivf is not None:
            print("imbalance_factor=", index_ivf.invlists.imbalance_factor())

        no_precomputed_tables  = index_params.get("no_precomputed_tables", True)
        if no_precomputed_tables:
            if isinstance(index_ivf, faiss.IndexIVFPQ):
                print("disabling precomputed table")
                index_ivf.use_precomputed_table = -1
                index_ivf.precomputed_table.clear()

        precomputed_table_size = 0
        if hasattr(index_ivf, 'precomputed_table'):
            precomputed_table_size = index_ivf.precomputed_table.size() * 4
        print("precomputed tables size:", precomputed_table_size)

        searchthreads  = index_params.get("searchthreads", -1)
        if searchthreads == -1:
            print("Search threads:", faiss.omp_get_max_threads())
        else:
            print("Setting nb of threads to", searchthreads)
            faiss.omp_set_num_threads(searchthreads)

        parallel_mode  = index_params.get("parallel_mode", 3)
        if parallel_mode != -1:
            print("setting IVF parallel mode to", parallel_mode)
            index_ivf.parallel_mode
            index_ivf.parallel_mode = parallel_mode

        # prep for the searches

        self.ps = faiss.ParameterSpace()
        self.ps.initialize(index)

        search_bs  = index_params.get("search_bs", 8192)
        index_wrap = IndexQuantizerOnGPU(index, search_bs)

        self.cpuindex = index
        self.index = index_wrap

    def load_index(self, dataset):

        index_params = self._index_params

        if not os.path.exists(self.index_name(dataset)):
            if 'url' not in self._index_params:
                return False

            print('Downloading index in background. This can take a while.')
            download_accelerated(self._index_params['url'], self.index_name(dataset), quiet=True)

        print("Loading index",self.index_name(dataset))

        index = faiss.read_index(self.index_name(dataset))

        index_ivf, vec_transform = unwind_index_ivf(index)
        if vec_transform is None:
            vec_transform = lambda x: x
        if index_ivf is not None:
            print("imbalance_factor=", index_ivf.invlists.imbalance_factor())

        no_precomputed_tables  = index_params.get("no_precomputed_tables", True)
        if no_precomputed_tables:
            if isinstance(index_ivf, faiss.IndexIVFPQ):
                print("disabling precomputed table")
                index_ivf.use_precomputed_table = -1
                index_ivf.precomputed_table.clear()

        precomputed_table_size = 0
        if hasattr(index_ivf, 'precomputed_table'):
            precomputed_table_size = index_ivf.precomputed_table.size() * 4
        print("precomputed tables size:", precomputed_table_size)

        # prep for the searches

        searchthreads  = index_params.get("searchthreads", -1)
        if searchthreads == -1:
            print("Search threads:", faiss.omp_get_max_threads())
        else:
            print("Setting nb of threads to", searchthreads)
            faiss.omp_set_num_threads(searchthreads)

        parallel_mode  = index_params.get("parallel_mode", 3)
        if parallel_mode != -1:
            print("setting IVF parallel mode to", parallel_mode)
            index_ivf.parallel_mode
            index_ivf.parallel_mode = parallel_mode

        self.ps = faiss.ParameterSpace()
        self.ps.initialize(index)

        search_bs  = index_params.get("search_bs", 8092)
        index_wrap = IndexQuantizerOnGPU(index, search_bs)

        self.cpuindex = index
        self.index = index_wrap

        return True

    def set_query_arguments(self, query_args):
        faiss.cvar.indexIVF_stats.reset()
        self.ps.set_index_parameters(self.cpuindex, query_args)
        self.qas = query_args

    # shall we return something interesting here?
    def get_additional(self):
        return {"dist_comps": faiss.cvar.indexIVF_stats.ndis}

    def __str__(self):
        return f'FaissIVFPQ({self.qas})'

    def query(self, X, n):
        self.res = self.index.search(X, n)

    def range_query(self, X, radius):
        self.res = self.index.range_search(X, radius)

    def get_results(self):
        D, I = self.res
        return I

    def get_range_results(self):
        return self.res