# Library Imports
import numpy as np
import tensorflow as tf

class ReplayBuffer:
    """Defines the Buffer dataset from which the agent learns"""
    def __init__(self, max_size, input_shape, dim_actions):
        self.mem_size = max_size
        self.mem_cntr = 0
        self.state_memory = np.zeros((self.mem_size, input_shape), dtype=np.float32)
        self.new_state_memory = np.zeros((self.mem_size, input_shape), dtype=np.float32)
        self.action_memory = np.zeros((self.mem_size, dim_actions), dtype=np.float32)
        self.reward_memory = np.zeros(self.mem_size, dtype=np.float32)
        self.terminal_memory = np.zeros(self.mem_size, dtype=np.bool)

    def store_transition(self, state, action, reward, new_state, done):
        index = self.mem_cntr % self.mem_size
        self.state_memory[index] = state
        self.new_state_memory[index] = new_state
        self.action_memory[index] = action
        self.reward_memory[index] = reward
        self.terminal_memory[index] = done
        self.mem_cntr += 1

    def sample_buffer(self, batch_size):
        max_mem = min(self.mem_cntr, self.mem_size)
        batch = np.random.choice(max_mem, batch_size, replace=False)
        states = self.state_memory[batch]
        _states = self.new_state_memory[batch]
        actions = self.action_memory[batch]
        rewards = self.reward_memory[batch]
        dones = self.terminal_memory[batch]
        return states, actions, rewards, _states, dones

class Critic(tf.keras.Model):
    """Defines a Critic Deep Learning Network"""
    def __init__(self, density=512, name='critic'):
        super(Critic, self).__init__()

        self.model_name = name
        self.checkpoint = self.model_name + '.h5'

        self.H1 = tf.keras.layers.Dense(density, activation='relu')
        self.H2 = tf.keras.layers.Dense(density, activation='relu')
        self.drop = tf.keras.layers.Dropout(0.2)
        self.H3 = tf.keras.layers.Dense(density, activation='relu')
        self.H4 = tf.keras.layers.Dense(density, activation='relu')
        self.Q = tf.keras.layers.Dense(1, activation=None)

    @tf.function()
    def call(self, state, action):
        action = self.H1(tf.concat([state, action], axis=1))
        action = self.H2(action)
        action = self.drop(action)
        action = self.H3(action)
        action = self.H4(action)
        Q = self.Q(action)
        return Q

class Actor(tf.keras.Model):
    """Defines a Actor Deep Learning Network"""
    def __init__(self, n_actions, density=512, name='actor'):
        super(Actor, self).__init__()

        self.model_name = name
        self.checkpoint = self.model_name + '.h5'

        self.H1 = tf.keras.layers.Dense(density, activation='relu')
        self.H2 = tf.keras.layers.Dense(density, activation='relu')
        self.drop = tf.keras.layers.Dropout(0.2)
        self.H3 = tf.keras.layers.Dense(density, activation='relu')
        self.H4 = tf.keras.layers.Dense(density, activation='relu')
        self.mu = tf.keras.layers.Dense(n_actions, activation='tanh')

    @tf.function()
    def call(self, state):
        state = self.H1(state)
        state = self.H2(state)
        state = self.drop(state)
        state = self.H3(state)
        state = self.H4(state)
        mu = self.mu(state)
        return mu

class Agent:
    """Defines a RL Agent based on Actor-Critc method"""
    def __init__(self, env, alpha=0.0001, beta=0.001,
                 gamma=0.99, max_size=250000, tau=0.005,
                 batch_size=64, noise=0.1):

        self.env = env
        self.gamma = gamma
        self.tau = tau
        self.n_actions = env.action_space.shape[0]
        self.obs_shape = env.observation_space['observation'].shape[0] + \
                         env.observation_space['achieved_goal'].shape[0] + \
                         env.observation_space['desired_goal'].shape[0]
        self.memory = ReplayBuffer(max_size, self.obs_shape, self.n_actions)

        self.batch_size = batch_size
        self.noise = noise
        self.max_action = env.action_space.high
        self.min_action = env.action_space.low

        self.actor = Actor(self.n_actions, name='actor')
        self.critic = Critic(name='critic')
        self.target_actor = Actor(self.n_actions, name='target_actor')
        self.target_critic = Critic(name='target_critic')

        self.actor.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=alpha))
        self.critic.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=beta))
        self.target_actor.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=alpha))
        self.target_critic.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=beta))
        self.update_networks(tau=1)

    def update_networks(self, tau=None):
        if tau is None:
            tau = self.tau

        weights = []
        targets = self.target_actor.weights
        for i, weight in enumerate(self.actor.weights):
            weights.append(weight * tau + targets[i] * (1 - tau))
        self.target_actor.set_weights(weights)

        weights = []
        targets = self.target_critic.weights
        for i, weight in enumerate(self.critic.weights):
            weights.append(weight * tau + targets[i] * (1 - tau))
        self.target_critic.set_weights(weights)

    def store(self, state, action, reward, new_state, done):
        self.memory.store_transition(state, action, reward, new_state, done)

    def save_models(self, path):
        self.actor.save_weights(path + self.actor.checkpoint)
        self.critic.save_weights(path + self.critic.checkpoint)
        self.target_actor.save_weights(path + self.target_actor.checkpoint)
        self.target_critic.save_weights(path + self.target_critic.checkpoint)

    def load_models(self, path):
        self.actor.load_weights(path + self.actor.checkpoint)
        self.critic.load_weights(path + self.critic.checkpoint)
        self.target_actor.load_weights(path + self.target_actor.checkpoint)
        self.target_critic.load_weights(path + self.target_critic.checkpoint)

    def choose_action(self, observation):
        state = tf.convert_to_tensor([observation], dtype=tf.float32)
        actions = self.actor(state)
        actions += tf.random.normal(shape=[self.n_actions], mean=0.0, stddev=self.noise)
        actions = tf.clip_by_value(actions, self.min_action, self.max_action)
        return actions[0]

    def optimize(self, steps):

        for i in range(steps):
            if self.memory.mem_cntr < self.batch_size:
                return

            state, action, reward, new_state, done = self.memory.sample_buffer(self.batch_size)

            with tf.GradientTape() as tape:
                target_actions = self.target_actor(new_state)
                critic_value_ = tf.squeeze(self.target_critic(new_state, target_actions), 1)
                critic_value = tf.squeeze(self.critic(state, action), 1)
                target = reward + self.gamma * critic_value_ * (1.0 - done)
                critic_loss = tf.keras.losses.MSE(target, critic_value)

            critic_network_gradient = tape.gradient(critic_loss, self.critic.trainable_variables)
            self.critic.optimizer.apply_gradients(zip(critic_network_gradient, self.critic.trainable_variables))

            with tf.GradientTape() as tape:
                new_policy_actions = self.actor(state)
                actor_loss = -self.critic(state, new_policy_actions)
                actor_loss = tf.math.reduce_mean(actor_loss)

            actor_network_gradient = tape.gradient(actor_loss, self.actor.trainable_variables)
            self.actor.optimizer.apply_gradients(zip(actor_network_gradient, self.actor.trainable_variables))
            self.update_networks()


from pathlib import Path
import os
import pickle

import pandas as pd
import numpy as np
from tensorflow.keras.utils import normalize, to_categorical
from sklearn.neighbors import BallTree, DistanceMetric
from tensorflow import keras

BASE = Path(os.path.realpath(__file__)).parent
DATA = BASE / '../backend/app/data'
node_data = pd.read_csv(DATA / 'edge_with_location.csv', index_col=0)
node_embeddings = pd.read_csv(DATA / "dgi_gs_final.csv", index_col=0)

# cors = []
# for i in tqdm(range(len(node_data))):
#     cors.append([node_data.lat.iloc[i], node_data.lon.iloc[i]])

# X = np.array([[np.radians(x) for x in list_] for list_ in cors])
# tree = BallTree(X, leaf_size=3, metric = DistanceMetric.get_metric('haversine'))

with open(DATA / 'balltree.pkl', 'rb') as f:
    tree = pickle.load(f)

# dist, idx = tree.query([cors[0]], k = 1)

def get_eta(samples, need_to_project_flag=False, model_name=DATA/"regression_prod.pt") -> float:
    def summ(first, second):
        return [x + y for x, y in zip(first, second)]

    def zerolistmaker(n):
        return list(np.array([0]*128))

    ids_to_remove = []

    if (need_to_project_flag == True):
        for i in tqdm(range(len(samples))):
            edges_id=[]
            ok = True
            for cor in samples[i]['cors']:
                dist, idx = tree.query([[cor[0]* np.pi / 180, cor[1]* np.pi / 180]], k = 1)
                if dist*6371000 > 100:
                    ids_to_remove.append(i)
                    ok = False
                    break
                edges_id.append(idx[0][0])
            if ok:
                samples[i]['edges'] = edges_id
        ids_to_remove = list(set(ids_to_remove))
    samples_copy = samples.copy()

    for i in ids_to_remove[::-1]:
        samples_copy.pop(i)

    if (need_to_project_flag == True):
        for samp in samples_copy:
            del samp['cors']

    emb_len = 128
    X_pre = samples_copy
    graph = node_data
    emb = []
    for i in tqdm(range(len(X_pre))):
        path_emb = zerolistmaker(emb_len)
        arr = X_pre[i]["edges"]
        for j in range(len(arr)):
            ind = graph[graph["edge_id"] == int(arr[j])].index
            if (len(ind) != 0):
                ind = graph[graph["edge_id"] == int(arr[j])].index[0]
            else:
