"""Tensorflow utility functions for training"""
# tensorboard --logdir=experiments/base_model/
# tensorboard --logdir=experiments/base_model/train_summaries
# tensorboard --logdir=experiments/base_model/eval_summaries

import logging
import os

from tqdm import trange
import tensorflow as tf
from tensorflow.python.tools.inspect_checkpoint import print_tensors_in_checkpoint_file

from model.utils import save_dict_to_json, load_best_metric, get_expaned_metrics, save_vars_to_file
from model.evaluation import evaluate_sess, evaluate_on_train_sess, take_train_samples_sess, evaluate_sess_sample

from model.mabp import rl, rl_weights
import tensorflow.contrib.slim as slim
import numpy as np

def train_initial_sess(sess, model_spec, num_steps, writer, params, \
    old_index, numbers_of_selections, sums_of_reward, arm_weights, max_upper_bound, old_loss_val):
    """Train the model on `num_steps` batches

    Args:
        sess: (tf.Session) current session
        model_spec: (dict) contains the graph operations or nodes needed for training
        num_steps: (int) train for this number of batches
        writer: (tf.summary.FileWriter) writer for summaries
        params: (Params) hyperparameters
    """
    # Get relevant graph operations or nodes needed for training

    loss = model_spec['loss']
    # train_op = model_spec['train_op']
    update_metrics = model_spec['update_metrics']
    metrics = model_spec['metrics']
    summary_op = model_spec['summary_op']
    # reward = model_spec['reward']
    global_step = tf.train.get_global_step()
    # Load the training dataset into the pipeline and initialize the metrics local variables
    # sess.run(model_spec['iterator_init_op'])
    sess.run(model_spec['metrics_init_op'])
    # Use tqdm for progress bar
    t = trange(int(num_steps))
    epoch_loss = []
    index = 0
    # for i in t:
    #     # Evaluate summaries for tensorboard only once in a while
    #     if i == params.save_summary_steps - 1:
    #     # if i % params.save_summary_steps == 0:
    #         # Perform a mini-batch update
    #         _, _, loss_val, summ, global_step_val = sess.run([train_op, update_metrics, loss,
    #                                                           summary_op, global_step])
    #         # Write summaries for tensorboard
    #         writer.add_summary(summ, global_step_val)
    #     else:
    #         _, _, loss_val = sess.run([train_op, update_metrics, loss])
    #     # Log the loss in the tqdm progress bar
    #     # t.set_postfix(loss='{:05.3f}'.format(loss_val))
    #     epoch_loss.append(loss_val)
    for i in t:
        # Evaluate summaries for tensorboard only once in a while
        if i == params.save_summary_steps - 1:
        # if i % params.save_summary_steps == 0:
            # Perform a mini-batch update
            _, loss_val, summ, global_step_val = sess.run([update_metrics, loss,
                                                              summary_op, global_step])
            # Write summaries for tensorboard
            writer.add_summary(summ, global_step_val)
        else:
            _, loss_val = sess.run([update_metrics, loss])
        # Log the loss in the tqdm progress bar
        # t.set_postfix(loss='{:05.3f}'.format(loss_val))
        epoch_loss.append(loss_val)
    #################################################################################
    if 'retrain' in params.loss_fn:
        # calculate the rewards of the (10) clusters --> pick and update the action variable
        global_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='model/rewards')
        action_var=[v for v in global_vars if 'action' in v.name][0]
        action_var_value = sess.run([action_var])
        # logging.info('****action_var at train_sess:\n {}'.format(action_var_value))
        # logging.info('*****old_index is {}, loss_val is {}'.format(old_index, loss_val))
        sums_of_reward[old_index] += loss_val - old_loss_val
        old_loss_val = loss_val
        index = old_index
        # index, numbers_of_selections, arm_weights, \
        # max_upper_bound =  rl_weights(params, numbers_of_selections, \
        #     sums_of_reward, max_upper_bound, \
        #     t, arm_weights)
        # logging.info('numbers_of_selections at i:\n {}'.format(numbers_of_selections))
        # total_reward += reward
        # index = 2
        assign_op = action_var.assign([index])
        sess.run(assign_op)
        action_var_value = sess.run([action_var])
        # logging.info('******updated action_var at train_sess:\n {}'.format(action_var_value))
    #################################################################################
    metrics_values = {k: v[0] for k, v in metrics.items()}
    metrics_val = sess.run(metrics_values)
    expanded_metrics_val = get_expaned_metrics(metrics_val)
    metrics_string = " ; ".join("{}: {:05.4f}".format(k, v) for k, v in expanded_metrics_val.items())
    # logging.info("- Train metrics: " + metrics_string)
    # logging.info("-- Reward: \n")
    # logging.info(reward)
    return epoch_loss, index, sums_of_reward, arm_weights, max_upper_bound, old_loss_val

def train_sess(sess, model_spec, num_steps, writer, params, \
    old_index, numbers_of_selections, sums_of_reward, arm_weights, max_upper_bound, old_loss_val):
    """Train the model on `num_steps` batches

    Args:
        sess: (tf.Session) current session
        model_spec: (dict) contains the graph operations or nodes needed for training
        num_steps: (int) train for this number of batches
        writer: (tf.summary.FileWriter) writer for summaries
        params: (Params) hyperparameters
    """
    # Get relevant graph operations or nodes needed for training

    gradients = model_spec['gradients']
    loss = model_spec['loss']
    train_op = model_spec['train_op']
    update_metrics = model_spec['update_metrics']
    metrics = model_spec['metrics']
    summary_op = model_spec['summary_op']
    # reward = model_spec['reward']
    global_step = tf.train.get_global_step()
    # Load the training dataset into the pipeline and initialize the metrics local variables
    # sess.run(model_spec['iterator_init_op'])
    sess.run(model_spec['metrics_init_op'])
    # Use tqdm for progress bar
    t = trange(int(num_steps))
    index = -1
    epoch_loss = []
    for i in t:
        # Evaluate summaries for tensorboard only once in a while
        if i == params.save_summary_steps - 1:
        # if i % params.save_summary_steps == 0:
            # Perform a mini-batch update
            _, _, loss_val, summ, global_step_val, gradients_vars = sess.run([train_op, update_metrics, loss,
                                                              summary_op, global_step, gradients])
            # Write summaries for tensorboard
            writer.add_summary(summ, global_step_val)
        else:
            _, _, loss_val, gradients_vars = sess.run([train_op, update_metrics, loss, gradients])
        #################################################################################
        # if params.group:
        #     # calculate the rewards of the (10) clusters --> pick and update the action variable
        #     global_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='model/rewards')
        #     action_var=[v for v in global_vars if 'action' in v.name][0]
        #     action_var_value = sess.run([action_var])
        #     logging.info('****action_var at train_sess:\n {}'.format(action_var_value))
        #     # logging.info('*****old_index is {}, loss_val is {}'.format(old_index, loss_val))
        #     sums_of_reward[old_index] += loss_val - old_loss_val
        #     old_loss_val = loss_val
        #     index, numbers_of_selections, arm_weights, max_upper_bound = rl(params, numbers_of_selections, sums_of_reward, max_upper_bound, int(num_steps), arm_weights)
        #     # logging.info('numbers_of_selections at i:\n {}'.format(numbers_of_selections))
        #     # total_reward += reward
        #     # index = 2
        #     assign_op = action_var.assign([index])
        #     sess.run(assign_op)
        #     action_var_value = sess.run([action_var])
        #     logging.info('******updated action_var at train_sess:\n {}'.format(action_var_value))
        ##############################jk ###################################################
        # Log the loss in the tqdm progress bar
        # t.set_postfix(loss='{:05.3f}'.format(loss_val))
        loss_val = 0
        for g in gradients_vars:
            loss_val += np.sum(np.square(g))
        epoch_loss.append(loss_val)
        # logging.info(gradients_vars)
        # #################################################################################
        if 'retrain' in params.loss_fn:
            # calculate the rewards of the (10) clusters --> pick and update the action variable
            global_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='model/rewards')
            action_var=[v for v in global_vars if 'action' in v.name][0]
            action_var_value = sess.run([action_var])
            logging.info('****action_var at train_sess:\n {}'.format(action_var_value))
            # logging.info('*****old_index is {}, loss_val is {}'.format(old_index, loss_val))
            ################################CHANGED TO L2-norm Gradient########################
            # reward = loss_val - old_loss_val
            reward = loss_val
            sums_of_reward[old_index] = reward
            # logging.info('--------- mini-batch reward:{}'.format(reward))
            # sums_of_reward[old_index] += loss_val - old_loss_val
            old_loss_val = loss_val
            index, numbers_of_selections, arm_weights, max_upper_bound = rl_weights(params, numbers_of_selections, sums_of_reward, max_upper_bound, int(num_steps), arm_weights)
            # logging.info('numbers_of_selections at i:\n {}'.format(numbers_of_selections))
            # total_reward += reward
            # index = 2
            if index == -1:
                index = params.num_clusters - 1
            assign_op = action_var.assign([index])
            sess.run(assign_op)
            action_var_value = sess.run([action_var])
            logging.info('******updated action_var at train_sess:\n {}'.format(action_var_value))
    # #################################################################################
    metrics_values = {k: v[0] for k, v in metrics.items()}
    metrics_val = sess.run(metrics_values)
    expanded_metrics_val = get_expaned_metrics(metrics_val)
    metrics_string = " ; ".join("{}: {:05.4f}".format(k, v) for k, v in expanded_metrics_val.items())
    logging.info("- Train metrics: " + metrics_string)
    # logging.info("-- Reward: \n")
    # logging.info(reward)
    return epoch_loss, index, numbers_of_selections, sums_of_reward, arm_weights, max_upper_bound, old_loss_val

def model_summary():
    model_vars = tf.trainable_variables()
    slim.model_analyzer.analyze_vars(model_vars, print_info=True)

def isSavingWeights(eval_metrics, best_eval_metrics):
    for i in range(len(eval_metrics)):
        if eval_metrics[i] > best_eval_metrics[i]:
            return True
        elif eval_metrics[i] < best_eval_metrics[i]:
            return False
        else:
            continue
    return False

def get_pretrained_include(params):
    pretrained_include = ['model/cnn']
    if not params.use_kfac:
        pretrained_include.append('model/c_cnn')
    if not params.collect and params.loss_fn != 'cnn' and params.loss_fn != 'retrain_regu_mine3':
        pretrained_include.append('model/mask')
    return pretrained_include