#                print(str(i) + " " + str(j) + " " + arr[j] + " "+ str(graph[graph["edge_id"] == int(arr[j])].index[0]))
            path_emb = summ(node_embeddings.iloc[ind], path_emb)
        emb.append(path_emb)

    X_pre_c = pd.DataFrame(samples_copy)

    X = X_pre_c.join(pd.DataFrame(emb)).dropna().drop(["edges"], axis = 1).reset_index()
    wind_dir_classes = pd.get_dummies(X.wind, prefix='wind_dir_class')
    day_classes = pd.get_dummies(X.day_period, prefix='day_class')
    X[['day_class_0','day_class_1','day_class_2','day_class_3','wind_dir_class_0','wind_dir_class_45','wind_dir_class_135','wind_dir_class_180','wind_dir_class_270','wind_dir_class_315']] = 0
    X[day_classes.columns[0]] = day_classes[day_classes.columns[0]].values[0]
#    print(day_classes[day_classes.columns[0]].values[0])
    X[wind_dir_classes.columns[0]] = wind_dir_classes[wind_dir_classes.columns[0]]
#    print(X)
    Y = X["time"]
    X = X.drop(["time", "index", "id", "start_timestamp", "day_period", "wind"], axis = 1)
#    print(X)

    model = keras.models.load_model(model_name)
    from tensorflow.keras.utils import normalize, to_categorical
    test_x = normalize(np.asarray(X), axis=1)
    print(model(test_x))
    return model(test_x)

if __name__ == '__main__':
    local_edges = [572896, 838883, 838884, 838885, 838886, 838887, 838879, 838880, 484740, 484741, 484742, 484743, 484781, 484720, 484721, 990772, 572909, 572908, 990773, 226536, 226537, 226538, 226539, 227505, 227506, 227810, 227811, 227812, 227813, 573338, 573342, 573341, 573340, 226706, 226689, 226690, 226691, 226692, 226693, 114049, 114048, 114047, 114046, 227507, 959442, 959443, 227332, 227331, 227458, 227459, 227460, 227461, 227462, 227463, 227464, 227465, 227466, 227467, 227468, 227469, 227470, 227471, 895542, 895541, 225474, 225475, 225432, 225433]
    # print(node_embeddings.iloc[572896])
    from get_samples import get_sample
    samples = get_sample(local_edges)
    # print(samples)
    print(get_eta(samples))
    # print(max(node_data.edge_id))
    # print(min(node_data.edge_id))


# --------------------------------------------------------
# Relation Networks for Object Detection
# Copyright (c) 2017 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Jiayuan Gu, Dazhi Cheng
# --------------------------------------------------------
"""
Nms Multi-thresh Target Operator selects foreground and background roi,
    and assigns label, bbox_transform to them.
"""

import mxnet as mx
import numpy as np

from bbox.bbox_transform import bbox_overlaps

class NmsMultiTargetOp(mx.operator.CustomOp):
    def __init__(self, target_thresh):
        super(NmsMultiTargetOp, self).__init__()
        self._target_thresh = target_thresh
        self._num_thresh = len(target_thresh)

    def forward(self, is_train, req, in_data, out_data, aux):
        # bbox, [first_n, num_fg_classes, 4]
        bbox = in_data[0].asnumpy()
        num_boxes = bbox.shape[0]
        num_fg_classes = bbox.shape[1]
        gt_box = in_data[1].asnumpy()
        # score, [first_n, num_fg_classes]
        score = in_data[2].asnumpy()

        batch_image, num_gt, code_size = gt_box.shape
        assert batch_image == 1, 'only support batch_image=1, but receive %d' % num_gt
        assert code_size == 5, 'code_size of gt should be 5, but receive %d' % code_size
        assert len(score.shape) == 2, 'shape of score is %d instead of 2.' % len(score.shape)
        assert score.shape[1] == num_fg_classes, 'number of fg classes should be same for boxes and scores'

        output_list = []
        for cls_idx in range(0, num_fg_classes):
            valid_gt_mask = (gt_box[0, :, -1].astype(np.int32) == (cls_idx + 1))
            valid_gt_box = gt_box[0, valid_gt_mask, :]
            num_valid_gt = len(valid_gt_box)

            if num_valid_gt == 0:
                output = np.zeros(shape=(num_boxes, self._num_thresh), dtype=np.float32)
                output_list.append(output)
            else:
                bbox_per_class = bbox[:, cls_idx, :]
                score_per_class = score[:, cls_idx:cls_idx + 1]
                overlap_mat = bbox_overlaps(bbox_per_class.astype(np.float),
                                            valid_gt_box[:, :-1].astype(np.float))

                eye_matrix = np.eye(num_valid_gt)
                output_list_per_class = []

                for thresh in self._target_thresh:
                    # following mAP metric
                    overlap_mask = (overlap_mat > thresh)
                    valid_bbox_indices = np.where(overlap_mask)[0]
                    # require score be 2-dim
                    overlap_score = np.tile(score_per_class, (1, num_valid_gt))
                    overlap_score *= overlap_mask
                    max_overlap_indices = np.argmax(overlap_mat, axis=1)
                    max_overlap_mask = eye_matrix[max_overlap_indices]
                    overlap_score *= max_overlap_mask
                    max_score_indices = np.argmax(overlap_score, axis=0)
                    output = np.zeros((num_boxes,))
                    output[np.intersect1d(max_score_indices, valid_bbox_indices)] = 1
                    output_list_per_class.append(output)
                output_per_class = np.stack(output_list_per_class, axis=-1)
                output_list.append(output_per_class)
        blob = np.stack(output_list, axis=1).astype(np.float32, copy=False)
        self.assign(out_data[0], req[0], blob)

    def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
        self.assign(in_grad[0], req[0], 0)
        self.assign(in_grad[1], req[1], 0)
        self.assign(in_grad[2], req[2], 0)

@mx.operator.register("nms_multi_target")
class NmsMultiTargetProp(mx.operator.CustomOpProp):
    def __init__(self, target_thresh):
        super(NmsMultiTargetProp, self).__init__(need_top_grad=False)
        self._target_thresh = np.fromstring(target_thresh[1:-1], dtype=float, sep=' ')
        self._num_thresh = len(self._target_thresh)

    def list_arguments(self):
        return ['bbox', 'gt_bbox', 'score']

    def list_outputs(self):
        return ['nms_multi_target']

    def infer_shape(self, in_shape):
        bbox_shape = in_shape[0]
        # gt_box_shape = in_shape[1]
        score_shape = in_shape[2]

        assert bbox_shape[0] == score_shape[0], 'ROI number should be same for bbox and score'

        num_boxes = bbox_shape[0]
        num_fg_classes = bbox_shape[1]
        output_shape = (num_boxes, num_fg_classes, self._num_thresh)

        return in_shape, [output_shape]

    def create_operator(self, ctx, shapes, dtypes):
        return NmsMultiTargetOp(self._target_thresh)

    def declare_backward_dependency(self, out_grad, in_data, out_data):
        return []


# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION.  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.
""" Finetuning the library models for question-answering on SQuAD (Bert, Roberta)."""
import argparse
import logging
import os
import random
import glob
import time
import pickle
import json
import gc
import math
from shutil import rmtree, copy
from collections import namedtuple
import tensorflow as tf
from tensorflow.keras.losses import sparse_categorical_crossentropy as sce

import numpy as np
from tqdm import tqdm; tqdm.monitor_interval = 0  # noqa

from transformers import BertConfig, BertTokenizer, RobertaConfig, RobertaTokenizer

import optimization
from utils_nq import read_nq_examples, convert_examples_to_crops, \
    write_predictions, get_add_tokens, convert_preds_to_df, \
    read_candidates
from models import TFBertForNaturalQuestionAnswering, TFRobertaForNaturalQuestionAnswering
from eval_server import main as evaluate_on_nq

logger = logging.getLogger(__name__)

ALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys())
                  for conf in (BertConfig, )), ())

MODEL_CLASSES = {
    'bert': (BertConfig, TFBertForNaturalQuestionAnswering, BertTokenizer),
    'roberta': (RobertaConfig, TFRobertaForNaturalQuestionAnswering, RobertaTokenizer),
}

RawResult = namedtuple("RawResult", ["unique_id", "start_logits", "end_logits",
    "long_logits"])

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    tf.random.set_seed(args.seed)

def train(args, model_class, tokenizer, config, strategy):
    train_dataset, _ = load_and_cache_crops(args, tokenizer, evaluate=False)
    num_replicas = strategy.num_replicas_in_sync
    args.train_batch_size = args.per_tpu_train_batch_size * num_replicas
    num_steps_per_epoch = len(train_dataset['input_ids']) // args.train_batch_size
    t_total = num_steps_per_epoch * args.num_train_epochs
    warmup_steps = int(args.warmup * t_total)

    opt = tf.data.Options()
    opt.experimental_deterministic = True
    train_ds = tf.data.Dataset.from_tensor_slices(train_dataset).with_options(opt)
    train_ds = train_ds.repeat()
    train_ds = train_ds.shuffle(buffer_size=100, seed=args.seed)
    train_ds = train_ds.batch(batch_size=args.train_batch_size, drop_remainder=True)
    train_ds = strategy.experimental_distribute_dataset(train_ds)
    train_ds = iter(train_ds)

    with strategy.scope():
        model = model_class.from_pretrained(args.model_name_or_path, config=config)
        if args.init_weights:
            logger.info(f"Initializing from {args.init_weights}")
            model.load_weights(args.init_weights)

        optimizer = optimization.create_optimizer(args.learning_rate,
            num_steps_per_epoch * args.num_train_epochs, warmup_steps)
        # optimizer = tf.keras.optimizers.Adam(learning_rate=args.learning_rate)

    print(model.summary())

    @tf.function
    def train_step(batch):
        with tf.GradientTape() as tape:
            outputs = model(batch, training=True)
            start_loss = sce(batch['start_positions'], outputs[0], from_logits=True)
            end_loss = sce(batch['end_positions'], outputs[1], from_logits=True)
            long_loss = sce(batch['long_positions'], outputs[2], from_logits=True)
            loss = ((tf.reduce_mean(start_loss) + tf.reduce_mean(end_loss) / 2.0) +
                tf.reduce_mean(long_loss)) / 2.0
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        return loss

    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset["input_ids"]))
    logger.info("  Num Epochs = %.1f", args.num_train_epochs)
    logger.info("  Instantaneous batch size per TPU = %d", args.per_tpu_train_batch_size)
    logger.info("  Total optimization steps = %d", t_total)

    global_step = 1
    num_samples = 0
    smooth = 0.99
    set_seed(args)
    for epoch in range(1, math.ceil(args.num_train_epochs) + 1):
        running_loss = 0.0
        epoch_iterator = tqdm(range(num_steps_per_epoch))
        for step in epoch_iterator:
            if global_step >= t_total:
                logger.info(f'Finished training at epoch {epoch}, step {global_step}')
                break
            batch = next(train_ds)
            loss = strategy.experimental_run_v2(train_step, args=(batch, ))
            loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, loss, axis=None)
            global_step += 1
            running_loss = smooth * running_loss + (1. - smooth) * float(loss)

            if args.save_steps > 0 and global_step % args.save_steps == 0:
                # Save model checkpoint
                step_str = '%06d' % global_step
                ckpt_dir = os.path.join(args.output_dir, 'checkpoint-{}'.format(step_str))
                os.makedirs(ckpt_dir, exist_ok=True)
                weights_fn = os.path.join(ckpt_dir, 'weights.h5')
                model.save_weights(weights_fn)
                tokenizer.save_pretrained(ckpt_dir)

                # remove too many checkpoints
                checkpoint_fns = sorted(glob.glob(os.path.join(args.output_dir, 'checkpoint-*')))
                for fn in checkpoint_fns[:-args.n_keep]:
                    rmtree(fn)

            epoch_iterator.set_postfix({'ep': '%d/%.1f' % (epoch, round(args.num_train_epochs, 1)),
                'samples': num_samples, 'trl': round(running_loss, 4)})
            num_samples += args.train_batch_size

    return global_step, running_loss