def evaluate_on_train(eval_model_spec,
    model_dir, params, restore_from, global_epoch=0):
    """Train the model and evaluate every epoch.
    Args:
        train_model_spec: (dict) contains the graph operations or nodes needed for training
        eval_model_spec: (dict) contains the graph operations or nodes needed for evaluation
        model_dir: (string) directory containing config, weights and log
        params: (Params) contains hyperparameters of the model.
                Must define: num_epochs, train_size, batch_size, eval_size, save_summary_steps
        restore_from: (string) directory or file containing weights to restore the graph
    """
    # Initialize tf.Saver instances to save weights during training
    last_saver = tf.train.Saver() # will keep last 5 epochs
    best_saver = tf.train.Saver(max_to_keep=1) # only keep 1 best checkpoint (best on eval)
    begin_at_epoch = 0
    with tf.Session() as sess:
        # Initialize model variables
        sess.run(eval_model_spec['variable_init_op'])
        best_json_path = os.path.join(model_dir, "metrics_eval_best_weights.json")
        # Reload weights from directory if specified
        # restor from the previous learner
        if restore_from is not None:
            save_path = os.path.join(model_dir, restore_from)
            if os.path.isdir(save_path):
                save_path = tf.train.latest_checkpoint(save_path)
                begin_at_epoch = int(save_path.split('-')[-1])
                global_epoch = begin_at_epoch + 1
            logging.info("Restoring parameters from {}".format(save_path))
            # last_saver = tf.train.import_meta_graph(save_path+".meta")
            pretrained_include = get_pretrained_include(params)
            pretrained_vars = tf.contrib.framework.get_variables_to_restore(include=pretrained_include)
            pretrained_saver = tf.train.Saver(pretrained_vars, name="pretrained_saver")
            pretrained_saver.restore(sess, save_path)
        model_summary()
        best_saver = tf.train.Saver(max_to_keep=1)
        # Run one epoch
        logging.info("Epoch {}/{}".format(begin_at_epoch + 1, \
            begin_at_epoch + 1))
        # Compute number of batches in one epoch (one full pass over the training set)
        # Evaluate for one epoch on validation set
        num_steps = (params.train_size + params.batch_size - 1) // params.batch_size
        metrics = evaluate_on_train_sess(sess, eval_model_spec, num_steps, params)
        # loss_evaluate_on_train = sess.run(eval_model_spec['metrics']['loss'])
        # logging.info('loss_evaluate_on_train')
        # print(loss_evaluate_on_train)
        if params.loss_fn == 'cnn' or params.loss_fn == 'retrain_regu':
            cnn_vars=[v for v in tf.trainable_variables() if 'model/cnn' in v.name]
            c_cnn_vars=tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='model/c_cnn')
            # c_cnn_vars=[v for v in tf.trainable_variables() if 'model/c_cnn' in v.name]
            update_weights = [tf.assign(c, old) for (c, old) in \
            zip(c_cnn_vars, cnn_vars)]
            sess.run(update_weights)
        # # Save latest eval metrics in a json file in the model directory
        eval_on_train_json_path = os.path.join(model_dir, "metrics_eval_on_train.json")
        save_dict_to_json(metrics, eval_on_train_json_path)
        best_save_path = os.path.join(model_dir, 'best_weights', 'after-epoch')
        best_save_path = best_saver.save(sess, best_save_path, global_step=global_epoch)
        logging.info("- Found new best metric score, saving in {}".format(best_save_path))
    return global_epoch

def save_var(sess, name, epoch):
    cnn_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="model/cnn/{}".format(name))[0]
    cnn_vars = sess.run(cnn_vars)
    cnn_vars = cnn_vars.flatten().tolist()
    # print(cnn_vars)
    save_vars_to_file(cnn_vars, "./weights/corr_{}_output_{}".format(name, epoch))