def evaluate(args, model, tokenizer, strategy):
    val = 'val' in args.predict_fn
    prefix = 'val' if val else 'test'
    raw_results_fn = os.path.join(args.output_dir, f'{prefix}_raw_results.pkl')
    prediction_fn = os.path.join(args.output_dir, f'{prefix}_predictions.json')
    nbest_fn = os.path.join(args.output_dir, f'{prefix}_nbest_predictions.json')
    csv_fn = os.path.join(args.output_dir, f'{prefix}_submission.csv')
    null_log_odds_fn = os.path.join(args.output_dir, f'{prefix}_null_odds.json')
    eval_dataset, crops = load_and_cache_crops(args, tokenizer, evaluate=True)
    if not os.path.exists(raw_results_fn):
        if not os.path.exists(args.output_dir):
            os.makedirs(args.output_dir)

        num_replicas = strategy.num_replicas_in_sync
        args.eval_batch_size = args.per_tpu_eval_batch_size * num_replicas

        # pad dataset to multiple of `args.eval_batch_size`
        eval_dataset_length = len(eval_dataset['input_ids'])
        padded_length = math.ceil(eval_dataset_length / args.eval_batch_size) * args.eval_batch_size
        num_pad = padded_length - eval_dataset_length
        for ti, t in eval_dataset.items():
            pad_tensor = tf.expand_dims(tf.zeros_like(t[0]), 0)
            pad_tensor = tf.repeat(pad_tensor, num_pad, 0)
            eval_dataset[ti] = tf.concat([t, pad_tensor], 0)

        # create eval dataset
        eval_dataset['example_index'] = tf.range(padded_length, dtype=tf.int32)
        eval_ds = tf.data.Dataset.from_tensor_slices(eval_dataset)
        eval_ds = eval_ds.batch(batch_size=args.eval_batch_size, drop_remainder=True)
        eval_ds = strategy.experimental_distribute_dataset(eval_ds)
        eval_ds = iter(eval_ds)

        # eval
        logger.info("***** Running evaluation *****")
        logger.info("  Num examples = %d", eval_dataset_length)
        logger.info("  Batch size = %d", args.eval_batch_size)

        @tf.function
        def predict_step(batch):
            outputs = model(batch, training=False)
            return outputs

        all_results = []
        tic = time.time()
        for batch_ind in tqdm(range(padded_length // args.eval_batch_size)):
            # if batch_ind > 40:
            #     break
            batch = next(eval_ds)
            example_indexes = tf.concat(batch['example_index'].values, 0)
            outputs = strategy.experimental_run_v2(predict_step, args=(batch, ))
            outputs = [tf.concat(o.values, 0).numpy() for o in outputs]

            for i, example_index in enumerate(example_indexes):
                # filter out padded samples
                if example_index >= eval_dataset_length:
                    continue

                eval_crop = crops[example_index]
                unique_id = int(eval_crop.unique_id)
                start_logits = outputs[0][i]
                end_logits = outputs[1][i]
                long_logits = outputs[2][i]

                start_logits = [float(x) for x in start_logits]
                end_logits = [float(x) for x in end_logits]
                long_logits = [float(x) for x in long_logits]

                result = RawResult(unique_id=unique_id,
                                   start_logits=start_logits,
                                   end_logits=end_logits,
                                   long_logits=long_logits)
                all_results.append(result)

        eval_time = time.time() - tic
        logger.info("  Evaluation done in total %f secs (%f sec per example)",
            eval_time, eval_time / padded_length)
        with open(raw_results_fn, 'wb') as f:
            pickle.dump(all_results, f)
        print(f'***** Done wrote raw results to {raw_results_fn} *****')

    else:
        print(f'***** Loading raw results from {raw_results_fn} *****')
        with open(raw_results_fn, 'rb') as f:
          all_results = pickle.load(f)

    examples_gen = read_nq_examples(args.predict_fn, is_training=False)
    preds = write_predictions(examples_gen, crops, all_results, args.n_best_size,
                              args.max_answer_length,
                              prediction_fn, nbest_fn, null_log_odds_fn,
                              args.verbose_logging,
                              args.short_null_score_diff_threshold, args.long_null_score_diff_threshold)
    del crops, all_results
    gc.collect()
    candidates = read_candidates(['simplified-nq-train.jsonl', 'simplified-nq-test.jsonl'])
    sub = convert_preds_to_df(preds, candidates).sort_values('example_id')
    sub.to_csv(csv_fn, index=False, columns=['example_id', 'PredictionString'])
    print(f'***** Wrote submission to {csv_fn} *****')
    result = {}
    if val:
        EvalArgs = namedtuple('EvalArgs', ['csv', 'fn'])
        eval_args = EvalArgs(csv=args.gt_csv, fn=csv_fn)
        result = evaluate_on_nq(eval_args)
    return result

def load_and_cache_crops(args, tokenizer, evaluate=False):
    # Load data crops from cache or dataset file
    input_file = args.predict_fn if evaluate else args.train_fn
    cache_prefix = 'train'
    if evaluate:
        if 'test' in input_file:
            cache_prefix = 'test'
        else:
            cache_prefix = 'val'

    cached_crops_fn = os.path.join(os.path.dirname(input_file), 'cached_{}_{}_{}.pkl'.format(
        cache_prefix, list(filter(None, args.output_dir.split('/'))).pop(),
        str(args.max_seq_length)))

    if os.path.exists(cached_crops_fn) and not args.overwrite_cache:
        logger.info("Loading crops from cached file %s", cached_crops_fn)
        with open(cached_crops_fn, "rb") as f:
            crops = pickle.load(f)
    else:
        if args.model_name_or_path.startswith('bert'):
          cls_token = '[CLS]'
          sep_token = '[SEP]'
          pad_id = tokenizer.convert_tokens_to_ids(['[PAD]'])[0]
          sep_token_extra = False
        elif args.model_name_or_path.startswith('roberta'):
          cls_token = '<s>'
          sep_token = '</s>'
          pad_id = tokenizer.convert_tokens_to_ids(['<pad>'])[0]
          sep_token_extra = True

        doc_stride = args.eval_doc_stride if evaluate else args.train_doc_stride
        logger.info(f"Creating crops of length {args.max_seq_length} with stride {doc_stride}"
            f" from dataset file {input_file} using cls_token {cls_token}, sep_token {sep_token}"
            f" and pad_id {pad_id}")

        examples_gen = read_nq_examples(input_file, is_training=not evaluate)
        crops = convert_examples_to_crops(examples_gen=examples_gen,
                                          tokenizer=tokenizer,
                                          max_seq_length=args.max_seq_length,
                                          doc_stride=doc_stride,
                                          max_query_length=args.max_query_length,
                                          is_training=not evaluate,
                                          cls_token_segment_id=0,
                                          pad_token_segment_id=0,
                                          cls_token=cls_token,
                                          sep_token=sep_token,
                                          pad_id=pad_id,
                                          p_keep_impossible=args.p_keep_impossible if not evaluate else 1.0,
                                          sep_token_extra=sep_token_extra)
        logger.info("Saving crops into cached file %s", cached_crops_fn)
        with open(cached_crops_fn, "wb") as f:
            pickle.dump(crops, f)

    # stack
    all_input_ids = tf.stack([c.input_ids for c in crops], 0)
    all_attention_mask = tf.stack([c.attention_mask for c in crops], 0)
    all_token_type_ids = tf.stack([c.token_type_ids for c in crops], 0)
    # all_p_mask = tf.stack([c.p_mask for c in crops], 0)
    all_attention_mask = tf.cast(all_attention_mask, tf.int32)

    if evaluate:
        dataset = {
            'input_ids': all_input_ids,
            'attention_mask': all_attention_mask,
            'token_type_ids': all_token_type_ids
        }

    else:
        all_start_positions = tf.convert_to_tensor([f.start_position for f in crops], dtype=tf.int32)
        all_end_positions = tf.convert_to_tensor([f.end_position for f in crops], dtype=tf.int32)
        all_long_positions = tf.convert_to_tensor([f.long_position for f in crops], dtype=tf.int32)
        dataset = {
            'input_ids': all_input_ids,
            'attention_mask': all_attention_mask,
            'token_type_ids': all_token_type_ids,
            'start_positions': all_start_positions,
            'end_positions': all_end_positions,
            'long_positions': all_long_positions,
        }

    # https://github.com/huggingface/transformers/blob/master/examples/run_squad.py#L206
    if args.model_name_or_path.startswith('roberta'):
      logger.info(f"Removing token_type_ids for {args.model_name_or_path}")
      dataset.pop('token_type_ids')

    return dataset, crops

def get_strategy():
    try:
        tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection
        print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
        tf.config.experimental_connect_to_cluster(tpu)
        tf.tpu.experimental.initialize_tpu_system(tpu)
        strategy = tf.distribute.experimental.TPUStrategy(tpu)
    except ValueError as e:
        print(e)
        print('No TPU detected')
        tpu = None
        strategy = tf.distribute.get_strategy()

    return strategy

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

    # Required parameters
    parser.add_argument("--train_fn", default="nq-train-v1.0.1.json", type=str)
    parser.add_argument("--predict_fn", default="nq-val-v1.0.1.json", type=str)
    parser.add_argument("--gt_csv", default="nq-val-v1.0.1.csv", type=str)
    parser.add_argument("--model_name_or_path",
        default="bert-large-uncased-whole-word-masking-finetuned-squad", type=str)
    parser.add_argument("--init_weights", default="", type=str)
    parser.add_argument("--id", default="0", type=str)

    # Other parameters
    parser.add_argument('--short_null_score_diff_threshold', type=float, default=0.0)
    parser.add_argument('--long_null_score_diff_threshold', type=float, default=0.0)
    parser.add_argument("--max_seq_length", default=512, type=int)
    parser.add_argument("--train_doc_stride", default=256, type=int)
    parser.add_argument("--eval_doc_stride", default=256, type=int)
    parser.add_argument("--max_query_length", default=64, type=int)
    parser.add_argument("--do_not_train", action='store_true')
    parser.add_argument("--do_not_eval", action='store_true')
    parser.add_argument("--per_tpu_train_batch_size", default=8, type=int)
    parser.add_argument("--per_tpu_eval_batch_size", default=8, type=int)
    parser.add_argument("--learning_rate", default=3e-5, type=float)
    parser.add_argument("--num_train_epochs", default=2.0, type=float)
    parser.add_argument("--warmup", default=0.1, type=float)
    parser.add_argument("--n_best_size", default=10, type=int)
    parser.add_argument("--max_answer_length", default=30, type=int)
    parser.add_argument("--verbose_logging", action='store_true')
    parser.add_argument('--logging_steps', type=int, default=1_000_000)
    parser.add_argument('--save_steps', type=int, default=200,
                        help="Save checkpoint every X updates steps.")
    parser.add_argument('--overwrite_cache', action='store_true',
                        help="Overwrite the cached training and evaluation sets")
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--n_keep', type=int, default=2, help="Number of checkpoints to keep.")
    parser.add_argument('--p_keep_impossible', type=float,
                        default=0.1, help="The fraction of impossible"
                        " samples to keep.")
    parser.add_argument('--do_enumerate', action='store_true')

    args = parser.parse_args()
    # Setup logging
    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                        datefmt='%m/%d/%Y %H:%M:%S',
                        level=logging.INFO)

    # Set seed
    set_seed(args)
    if args.model_name_or_path.startswith('bert'):
        do_lower_case = 'uncased' in args.model_name_or_path
    elif args.model_name_or_path.startswith('roberta'):
        # https://github.com/huggingface/transformers/pull/1386/files
        do_lower_case = False

    args.output_dir = f"nq_{args.model_name_or_path.split('-')[0]}_" \
        f"{'uncased' if do_lower_case else 'cased'}_{args.id}"
    strategy = get_strategy()
    logger.info("Running on %s replicas" % strategy.num_replicas_in_sync)
    logger.info("Training / evaluation parameters %s", args)
    args.model_type = args.model_name_or_path.split('-')[0].lower()
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    config = config_class.from_pretrained(args.model_name_or_path, cache_dir='transformers_cache')
    logger.info(f"Using {'uncased' if do_lower_case else 'cased'} tokenizer")
    tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path,
        do_lower_case=do_lower_case, cache_dir='transformers_cache')
    add_tokens = get_add_tokens(do_enumerate=args.do_enumerate)
    if args.model_name_or_path.startswith('bert'):
      offset = 1
    elif args.model_name_or_path.startswith('roberta'):
      offset = 50100
      # https://github.com/huggingface/transformers/issues/1234
      # config.type_vocab_size = 2

    num_added = tokenizer.add_tokens(add_tokens, offset=offset)
    logger.info(f"Added {num_added} tokens")
    # Training
    if not args.do_not_train:
        # Create output directory if needed
        if not os.path.exists(args.output_dir):
            os.makedirs(args.output_dir)

        # Save the code
        src_fns = glob.glob("*.py")
        for src_fn in src_fns:
            dst = os.path.join(args.output_dir, src_fn)
            copy(src_fn, dst)

        # Save the run arguments
        with open(os.path.join(args.output_dir, 'args.json'), 'w') as f:
            args_ = vars(args).copy()
            args_['device'] = ''
            json.dump(args_, f, indent=2)

        global_step, tr_loss = train(args, model_class, tokenizer, config, strategy)
        logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)

    if not args.do_not_eval:
        # Reload the model
        glob_str = os.path.join(args.output_dir, 'checkpoint-*')
        checkpoint_fns = sorted(glob.glob(glob_str))
        if len(checkpoint_fns) == 0:
          raise ValueError(f'No weights for {glob_str}')

        checkpoint = checkpoint_fns[-1]
        logger.info("Evaluate the following checkpoint: %s", checkpoint)
        weights_fn = os.path.join(checkpoint, 'weights.h5')
        with strategy.scope():
            model = model_class(config)
            model(model.dummy_inputs, training=False)
            model.load_weights(weights_fn)

        # Evaluate
        result = evaluate(args, model, tokenizer, strategy)
        logger.info("Result: {}".format(result))

if __name__ == "__main__":
    main()


import os
import gym
import retro
import retrowrapper
import numpy as np
import tensorflow as tf
from Actions import ACTIONS
from model import ActorCritic
from env_wrap import make_sonic_env

weights_path = os.getcwd()+'/weights/GreenHillZone_Act1/sonic_weight'
env = make_sonic_env(game='SonicTheHedgehog-Genesis', state='GreenHillZone.Act1', test=True, record=True)
model = ActorCritic(len(ACTIONS), 512).model

model.load_weights(weights_path)
state = env.reset()
state = tf.cast(np.expand_dims(state, axis=0), dtype=tf.float32)

ep_reward = 0
done = False
new_ep = True
while not done:
    env.render()
    if new_ep:
        ht = tf.zeros((1, 512))
        ct = tf.zeros((1, 512))
        new_ep = False

    act_score, value, (ht, ct) = model((state, (ht, ct)))
    action = tf.math.argmax(tf.math.softmax(act_score[0]))

    state, reward, done, info = env.step(action)
    state = tf.cast(np.expand_dims(state, axis=0), dtype=tf.float32)
    print(info['x'])
    ep_reward += reward

env.close()


# -*- coding: utf-8 -*-
from __future__ import absolute_import
import os

os.system(
    "pip install --extra-index-url https://developer.download.nvidia.com/compute/redist/cuda/10.0 nvidia-dali && pip3 install torch torchvision")

import threading
import random

import tensorflow as tf
import torch
import torchvision as tv
import numpy as np

import skeleton
from architectures.mc3 import ResNet as MC3
from skeleton.projects import LogicModel, get_logger
from skeleton.projects.others import NBAC, AUC
import time
from collections import OrderedDict
from skeleton.utils.log_utils import timeit

import pdb

torch.backends.cudnn.benchmark = True
threads = [
    threading.Thread(target=lambda: torch.cuda.synchronize()),
    threading.Thread(target=lambda: tf.Session())
]
[t.start() for t in threads]
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

LOGGER = get_logger(__name__)

def set_random_seed_all(seed, deterministic=False):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    tf.random.set_random_seed(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

def get_top_players(data, sort_keys, reverse=True, n=2, order=True):
    top = sorted(data.items(), key=lambda x: x[1][sort_keys], reverse=reverse)[:n]
    if order:
        return OrderedDict(top)
    return dict(top)

class Model(LogicModel):
    def __init__(self, metadata):
        set_random_seed_all(0xC0FFEE,True)
        super(Model, self).__init__(metadata)

    def build(self):
        base_dir = os.path.dirname(os.path.abspath(__file__))
        in_channels = self.base_info['dataset']['shape'][-1]
        num_class = self.base_info['dataset']['num_class']
        # torch.cuda.synchronize()

        LOGGER.info('[init] session')
        [t.join() for t in threads]

        self.device = torch.device('cuda', 0)
        self.session = tf.Session()

        LOGGER.info('[init] Model')

        model_space = [MC3,MC3,MC3,MC3]
        self.model_space =[]
        self.ensembleconfig.MODEL_LENGTH = len(self.model_space)
        for key,model in enumerate(model_space):
            if model in [MC3]:
                Network = model(in_channels, num_class)
                model_path = os.path.join(base_dir, 'models')
                LOGGER.info('model path: %s', model_path)
                Network.init(model_dir=model_path, gain=1.0)
            else:
                Network = model(in_channels, num_class)
                Network.init(model_dir=model_path, gain=1.0)
            self.model_space.append(Network)

        self.model = self.model_space[self.ensembleconfig.MODEL_INDEX]
        self.model_pred = self.model_space[self.ensembleconfig.MODEL_INDEX].eval()

        LOGGER.info('[init] copy to device')
        self.model = self.model.to(device=self.device, non_blocking=True)  # .half()
        self.model_pred = self.model_pred.to(device=self.device, non_blocking=True)  # .half()self。
        LOGGER.info('[init] done.')

    def update_model(self):
        self.is_half = self.model._half
        num_class = self.base_info['dataset']['num_class']

        epsilon = min(0.1, max(0.001, 0.001 * pow(num_class / 10, 2)))
        if self.is_multiclass():
            self.model.loss_fn = torch.nn.BCEWithLogitsLoss(reduction='none')
            self.tau = 8.0
            LOGGER.info('[update_model] %s (tau:%f, epsilon:%f)', self.model.loss_fn.__class__.__name__, self.tau,
                        epsilon)
        else:
            self.model.loss_fn = torch.nn.CrossEntropyLoss(reduction='none')
            self.tau = 8.0
            LOGGER.info('[update_model] %s (tau:%f, epsilon:%f)', self.model.loss_fn.__class__.__name__, self.tau,
                        epsilon)
        self.model_pred.loss_fn = self.model.loss_fn

        self.model.hyper_params['conditions']['use_fast_auto_aug'] = False
        times = self.model.hyper_params['dataset']['input'][0]
        self.model.set_video(times=times)
        self.model_pred.set_video(times=times)

        self.init_opt()
        LOGGER.info('[update] done.')

    def init_opt(self):
        steps_per_epoch = self.model.hyper_params['dataset']['steps_per_epoch']
        batch_size = self.model.hyper_params['dataset']['batch_size']
        init_lr = self.model.hyper_params['optimizer']['lr']
        warmup_multiplier = self.model.hyper_params['optimizer']['warmup_multiplier']
        warmup_epoch = self.model.hyper_params['optimizer']['warmup_epoch']
        self.model.init_opt(steps_per_epoch,batch_size,init_lr,warmup_multiplier,warmup_epoch)
        LOGGER.info('[optimizer] %s (batch_size:%d)', self.model.optimizer._optimizer.__class__.__name__, batch_size)

    @timeit
    def adapt(self, remaining_time_budget=None):
        epoch = self.model.info['loop']['epoch']
        input_shape = self.model.hyper_params['dataset']['input']
        height, width = input_shape[:2]
        batch_size = self.model.hyper_params['dataset']['batch_size']

        train_score = np.average([c['train']['score'] for c in self.model.checkpoints[-5:]])
        valid_score = np.average([c['valid']['score'] for c in self.model.checkpoints[-5:]])
        LOGGER.info('[adapt] [%04d/%04d] train:%.3f valid:%.3f',
                    epoch, self.model.hyper_params['dataset']['max_epoch'],
                    train_score, valid_score)

        self.model.use_test_time_augmentation = self.model.info['loop']['test'] > 1

        if self.model.hyper_params['conditions']['use_fast_auto_aug']:
            self.model.hyper_params['conditions']['use_fast_auto_aug'] = valid_score < 0.995

        # Adapt Apply Fast auto aug
        if self.model.hyper_params['conditions']['use_fast_auto_aug'] and \
                (train_score > 0.995 or self.model.info['terminate']) and \
                remaining_time_budget > 120 and \
                valid_score > 0.01 and \
                self.dataloaders['valid'] is not None and \
                not self.model.update_transforms:
            LOGGER.info('[adapt] search fast auto aug policy')
            self.update_transforms = True

    def activation(self, logits):
        if self.is_multiclass():
            logits = torch.sigmoid(logits)
            prediction = (logits > 0.5).to(logits.dtype)
        else:
            logits = torch.softmax(logits, dim=-1)
            _, k = logits.max(-1)
            prediction = torch.zeros(logits.shape, dtype=logits.dtype, device=logits.device).scatter_(-1, k.view(-1, 1),
                                                                                1.0)
        return logits, prediction

    def decision_if_single_ensemble(self,model):
        pre_ens_round_start = 3
        if model.round_idx > pre_ens_round_start:
            model.single_ensemble = True
            return model.single_ensemble
        else:
            model.single_ensemble = False
            return model.single_ensemble

    @timeit
    def epoch_train(self, epoch, train, model=None):
        model = model if model is not None else self.model
        model.round_idx += 1
        if epoch < 0:
            optimizer = model.optimizer_fc
        else:
            optimizer = model.optimizer

        model.train()
        model.zero_grad()
        if model.info['condition']['first']['train']:
            num_steps = 10000
        else:
            num_steps = model.hyper_params['dataset']['steps_per_epoch']
        # train
        metrics = []
        scores = []
        score = 0
        step = 0

        for step, (examples, labels) in zip(range(num_steps), train):
            if examples.shape[0] == 1:
                examples = examples[0]
                labels = labels[0]
            original_labels = labels
            if not self.is_multiclass():
                labels = labels.argmax(dim=-1)
            # LOGGER.info('*'*30+str(self.is_multiclass()) + '*'*30)
            skeleton.nn.MoveToHook.to((examples, labels), self.device, self.is_half)
            logits, loss = model(examples, labels, tau=self.tau, reduction='avg')
            loss = loss.sum()
            loss.backward()

            max_epoch =model.hyper_params['dataset']['max_epoch']
            optimizer.update(maximum_epoch=max_epoch)
            optimizer.step()
            model.zero_grad()
            if model.info['condition']['first']['train']:
                logits, prediction = self.activation(logits.float())
                score = AUC(logits, original_labels.float())
                scores.append(score)
                if step > 10 and sum(scores[-10:]) > 2.:
                    break

            if step == num_steps - 1:
                logits, prediction = self.activation(logits.float())
                score = AUC(logits, original_labels.float())

            metrics.append({
                'loss': loss.detach().float().cpu(),
                'score': 0,
            })

        train_loss = np.average([m['loss'] for m in metrics])
        train_score = score
        optimizer.update(train_loss=train_loss)

        self.train_loss = train_loss
        self.train_score = train_score

        return {
            'loss': train_loss,
            'score': train_score,
            'run_steps': step
        }

    @timeit
    def epoch_valid(self, epoch, valid, model=None,reduction='avg'):
        test_time_augmentation = False
        model = model if model is not None else self.model
        model.train()
        num_steps = len(valid)
        metrics = []
        tau = self.tau

        with torch.no_grad():
            for step, (examples, labels) in zip(range(num_steps), valid):
                original_labels = labels
                if not self.is_multiclass():
                    labels = labels.argmax(dim=-1)
                batch_size = examples.size(0)

                if model.use_test_time_augmentation and test_time_augmentation:
                    examples = torch.cat([examples, torch.flip(examples, dims=[-1])], dim=0)
                    labels = torch.cat([labels, labels], dim=0)

                skeleton.nn.MoveToHook.to((examples, labels), self.device, self.is_half)
                logits, loss = model(examples, labels, tau=tau, reduction=reduction)

                if model.use_test_time_augmentation and test_time_augmentation:
                    logits1, logits2 = torch.split(logits, batch_size, dim=0)
                    logits = (logits1 + logits2) / 2.0

                logits, prediction = self.activation(logits.float())
                tpr, tnr, nbac = NBAC(prediction, original_labels.float())
                if reduction == 'avg':
                    auc = AUC(logits, original_labels.float())
                else:
                    auc = max([AUC(logits[i:i + 16], original_labels[i:i + 16].float()) for i in
                               range(int(len(logits)) // 16)])

                score = auc if model.hyper_params['conditions']['score_type'] == 'auc' else float(nbac.detach().float())
                metrics.append({
                    'loss': loss.detach().float().cpu(),
                    'score': score,
                })

                LOGGER.debug(
                    '[valid] [%02d] [%03d/%03d] loss:%.6f AUC:%.3f NBAC:%.3f tpr:%.3f tnr:%.3f, lr:%.8f',
                    epoch, step, num_steps, loss, auc, nbac, tpr, tnr,
                    model.optimizer.get_learning_rate()
                )

            if reduction == 'avg':
                valid_loss = np.average([m['loss'] for m in metrics])
                valid_score = np.average([m['score'] for m in metrics])
            elif reduction in ['min', 'max']:
                valid_loss = np.min([m['loss'] for m in metrics])
                valid_score = np.max([m['score'] for m in metrics])
            else:
                raise Exception('not support reduction method: %s' % reduction)

        model.optimizer.update(valid_loss=np.average(valid_loss))

        self.valid_loss = valid_loss
        self.valid_score = valid_score

        model.g_his_eval_dict[model.round_idx] = {
            "t_loss": self.train_loss,
            "t_acc": self.train_score,
            "v_loss": self.valid_loss,
            "v_acc": self.valid_score
        }

        return {
            'loss': valid_loss,
            'score': valid_score,
        }

    def skip_valid(self, model):

        model.g_his_eval_dict[model.round_idx] = {
            "t_loss": self.train_loss,
            "t_acc": self.train_score,
            "v_loss": 99.9,
            "v_acc": model.info['loop']['epoch'] * 1e-4
        }

        return {
            'loss': 99.9,
            'score': model.info['loop']['epoch'] * 1e-4,
        }

    @timeit
    def prediction(self, dataloader, model, checkpoints,test_time_augmentation=True, detach=True, num_step=None):
        tau = self.tau
        best_idx = np.argmax(np.array([c['valid']['score'] for c in checkpoints]))
        best_loss = checkpoints[best_idx]['valid']['loss']
        best_score = checkpoints[best_idx]['valid']['score']

        states = checkpoints[best_idx]['model']
        model.load_state_dict(states)
        LOGGER.info('best checkpoints at %d/%d (valid loss:%f score:%f) tau:%f',
                    best_idx + 1, len(checkpoints), best_loss, best_score, tau)

        num_step = len(dataloader) if num_step is None else num_step
        model.train()

        with torch.no_grad():
            predictions = []
            for step, (examples, labels) in enumerate(dataloader):

                batch_size = examples.size(0)
                if model.use_test_time_augmentation and test_time_augmentation:
                    examples = torch.cat([examples, torch.flip(examples, dims=[-1])], dim=0)

                # skeleton.nn.MoveToHook.to((examples, labels), self.device, self.is_half)
                logits = model(examples, tau=tau)

                if model.use_test_time_augmentation and test_time_augmentation:
                    logits1, logits2 = torch.split(logits, batch_size, dim=0)
                    logits = (logits1 + logits2) / 2.0

                logits, prediction = self.activation(logits)

                if detach:
                    predictions.append(logits.detach().float().cpu().numpy())
                else:
                    predictions.append(logits)

            if detach:
                predictions = np.concatenate(predictions, axis=0).astype(np.float)
            else:
                predictions = torch.cat(predictions, dim=0)
        return predictions


import argparse,logging
import numpy as np
import mxnet as mx
import cv2, dlib
from lightened_moon import lightened_moon_feature

logger = logging.getLogger()
logger.setLevel(logging.INFO)
import pdb

def main():
    symbol = lightened_moon_feature(num_classes=40, use_fuse=True)
    devs = mx.cpu() if args.gpus is None else [mx.gpu(int(i)) for i in args.gpus.split(',')]
    _, arg_params, aux_params = mx.model.load_checkpoint(args.model_load_prefix, args.model_load_epoch)
    detector = dlib.get_frontal_face_detector()
    face_cascade = cv2.CascadeClassifier(args.opencv)
    # read img and drat face rect
    img = cv2.imread(args.img)
    faces = detector(img, 1)
    gray = np.zeros(img.shape[0:2])
    if len(faces) == 0:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        opencv_faces = face_cascade.detectMultiScale(gray, 1.3, 5)
        for (x,y,w,h) in opencv_faces:
            faces.append(dlib.rectangle(int(x), int(y), int(x+w), int(y+h)))
    max_face = faces[0]
    if len(faces) > 0:
            max_face = max(faces, key=lambda rect: rect.width() * rect.height())
    for f in faces:
        if f == max_face:
            cv2.rectangle(img, (f.left(), f.top()), (f.right(), f.bottom()), (0,0,255), 2)
        else:
            cv2.rectangle(img, (f.left(), f.top()), (f.right(), f.bottom()), (255,0,0), 2)
    cv2.imwrite(args.img.replace('jpg', 'png'), img)
    # crop face area
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    pad = [0.25, 0.25, 0.25, 0.25] if args.pad is None else args.pad
    left = int(max(0, max_face.left() - max_face.width()*float(pad[0])))
    top = int(max(0, max_face.top() - max_face.height()*float(pad[1])))
    right = int(min(gray.shape[1], max_face.right() + max_face.width()*float(pad[2])))
    bottom = int(min(gray.shape[0], max_face.bottom()+max_face.height()*float(pad[3])))
    gray = gray[left:right, top:bottom]
    gray = cv2.resize(gray, (args.size, args.size))/255.0
    img = np.expand_dims(np.expand_dims(gray, axis=0), axis=0)
    # get pred
    arg_params['data'] = mx.nd.array(img, devs)
    exector = symbol.bind(devs, arg_params ,args_grad=None, grad_req="null", aux_states=aux_params)
    exector.forward(is_train=False)
    exector.outputs[0].wait_to_read()
    output = exector.outputs[0].asnumpy()
    text = ["5_o_Clock_Shadow","Arched_Eyebrows","Attractive","Bags_Under_Eyes","Bald", "Bangs","Big_Lips","Big_Nose",
            "Black_Hair","Blond_Hair","Blurry","Brown_Hair","Bushy_Eyebrows","Chubby","Double_Chin","Eyeglasses","Goatee",
            "Gray_Hair", "Heavy_Makeup","High_Cheekbones","Male","Mouth_Slightly_Open","Mustache","Narrow_Eyes","No_Beard",
            "Oval_Face","Pale_Skin","Pointy_Nose","Receding_Hairline","Rosy_Cheeks","Sideburns","Smiling","Straight_Hair",
            "Wavy_Hair","Wearing_Earrings","Wearing_Hat","Wearing_Lipstick","Wearing_Necklace","Wearing_Necktie","Young"]
    pred = np.ones(40)
    print("attribution is:")
    for i in range(40):
        print text[i].rjust(20)+" : \t",
        if output[0][i] < 0:
            pred[i] = -1
            print "No"
        else:
            pred[i] = 1
            print "Yes"

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="predict the face attribution of one input image")
    parser.add_argument('--gpus', type=str, help='the gpus will be used, e.g "0,1,2,3"')
    parser.add_argument('--img', type=str, default='./test.jpg', help='the input img path')
    parser.add_argument('--size', type=int, default=128,
                        help='the image size of lfw aligned image, only support squre size')
    parser.add_argument('--opencv', type=str, default='../model/opencv/cascade.xml',
                        help='the opencv model path')
    parser.add_argument('--pad', type=float, nargs='+',
                                 help="pad (left,top,right,bottom) for face detection region")
    parser.add_argument('--model-load-prefix', type=str, default='../model/lightened_moon/lightened_moon_fuse',
                        help='the prefix of the model to load')
    parser.add_argument('--model-load-epoch', type=int, default=82,
                        help='load the model on an epoch using the model-load-prefix')
    args = parser.parse_args()
    logging.info(args)
    main()


import numpy as np
from sklearn.neighbors import BallTree, KDTree
import os
import gc
import pickle as pkl
import copy

def array_min2d(x):
    x = np.array(x)
    if x.ndim >= 2:
        return x
    return x.reshape(-1, 1)

class EpisodicMemory(object):
    def __init__(self, buffer_size, state_dim, action_dim, action_shape, obs_shape, gamma=0.99,w_q=0.1):
        self.ec_buffer = []
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.capacity = buffer_size
        self.curr_capacity = 0

        self.latent_buffer = np.zeros((buffer_size, state_dim + action_dim))
        self.q_values = -np.inf * np.ones(buffer_size + 1)
        self.replay_buffer = np.empty((buffer_size,) + obs_shape, np.float32)
        self.action_buffer = np.empty((buffer_size,) + action_shape, np.float32)
        self.reward_buffer = np.empty((buffer_size,), np.float32)
        self.steps = np.empty((buffer_size,), np.int)
        self.done_buffer = np.empty((buffer_size,), np.bool)
        self.next_id = -1 * np.ones(buffer_size)
        self.prev_id = [[] for _ in range(buffer_size)]
        self.ddpg_q_values = -np.inf * np.ones(buffer_size)
        self.lru = np.zeros(buffer_size)
        self.time = 0
        self.gamma = gamma
        self.hashes = dict()
        self.reward_mean = None
        self.kd_tree = None
        self.state_kd_tree = None
        self.build_tree = False
        self.build_tree_times = 0
        self.w_q = w_q

    def save(self, filedir):
        pkl.dump(self, open(os.path.join(filedir, "episodic_memory.pkl"), "wb"))

    # def capacity(self):
    #     return [buffer.curr_capacity for buffer in self.ec_buffer]

    def add(self, obs, action, state, encoded_action, sampled_return, next_id=-1):
        if state is not None and encoded_action is not None:
            state, encoded_action = np.squeeze(state), np.squeeze(encoded_action)
            if len(encoded_action.shape) == 0:
                encoded_action = encoded_action[np.newaxis, ...]
        if self.curr_capacity >= self.capacity:
            # find the LRU entry
            # priority = self.w_q *(self.q_values[:self.capacity]) + self.lru
            # priority = self.q_values[:self.capacity]
            priority = self.lru
            index = int(np.argmin(priority))
            # index = int(np.argmin(self.q_values))
            self.prev_id[index] = []
            self.next_id[index] = -1
            self.q_values[index] = -np.inf
            old_key = tuple(
                np.squeeze(np.concatenate([self.replay_buffer[index], self.action_buffer[index]])).astype('float32'))
            self.hashes.pop(old_key, None)

        else:
            index = self.curr_capacity
            self.curr_capacity += 1
        self.replay_buffer[index] = obs
        self.action_buffer[index] = action
        if state is not None and encoded_action is not None:
            self.latent_buffer[index] = np.concatenate([state, encoded_action])
        self.q_values[index] = sampled_return
        self.lru[index] = self.time
        new_key = tuple(np.squeeze(np.concatenate([obs, action])).astype('float32'))
        self.hashes[new_key] = index

        if next_id >= 0:
            self.next_id[index] = next_id
            if index not in self.prev_id[next_id]:
                self.prev_id.append(index)
        self.time += 0.01
        return index

    def peek(self, action, state, value_decay, modify, next_id=-1, allow_decrease=False, knn=1):
        state = np.squeeze(state)
        if len(state.shape) == 1:
            state = state[np.newaxis, ...]
        action = np.squeeze(action)
        if len(action.shape) == 0:
            action = action[np.newaxis, ...]
        if len(action.shape) == 1:
            action = action[np.newaxis, ...]
        h = np.concatenate([state, action], axis=1)
        if self.curr_capacity == 0 or self.build_tree == False:
            return None, None
        # print(h.shape)
        dist, ind = self.kd_tree.query(h, k=knn)
        dist, ind = dist[0], ind[0]
        if dist[0] < 1e-5:
            self.lru[ind] = self.time
            self.time += 0.01
            if modify:
                # if value_decay > self.q_values[ind]:
                if allow_decrease:
                    self.q_values[ind] = value_decay
                else:
                    self.q_values[ind] = max(value_decay, self.q_values[ind])
                if next_id >= 0:
                    self.next_id[ind] = next_id
                    if ind not in self.prev_id[next_id]:
                        self.prev_id[next_id].append(ind)

            return self.q_values[ind], ind
        else:
            queried_q_value = 0.0
            dist = dist / (1e-12 + np.max(dist)) + 1e-13
            coeff = -np.log(dist)
            coeff = coeff / np.sum(coeff)
            for i, index in enumerate(ind):
                queried_q_value += self.q_values[index] * coeff[i]

            return queried_q_value, None

    # def sample(self, batch_size, num_neg=1):
    #     capacity = sum([buffer.curr_capacity for buffer in self.ec_buffer])
    #     assert 0 < batch_size < capacity, "can't sample that much!"
    #     anchor_idxes = []
    #     pos_idxes = []
    #     neg_idxes = []
    #     loop_count = 0
    #     while len(anchor_idxes) < batch_size:
    #         loop_count += 1
    #         if loop_count > self.curr_capacity * 100:
    #             return None
    #         rand_idx = np.random.randint(0, self.curr_capacity)
    #         if self.next_id[rand_idx] > 0:
    #             anchor_idxes.append(rand_idx)
    #             rand_pos = np.random.randint(0, len(self.next_id[rand_idx]))
    #             pos_idxes.append(self.next_id[rand_idx][rand_pos])
    #             prev_place = self.prev_id[rand_idx]
    #             next_place = [self.next_id[rand_idx]]
    #             neg_action, neg_idx = self.sample_neg_keys(
    #                 [rand_idx] + prev_place + next_place, num_neg)
    #             neg_idxes.append(neg_idx)
    #     neg_idxes = np.array(neg_idxes).reshape(-1)
    #     # anchor_obses = [self.ec_buffer[action].obses[id] for id, action in zip(anchor_idxes, anchor_actions)]
    #     # anchor_keys = [self.ec_buffer[action].hashes[id] for id, action in zip(anchor_idxes, anchor_actions)]
    #     # pos_keys = [self.ec_buffer[action].hashes[id] for id, action in zip(pos_idxes, pos_actions)]
    #     # neg_keys = [[self.ec_buffer[action].hashes[id] for id, action in zip(neg_idxes[i], neg_actions[i])] for i in
    #     #             range(len(neg_idxes))]
    #
    #     anchor_obs = [self.replay_buffer[ind] for ind in anchor_idxes]
    #     pos_obs = [self.replay_buffer[ind] for ind in pos_idxes]
    #     neg_obs = [self.replay_buffer[ind] for ind in neg_idxes]
    #     anchor_values = [self.q_values[index] for index in anchor_idxes]
    #     anchor_actions = [self.action_buffer[index] for index in anchor_idxes]
    #     return anchor_obs, pos_obs, neg_obs, anchor_values, anchor_actions

    def sample_neg_keys(self, avoids, batch_size):
        # sample negative keys
        assert batch_size + len(
            avoids) <= self.capacity, "can't sample that much neg samples from episodic memory!"
        places = []
        while len(places) < batch_size:
            ind = np.random.randint(0, self.curr_capacity)
            if ind not in places:
                places.append(ind)
        return places

    def update(self, idxes, reprs):
        self.latent_buffer[idxes] = reprs

    def state_knn_max_value(self, state, knn=4):
        knn = min(self.curr_capacity, knn)
        state = np.squeeze(state)
        if len(state.shape) <= 1:
            state = state[np.newaxis, ...]
        if not self.build_tree:
            return -np.inf
        dist, inds = self.state_kd_tree.query(state, k=knn)
        inds = inds[0]
        q_values = [self.q_values[i] for i in inds]
        return np.max(q_values)

    def state_knn_action(self, state_key, knn):
        knn = min(self.curr_capacity, knn)
        if self.curr_capacity < 2000 or self.build_tree == False:
            return None, None
        key = np.squeeze(state_key)
        while len(key.shape) <= 1:
            key = key[np.newaxis, ...]
        dist, ind = self.state_kd_tree.query(key, k=knn)
        ind = ind[0]
        q_values = self.q_values[ind]
        # temperature = 0.001
        # coeff = np.sum(np.square(self.latent_buffer[ind, :self.state_dim] - state_key), axis=1)
        # coeff = coeff / np.max(coeff)
        # print(coeff.shape,knn)
        # coeff = np.exp(-coeff / temperature)
        # coeff = coeff / np.sum(coeff)
        # q_values = q_values * coeff
        ind_max = ind[np.argmax(q_values)]
        max_action = self.action_buffer[ind_max]
        max_q = self.q_values[ind_max]
        return max_action, max_q

    def knn_value(self, state_key, action_key, knn):
        knn = min(self.curr_capacity, knn)
        if self.curr_capacity == 0 or self.build_tree == False:
            return None
        combined_key = np.concatenate((state_key, action_key), axis=1)
        key = np.squeeze(combined_key)
        if len(key.shape) <= 1:
            key = key[np.newaxis, ...]
        dist, ind = self.kd_tree.query(key, k=knn)
        # dist, ind = dist[0], ind[0]
        # dist = dist / (1e-12 + np.max(dist, axis=1, keepdims=True)) + 1e-13
        dist = dist + 1e-13
        # coeff = -np.log(dist)
        # coeff = np.exp(-dist / 0.1)
        coeff = 1. / dist ** 2
        coeff = coeff / np.sum(coeff, axis=1, keepdims=True)
        # value = 0.0
        queried_q_value = []
        # print("here",ind.shape,coeff.shape)
        for i in range(len(ind)):
            queried_q_value.append(0.0)
            for j, index in enumerate(ind[i]):
                queried_q_value[i] += self.q_values[index] * coeff[i, j]
                # self.lru[index] = self.time
                # self.time += 0.01

        return np.array(queried_q_value)[:, np.newaxis]

    def hash_value(self, state, action):
        if not self.build_tree:
            return None
        index = []
        for i in range(len(state)):
            key = tuple(np.squeeze(np.concatenate([state[i], action[i]])).astype('float32'))
            ind = self.hashes.get(key, self.capacity)
            # if ind == self.capacity:
            #     print("hash failure in episodic memory",self.curr_capacity)
            # else:
            #     print("success hash",self.capacity)
            index.append(ind)
        return (self.q_values[index]).reshape(-1)

    def knn(self, state_key, action_key, knn):
        combined_key = np.concatenate((state_key, action_key), axis=1)
        knn = min(self.curr_capacity, knn)
        if self.curr_capacity == 0 or self.build_tree == False:
            return None
        key = np.squeeze(combined_key)
        if len(key.shape) == 1:
            key = key[np.newaxis, ...]
        dist, ind = self.kd_tree.query(key, k=knn)
        # value = 0.0
        queried_q_values = []
        queried_obs = []
        queried_actions = []
        for i in range(len(ind)):
            queried_q_values.append([])
            queried_obs.append([])
            queried_actions.append([])
            for j, index in enumerate(ind[0]):
                queried_q_values[i].append(self.q_values[index])
                queried_obs[i].append(self.latent_buffer[index, :self.state_dim])
                queried_actions[i].append(self.latent_buffer[index, self.state_dim:])
                # self.lru[index] = self.time

        # self.time += 0.01

        return np.array(queried_obs), np.array(queried_actions), np.array(queried_q_values)

    def knn_with_intrinsic(self, state_key, action_key, knn):
        combined_key = np.concatenate((state_key, action_key), axis=1)
        knn = min(self.curr_capacity, knn)
        if self.curr_capacity == 0 or self.build_tree == False:
            return None
        key = np.squeeze(combined_key)
        if len(key.shape) == 1:
            key = key[np.newaxis, ...]
        dist, ind = self.kd_tree.query(key, k=knn)
        # value = 0.0
        queried_q_values = []
        queried_obs = []
        queried_actions = []
        for i in range(len(ind)):
            queried_q_values.append([])
            queried_obs.append([])
            queried_actions.append([])
            for j, index in enumerate(ind[0]):
                queried_q_values[i].append(self.q_values[index])
                queried_obs[i].append(self.latent_buffer[index, :self.state_dim])
                queried_actions[i].append(self.latent_buffer[index, self.state_dim:])
                self.lru[index] = self.time

        self.time += 0.01

        return np.array(queried_obs), np.array(queried_actions), np.array(queried_q_values)

    def update_sequence_horizon(self, sequence, horizen=100):
        # print(sequence)
        if self.reward_mean is None:
            reward_sum = 0
            for _, _, _, _, r, _, _ in sequence:
                reward_sum += r
            self.reward_mean = reward_sum / (len(sequence) - 1)
        next_id = -1
        Rtd = 0
        time_step = 0
        returns = []
        for obs, a, z, encoded_a, r, q_tp1, done in reversed(sequence):
            time_step += 1
            Rtd = self.gamma * Rtd + r
            returns.append(Rtd)
            # if self.gamma == 1:
            #     corrected_Rtd = Rtd + self.reward_mean * time_step
            # else:
            #     corrected_Rtd = Rtd + self.reward_mean * self.gamma * (1 - self.gamma ** time_step) / (1 - self.gamma)
            if time_step > horizen:
                corrected_Rtd = Rtd - self.gamma ** horizen * returns[time_step - 100]
                qd, current_id = self.peek(encoded_a, z, corrected_Rtd, True)
                if current_id is None:  # new action
                    current_id = self.add(obs, a, z, encoded_a, corrected_Rtd, next_id)

                self.replay_buffer[current_id] = obs
                self.reward_buffer[current_id] = r
            else:
                current_id = -1
            next_id = current_id
        self.reward_mean = np.mean(self.reward_buffer[:self.curr_capacity])
        return

    def update_sequence_corrected(self, sequence):
        # print(sequence)
        next_id = -1
        Rtd = 0
        for obs, a, z, encoded_a, r, q_tp1, done in reversed(sequence):
            # print(np.mean(z))
            if done and q_tp1 is None:
                continue
            if done and q_tp1 is not None:
                Rtd = float(np.squeeze(q_tp1))
            else:
                Rtd = self.gamma * Rtd + r

            # qd, current_id = self.peek(encoded_a, z, Rtd, True)
            # if current_id is None:  # new action
            current_id = self.add(obs, a, z, encoded_a, Rtd, next_id)

            self.replay_buffer[current_id] = obs
            self.reward_buffer[current_id] = r
            next_id = int(current_id)
        return

    def update_sequence_with_critic(self, sequence, q_values):
        next_id = -1
        Rtd = 0
        steps = len(q_values)
        for q, experience in zip(reversed(q_values), reversed(sequence)):
            obs, a, z, encoded_a, r, done = experience
            Rtd = self.gamma * Rtd + r
            steps -= 1
            qd, current_id = self.peek(encoded_a, z, q, True)
            if current_id is None:
                current_id = self.add(obs, a, z, encoded_a, q, next_id)

            self.replay_buffer[current_id] = obs
            self.reward_buffer[current_id] = r
            self.done_buffer[current_id] = done
            # self.ddpg_q_values[current_id] = q
            self.steps[current_id] = steps
            next_id = current_id

        return

    def update_sequence(self, sequence):
        # print(sequence)
        next_id = -1
        Rtd = 0
        for obs, a, z, encoded_a, r, done in reversed(sequence):
            # print(np.mean(z))

            Rtd = self.gamma * Rtd + r

            qd, current_id = self.peek(encoded_a, z, Rtd, True)
            if current_id is None:  # new action
                current_id = self.add(obs, a, z, encoded_a, Rtd, next_id)

            self.replay_buffer[current_id] = obs
            self.reward_buffer[current_id] = r
            self.done_buffer[current_id] = done

            next_id = current_id
        return

    def update_sequence_iterate(self, sequence, knn):
        # print(sequence)
        next_id = -1

        delta = np.inf
        last_value, last_key = None, None
        while delta > 1e6:
            Rtd = 0
            for obs, a, z, encoded_a, r, optimal_encoded_a, done in reversed(sequence):
                # print(np.mean(z))
                if done:
                    last_key = (z, encoded_a, knn)
                    Rtd = self.knn_value(z, encoded_a, knn=knn)

                    if Rtd is None:
                        Rtd = 0
                    else:
                        Rtd = float(Rtd)
                    last_value = Rtd
                    continue
                else:
                    Rtd = self.gamma * Rtd + r
                estimated_Rtd = -np.inf
                # estimated_Rtd = self.state_knn_max_value(obs, 4)
                # estimated_Rtd = -np.inf if estimated_Rtd is None else estimated_Rtd
                Rtd = max(estimated_Rtd, Rtd)
                qd, current_id = self.peek(encoded_a, z, Rtd, True, allow_decrease=done)
                if current_id is None:  # new action
                    current_id = self.add(obs, a, z, encoded_a, Rtd, next_id)

                self.replay_buffer[current_id] = obs
                self.reward_buffer[current_id] = r
                next_id = current_id

            self.update_kdtree()
            Rtd_knn = float(self.knn_value(*last_key))
            delta = abs(last_value - Rtd_knn)
            print(delta)
        return

    def update_kdtree(self, use_repr=True):
        if self.build_tree:
            del self.kd_tree
            del self.state_kd_tree
            # del self.hash_tree
        if self.curr_capacity <= 0:
            return
        # print("build tree", self.curr_capacity)
        # self.tree = KDTree(self.states[:self.curr_capacity])
        self.kd_tree = KDTree(self.latent_buffer[:self.curr_capacity])
        if use_repr:
            self.state_kd_tree = KDTree(self.latent_buffer[:self.curr_capacity, :self.state_dim])
        else:
            self.state_kd_tree = KDTree(self.replay_buffer[:self.curr_capacity])
        # self.hash_tree = KDTree(self.hashes[:self.curr_capacity])
        self.build_tree = True
        self.build_tree_times += 1
        if self.build_tree_times == 50:
            self.build_tree_times = 0
            gc.collect()

    def sample_negative(self, batch_size, batch_idxs, batch_idxs_next, batch_idx_pre):
        neg_batch_idxs = []
        i = 0
        while i < batch_size:
            neg_idx = np.random.randint(0, self.curr_capacity - 2)
            if neg_idx != batch_idxs[i] and neg_idx != batch_idxs_next[i] and neg_idx not in batch_idx_pre[i]:
                neg_batch_idxs.append(neg_idx)
                i += 1
        neg_batch_idxs = np.array(neg_batch_idxs)
        return self.replay_buffer[neg_batch_idxs]

    def switch_first_half(self, obs0, obs1, batch_size):
        tmp = copy.copy(obs0[:batch_size // 2, ...])
        obs0[:batch_size // 2, ...] = obs1[:batch_size // 2, ...]
        obs1[:batch_size // 2, ...] = tmp
        return obs0, obs1

    def sample(self, batch_size, mix=False):
        # Draw such that we always have a proceeding element
        if self.curr_capacity + 2 < batch_size:
            return None
        batch_idxs = []
        batch_idxs_next = []
        while len(batch_idxs) < batch_size:
            rnd_idx = np.random.randint(0, self.curr_capacity)
            if self.next_id[rnd_idx] == -1:
                continue
            batch_idxs.append(rnd_idx)
            batch_idxs_next.append(self.next_id[rnd_idx])

        batch_idxs = np.array(batch_idxs).astype(np.int)
        batch_idxs_next = np.array(batch_idxs_next).astype(np.int)
        batch_idx_pre = [self.prev_id[id] for id in batch_idxs]

        obs0_batch = self.replay_buffer[batch_idxs]
        obs1_batch = self.replay_buffer[batch_idxs_next]
        obs2_batch = self.sample_negative(batch_size, batch_idxs, batch_idxs_next, batch_idx_pre)
        action_batch = self.action_buffer[batch_idxs]
        reward_batch = self.reward_buffer[batch_idxs]
        terminal1_batch = self.done_buffer[batch_idxs]
        q_batch = self.q_values[batch_idxs]

        if mix:
            obs0_batch, obs1_batch = self.switch_first_half(obs0_batch, obs1_batch, batch_size)
        result = {
            'obs0': array_min2d(obs0_batch),
            'obs1': array_min2d(obs1_batch),
            'obs2': array_min2d(obs2_batch),
            'rewards': array_min2d(reward_batch),
            'actions': array_min2d(action_batch),
            'terminals1': array_min2d(terminal1_batch),
            'return': array_min2d(q_batch),
        }
        return result