def train_and_evaluate(train_model_spec, eval_model_spec,
    model_dir, params, restore_from=None, global_epoch=1):
    """Train the model and evaluate every epoch.

    Args:
        train_model_spec: (dict) contains the graph operations or nodes needed for training
        eval_model_spec: (dict) contains the graph operations or nodes needed for evaluation
        model_dir: (string) directory containing config, weights and log
        params: (Params) contains hyperparameters of the model.
                Must define: num_epochs, train_size, batch_size, eval_size, save_summary_steps
        restore_from: (string) directory or file containing weights to restore the graph
    """
    # Initialize tf.Saver instances to save weights during training
    last_saver = tf.train.Saver() # will keep last 5 epochs
    best_saver = tf.train.Saver(max_to_keep=1)  # only keep 1 best checkpoint (best on eval)
    begin_at_epoch = 0
    # MAB weight sampling
    num_clusters = params.num_clusters#10
    rewards = [0] * num_clusters
    weight_numbers_of_selections = [0] * num_clusters
    weight_sums_of_reward = [0] * num_clusters
    weight_arm_weights = [1] * num_clusters
    weight_max_upper_bound = 0
    old_index = 0
    old_loss_val = 0
    with tf.Session() as sess:
        # Initialize model variables
        sess.run(train_model_spec['variable_init_op'])
        # For tensorboard (takes care of writing summaries to files)
        train_writer = tf.summary.FileWriter(os.path.join(model_dir, 'train_summaries'), sess.graph)
        eval_writer = tf.summary.FileWriter(os.path.join(model_dir, 'vali_summaries'), sess.graph)
        best_json_path = os.path.join(model_dir, "metrics_eval_best_weights.json")
        best_eval_metrics = [0.0, -float('inf')]
        global_epoch = 0
        # Reload weights from directory if specified
        # restor from the previous learner
        if restore_from is not None:
            save_path = os.path.join(model_dir, restore_from)
            if os.path.isdir(save_path):
                save_path = tf.train.latest_checkpoint(save_path)
                begin_at_epoch = int(save_path.split('-')[-1])
                global_epoch = begin_at_epoch
            logging.info("Restoring parameters from {}".format(save_path))
            pretrained_include = get_pretrained_include(params)
            pretrained_vars = tf.contrib.framework.get_variables_to_restore(include=pretrained_include)
            pretrained_saver = tf.train.Saver(pretrained_vars, name="pretrained_saver")
            pretrained_saver.restore(sess, save_path)
            # last_best_eval_metric = load_best_metric(best_json_path)
            # best_eval_metrics = [last_best_eval_metric['accuracy'], -last_best_eval_metric['loss']]
            logging.info(best_eval_metrics)
        model_summary()
        # for each learner
        num_train_steps = (params.train_size + params.batch_size - 1) // params.batch_size
        num_train_steps = int(num_train_steps)
        if params.finetune:
            # initial rewards for all arms
            for i in range(num_clusters):
                old_index = i
                _, _, weight_sums_of_reward, weight_arm_weights, weight_max_upper_bound, old_loss_val = train_initial_sess(sess, train_model_spec, num_train_steps, \
                train_writer, params, old_index, weight_numbers_of_selections, weight_sums_of_reward, weight_arm_weights, weight_max_upper_bound, old_loss_val)
        # now real rl
        early_stopping_count = 0
        epoch_cut_off = int((begin_at_epoch + params.num_epochs) * params.epoch_cutoff)
        for epoch in range(begin_at_epoch, begin_at_epoch + params.num_epochs):
            if early_stopping_count == int(params.early_stoping_epochs):
                logging.info("Early stopping at epoch {}/{}".format(epoch + 1, \
                    begin_at_epoch + params.num_epochs))
                break
            # Run one epoch
            logging.info("Epoch {}/{}".format(epoch + 1, \
                begin_at_epoch + params.num_epochs))
            # Compute number of batches in one epoch (one full pass over the training set)

            # MAB data sampling
            sum_loss = [0] * num_train_steps
            numbers_of_selections = [0] * num_train_steps
            # UCB specific
            sums_of_reward = [0] * num_train_steps
            arm_weights = [1] * num_train_steps
            # UCB specific
            max_upper_bound = 0
            total_reward = 0
            batch_loss, old_index, weight_numbers_of_selections, weight_sums_of_reward, weight_arm_weights, weight_max_upper_bound, old_loss_val = train_sess(sess, train_model_spec, num_train_steps, \
                train_writer, params, old_index, weight_numbers_of_selections, weight_sums_of_reward, weight_arm_weights, weight_max_upper_bound, old_loss_val)

            # sum_loss = batch_loss
            # # sum_loss = [s+n for (s, n) in zip(batch_loss, sum_loss)]
            # sum_loss = [float(v) for v in sum_loss]
            # # logging.info('sum_loss :\n {}'.format(sum_loss))
            # consk = int(params.consk)
            # for i in range(num_train_steps):
            #     index, reward, numbers_of_selections, sums_of_reward, \
            #     max_upper_bound = rl(params, sum_loss, numbers_of_selections, \
            #         sums_of_reward, max_upper_bound, \
            #         (epoch - begin_at_epoch + 1) / consk, arm_weights)
            #     if params.rl == 'EXP3':
            #         arm_weights = sums_of_reward
            #     # logging.info('numbers_of_selections at i:\n {}'.format(numbers_of_selections))
            #     total_reward += reward

            # Save weights
            # if epoch >= epoch_cut_off:
            #     # cnn_vars=[v for v in tf.trainable_variables() if 'model/cnn/weights1_1' in v.name]
            #     # cnn_vars = tf.get_variable('model/cnn/weights1_1')
            #     save_var(sess, 'weights1_1', epoch)
            #     save_var(sess, 'weights1_2', epoch)
            #     # save_var(sess, 'weights3_1', epoch)
            #     save_var(sess, 'weights3_2', epoch)
            save_var(sess, 'weights1_1', epoch)
            save_var(sess, 'weights1_2', epoch)
            save_var(sess, 'weights3_2', epoch)
            last_save_path = os.path.join(model_dir, 'last_weights', 'after-epoch')
            last_saver.save(sess, last_save_path, global_step=global_epoch)
            # # Evaluate for one epoch on validation set
            num_vali_steps = (params.vali_size + params.batch_size - 1) // params.batch_size
            num_vali_steps = int(num_vali_steps)
            metrics = evaluate_sess(sess, eval_model_spec, num_vali_steps, eval_writer, params)
            # If best_eval, best_save_path
            accuracy_metric = round(metrics['accuracy'], 6)
            loss_metric = -round(metrics['loss'], 6)
            # save_batch()
            eval_metrics = [accuracy_metric, loss_metric]
            # logging.info('global_epoch: {}, best_eval_metrics: {}, \
            #     eval_metric: {}', global_epoch, best_eval_metrics, eval_metric)
            # logging.info('isSavingWeights(eval_metrics, best_eval_metrics) {}'.\
            # format(isSavingWeights(eval_metrics, best_eval_metrics)))
            if isSavingWeights(eval_metrics, best_eval_metrics):
                # rest early_stopping_count
                early_stopping_count = 0
                # and isSavingWeights
                best_eval_metrics = eval_metrics
                # Save weights
                if params.loss_fn == 'cnn' and not params.use_kfac:
                    cnn_vars=[v for v in tf.trainable_variables() if 'model/cnn' in v.name]
                    c_cnn_vars=[v for v in tf.trainable_variables() if 'model/c_cnn' in v.name]
                    update_weights = [tf.assign(c, old) for (c, old) in \
                    zip(c_cnn_vars, cnn_vars)]
                    sess.run(update_weights)
                    best_save_path = os.path.join(model_dir, 'best_weights', 'after-epoch')
                    best_save_path = best_saver.save(sess, best_save_path, global_step=global_epoch)
                    logging.info("- Make a copy of cnn vars, saving in {}".format(best_save_path))
                elif params.loss_fn == 'retrain_regu_mine3':
                    # c_cnn_vars=tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='model/cnn')
                    c_cnn_vars=[v for v in tf.trainable_variables() if 'model/cnn' in v.name]
                    cnn_vars=[v for v in tf.trainable_variables() if 'model/mask' in v.name]
                    update_weights = [tf.assign(c, tf.multiply(old, c)) for (c, old) in \
                    zip(c_cnn_vars, cnn_vars)]
                    sess.run(update_weights)
                    best_save_path = os.path.join(model_dir, 'best_weights', 'after-epoch')
                    best_save_path = best_saver.save(sess, best_save_path, global_step=global_epoch)
                    logging.info("- Updated cnn vars, saving in {}".format(best_save_path))
                best_save_path = os.path.join(model_dir, 'best_weights', 'after-epoch')
                best_save_path = best_saver.save(sess, best_save_path, global_step=global_epoch)
                logging.info("- Found new best metric score, saving in {}".format(best_save_path))
                # Save best eval metrics in a json file in the model directory
                save_dict_to_json(metrics, best_json_path)
            else:
                early_stopping_count = early_stopping_count + 1
            # Save latest eval metrics in a json file in the model directory
            last_json_path = os.path.join(model_dir, "metrics_eval_last_weights.json")
            save_dict_to_json(metrics, last_json_path)
            global_epoch += 1
        # update in the end is wrong as not the best weights are copied
        '''
        if params.loss_fn == 'cnn' and not params.use_kfac:
            cnn_vars=[v for v in tf.trainable_variables() if 'model/cnn' in v.name]
            c_cnn_vars=[v for v in tf.trainable_variables() if 'model/c_cnn' in v.name]
            update_weights = [tf.assign(c, old) for (c, old) in \
            zip(c_cnn_vars, cnn_vars)]
            sess.run(update_weights)
            best_save_path = os.path.join(model_dir, 'best_weights', 'after-epoch')
            best_save_path = best_saver.save(sess, best_save_path, global_step=global_epoch)
            logging.info("- Make a copy of cnn vars, saving in {}".format(best_save_path))
        elif params.loss_fn == 'retrain_regu_mine3':
            # c_cnn_vars=tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='model/cnn')
            c_cnn_vars=[v for v in tf.trainable_variables() if 'model/cnn' in v.name]
            cnn_vars=[v for v in tf.trainable_variables() if 'model/mask' in v.name]
            update_weights = [tf.assign(c, tf.multiply(old, c)) for (c, old) in \
            zip(c_cnn_vars, cnn_vars)]
            sess.run(update_weights)
            best_save_path = os.path.join(model_dir, 'best_weights', 'after-epoch')
            best_save_path = best_saver.save(sess, best_save_path, global_step=global_epoch)
            logging.info("- Updated cnn vars, saving in {}".format(best_save_path))
        '''
        begin_at_epoch = global_epoch
        early_stopping_count = 0
        for epoch in range(begin_at_epoch, begin_at_epoch + params.num_epochs2):
            # if early_stopping_count == int(params.early_stoping_epochs):
            #     logging.info("Early stopping at epoch {}/{}".format(epoch + 1, \
            #         begin_at_epoch + params.num_epochs))
            #     break
            # Run one epoch
            logging.info("Epoch {}/{}".format(epoch + 1, \
                begin_at_epoch + params.num_epochs2))

            # MAB data sampling
            sum_loss = [0] * num_train_steps
            numbers_of_selections = [0] * num_train_steps
            # UCB specific
            sums_of_reward = [0] * num_train_steps
            arm_weights = [1] * num_train_steps
            # UCB specific
            max_upper_bound = 0
            total_reward = 0
            batch_loss, old_index, weight_numbers_of_selections, weight_sums_of_reward, weight_arm_weights, weight_max_upper_bound, old_loss_val = train_sess(sess, train_model_spec, num_train_steps, \
                train_writer, params, old_index, weight_numbers_of_selections, weight_sums_of_reward, weight_arm_weights, weight_max_upper_bound, old_loss_val)

            sum_loss = batch_loss
            # sum_loss = [s+n for (s, n) in zip(batch_loss, sum_loss)]
            sum_loss = [float(v) for v in sum_loss]
            # logging.info('sum_loss :\n {}'.format(sum_loss))
            consk = int(params.consk)
            for i in range(num_train_steps):
                index, reward, numbers_of_selections, sums_of_reward, \
                max_upper_bound = rl(params, sum_loss, numbers_of_selections, \
                    sums_of_reward, max_upper_bound, \
                    (epoch - begin_at_epoch + 1) / consk, arm_weights)
                if params.rl == 'EXP3':
                    arm_weights = sums_of_reward
                # logging.info('numbers_of_selections at i:\n {}'.format(numbers_of_selections))
                total_reward += reward
            # logging.info('len of sum_loss: {}'.format(len(sum_loss)))
            save_var(sess, 'weights1_1', epoch)
            save_var(sess, 'weights1_2', epoch)
            save_var(sess, 'weights3_2', epoch)
            # Save weights
            last_save_path = os.path.join(model_dir, 'last_weights', 'after-epoch')
            # global_epoch = int(params.num_learners) * int(params.num_epochs) + epoch + 1
            last_saver.save(sess, last_save_path, global_step=global_epoch)
            metrics = evaluate_sess(sess, eval_model_spec, num_vali_steps, eval_writer, params)
            # If best_eval, best_save_path
            accuracy_metric = round(metrics['accuracy'], 6)
            loss_metric = -round(metrics['loss'], 6)
            # save_batch()
            eval_metrics = [accuracy_metric, loss_metric]
            # logging.info('global_epoch: {}, best_eval_metrics: {}, \
            #     eval_metric: {}', global_epoch, best_eval_metrics, eval_metric)
            if isSavingWeights(eval_metrics, best_eval_metrics):
                # rest early_stopping_count
                early_stopping_count = 0
                # and isSavingWeights
                best_eval_metrics = eval_metrics
                # Save weights
                # trainalbe_vars = {v.name: v for v in tf.trainable_variables() if 'model' in v.name}
                # print(trainalbe_vars.keys())
                if params.loss_fn == 'cnn' and not params.use_kfac:
                    cnn_vars=[v for v in tf.trainable_variables() if 'model/cnn' in v.name]
                    c_cnn_vars=[v for v in tf.trainable_variables() if 'model/c_cnn' in v.name]
                    update_weights = [tf.assign(c, old) for (c, old) in \
                    zip(c_cnn_vars, cnn_vars)]
                    sess.run(update_weights)
                best_save_path = os.path.join(model_dir, 'best_weights', 'after-epoch')
                # global_epoch = int(params.num_learners) * int(params.num_epochs) + epoch + 1
                best_save_path = best_saver.save(sess, best_save_path, global_step=global_epoch)
                logging.info("- Found new best metric score, saving in {}".format(best_save_path))
                # Save best eval metrics in a json file in the model directory
                save_dict_to_json(metrics, best_json_path)
            else:
                early_stopping_count = early_stopping_count + 1
            # Save latest eval metrics in a json file in the model directory
            last_json_path = os.path.join(model_dir, "metrics_eval_last_weights.json")
            save_dict_to_json(metrics, last_json_path)
            global_epoch += 1

        logging.info('train num_steps: {}'.format((params.train_size + params.batch_size - 1) // params.batch_size))
        logging.info('weight_sums_of_reward: {}'.format(weight_sums_of_reward))
        # logging.info('numbers_of_selections: {}'.format(numbers_of_selections))

        # logging.info('numbers_of_selections:\n {}'.format(numbers_of_selections))
        sorted_index = sorted(range(num_train_steps), key=lambda k: numbers_of_selections[k], reverse=True)
        # top_sorted_index = sorted_index[0: int(num_train_steps*params.top_ratio)+1]
        sample_batchs = (params.sample_size + params.batch_size - 1) // params.batch_size
        top_sorted_index = sorted_index[0: int(sample_batchs)+1]
        logging.info('len(top_sorted_index) in training: {}'.format(len(top_sorted_index)))
        take_train_samples_sess(sess, eval_model_spec, num_train_steps, params, top_sorted_index)
    return global_epoch


#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-
import os

import rospy

from sensor_msgs.msg import PointCloud2, Image, PointCloud, PointField, JointState
from geometry_msgs.msg import Point32, Point, PoseStamped, Pose, Twist
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint

import numpy as np
import time
from std_msgs.msg import String, Header, Int64, Bool
import sensor_msgs.point_cloud2 as pc2
import std_msgs.msg
import tf
import message_filters
from robobreizh.msg import CloudIndexed, CloudSources, GraspConfigList, BoundingBoxCoord, GraspConfig
from robobreizh.srv import detect_grasps, object_detection, grasping
import moveit_commander
import geometry_msgs.msg
from math import pi
from std_msgs.msg import String
from moveit_commander.conversions import pose_to_list
from utils import *
from scipy.spatial.transform import Rotation as R
import sys
from tf import TransformListener
import tf2_ros
import tf2_geometry_msgs
import threading
from tf.transformations import euler_from_quaternion, quaternion_from_euler
from move_base_msgs.msg import MoveBaseAction, MoveBaseGoal
import copy

class Grasping():
	def __init__(self):
		#rospy.init_node('GraspingNode', anonymous=False)

		self.object_pose = [0,0,0,0]
		self.cloud_points = []

		# self.object_pc_pub = rospy.Publisher("crop_pointcloud", PointCloud2, queue_size=10)
		# self.cloud_index_pub = rospy.Publisher("gpd_cloud_indexed", CloudIndexed, queue_size=10)

		# MOVEIT CONFIG
		self.listener = tf.TransformListener()

		moveit_commander.roscpp_initialize(sys.argv)
		self.robot = moveit_commander.RobotCommander()
		self.whole_body = moveit_commander.MoveGroupCommander("whole_body_light")
		self.scene = moveit_commander.PlanningSceneInterface()
		self.arm = moveit_commander.MoveGroupCommander('arm')
		self.object_coord = ""
		self.cam_frame = "head_rgbd_sensor_rgb_frame"
	# 	t = threading.Thread(target=self.publish_frame)
	# 	t.start()

	def publish_frame(self):
		rate = rospy.Rate(10)
		while(not rospy.is_shutdown()):
			if self.object_coord != "":
				self.tb = tf.TransformBroadcaster()
				pose = [self.object_coord.position.x, self.object_coord.position.y, self.object_coord.position.z]
				orientation = [self.object_coord.orientation.x, self.object_coord.orientation.y, self.object_coord.orientation.z, self.object_coord.orientation.w]
				self.tb.sendTransform(pose, orientation,  rospy.Time.now(), "current_grasp", self.cam_frame)
				rate.sleep()

	def continuous_node(self):
		# PUBLISHERS
		self.grasp_pub = rospy.Publisher("best_grasp", GraspConfig, queue_size=10)

		# SUBSCRIBERS
		self.detected_obj_sub = message_filters.Subscriber('/detected_object', BoundingBoxCoord)
		self.pc_sub = message_filters.Subscriber('/hsrb/head_rgbd_sensor/depth_registered/rectified_points', PointCloud2)

		while not rospy.is_shutdown():

			ts = message_filters.ApproximateTimeSynchronizer([self.detected_obj_sub, self.pc_sub], 10)
			ts.registerCallback(self.callback)

			rospy.spin()

	def callback(self, detected_obj_sub, pc_sub):
		obj = detected_obj_sub
		msg = self.create_cloud_indexed(obj, pc_sub)

		rospy.wait_for_service('/detect_grasps_server/detect_grasps')
		grasp_service = rospy.ServiceProxy('/detect_grasps_server/detect_grasps', detect_grasps)

		try:
			resp = grasp_service(msg)
		except rospy.ServiceException as exc:
			print("Service did not process request: " + str(exc))

		best_res = GraspConfig()
		best_res = resp.grasp_configs.grasps[0]
		self.grasp_pub.publish(best_res)

	def service_node(self):
		s = rospy.Service('grasping', grasping, self.handle_grasping)
		rospy.loginfo("Grasping Node: Waiting for Request...")
		rospy.spin()

	def handle_grasping(self, req):
		pose = req.best_grasp
		self.object_coord = pose
		res = Bool()
		res.data = self.grasp_gound(pose)
		res.data = self.grasp_table1(pose)

		return res

	def move_base_goal2(self, x, y):
		navclient = actionlib.SimpleActionClient('/move_base', MoveBaseAction)
		goal = MoveBaseGoal()

		goal.target_pose.header.frame_id = "map"

		goal.target_pose.pose.position.x = x
		goal.target_pose.pose.position.y = y

		goal.target_pose.pose.orientation = quaternion_from_euler(0, 0, 90)

		navclient.send_goal(goal)
		navclient.wait_for_result()
		state = navclient.get_state()
		return True if state == 3 else False

	def grasp_ground(self, grasp_pose):
		time.sleep(1)

		# Move toward the object
		# dist = 0.6

		# obj_pose = self.transform_frame(grasp_pose.actual_pose, "map", "odom").pose
		# #obj_pose = grasp_pose.pre_pose
		# print(obj_pose)
		# (trans,rot) = self.listener.lookupTransform('/map', '/base_link', rospy.Time(0))
		# print(trans)
		# x = obj_pose.position.x - trans[0]
		# y = obj_pose.position.y - trans[1]
		# v = [x,y]
		# dist_actual = np.sqrt(np.square(x) + np.square(y))
		# print(dist_actual)
		# if dist_actual > dist:
		# 	u = np.array(v/np.linalg.norm(v))
		# 	print(u)
		# 	dest = []
		# 	dest.append(obj_pose.position.x - dist*u[0])
		# 	dest.append(obj_pose.position.x - dist*u[1])

		# 	print(dest)
		# 	move_head_tilt(0.0)
		# 	move_base_goal(dest[0], dest[1], 90)

		move_hand(1)

		#self.visualize_pose(grasp_pose.pre_pose)
		#grasp_pose.pre_pose.position.z = grasp_pose.pre_pose.position.z + 0.1
		self.move_to(grasp_pose.pre_pose, 'arm')

		#time.sleep(1)
		# Check if the grasp pose z value is above the hand depth to avoid htting the ground while trying to grasp
		if grasp_pose.actual_pose.position.z <= 0.066:
			grasp_pose.actual_pose.position.z = 0.05

		# group = moveit_commander.MoveGroupCommander("arm")
		# group.allow_replanning(True)
		# group.set_num_planning_attempts(5)
		# group.set_workspace([-3.0, -3.0, 3.0, 3.0])
		# group.set_planning_time(10)

		# waypoints = []
		# pose_to_perform = group.get_current_pose().pose

		# pose_to_perform.position = grasp_pose.actual_pose.position
		# pose_to_perform.orientation = grasp_pose.actual_pose.orientation
		# waypoints.append(copy.deepcopy(pose_to_perform))

		# print(pose_to_perform)

		# (plan, fraction) = group.compute_cartesian_path(waypoints, 0.01, 0.0)

		# group.execute(plan)
		# # Calling `stop()` ensures that there is no residual movement
		# group.stop()
		# # It is always good to clear your targets after planning with poses.
		# # Note: there is no equivalent function for clear_joint_value_targets()
		# group.clear_pose_targets()
		self.move_to(grasp_pose.actual_pose, 'arm')

		#move_hand(0)

		return True

	def grasp_table2(self, grasp_pose):

		move_hand(1)

		self.move_to(grasp_pose.pre_pose, 'arm')

		rospy.sleep(1.)

		self.move_to(grasp_pose.actual_pose, 'arm')

		return True

	def grasp_table1(self, grasp_pose):
		move_hand(1)

		self.move_to(grasp_pose.pre_pose, 'arm')

		self.move_to(grasp_pose.actual_pose, 'arm')

		return True

	def move_arm_vision(self):
		group_name = "arm"
		move_group = moveit_commander.MoveGroupCommander(group_name)
		move_group.allow_replanning(True)
		#move_group.set_workspace([-2.0, -2.0, 2.0, 2.0])
		joint_goal = move_group.get_current_joint_values()

		joints_index = {'arm_lift_joint':0, 'arm_flex_joint':1, 'arm_roll_joint':2, 'wrist_flex_joint':3, 'wrist_roll_joint':4, 'wrist_ft_sensor_frame_joint':5, 'wrist_ft_sensor_frame_inverse_joint':6, 'hand_palm_joint':7,}

		joint_goal[3] = -2.0
		#joint_goal[2] = 1.4
		#joint_goal[1] = -0.1

		try:
			move_group.go(joint_goal, wait=True)
		except moveit_commander.MoveItCommanderException as exc:
			print("")
		move_group.stop()
		move_hand(0)

	def move_arm_depose(self):
		group_name = "arm"
		move_group = moveit_commander.MoveGroupCommander(group_name)
		move_group.allow_replanning(True)
		#move_group.set_workspace([-2.0, -2.0, 2.0, 2.0])
		joint_goal = move_group.get_current_joint_values()

		joints_index = {'arm_lift_joint':0, 'arm_flex_joint':1, 'arm_roll_joint':2, 'wrist_flex_joint':3, 'wrist_roll_joint':4, 'wrist_ft_sensor_frame_joint':5, 'wrist_ft_sensor_frame_inverse_joint':6, 'hand_palm_joint':7,}

		joint_goal[1] = -0.4

		try:
			move_group.go(joint_goal, wait=True)
		except moveit_commander.MoveItCommanderException as exc:
			print("")
		move_group.stop()
		move_hand(0)

	def transform_frame(self, pose, frame_dest, frame_source):
		tf2_buffer = tf2_ros.Buffer(rospy.Duration(1200.0)) #tf buffer length
		tf2_listener = tf2_ros.TransformListener(tf2_buffer)

		transform = tf2_buffer.lookup_transform(frame_dest,
									   frame_source, #source frame
									   rospy.Time(0), #get the tf at first available time
									   rospy.Duration(2.0)) #wait for 1 second

		pose_stamped_to_transform = PoseStamped()
		pose_stamped_to_transform.pose.position = pose.position
		pose_stamped_to_transform.pose.orientation = pose.orientation

		return(tf2_geometry_msgs.do_transform_pose(pose_stamped_to_transform, transform))

	def move_to(self, pose, target):
		group = moveit_commander.MoveGroupCommander(target)
		group.allow_replanning(True)
		group.set_num_planning_attempts(5)
		group.set_workspace([-3.0, -3.0, 3.0, 3.0])
		group.set_planning_time(10)
		pose_to_perform = group.get_current_pose()

		print(group.get_current_pose())
		pose_to_perform.header.frame_id = "/odom"
		pose_to_perform.pose.position = pose.position
		pose_to_perform.pose.orientation = pose.orientation
		print(pose_to_perform)
		#group.set_pose_target(pose_to_perform)
		group.set_joint_value_target(pose_to_perform, "hand_palm_link", True)

		plan = group.go()
		# Calling `stop()` ensures that there is no residual movement
		group.stop()
		# It is always good to clear your targets after planning with poses.
		# Note: there is no equivalent function for clear_joint_value_targets()
		group.clear_pose_targets()

	def visualize_pose(self, pose):

		pub = rospy.Publisher("/visualize_pose", PointCloud2, queue_size=10)
		points = []
		points.append([pose.position.x, pose.position.y, pose.position.z])

		fields = [PointField('x', 0, PointField.FLOAT32, 1),
		  PointField('y', 4, PointField.FLOAT32, 1),
		  PointField('z', 8, PointField.FLOAT32, 1),
		  ]
		header = Header()
		header.frame_id = "odom"
		header.stamp = rospy.Time.now()
		poses_msg = pc2.create_cloud(header, fields, points)

		pub.publish(poses_msg)

	def trac_IK(self, grasp_pose):
		urdfstring = ''.join(open('/opt/ros/melodic/share/hsrb_description/robots/hsrb4s.urdf', 'r').readlines())
		#urdfstring = rospy.get_param('/robot_description')
		joint_values = rospy.wait_for_message("/hsrb/joint_states", JointState)
		joint_pose = []
		ik = IK("arm_lift_link", "hand_palm_link", urdf_string=urdfstring)

		for name in ik.joint_names:
			for i in range(0, len(joint_values.name)):
				if joint_values.name[i] == name:
					joint_pose.append(joint_values.position[i])

		joint_pose.append(0.0)

		print("TRAC IK")
		print(ik.get_ik(joint_pose, grasp_pose.position.x, grasp_pose.position.y, grasp_pose.position.z, grasp_pose.orientation.x, grasp_pose.orientation.y, grasp_pose.orientation.z, grasp_pose.orientation.w))

# 	seed_state = [0.0] * ik_solver.number_of_joints

# ik_solver.get_ik(seed_state,
#                 0.45, 0.1, 0.3,  # X, Y, Z
#                 0.0, 0.0, 0.0, 1.0)  # QX, QY, QZ, QW

	def move_pose_joint(self, pose, joint):
		group_name = "arm"
		move_group = moveit_commander.MoveGroupCommander(group_name)
		move_group.allow_replanning(True)
		move_group.set_workspace([-2.0, -2.0, 2.0, 2.0])
		joint_goal = move_group.get_current_joint_values()

		# 0 = arm_lift_joint
		# 1 = arm_flex_joint
		# 2 = arm_roll_joint
		# 3 = wrist_flex_joint
		# 4 = wrist_roll_joint
		# 5 = wrist_ft_sensor_frame_joint
		# 6 = wrist_ft_sensor_frame_inverse_joint
		# 7 = hand_palm_joint

		joints_index = {'arm_lift_joint':0, 'arm_flex_joint':1, 'arm_roll_joint':2, 'wrist_flex_joint':3, 'wrist_roll_joint':4, 'wrist_ft_sensor_frame_joint':5, 'wrist_ft_sensor_frame_inverse_joint':6, 'hand_palm_joint':7,}

		joint_goal[joints_index.get(joint)] = joint_goal[joints_index.get(joint)] + pose

		move_group.go(joint_goal, wait=True)
		move_group.stop()

	def move_base(self, angle):
		PI = 3.1415926535897
		velocity_publisher = rospy.Publisher("/hsrb/cmd_vel_controller/cmd_vel", Twist, queue_size=10)
		vel_msg = Twist()

		#Converting from angles to radians
		speed = 10 # degree/s
		angular_speed = speed*2*PI/360

		#We wont use linear components
		vel_msg.linear.x=0
		vel_msg.linear.y=0
		vel_msg.linear.z=0
		vel_msg.angular.y = 0

		# Checking if our movement is CW or CCW
		if angle < 0:
			vel_msg.angular.z = -abs(angular_speed)
		else:
			vel_msg.angular.z = abs(angular_speed)

		# Setting the current time for distance calculus
		t0 = rospy.Time.now().to_sec()
		current_angle = 0

		while(current_angle < angle):
			velocity_publisher.publish(vel_msg)
			t1 = rospy.Time.now().to_sec()
			current_angle = angular_speed*(t1-t0)

		#Forcing our robot to stop
		vel_msg.angular.z = 0
		velocity_publisher.publish(vel_msg)

	def compute_inverse_kinematics(self, grasp_pose):
		#self.trac_IK(grasp_pose)
		group_joints = ["odom_r", "arm_lift_joint", "arm_flex_joint", "arm_roll_joint", ]
		tf2_buffer = tf2_ros.Buffer(rospy.Duration(1200.0)) #tf buffer length
		tf2_listener = tf2_ros.TransformListener(tf2_buffer)

		transform = tf2_buffer.lookup_transform("base_link",
									   "head_rgbd_sensor_rgb_frame", #source frame
									   rospy.Time(0), #get the tf at first available time
									   rospy.Duration(2.0)) #wait for 1 second

		pose_stamped_to_transform = PoseStamped()
		pose_stamped_to_transform.pose.position = grasp_pose.position
		pose_stamped_to_transform.pose.orientation = grasp_pose.orientation

		#grasp_pose = tf2_geometry_msgs.do_transform_pose(pose_stamped_to_transform, transform)

		pose_target = [grasp_pose.position.x, grasp_pose.position.y, grasp_pose.position.z] # End effector position = Te.
		orient_target = [grasp_pose.orientation.x, grasp_pose.orientation.y, grasp_pose.orientation.z, grasp_pose.orientation.w] # End effector orientation as a quaternion.

		# R0_g = end-effector(gripper) rotation transformation(4X4)
		R0_g = tf.transformations.quaternion_matrix(orient_target)
		# D0_g = end-effector(gripper) translation transformation(3X3)
		D0_g = tf.transformations.translation_matrix(pose_target)
		print(D0_g)

		origin, xaxis, yaxis, zaxis = (0, 0, 0), (1, 0, 0), (0, 1, 0), (0, 0, 1)

		(trans,rot) = self.listener.lookupTransform('/arm_lift_link', '/odom_xr_link', rospy.Time(0))
		dlift = np.linalg.norm(trans)

		(trans,rot) = self.listener.lookupTransform('/hand_palm_link', '/wrist_flex_link', rospy.Time(0))
		dwrist = np.linalg.norm(trans)

		(trans,rot) = self.listener.lookupTransform('/arm_flex_link', '/wrist_flex_link', rospy.Time(0))
		darm = np.linalg.norm(trans)

		approach = [grasp_pose.approach.x, grasp_pose.approach.y, grasp_pose.approach.z]

		binormal = [grasp_pose.binormal.x, grasp_pose.binormal.y, grasp_pose.binormal.z]

		hand_axis = [grasp_pose.axis.x, grasp_pose.axis.y, grasp_pose.axis.z]

		R = [approach, binormal, hand_axis]

		# 1. Compute the base yaw angle
		#(trans,rot) = self.listener.lookupTransform('/base_link', '/hand_palm_link', rospy.Time(0))
		#trans = tf.transformations.translation_matrix([0.0,0.0,-dwrist])

		(trans,rot) = self.listener.lookupTransform('/base_link', '/current_grasp', rospy.Time(0))
		print(R)
		wristX = trans[0] - (dwrist*R[0][2])
		wristY = trans[1] - (dwrist*R[1][2])
		wristZ = trans[2] - (dwrist*R[2][2])
		wrist_pos = [wristX,wristY,wristZ]

		print("wrist_pos: {}".format(wrist_pos))
		base_yaw = np.arctan2(wrist_pos[1], wrist_pos[0])
		print("base_yaw: {}".format(base_yaw))

		tf2_buffer = tf2_ros.Buffer(rospy.Duration(1200.0)) #tf buffer length
		tf2_listener = tf2_ros.TransformListener(tf2_buffer)

		transform = tf2_buffer.lookup_transform("current_grasp",
									   "base_link", #source frame
									   rospy.Time(0), #get the tf at first available time
									   rospy.Duration(2.0)) #wait for 1 second

		pose_stamped_to_transform = PoseStamped()
		pose_stamped_to_transform.pose.position.x = wristX
		pose_stamped_to_transform.pose.position.y = wristY
		pose_stamped_to_transform.pose.position.z = wristZ

		wrist_pos = tf2_geometry_msgs.do_transform_pose(pose_stamped_to_transform, transform).pose.position
		print(wrist_pos)
		# 2. Compute torso lift height
		lift_heigth = wrist_pos.z - np.sqrt(np.square(darm) - np.square(wrist_pos.x) - np.square(wrist_pos.y))

		#lift_heigth = wrist_pos[2] - np.sqrt((darm**2) - (wrist_pos[0]**2) - ((wrist_pos[1])**2))
		print("lift_heigth: {}".format(lift_heigth))

		# 3. Compute shoulder pitch angle
		shoulder_angle = np.arcsin(np.sqrt(np.square(wrist_pos.x) + np.square(wrist_pos.y))/darm)
		print("shoulder_angle: {}".format(shoulder_angle))

		# 4. Compute the wirst joint angles
		trans_arm = tf.transformations.translation_matrix([0.0, 0.0, -darm])
		trans_lift = tf.transformations.translation_matrix([0.0, 0.0, -dlift])
		trans_wrist = tf.transformations.translation_matrix([0.0, 0.0, -dwrist])

		Ry = tf.transformations.rotation_matrix(-shoulder_angle, yaxis)
		print(Ry)
		Rz = tf.transformations.rotation_matrix(-base_yaw, zaxis)
		print(Rz)
		R = np.dot(np.dot(np.dot(np.dot(trans_arm, Ry), trans_lift), Rz), trans)

		print(R)
		roll_arm = np.arctan2(R[1][2], R[0][2])
		pitch_arm = np.arccos(R[2][2])
		roll_gripper = np.arctan2(R[2][1], -R[2][0])
		print("roll_arm: {}".format(roll_arm))
		print("pitch_arm: {}".format(pitch_arm))
		print("roll_gripper: {}".format(roll_gripper))

		#arm_pub = rospy.Publisher()arm_lift_joint

		return base_yaw, lift_heigth, shoulder_angle, roll_arm, pitch_arm, roll_gripper

	def grasp_object2(self, grasp_pose):
		self.object_coord = grasp_pose

		approach = [grasp_pose.approach.x, grasp_pose.approach.y, grasp_pose.approach.z]

		binormal = [grasp_pose.binormal.x, grasp_pose.binormal.y, grasp_pose.binormal.z]

		hand_axis = [grasp_pose.axis.x, grasp_pose.axis.y, grasp_pose.axis.z]

		aperture = grasp_pose.width

		orientation = grasp_pose.orientation

		pose_stamped_to_transform = geometry_msgs.msg.PoseStamped()
		pose_stamped_to_transform.header.frame_id = "head_rgbd_sensor_rgb_frame"
		pose_stamped_to_transform.pose.orientation = grasp_pose.orientation
		pose_stamped_to_transform.pose.position = grasp_pose.position

		# arm.set_joint_value_target([0.5, 0, 0, 0, 0, 0])
		# arm.go()

		current_pose = get_relative_coordinate("map", "hand_palm_link")

		tf2_buffer = tf2_ros.Buffer(rospy.Duration(1200.0)) #tf buffer length
		tf2_listener = tf2_ros.TransformListener(tf2_buffer)

		transform = tf2_buffer.lookup_transform("odom",
									   pose_stamped_to_transform.header.frame_id, #source frame
									   rospy.Time(0), #get the tf at first available time
									   rospy.Duration(2.0)) #wait for 1 second

		pose_transformed = tf2_geometry_msgs.do_transform_pose(pose_stamped_to_transform, transform)
		print "New coordinates in the map frame 1: "
		print(pose_transformed)

		# if self.listener.frameExists("/head_rgbd_sensor_rgb_frame") and self.listener.frameExists("/map"):
		# 	t = self.listener.getLatestCommonTime("/head_rgbd_sensor_rgb_frame", "/map")
		# 	p1 = geometry_msgs.msg.PoseStamped()
		# 	p1.header.frame_id = "head_rgbd_sensor_rgb_frame"
		# 	p1.pose.orientation = grasp_pose.orientation
		# 	p1.pose.position = grasp_pose.position
		# 	p_in_end_effector = PoseStamped()
		# 	p_in_end_effector = self.listener.transformPose("/map", p1)
		# 	print "Position of the gripper in the map frame:"
		# 	print p_in_end_effector
		# else:
		# 	print "no transform available"

		#r = R.from_euler('xyz', [approach, binormal, hand_axis])
		# r = R.from_rotvec([approach, binormal, hand_axis])

		# quat = r.as_quat()
		# print(r.as_euler('xyz', degrees=True))
		# roll, pitch, yaw = r.as_euler('xyz', degrees=True)

		#p_in_end_effector.pose.orientation = orientation

		#p_in_end_effector.pose.position = grasp_pose.position
		current_pose = get_relative_coordinate("map", "hand_palm_link")

		move_arm_neutral()
		move_hand(1)

		try:
			pose_to_perform = geometry_msgs.msg.PoseStamped()
			pose_to_perform.position = current_pose.position
			pose_to_perform.header = current_pose.header
			pose_to_perform.header.stamp = rospy.Time.now()
			pose_to_perform.orientation = pose_transformed.orientation
			self.move_to(pose_transformed, 'whole_body_light')
		except:
			rospy.logerr('fail to grasp')
			return False

		rospy.sleep(1)
		current_pose = get_relative_coordinate("map", "hand_palm_link")

		try:
			pose_to_perform = geometry_msgs.msg.PoseStamped()
			pose_to_perform.position = pose_transformed.position
			pose_to_perform.header = current_pose.header
			pose_to_perform.header.stamp = rospy.Time.now()

			pose_to_perform.orientation = grasp_pose.orientation
			self.move_to(pose_to_perform, 'whole_body_light')
		except:
			rospy.logerr('fail to grasp')
			return False

		move_hand(0)

		return True

	def transform_Pose(self, pose_array):
		pose = geometry_msgs.msg.Pose()
		pose.position.x = pose_array.x
		pose.position.y = pose_array.y
		pose.position.z = pose_array.z
		return pose

	def rotate(self, pose, target):
		pose_goal = geometry_msgs.msg.Pose()
		pose_goal.pose.orientation = quaternion_from_euler(pose[0], pose[1], pose[2])
		group = moveit_commander.MoveGroupCommander(target)
		group.allow_replanning(True)
		group.set_pose_target(pose_goal)
		plan = group.go(wait=True)
		# Calling `stop()` ensures that there is no residual movement
		group.stop()
		# It is always good to clear your targets after planning with poses.
		# Note: there is no equivalent function for clear_joint_value_targets()
		group.clear_pose_targets()

	def move_arm(self, pose):
		pose_goal = geometry_msgs.msg.Pose()
		pose_goal.pose.orientation = quaternion_from_euler(pose[0], pose[1], pose[2])

	def attached_object(self):
		grasping_group = 'hand'
		touch_links = self.robot.get_link_names(group=grasping_group)
		self.scene.attach_box(eef_link, box_name, touch_links=touch_links)

	def create_cloud_indexed(self, object_pose, pointcloud_sub):
		points = []
		cam_source = []

		pts = list(pc2.read_points(pointcloud_sub, field_names=("x", "y", "z", "rgb"), skip_nans=True))
		for pt in pts:
			cam_source.append(Int64(0))

		for i in range(object_pose.x_min.data, object_pose.x_max.data):
			for j in range(object_pose.y_min.data, object_pose.y_max.data):
				points.append(Int64(i + (j*640)))

		self.listener.waitForTransform('/map', '/head_rgbd_sensor_rgb_frame', rospy.Time(), rospy.Duration(2.0))
		(trans,rot) = self.listener.lookupTransform('/map', '/head_rgbd_sensor_rgb_frame', rospy.Time(0))

		point_msg = Point()
		point_msg.x = trans[0]
		point_msg.y = trans[1]
		point_msg.z = trans[2]

		cloud_source_msg = CloudSources()
		cloud_source_msg.cloud = pointcloud_sub
		cloud_source_msg.camera_source = cam_source
		cloud_source_msg.view_points = [point_msg]

		cloud_index_msg = CloudIndexed()
		cloud_index_msg.cloud_sources = cloud_source_msg
		cloud_index_msg.indices = points

		return cloud_index_msg

	def estimate_pose(self, points, pointcloud_sub):
		return list(pc2.read_points(pointcloud_sub, field_names=("x", "y", "z"), skip_nans=True, uvs=points))

if __name__ == "__main__":

	grasping_node = Grasping()
	# Here you can switch between two mode:
	# 1. Continuous detection by ROS subscriber/callback (asynchronous)
	# 2. Synchronous detection via ROS Service (Server/Client-like)

	#grasping_node.continuous_node()
	grasping_node.service_node()

# Ros_goal cmd;
#     int nb_pos = mat_cmd.n_cols;
#     cmd.trajectory.points.resize(nb_pos);
#     cmd.trajectory.joint_names = jointnames;
#     for(int pos=0; pos<nb_pos; pos++)
#     {
#         cmd.trajectory.points[pos].time_from_start = ros::Duration(mat_cmd(0,pos)/2.5+1.0);// /1.5 augmente la vitesse
#         for(int j=0; j<rob.nb_joints(); j++)
#         {
#             cmd.trajectory.points[pos].positions.emplace_back((float)mat_cmd(j+1,pos));
#         }
#     }
#     return std::move(cmd);

# rostopic pub /hsrb/arm_trajectory_controller/follow_joint_trajectory/goal control_msgs/FollowJointTrajectoryActionGoal "header:
#   seq: 0
#   stamp:
#     secs: 0
#     nsecs: 0
#   frame_id: ''
# goal_id:
#   stamp:
#     secs: 0
#     nsecs: 0
#   id: ''
# goal:
#   trajectory:
#     header:
#       seq: 0
#       stamp:
#         secs: 0
#         nsecs: 0
#       frame_id: ''
#     joint_names:
#     - '"arm_lift_joint", "arm_flex_joint", "arm_roll_joint", "wrist_flex_joint", "wrist_roll_joint"'
#   goal_time_tolerance: {secs: 0, nsecs: 0}"  acceleration: 0.0}


"""
This file contains a function to decode a single synthetic tree
"""
import pandas as pd
import numpy as np
from syn_net.utils.data_utils import ReactionSet
from dgllife.model import load_pretrained
from syn_net.utils.predict_utils import tanimoto_similarity, load_modules_from_checkpoint, mol_fp
from syn_net.utils.predict_beam_utils import synthetic_tree_decoder

# define some constants (here, for the Hartenfeller-Button test set)
nbits = 4096
out_dim = 300
rxn_template = 'hb'
featurize = 'fp'
param_dir = 'hb_fp_2_4096'
ncpu = 16

# define model to use for molecular embedding
model_type = 'gin_supervised_contextpred'
device = 'cpu'
mol_embedder = load_pretrained(model_type).to(device)
mol_embedder.eval()

# load the purchasable building block embeddings
bb_emb = np.load('/pool001/whgao/data/synth_net/st_hb/enamine_us_emb_fp_256.npy')

# define path to the reaction templates and purchasable building blocks
path_to_reaction_file = '/pool001/whgao/data/synth_net/st_' + rxn_template + '/reactions_' + rxn_template + '.json.gz'
path_to_building_blocks = '/pool001/whgao/data/synth_net/st_' + rxn_template + '/enamine_us_matched.csv.gz'

# define paths to pretrained modules
param_path = '/home/whgao/scGen/synth_net/synth_net/params/' + param_dir + '/'
path_to_act = param_path + 'act.ckpt'
path_to_rt1 = param_path + 'rt1.ckpt'
path_to_rxn = param_path + 'rxn.ckpt'
path_to_rt2 = param_path + 'rt2.ckpt'

# load the purchasable building block SMILES to a dictionary
building_blocks = pd.read_csv(path_to_building_blocks, compression='gzip')['SMILES'].tolist()
bb_dict = {building_blocks[i]: i for i in range(len(building_blocks))}

# load the reaction templates as a ReactionSet object
rxn_set = ReactionSet()
rxn_set.load(path_to_reaction_file)
rxns = rxn_set.rxns

# load the pre-trained modules
act_net, rt1_net, rxn_net, rt2_net = load_modules_from_checkpoint(
    path_to_act=path_to_act,
    path_to_rt1=path_to_rt1,
    path_to_rxn=path_to_rxn,
    path_to_rt2=path_to_rt2,
    featurize=featurize,
    rxn_template=rxn_template,
    out_dim=out_dim,
    nbits=nbits,
    ncpu=ncpu,
)

def func(smi):
    """
    Generates the synthetic tree for the input moleular string.

    Args:
        smi (str): Molecule (SMILES) to decode.

    Returns:
        np.ndarray or None: State of the generated synthetic tree.
        float: The best score.
        SyntheticTree: The generated synthetic tree.
    """
    emb = mol_fp(smi)
    try:
        tree, action = synthetic_tree_decoder(emb, building_blocks, bb_dict, rxns, mol_embedder, act_net, rt1_net, rxn_net, rt2_net, bb_emb, beam_width=10, rxn_template=rxn_template, n_bits=nbits, max_step=15)
    except Exception as e:
        print(e)
        action = -1

    # tree, action = synthetic_tree_decoder(emb, building_blocks, bb_dict, rxns, mol_embedder, act_net, rt1_net, rxn_net, rt2_net, max_step=15)

    # import ipdb; ipdb.set_trace(context=9)
    # tree._print()
    # print(action)
    # print(np.max(oracle(tree.get_state())))
    # print()

    if action != 3:
        return None, 0, None
    else:
        scores = tanimoto_similarity(emb, tree.get_state())
        max_score_idx = np.where(scores == np.max(scores))[0][0]
        return tree.get_state()[max_score_idx], np.max(scores), tree


import dataclasses
import os
import json
import time

import torch
from torch import device as TorchDevice, nn
import numpy as np

import hifigan
import fregan
from conformer.Models import ModelConfig
from dataset import PreProcessConfig, TrainConfig
from model import FastSpeech2, ScheduledOptim, VoiceVoxScheduledOptim

from typing import TypedDict, Tuple, Union, Optional, List

from model import VoiceVoxVariancePredictor, VoiceVoxDecoder

@dataclasses.dataclass
class ArgsDict:
    restore_step: int
    preprocess_config: str
    model_config: str
    train_config: str

def get_model(
    args: ArgsDict,
    configs: Tuple[PreProcessConfig, ModelConfig, TrainConfig],
    device: TorchDevice,
    train: bool = False
) -> Tuple[FastSpeech2, Union[ScheduledOptim, None]]:
    preprocess_config, model_config, train_config = configs

    model = FastSpeech2(preprocess_config, model_config).to(device)
    if args.restore_step:
        ckpt_path = os.path.join(
            train_config["path"]["ckpt_path"],
            "{}.pth.tar".format(args.restore_step),
        )
        ckpt = torch.load(ckpt_path, map_location=device)
        model.load_state_dict(ckpt["model"])

    if train:
        scheduled_optim = ScheduledOptim(
            model, train_config, model_config, args.restore_step
        )
        if args.restore_step:
            scheduled_optim.load_state_dict(ckpt["optimizer"])
        model.train()
        return model, scheduled_optim

    model.eval()
    model.requires_grad_ = False
    return model, None

def get_vv_model(
    args: ArgsDict,
    configs: Tuple[PreProcessConfig, ModelConfig, TrainConfig],
    device: TorchDevice,
    train: bool = False
) -> Tuple[VoiceVoxVariancePredictor, VoiceVoxDecoder, Union[VoiceVoxScheduledOptim, None]]:
    preprocess_config, model_config, train_config = configs

    variance_model = VoiceVoxVariancePredictor(preprocess_config, model_config).to(device)
    decoder_model = VoiceVoxDecoder(preprocess_config, model_config).to(device)
    if args.restore_step:
        ckpt_path = os.path.join(
            train_config["path"]["ckpt_path"],
            "{}.pth.tar".format(args.restore_step),
        )
        ckpt = torch.load(ckpt_path, map_location=device)
        variance_model.load_state_dict(ckpt["variance_model"])
        decoder_model.load_state_dict(ckpt["decoder_model"])

    if train:
        scheduled_optim = VoiceVoxScheduledOptim(
            variance_model, decoder_model, train_config, model_config, args.restore_step
        )
        if args.restore_step:
            scheduled_optim.load_state_dict(ckpt["optimizer"])
        variance_model.train()
        decoder_model.train()
        return variance_model, decoder_model, scheduled_optim

    variance_model.eval()
    decoder_model.eval()
    variance_model.requires_grad_ = False
    decoder_model.requires_grad_ = False
    return variance_model, decoder_model, None

def get_param_num(model: nn.DataParallel) -> int:
    num_param = sum(param.numel() for param in model.parameters())
    return num_param

def get_vocoder(config: ModelConfig, device: TorchDevice) -> Union[hifigan.Generator, fregan.Generator]:
    name = config["vocoder"]["model"]
    speaker = config["vocoder"]["speaker"]

    if name == "HiFi-GAN" or name == "Fre-GAN":
        folder_name = "hifigan" if name == "HiFi-GAN" else "fregan"
        if name == "HiFi-GAN":
            with open(f"{folder_name}/config.json", "r") as f:
                config = json.load(f)
            config = hifigan.AttrDict(config)
            vocoder = hifigan.Generator(config)
        else:
            config = fregan.Config()
            vocoder = fregan.Generator(config)
        if speaker == "LJSpeech":
            ckpt = torch.load(f"{folder_name}/generator_LJSpeech.pth.tar", map_location=device)
        elif speaker == "universal":
            ckpt = torch.load(f"{folder_name}/generator_universal.pth.tar", map_location=device)
        vocoder.load_state_dict(ckpt["generator"])
        vocoder.eval()
        vocoder.remove_weight_norm()
        vocoder.to(device)
        return vocoder
    else:
        raise RuntimeError(f"{name} vocoder is not supported.")

def vocoder_infer(
    mels: torch.Tensor,
    vocoder: Union[hifigan.Generator, fregan.Generator],
    model_config: ModelConfig,
    preprocess_config: PreProcessConfig,
    lengths: Optional[int] = None
) -> List[np.ndarray]:
    name = model_config["vocoder"]["model"]
    with torch.no_grad():
        start = time.time()
        wavs = vocoder(mels).squeeze(1)
        print("RTF:", f"{time.time() - start}s")

    wavs = (
        wavs.cpu().numpy()
        * preprocess_config["preprocessing"]["audio"]["max_wav_value"]
    ).astype("int16")
    wavs: List[np.ndarray] = [wav for wav in wavs]

    for i in range(len(mels)):
        if lengths is not None:
            wavs[i] = wavs[i][: lengths[i]]

    return wavs


# The following code is largely borrowed from:
# https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail/blob/master/a2c_ppo_acktr/storage.py

from collections import namedtuple
# from multiprocessing import Pool

import numpy as np
import torch
from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler

def _flatten_helper(T, N, _tensor):
    return _tensor.view(T * N, *_tensor.size()[2:])

class RolloutStorage(object):

    def __init__(self, num_steps, num_processes, obs_shape, action_space,
                 rec_state_size, device):

        if action_space.__class__.__name__ == 'Discrete':
            self.n_actions = 1
            action_type = torch.long
        else:
            self.n_actions = action_space.shape[0]
            action_type = torch.float32

        self.obs = torch.zeros((num_steps + 1, num_processes, *obs_shape), device=device)
        self.rec_states = torch.zeros((num_steps + 1, num_processes,
                                      rec_state_size), device=device)
        self.rewards = torch.zeros((num_steps, num_processes), device=device)
        self.value_preds = torch.zeros((num_steps + 1, num_processes), device=device)
        self.returns = torch.zeros((num_steps + 1, num_processes), device=device)
        self.action_log_probs = torch.zeros((num_steps, num_processes), device=device)
        self.actions = torch.zeros((num_steps, num_processes, self.n_actions),
                                   dtype=action_type, device=device)
        self.masks = torch.ones((num_steps + 1, num_processes), device=device)

        self.num_steps = num_steps
        self.step = 0
        self.has_extras = False
        self.extras_size = None

    def to(self, device):
        self.obs = self.obs.to(device)
        self.rec_states = self.rec_states.to(device)
        self.rewards = self.rewards.to(device)
        self.value_preds = self.value_preds.to(device)
        self.returns = self.returns.to(device)
        self.action_log_probs = self.action_log_probs.to(device)
        self.actions = self.actions.to(device)
        self.masks = self.masks.to(device)
        if self.has_extras:
            self.extras = self.extras.to(device)
        return self

    def insert(self, obs, rec_states, actions, action_log_probs, value_preds,
               rewards, masks):
        self.obs[self.step + 1].copy_(obs)
        self.rec_states[self.step + 1].copy_(rec_states)
        self.actions[self.step].copy_(actions.view(-1, self.n_actions))
        self.action_log_probs[self.step].copy_(action_log_probs)
        self.value_preds[self.step].copy_(value_preds)
        self.rewards[self.step].copy_(rewards)
        self.masks[self.step + 1].copy_(masks)

        self.step = (self.step + 1) % self.num_steps

    def after_update(self):
        self.obs[0].copy_(self.obs[-1])
        self.rec_states[0].copy_(self.rec_states[-1])
        self.masks[0].copy_(self.masks[-1])
        if self.has_extras:
            self.extras[0].copy_(self.extras[-1])

    def compute_returns(self, next_value, use_gae, gamma, tau):
        if use_gae:
            self.value_preds[-1] = next_value
            gae = 0
            for step in reversed(range(self.rewards.size(0))):
                delta = self.rewards[step] + gamma \
                        * self.value_preds[step + 1] * self.masks[step + 1] \
                        - self.value_preds[step]
                gae = delta + gamma * tau * self.masks[step + 1] * gae
                self.returns[step] = gae + self.value_preds[step]
        else:
            self.returns[-1] = next_value
            for step in reversed(range(self.rewards.size(0))):
                self.returns[step] = self.returns[step + 1] * gamma \
                                     * self.masks[step + 1] + self.rewards[step]

    def feed_forward_generator(self, advantages, num_mini_batch):

        num_steps, num_processes = self.rewards.size()[0:2]
        batch_size = num_processes * num_steps
        mini_batch_size = batch_size // num_mini_batch
        assert batch_size >= num_mini_batch, (
            "PPO requires the number of processes ({}) "
            "* number of steps ({}) = {} "
            "to be greater than or equal to the number of PPO mini batches ({})."
            "".format(num_processes, num_steps, num_processes * num_steps,
                      num_mini_batch))

        sampler = BatchSampler(SubsetRandomSampler(range(batch_size)),
                               mini_batch_size, drop_last=False)

        for indices in sampler:
            yield {
                'obs': self.obs[:-1].view(-1, *self.obs.size()[2:])[indices],
                'rec_states': self.rec_states[:-1].view(-1,
                                                        self.rec_states.size(-1))[indices],
                'actions': self.actions.view(-1, self.n_actions)[indices],
                'value_preds': self.value_preds[:-1].view(-1)[indices],
                'returns': self.returns[:-1].view(-1)[indices],
                'masks': self.masks[:-1].view(-1)[indices],
                'old_action_log_probs': self.action_log_probs.view(-1)[indices],
                'adv_targ': advantages.view(-1)[indices],
                'extras': self.extras[:-1].view(-1, self.extras_size)[indices] \
                    if self.has_extras else None,
            }

    def recurrent_generator(self, advantages, num_mini_batch):

        num_processes = self.rewards.size(1)
        assert num_processes >= num_mini_batch, (
            "PPO requires the number of processes ({}) "
            "to be greater than or equal to the number of "
            "PPO mini batches ({}).".format(num_processes, num_mini_batch))
        num_envs_per_batch = num_processes // num_mini_batch
        perm = torch.randperm(num_processes)
        T, N = self.num_steps, num_envs_per_batch

        for start_ind in range(0, num_processes, num_envs_per_batch):

            obs = []
            rec_states = []
            actions = []
            value_preds = []
            returns = []
            masks = []
            old_action_log_probs = []
            adv_targ = []
            if self.has_extras:
                extras = []

            for offset in range(num_envs_per_batch):

                ind = perm[start_ind + offset]
                obs.append(self.obs[:-1, ind])
                rec_states.append(self.rec_states[0:1, ind])
                actions.append(self.actions[:, ind])
                value_preds.append(self.value_preds[:-1, ind])
                returns.append(self.returns[:-1, ind])
                masks.append(self.masks[:-1, ind])
                old_action_log_probs.append(self.action_log_probs[:, ind])
                adv_targ.append(advantages[:, ind])
                if self.has_extras:
                    extras.append(self.extras[:-1, ind])

            # These are all tensors of size (T, N, ...)
            obs = torch.stack(obs, 1)
            actions = torch.stack(actions, 1)
            value_preds = torch.stack(value_preds, 1)
            returns = torch.stack(returns, 1)
            masks = torch.stack(masks, 1)
            old_action_log_probs = torch.stack(old_action_log_probs, 1)
            adv_targ = torch.stack(adv_targ, 1)
            if self.has_extras:
                extras = torch.stack(extras, 1)

            yield {
                'obs': _flatten_helper(T, N, obs),
                'actions': _flatten_helper(T, N, actions),
                'value_preds': _flatten_helper(T, N, value_preds),
                'returns': _flatten_helper(T, N, returns),
                'masks': _flatten_helper(T, N, masks),
                'old_action_log_probs': _flatten_helper(T, N, old_action_log_probs),
                'adv_targ': _flatten_helper(T, N, adv_targ),
                'extras': _flatten_helper(T, N, extras) if self.has_extras else None,
                'rec_states': torch.stack(rec_states, 1).view(N, -1),
            }

class GlobalRolloutStorage(RolloutStorage):

    def __init__(self, num_steps, num_processes, obs_shape, action_space,
                 rec_state_size, extras_size, device):
        super(GlobalRolloutStorage, self).__init__(num_steps, num_processes,
                                                   obs_shape, action_space, rec_state_size, device)
        self.extras = torch.zeros((num_steps + 1, num_processes, extras_size),
                                  dtype=torch.long, device=device)
        self.has_extras = True
        self.extras_size = extras_size

    def insert(self, obs, rec_states, actions, action_log_probs, value_preds,
               rewards, masks, extras):
        self.extras[self.step + 1].copy_(extras)
        super(GlobalRolloutStorage, self).insert(obs, rec_states, actions,
                                                 action_log_probs, value_preds, rewards, masks)

Datapoint = namedtuple('Datapoint',
                       ('input', 'target'))

class FIFOMemory(object):

    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = []
        self.position = 0

    def push(self, *args):
        """Saves a datapoint."""
        if len(self.memory) < self.capacity:
            self.memory.append(None)

        self.memory[self.position] = Datapoint(*args)
        if self.position == 0:
            x = self.memory[0][0]
            y = self.memory[0][1]
            self.batch_in_sizes = {}
            self.n_inputs = len(x)
            for dim in range(len(x)):
                self.batch_in_sizes[dim] = x[dim].size()

            self.batch_out_sizes = {}
            self.n_outputs = len(y)
            for dim in range(len(y)):
                self.batch_out_sizes[dim] = y[dim].size()

        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        """Samples a batch"""

        batch = {}
        inputs = []
        outputs = []

        for dim in range(self.n_inputs):
            inputs.append(torch.cat(batch_size *
                                    [torch.zeros(
                                        self.batch_in_sizes[dim]
                                    ).unsqueeze(0)]))

        for dim in range(self.n_outputs):
            outputs.append(torch.cat(batch_size *
                                     [torch.zeros(
                                         self.batch_out_sizes[dim]
                                     ).unsqueeze(0)]))

        indices = np.random.choice(len(self.memory), batch_size, replace=False)

        count = 0
        for i in indices:
            x = self.memory[i][0]
            y = self.memory[i][1]

            for dim in range(len(x)):
                inputs[dim][count] = x[dim]

            for dim in range(len(y)):
                outputs[dim][count] = y[dim]

            count += 1

        return (inputs, outputs)

    # def sample_loader(self, batch_size, num_samples, num_processes, device=None):
    #     with Pool(processes=num_processes) as pool:
    #         # multiple_results caches the next num_processes results if sampled
    #         multiple_results = [pool.apply_async(self.sample, (batch_size,))
    #                                 for i in range(num_samples)]
    #         for res in multiple_results:
    #             if device is None:
    #                 yield res.get()
    #             else:
    #                 inputs, outputs = res.get()
    #                 inputs = [x.to(device) for x in inputs]
    #                 outputs = [x.to(device) for x in outputs]
    #                 yield (inputs, outputs)

    def __len__(self):
        return len(self.memory)


import copy
import numpy as np
import os
import multiprocessing as mp
import pickle
import pdb
import yaml
import matplotlib.pyplot as plt
from pyvista import imred, image, spectra, tv
from tools import plots, html
from astropy import units as u
from astropy.nddata import CCDData, StdDevUncertainty
import scipy.signal

ROOT = os.path.dirname(os.path.abspath(__file__)) + '/../../'

def all(ymlfile,display=None,plot=None,verbose=True,clobber=True,wclobber=None,
        groups='all',solve=False,htmlfile='index.html',threads=0) :
    """ Reduce full night(s) of data given input configuration file
    """

    # read input configuration file for reductions
    f=open(ymlfile,'r')
    d=yaml.load(f, Loader=yaml.FullLoader)
    f.close()

    if display is None : plt.ioff()
    if display and threads > 0 :
        raise ValueError('no multiprocessing with display!')

    if type(groups) is not list : groups = [groups]

    # loop over multiple groups in input file
    for group in d['groups'] :

        if 'skip' in group:
            if group['skip']  : continue
        if groups[0] != 'all' and group['name'] not in groups  : continue

        # clear displays if given
        if display is not None : display.clear()
        if plot is not None : plot.clf()

        # set up Reducer, Combiner, and output directory
        inst = group['inst']
        try : conf = group['conf']
        except : conf=''
        print('Instrument: {:s}'.format(inst))
        try : red = imred.Reducer(inst=inst,conf=conf,dir=group['rawdir'],
                                  verbose=verbose,nfowler=group['nfowler'])
        except KeyError : red = imred.Reducer(inst=group['inst'],conf=conf,
                                  dir=group['rawdir'],verbose=verbose)
        reddir = group['reddir']+'/'
        try: os.makedirs(reddir)
        except FileExistsError : pass
        if htmlfile is not None : fhtml = html.head(reddir+'/'+htmlfile)
        else : fhtml = None

        #create superbiases if biases given
        if 'biases' in group :
            sbias = mkcal(group['biases'],'bias',red,reddir,
                          clobber=clobber,display=display,html=fhtml)
        else:
            print('no bias frames given')
            sbias = None

        #create superdarks if darks given
        if 'darks' in group :
            sdark = mkcal(group['darks'],'dark',red,reddir,
                          clobber=clobber,display=display,sbias=sbias,
                          html=fhtml)
        else:
            print('no dark frames given')
            sdark = None

        #create superflats if darks given
        if 'flats' in group :
            sflat = mkcal(group['flats'],'flat',red,reddir,
                          clobber=clobber,display=display,
                          sbias=sbias,sdark=sdark,html=fhtml)
        else:
            print('no flat frames given')
            sflat = None

        # create wavecals if arcs given
        if 'arcs' in group :
            # existing trace template
            traces=pickle.load(open(ROOT+'/data/'+inst+'/'+inst+'_traces.pkl','rb'))

            if wclobber is None : wclobber = clobber
            wavedict={}
            wavecals=group['arcs']
            for wavecal in wavecals :
                print('create wavecal : {:s}'.format(wavecal['id']))
                # existing wavecal template
                waves=pickle.load(open(ROOT+'/data/'+inst+'/'+wavecal['wref']+'.pkl','rb'))
                if wclobber :
                    make = True
                else :
                    make = False
                    try: waves_all = pickle.load(open(reddir+wavecal['id']+'.pkl','rb'))
                    except FileNotFoundError : make=True
                if make :
                    # combine frames
                    try : superbias = sbias[wavecal['bias']]
                    except KeyError: superbias = None
                    arcs=red.sum(wavecal['frames'],return_list=True,
                                 bias=superbias, crbox=[5,1],
                                 display=display)

                    print('  extract wavecal')
                    # loop over channels
                    waves_all=[]
                    for arc,wave,trace in zip(arcs,waves,traces) :

                        # loop over windows
                        waves_channel=[]
                        for iwind,(wcal,wtrace) in enumerate(zip(wave,trace)) :
                            try:
                                delattr(wcal,'ax')
                                delattr(wcal,'fig')
                            except: pass
                            try : file = wavecal['file']
                            except KeyError : file = None
                            # extract and ID lines
                            if wavecal['wavecal_type'] == 'echelle' :
                                shift=wtrace.find(arc,plot=display)
                                arcec=wtrace.extract(arc,plot=display)
                                wcal.identify(spectrum=arcec, rad=3,
                                          display=display, plot=plot,file=file)
                            elif wavecal['wavecal_type'] == 'longslit' :
                                r0=wtrace.rows[0]
                                r1=wtrace.rows[1]
                                # 1d for inspection
                                wtrace.pix0 +=30
                                if red.transpose : arc=red.imtranspose(arc)
                                arcec=wtrace.extract(arc,plot=display,rad=20)
                                arcec.data = arcec.data - \
                                   scipy.signal.medfilt(arcec.data,kernel_size=[1,101])
                                wcal.identify(spectrum=arcec, rad=3, plot=plot,
                                              display=display,
                                              lags=range(-500,500),file=file)
                                wcal.fit()

                                print("doing 2D wavecal...")
                                arcec=wtrace.extract2d(arc)
                                print(" remove continuum and smooth in rows")
                                arcec.data=arcec.data - scipy.signal.medfilt(arcec.data,kernel_size=[1,101])
                                # smooth vertically for better S/N, then sample accordingly
                                image.smooth(arcec,[5,1])
                                wcal.identify(spectrum=arcec, rad=3,
                                              display=display, plot=plot,
                                              nskip=5,lags=range(-50,50))

                            wcal.fit()
                            if plot is not None :
                                delattr(wcal,'ax')
                                delattr(wcal,'fig')
                            waves_channel.append(wcal)
                        waves_all.append(waves_channel)
                        if display is not None : display.clear()
                    pickle.dump(waves_all,open(reddir+wavecal['id']+'.pkl','wb'))
                    if plot is not None : plot.clf()
                else :
                    print('  already made!')
                wavedict[wavecal['id']] = waves_all
        else :
            print('no wavecal frames given')
            w=None

        # reduce objects
        if 'objects' in group :
            objects = group['objects']
            if 'image' in objects :
                # images
                for obj in objects['image']:
                    if html is not None :
                        try : fhtml.write('

{:s}
\n'.format(obj['id']))
                        except KeyError : pass
                        fhtml.write('

\n')
                    try : superbias = sbias[obj['bias']]
                    except KeyError: superbias = None
                    try : superdark = sdark[obj['dark']]
                    except KeyError: superdark = None
                    try : superflat = sflat[obj['flat']]
                    except KeyError: superflat = None
                    try : trim = obj['trim']
                    except KeyError: trim = False

                    if threads > 0  :
                        # if multiprocessing, do all frames in this object in parallel
                        pars=[]
                        for id in obj['frames'] :
                            pars.append((red,id,superbias,superdark,superflat,red.scat,red.crbox,trim,solve,reddir))
                        pool = mp.Pool(threads)
                        output = pool.map_async(process_thread, pars).get()
                        pool.close()
                        pool.join()
                    else :
                        output = obj['frames']

                    for iframe,(id,frames) in enumerate(zip(obj['frames'],output)) :
                        if display is not None : display.tvclear()
                        if threads == 0 :
                            # if not multiprocessing, then do the reduction
                            frames=red.reduce(id,bias=superbias,
                                             dark=superdark,
                                             flat=superflat,
                                             scat=red.scat,
                                             crbox=red.crbox,
                                             trim=trim,
                                             solve=solve,
                                             return_list=True,
                                             display=display)
                            name=frames[0].header['FILE']
                            red.write(frames,reddir+name,overwrite=True,png=True)
                        else : name=frames[0].header['FILE']
                        if html is not None :
                            fhtml.write('
{:s}'.format(name))
                            for frame in frames :
                                fhtml.write((''+
                                '\n').format(name,name))
                    if html is not None : fhtml.write('
')
                if html is not None : html.tail(fhtml)
                plt.ion()
            elif 'extract1d' in objects :
                # 1D spectra
                traces=pickle.load(open(ROOT+'/data/'+inst+'/'+inst+'_traces.pkl','rb'))
                for obj in objects['extract1d'] :
                    try : superbias = sbias[obj['bias']]
                    except KeyError: superbias = None
                    try : superdark = sdark[obj['dark']]
                    except KeyError: superdark = None
                    try : superflat = sflat[obj['flat']]
                    except KeyError: superflat = None
                    waves_all = wavedict[obj['wavecal']]
                    try : trim = obj['trim']
                    except KeyError: trim = False

                    if obj['flat_type'] == '1d' :
                        print('extracting 1d flat')
                        ecflat=[]
                        for trace in traces :
                            tmp=[]
                            for wtrace in trace :
                                shift=wtrace.find(superflat,plot=display)
                                tmp.append(wtrace.extract(superflat,plot=display))
                                #tmp.append(wtrace.extract(superflat,plot=display,medfilt=101))
                            ecflat.append(tmp)
                        superflat = None

                    if threads > 0  :
                        # if multiprocessing, do all frames in this object in parallel
                        pars=[]
                        for id in obj['frames'] :
                            pars.append((red,id,superbias,superdark,superflat,red.scat,red.crbox,trim,solve,reddir))
                        pool = mp.Pool(threads)
                        output = pool.map_async(process_thread, pars).get()
                        pool.close()
                        pool.join()
                    else :
                        output = obj['frames']

                    # now frames
                    for iframe,(id,frames) in enumerate(zip(obj['frames'],output)) :
                        if display is not None : display.clear()
                        print("extracting object {}".format(id))
                        if threads == 0 :
                            frames=red.reduce(id,bias=superbias,dark=superdark,
                                          flat=superflat,scat=red.scat,
                                          return_list=True,crbox=red.crbox,display=display)
                        if 'skyframes' in obj :
                            id = obj['skyframes'][iframe]
                            skyframes=red.reduce(id,bias=superbias,dark=superdark,
                                                 flat=superflat,scat=red.scat,
                                                 return_list=True,crbox=red.crbox,display=display)
                            for iframe,(frame,skyframe) in enumerate(zip(frames,skyframes)) :
                                header = frame.header
                                frames[iframe]= frame.subtract(skyframe)
                                frames[iframe].header = header

                        # extraction radius
                        try : rad = obj['rad']
                        except KeyError : rad = None

                        # retrace?
                        try : retrace = obj['retrace']
                        except KeyError : retrace = True

                        # initialize plots
                        if plot is not None :
                            plot.clf()
                            ax=[]
                            ax.append(plot.add_subplot(2,1,1))
                            ax.append(plot.add_subplot(2,1,2,sharex=ax[0]))
                            plot.subplots_adjust(hspace=0.001)

                        # loop over channels
                        max=0
                        for ichannel,(frame,wave,trace) in enumerate(zip(frames,waves_all,traces)) :
                            # loop over windows
                            for iwind,(wcal,wtrace) in enumerate(zip(wave,trace)) :
                                if red.transpose : frame=red.imtranspose(frame)
                                tmptrace=copy.deepcopy(wtrace)
                                if retrace :
                                    print('  retracing ....')
                                    shift=tmptrace.retrace(frame,plot=display,thresh=10)
                                else : shift=tmptrace.find(frame,plot=display)
                                ec=tmptrace.extract(frame,plot=display,rad=rad)
                                w=wcal.wave(image=np.array(ec.data.shape))
                                if obj['flat_type'] == '1d' :
                                    header=ec.header
                                    ec=ec.divide(ecflat[ichannel][iwind])
                                    ec.header=header
                                if plot is not None :
                                    gd=np.where(ec.mask == False)
                                    med=np.median(ec.data[gd[0],gd[1]])
                                    max=np.max([max,scipy.signal.medfilt(ec.data,[1,101]).max()])
                                    for row in range(ec.data.shape[0]) :
                                        gd=np.where(ec.mask[row,:] == False)[0]
                                        plots.plotl(ax[0],w[row,gd],ec.data[row,gd],yr=[0,1.2*max],xt='Wavelength',yt='Flux')
                                        plots.plotl(ax[1],w[row,gd],ec.data[row,gd]/ec.uncertainty.array[row,gd],xt='Wavelength',yt='S/N')
                                    plot.suptitle(ec.header['OBJNAME'])
                                    plt.draw()
                            wnew = 10.**np.arange(3.5,4.0,5.5e-6)
                            comb=wcal.scomb(ec,wnew,average=True,usemask=True)
                            if plot is not None :
                                plots.plotl(ax[0],wnew,comb.data,color='k')
                                plots.plotl(ax[1],wnew,comb.data/comb.uncertainty.array,color='k')
                                plt.draw()
                                plot.canvas.draw_idle()
                                plt.pause(0.1)
                                input("  hit a key to continue")
                            out = spectra.SpecData(ec,wave=w)
                            out.write(reddir+ec.header['FILE'].replace('.fits','.ec.fits'),png=True)
                            out = spectra.SpecData(comb,wave=wnew)
                            out.write(reddir+comb.header['FILE'],png=True)
                        if html is not None :
                            for frame in frames :
                                name=frames[0].header['FILE']
                                fhtml.write('
{:s}'.format(name))
                                fhtml.write((''+
                                '\n').format(name,name))
                                fhtml.write((''+
                                '\n').format(name,name.replace('.fits','.ec.fits').replace('.fits','.ec.fits')))
            elif 'extract2d' in objects :
                # 2D spectra
                print('extract2d')
#    traces=pickle.load(open(group['inst']+'_trace.pkl','rb'))
#    if type(traces) is not list : traces = [traces]
#    pdb.set_trace()
#    for id in group['objects']['extract2d']['frames'] :
#        frames=red.reduce(id,bias=sbias,dark=sdark,flat=sflat,scat=red.scat,return_list=True)
#        for frame,trace in zip(frames,traces) :
#            if red.transpose : frame=red.imtranspose(frame)
#            out=trace.extract2d(frame,plot=t)

def mkcal(cals,caltype,reducer,reddir,sbias=None,sdark=None,clobber=False,
          html=None,**kwargs) :
    """ Make calibration frames given input lists

        Args :
            cals : list of different sets of given calibration type, as dictionaries
            caltype : gives caltype, of 'bias', 'dark', 'flat'
            reddir : directory for cal frames
            clobber= : set to True to force construction even if cal frame already exists
    """

    # we will loop over (possibly) multiple individual cal products of this type
    # These may or may not be combined, depending on "use" tag
    outcal={}
    if html is not None :
        html.write('

{:s}

\n'.format(caltype))
    for cal in cals :
        calname = cal['id']
        try : superbias = sbias[cal['bias']]
        except KeyError: superbias = None
        try : superdark = sdark[cal['dark']]
        except KeyError: superdark = None
        try :
            print('create {:s} : {:s}'.format(caltype,calname))
            # if not clobber, try to read existing frames
            if clobber :
                make=True
            else :
                make=False
                if len(reducer.channels)==1 :
                    try : scal= CCDData.read(reddir+calname+'.fits',
                                             unit=u.dimensionless_unscaled)
                    except FileNotFoundError : make=True
                else :
                    scal=[]
                    for channel in reducer.channels :
                        try : scal.append(
                                CCDData.read(reddir+calname+'_'+channel+'.fits',
                                             unit=u.dimensionless_unscaled))
                        except FileNotFoundError : make=True
            if make :
                try :
                    # see if we are requested to make product from previous products by seeing if dictionary entries exist for frames
                    scal=[]
                    tot=[]
                    for frame in cal['frames'] :
                        out = outcal[frame]
                        if type(out) is not list : out=[out]
                        for i in range(len(out)) :
                            print('combining: {:s}'.format(frame))
                            try:
                                scal[i] = scal[i].add(out[i].multiply(out[i].header['MEANNORM']))
                                tot[i]+=out[i].header['MEANNORM']
                            except:
                                scal.append( copy.deepcopy(out[i].multiply(out[i].header['MEANNORM'])) )
                                tot.append(out[i].header['MEANNORM'])
                    for i in range(len(scal)) : scal[i] = scal[i].divide(tot[i])
                    if len(scal) is 1 : scal= scal[0]
                except :
                    # make calibration product from raw data frames
                    if caltype == 'bias' :
                        scal = reducer.mkbias(cal['frames'],**kwargs)
                    elif caltype == 'dark' :
                        scal = reducer.mkdark(cal['frames'],bias=superbias,**kwargs)
                    elif caltype == 'flat' :
                        scal = reducer.mkflat(cal['frames'],bias=superbias,dark=superdark,**kwargs)
                        try:
                            if cal['specflat'] : scal = reducer.mkspecflat(scal)
                        except: pass
                        reducer.scatter(scal,scat=reducer.scat,**kwargs)
                reducer.write(scal,reddir+calname+'.fits',overwrite=True,png=True)
                if html is not None :
                    html.write(
                      '
{:s}\n'.
                      format(calname,calname,calname))
            else : print('  already made!')
        except RuntimeError :
            print('error processing {:s} frames'.format(caltype))
        except KeyError:
            print('no {:s} frames given'.format(caltype))
            scal=None

        outcal[calname] = scal

    if html is not None : html.write('
\n')

    return outcal

def process_thread(pars) :

    red,id,superbias,superdark,superflat,scat,crbox,trim,solve,reddir = pars
    frames= red.reduce(id, bias=superbias,
                          dark=superdark,
                          flat=superflat,
                          scat=red.scat,
                          crbox=red.crbox,
                          trim=trim,
                          solve=solve,
                          return_list=True,display=None)
    name=frames[0].header['FILE']
    red.write(frames,reddir+name,overwrite=True,png=True)
    return frames