# coding: utf-8
import numpy as np
import random
import tensorflow as tf
import logging
import imageio
import read_data
from keras.utils import to_categorical

# from data_generator import DataGenerator
from place_pick_mil import MIL
# from evaluation.eval_reach import evaluate_vision_reach
# from evaluation.eval_push import evaluate_push
from tensorflow.python.platform import flags
import os

from functools import reduce
from operator import mul
def get_num_params():
    nums=0
    for variable in tf.trainable_variables():
        shape= variable.get_shape()
        nums+=reduce(mul, [dim.value for dim in shape], 1)
    return nums

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
FLAGS = flags.FLAGS
LOGGER = logging.getLogger(__name__)

## Dataset/method options
flags.DEFINE_string('experiment', 'pick_place', 'sim_vision_reach or sim_push')
flags.DEFINE_string('data_path', './pick_dataset_origin/human_robot_dataset/',
                    'path to the directory where demo files that containing robot states and actions are stored')
flags.DEFINE_string('demo_gif_dir', 'data', 'path to the videos of demonstrations')
flags.DEFINE_string('gif_prefix', 'object', 'prefix of the video directory for each task, e.g. object_0 for task 0')
flags.DEFINE_integer('im_width', 264,
                     'width of the images in the demo videos,  125 for sim_push, and 80 for sim_vision_reach')
flags.DEFINE_integer('im_height', 196,
                     'height of the images in the demo videos, 125 for sim_push, and 64 for sim_vision_reach')
flags.DEFINE_integer('num_channels', 3, 'number of channels of the images in the demo videos')
flags.DEFINE_integer('T', 3, 'time horizon of the demo videos, 50 for reach, 100 for push')
flags.DEFINE_bool('hsv', False, 'convert the image to HSV format')
flags.DEFINE_bool('use_noisy_demos', False, 'use noisy demonstrations or not (for domain shift)')
flags.DEFINE_string('noisy_demo_gif_dir', None, 'path to the videos of noisy demonstrations')
flags.DEFINE_string('noisy_demo_file', None,
                    'path to the directory where noisy demo files that containing robot states and actions are stored')
flags.DEFINE_bool('no_action', True, 'do not include actions in the demonstrations for inner update')
flags.DEFINE_bool('no_state', False, 'do not include states in the demonstrations during training')
flags.DEFINE_bool('no_final_eept', False, 'do not include final ee pos in the demonstrations for inner update')
flags.DEFINE_bool('zero_state', True,
                  'zero-out states (meta-learn state) in the demonstrations for inner update (used in the paper with video-only demos)')
flags.DEFINE_bool('two_arms', False, 'use two-arm structure when state is zeroed-out')
flags.DEFINE_integer('training_set_size', -1, 'size of the training set, 1500 for sim_reach, 693 for sim push, anzero_stated \
                                                -1 for all data except those in validation set')
flags.DEFINE_integer('val_set_size', 150, 'size of the training set, 150 for sim_reach and 76 for sim push')

## Training options
flags.DEFINE_integer('metatrain_iterations', 30000,  'number of metatraining iterations.')  # 30k for pushing, 50k for reaching and placing
flags.DEFINE_integer('meta_batch_size', 16,  'number of tasks sampled per meta-update')  # 15 for reaching, 15 for pushing, 12 for placing
flags.DEFINE_integer('meta_test_batch_size', 1,  'number of tasks sampled per meta-update')  # 15 for reaching, 15 for pushing, 12 for placing
flags.DEFINE_float('meta_lr', 1e-4, 'the base learning rate of the generator')
flags.DEFINE_integer('update_batch_size', 1,
                     'number of examples used for inner gradient update (K for K-shot learning).')
flags.DEFINE_float('train_update_lr', 1e-4,
                   'step size alpha for inner gradient update.')  # 0.001 for reaching, 0.01 for pushing and placing
flags.DEFINE_integer('num_updates', 1, 'number of inner gradient updates during training.')  # 5 for placing
flags.DEFINE_bool('clip', True, 'use gradient clipping for fast gradient')
flags.DEFINE_float('clip_max', 100.0, 'maximum clipping value for fast gradient')
flags.DEFINE_float('clip_min', -100.0, 'minimum clipping value for fast gradient')
# flags.DEFINE_float('clip_max', 20.0, 'maximum clipping value for fast gradient')
# flags.DEFINE_float('clip_min', -20.0, 'minimum clipping value for fast gradient')
flags.DEFINE_bool('fc_bt', True, 'use bias transformation for the first fc layer')
flags.DEFINE_bool('all_fc_bt', False, 'use bias transformation for all fc layers')
flags.DEFINE_bool('conv_bt', False, 'use bias transformation for the first conv layer, N/A for using pretraining')
flags.DEFINE_integer('bt_dim', 10, 'the dimension of bias transformation for FC layers')
flags.DEFINE_string('pretrain_weight_path', 'N/A', 'path to pretrained weights')
flags.DEFINE_bool('train_pretrain_conv1', False, 'whether to finetune the pretrained weights')
flags.DEFINE_bool('two_head', True, 'use two-head architecture')
flags.DEFINE_bool('learn_final_eept', False, 'learn an auxiliary loss for predicting final end-effector pose')
flags.DEFINE_bool('learn_final_eept_whole_traj', False, 'learn an auxiliary loss for predicting final end-effector pose \
                                                         by passing the whole trajectory of eepts (used for video-only models)')
flags.DEFINE_bool('stopgrad_final_eept', True,
                  'stop the gradient when concatenate the predicted final eept with the feature points')
flags.DEFINE_integer('final_eept_min', 6, 'first index of the final eept in the action array')
flags.DEFINE_integer('final_eept_max', 8, 'last index of the final eept in the action array')
flags.DEFINE_float('final_eept_loss_eps', 0.1, 'the coefficient of the auxiliary loss')
flags.DEFINE_float('act_loss_eps', 1.0, 'the coefficient of the action loss')
flags.DEFINE_float('loss_multiplier', 100.0,
                   'the constant multiplied with the loss value, 100 for reach and 50 for push')
flags.DEFINE_bool('use_l1_l2_loss', False, 'use a loss with combination of l1 and l2')
flags.DEFINE_float('l2_eps', 0.01, 'coeffcient of l2 loss')
flags.DEFINE_bool('shuffle_val', False, 'whether to choose the validation set via shuffling or not')

## Model options
flags.DEFINE_integer('random_seed', 0, 'random seed for training')
flags.DEFINE_bool('fp', True, 'use spatial soft-argmax or not')
flags.DEFINE_string('norm', 'layer_norm', 'batch_norm, layer_norm, or None')
flags.DEFINE_bool('dropout', False, 'use dropout for fc layers or not')
flags.DEFINE_float('keep_prob', 0.5, 'keep probability for dropout')
flags.DEFINE_integer('num_filters', 64,
                     'number of filters for conv nets -- 64 for placing, 16 for pushing, 40 for reaching.')
flags.DEFINE_integer('filter_size', 3, 'filter size for conv nets -- 3 for placing, 5 for pushing, 3 for reaching.')
flags.DEFINE_integer('num_conv_layers', 5, 'number of conv layers -- 5 for placing, 4 for pushing, 3 for reaching.')
flags.DEFINE_integer('num_strides', 3,
                     'number of conv layers with strided filters -- 3 for placing, 4 for pushing, 3 for reaching.')
flags.DEFINE_bool('conv', True, 'whether or not to use a convolutional network, only applicable in some cases')
flags.DEFINE_integer('num_fc_layers', 3, 'number of fully-connected layers')
flags.DEFINE_integer('layer_size', 200, 'hidden dimension of fully-connected layers')
flags.DEFINE_bool('temporal_conv_2_head', True,
                  'whether or not to use temporal convolutions for the two-head architecture in video-only setting.')
flags.DEFINE_bool('temporal_conv_2_head_ee', False, 'whether or not to use temporal convolutions for the two-head architecture in video-only setting \
                for predicting the ee pose.')
flags.DEFINE_integer('temporal_filter_size', 10, 'filter size for temporal convolution')
flags.DEFINE_integer('temporal_num_filters', 32, 'number of filters for temporal convolution')
flags.DEFINE_integer('temporal_num_filters_ee', 32, 'number of filters for temporal convolution for ee pose prediction')
flags.DEFINE_integer('temporal_num_layers', 3, 'number of layers for temporal convolution for ee pose prediction')
flags.DEFINE_integer('temporal_num_layers_ee', 3, 'number of layers for temporal convolution for ee pose prediction')
flags.DEFINE_string('init', 'xavier', 'initializer for conv weights. Choose among random, xavier, and he')
flags.DEFINE_bool('max_pool', False, 'Whether or not to use max pooling rather than strided convolutions')
flags.DEFINE_bool('stop_grad', False, 'if True, do not use second derivatives in meta-optimization (for axis_angle)')

## Logging, saving, and testing options
flags.DEFINE_bool('log', True, 'if false, do not log summaries, for debugging code.')
flags.DEFINE_string('save_dir', './atmaml_pick_logs', 'directory for summaries and checkpoints.')
# flags.DEFINE_string('save_dir', './amaml_human_pick_logs', 'directory for summaries and checkpoints.')
# flags.DEFINE_bool('resume', True, 'resume training if there is a model available')
flags.DEFINE_bool('resume', False, 'resume training if there is a model available')
flags.DEFINE_bool('train', True, 'True to train, False to test.')
flags.DEFINE_integer('restore_iter', -1, 'iteration to load model (-1 for latest model)')
# flags.DEFINE_integer('restore_iter', 20000, 'iteration to load model (-1 for latest model)')
flags.DEFINE_integer('train_update_batch_size', -1, 'number of examples used for gradient update during training \
                    (use if you want to test with a different number).')
flags.DEFINE_integer('test_update_batch_size', 1, 'number of demos used during test time')
flags.DEFINE_float('gpu_memory_fraction', 0.8, 'fraction of memory used in gpu')
flags.DEFINE_bool('record_gifs', True, 'record gifs during evaluation')
flags.DEFINE_integer('color_num', 4, '')
flags.DEFINE_integer('object_num', 4, '')
flags.DEFINE_integer('train_task_num', 6, '')
flags.DEFINE_integer('task_num', 8, '')
flags.DEFINE_integer('demo_num', 5, '')
flags.DEFINE_integer('index_range', 20, '')
flags.DEFINE_integer('index_train_range', 20, '')
# flags.DEFINE_string('demo_type', 'robot', 'robot or human')
flags.DEFINE_string('demo_type', 'human', 'robot or human')
flags.DEFINE_string('extra_type', 'robot', 'robot or human') #opposite to demo_type
flags.DEFINE_string('compare_type', 'robot', 'robot or human')
flags.DEFINE_string('target_type', 'robot', '')
# flags.DEFINE_float('weight_xy', 0.999, '')
# flags.DEFINE_float('weight_z', 0.001, '')
# flags.DEFINE_float('weight_rxyz', 0.001, '')
flags.DEFINE_integer('clip_action_size', 2, 'size of embedding')
flags.DEFINE_integer('action_size', 2, 'size of embedding')
flags.DEFINE_integer('state_size', 6, 'size of embedding')
flags.DEFINE_integer('output_data', 6, '')
# flags.DEFINE_float('margin', 1.0, 'margin of loss')
# flags.DEFINE_float('pre_margin', 1.0, 'margin of loss')
flags.DEFINE_float('pre_margin_clip', 1.0, '')
flags.DEFINE_float('pre_margin', 2.0, 'pre_margin of loss')
flags.DEFINE_float('margin', 2.0, 'margin of loss')
flags.DEFINE_float('pre_margin_coefficient', 10.0, 'margin of pre-contrastive conloss')
flags.DEFINE_float('margin_coefficient', 10.0, 'margin of post-contrastive loss')
flags.DEFINE_float('pre_tar_coefficient', 0.0, '')
flags.DEFINE_float('post_tar_coefficient', 0.0, '')
flags.DEFINE_bool('norm_pre_contrastive', True, 'norm for pre_task contrastive ')
flags.DEFINE_bool('norm_post_contrastive', True, 'norm for post_task contrastive ')
flags.DEFINE_bool('all_frame_contrastive', True, '')
flags.DEFINE_bool('cos_contrastive_loss', True, 'True for cos loss, False for kl loss ')
flags.DEFINE_bool('pre_contrastive', False, 'task contrastive for pre-update ')
flags.DEFINE_bool('post_contrastive', False, 'task contrastive for post-update ')
# flags.DEFINE_bool('amaml', False, 'true for amaml')
flags.DEFINE_bool('amaml', True, 'true for amaml')
flags.DEFINE_bool('amaml_extra_domain', True, 'true for extra_domain')
flags.DEFINE_bool('cross_domain', True, 'use human and robot ')
flags.DEFINE_bool('cross_adapt_domain', True, 'adpat human/robot demos')
flags.DEFINE_bool('zero_robot_domain', True, 'adpat human/robot demos')
flags.DEFINE_bool('random_adapt_domain', True, 'randomly adpat human/robot demos')
flags.DEFINE_bool('random_single_domain', True, 'randomly adpat human/robot demos')
flags.DEFINE_integer('random_single_domain_margin', 50, 'random single domain margin')
flags.DEFINE_integer('embed_size', 0, 'size of embedding')
flags.DEFINE_bool('tar_mil', False, 'with target loss')
flags.DEFINE_bool('record_loss', False, 'true for amaml')
flags.DEFINE_bool('contrastive_fc', False, '')
# flags.DEFINE_bool('amaml', False, 'true for amaml')

def generate_data(if_train=True):
    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size

    color_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.color_num
    print('color_list', color_list)

    object_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.object_num
    print('object_list', object_list)
    if if_train:
        task_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.train_task_num
    else:
        task_list = np.random.randint(FLAGS.train_task_num, FLAGS.task_num, size=batch_size)
    print('task_list', task_list)

    demo_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.demo_num
    print('demo_list', demo_list)

    target_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.demo_num
    print('target_list', target_list)

    obsas = []
    obsbs = []
    stateas = []
    statebs = []
    actionas = []
    actionbs = []

    color_num = ['color_blue', 'color_green', 'color_orange', 'color_yellow']
    # color_num = ['color_blue', 'color_green', 'color_orange']
    object_num = ['object_type_animal', 'object_type_car', 'object_type_dinosaur', 'object_type_tool']

    for element in range(0, batch_size):
        demo_path = '%s/%s/%s/%s/task_%d/demo_%d' % (
            FLAGS.data_path, color_num[color_list[element]], object_num[object_list[element]], FLAGS.demo_type,
            task_list[element], demo_list[element])
        target_path = '%s/%s/%s/%s/task_%d/demo_%d' % (
            FLAGS.data_path, color_num[color_list[element]], object_num[object_list[element]], FLAGS.target_type,
            task_list[element], target_list[element])

        print('demo_path', demo_path)
        print('target_path', target_path)

        index = np.random.randint(0, FLAGS.index_train_range)

        # if if_train:
        #     index = np.random.randint(0, FLAGS.index_train_range)
        # else:
        #     index = np.random.randint(FLAGS.index_train_range, FLAGS.index_range)

        if FLAGS.demo_type == 'robot':
            obsa, statea, actiona = read_data.Read_Robot_Data2(demo_path, FLAGS.T, index)
        elif FLAGS.demo_type == 'human':
            obsa, statea, actiona = read_data.Read_Human_Data2(demo_path, FLAGS.T, index)
        obsb, stateb, actionb = read_data.Read_Robot_Data2(target_path, FLAGS.T, index)

        obsas.append(obsa)
        obsbs.append(obsb)
        stateas.append(statea)
        statebs.append(stateb)
        actionas.append(actiona)
        actionbs.append(actionb)


    obsas = np.reshape(obsas,  [batch_size, FLAGS.T, FLAGS.im_width * FLAGS.im_height * FLAGS.num_channels])
    obsbs = np.reshape(obsbs,    [batch_size, FLAGS.T, FLAGS.im_width * FLAGS.im_height * FLAGS.num_channels])

    actionas = np.reshape(actionas, [batch_size, FLAGS.T, FLAGS.output_data])
    actionbs = np.reshape(actionbs, [batch_size, FLAGS.T, FLAGS.output_data])

    stateas = np.zeros([batch_size, FLAGS.T, FLAGS.output_data])
    statebs = np.zeros([batch_size, FLAGS.T, FLAGS.output_data])

    return obsas, obsbs, actionas, actionbs, stateas, statebs


def generate_atmaml_data(if_train=True):
    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size

    color_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.color_num
    # print('color_list', color_list)
    compare_color_list = (color_list + np.random.randint(1, FLAGS.color_num-1)) %  FLAGS.color_num
    # print('compare_color_list', compare_color_list)

    object_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.object_num
    # print('object_list', object_list)
    compare_object_list = (object_list + np.random.randint(1, FLAGS.object_num-1)) % FLAGS.object_num
    # print('compare_object_list', compare_object_list)

    if if_train:
        task_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.train_task_num
        compare_task_list = (task_list + np.random.randint(1, FLAGS.train_task_num)) % FLAGS.train_task_num
        # compare_task_list = (task_list + np.random.randint(1, FLAGS.task_num - 1)) % FLAGS.task_num
    else:
        task_list = np.random.randint(FLAGS.train_task_num, FLAGS.task_num, size=batch_size)
        compare_task_list = np.random.randint(FLAGS.train_task_num, FLAGS.task_num,  size=batch_size)

    # print('task_list', task_list)
    # print('compare_task_list', compare_task_list)


    demo_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.demo_num
    # print('demo_list', demo_list)
    target_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.demo_num
    # print('target_list', target_list)
    compare_demo_list = (task_list +  np.random.randint(1, FLAGS.demo_num-1)) %  FLAGS.demo_num
    # print('compare_demo_list', compare_demo_list)



    obsas = []
    obsbs = []
    obscs = []
    extra_obses = []

    stateas = []
    statebs = []
    statecs = []
    extra_states = []

    actionas = []
    actionbs = []
    actioncs = []
    extra_actions = []

    labelabs = []
    labelcs = []


    color_num = ['color_blue', 'color_green', 'color_orange', 'color_yellow']
    # color_num = ['color_blue', 'color_green', 'color_orange']
    object_num = ['object_type_animal', 'object_type_car', 'object_type_dinosaur', 'object_type_tool']

    for element in range(0, batch_size):

        demo_path = '%s/%s/%s/%s/task_%d/demo_%d' % (
            FLAGS.data_path, color_num[color_list[element]], object_num[object_list[element]], FLAGS.demo_type,
            task_list[element], demo_list[element])

        target_path = '%s/%s/%s/%s/task_%d/demo_%d' % (
            FLAGS.data_path, color_num[color_list[element]], object_num[object_list[element]], FLAGS.target_type,
            task_list[element], target_list[element])

        compare_path = '%s/%s/%s/%s/task_%d/demo_%d' % (
            FLAGS.data_path, color_num[compare_color_list[element]], object_num[compare_object_list[element]],
            FLAGS.compare_type, compare_task_list[element], compare_demo_list[element])

        if FLAGS.cross_domain:
                extra_path = '%s/%s/%s/%s/task_%d/demo_%d' % (FLAGS.data_path, color_num[color_list[element]], object_num[object_list[element]],
                                                              FLAGS.extra_type, task_list[element], demo_list[element])
        if  element==0:
            if FLAGS.cross_domain:
                print('extra_path', extra_path)
            print('demo_path', demo_path)
            print('target_path', target_path)
            print('compare_path', compare_path)


        index = np.random.randint(0, FLAGS.index_train_range)

        # if if_train:
        #     index = np.random.randint(0, FLAGS.index_train_range)
        # else:
        #     index = np.random.randint(FLAGS.index_train_range, FLAGS.index_range)

        if FLAGS.demo_type == 'robot':
            obsa, statea, actiona = read_data.Read_Robot_Data2(demo_path, FLAGS.T, index)
        else :
            obsa, statea, actiona = read_data.Read_Human_Data2(demo_path, FLAGS.T, index)

        if FLAGS.compare_type == 'robot':
            obsc, statec, actionc = read_data.Read_Robot_Data2(compare_path, FLAGS.T, index)
        else:
            obsc, statec, actionc = read_data.Read_Human_Data2(compare_path, FLAGS.T, index)

        if FLAGS.cross_domain:
            if FLAGS.extra_type == 'robot':
                extra_obs, extra_state, extra_action = read_data.Read_Robot_Data2(extra_path, FLAGS.T, index)
            else:
                extra_obs, extra_state, extra_action = read_data.Read_Human_Data2(extra_path, FLAGS.T, index)

        obsb, stateb, actionb = read_data.Read_Robot_Data2(target_path, FLAGS.T, index)
        if FLAGS.tar_mil:
            labelab  = to_categorical(color_list[element] * FLAGS.index_train_range + index, FLAGS.color_num*FLAGS.index_train_range)
            labelc = to_categorical(compare_color_list[element] * FLAGS.index_train_range + index, FLAGS.color_num*FLAGS.index_train_range)
            labelabs.append(labelab)
            labelcs.append(labelc)
            # print('element', element, 'labela and labelb', labelab, 'labelc', labelc)

        # print('----------------------------------------------------------------')

        obsas.append(obsa)
        obsbs.append(obsb)
        obscs.append(obsc)
        stateas.append(statea)
        statebs.append(stateb)
        statecs.append(statec)
        actionas.append(actiona)
        actionbs.append(actionb)
        actioncs.append(actionc)
        if FLAGS.cross_domain:
            extra_obses.append(extra_obs)
            extra_states.append(extra_state)
            extra_actions.append(extra_action)


    obsas = np.reshape(obsas,  [batch_size, FLAGS.T, FLAGS.im_width * FLAGS.im_height * FLAGS.num_channels])
    obsbs = np.reshape(obsbs,    [batch_size, FLAGS.T, FLAGS.im_width * FLAGS.im_height * FLAGS.num_channels])
    obscs = np.reshape(obscs, [batch_size, FLAGS.T, FLAGS.im_width * FLAGS.im_height * FLAGS.num_channels])


    actionas = np.reshape(actionas, [batch_size, FLAGS.T, FLAGS.output_data])
    actionbs = np.reshape(actionbs, [batch_size, FLAGS.T, FLAGS.output_data])
    actioncs = np.reshape(actioncs, [batch_size, FLAGS.T, FLAGS.output_data])


    stateas = np.zeros([batch_size, FLAGS.T, FLAGS.output_data])
    statebs = np.zeros([batch_size, FLAGS.T, FLAGS.output_data])
    statecs = np.zeros([batch_size, FLAGS.T, FLAGS.output_data])


    if FLAGS.tar_mil:
        labelabs = np.reshape(labelabs, [batch_size, 1, -1])
        labelcs = np.reshape(labelcs, [batch_size, 1, -1])
        # print('labelabs', labelabs.shape, 'labelcs', labelcs)
    else:
        labelabs = np.zeros([batch_size, 1, FLAGS.color_num*FLAGS.index_train_range])
        labelcs = np.zeros( [batch_size, 1, FLAGS.color_num*FLAGS.index_train_range])

    if FLAGS.cross_domain:
        extra_obses = np.reshape(extra_obses, [batch_size, FLAGS.T, FLAGS.im_width * FLAGS.im_height * FLAGS.num_channels])
        extra_actions = np.reshape(extra_actions, [batch_size, FLAGS.T, FLAGS.output_data])
        extra_states = np.zeros([batch_size, FLAGS.T, FLAGS.output_data])

        return obsas, obsbs, obscs, extra_obses, \
               actionas, actionbs, actioncs, extra_actions, \
               stateas, statebs, statecs, extra_states, \
               labelabs, labelcs
    else:
        return obsas, obsbs, obscs, actionas, actionbs, actioncs, stateas, statebs, statecs, labelabs, labelcs

def generate_single_place_data(obsas, obsbs, actionas, actionbs, stateas, statebs, if_train=True):

    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size

    # print('obsas, obsbs, actionas, actionbs, stateas, statebs', obsas.shape, obsbs.shape, actionas.shape, actionbs.shape, stateas.shape, statebs.shape)
    obsbs = np.reshape(obsbs[:, 0, :], [batch_size , 1, -1])
    actionbs = np.reshape(actionbs[:, 1, :], [batch_size , 1, -1])
    statebs = np.reshape(statebs[:, -1, :], [batch_size , 1, -1])
    # print('obsas, obsbs, actionas, actionbs, stateas, statebs of generate_single_place_data', obsas.shape, obsbs.shape, actionas.shape, actionbs.shape,
    # stateas.shape, statebs.shape)
    return obsas, obsbs, actionas, actionbs, stateas, statebs


def generate_single_atmaml_place_data(obsbs, actionbs, statebs, if_train=True):

    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size



    # print('obsas, obsbs, actionas, actionbs, stateas, statebs', obsbs.shape, actionbs.shape, statebs.shape)
    obsbs = np.reshape(obsbs[:, 0, :], [batch_size , 1, -1])
    actionbs = np.reshape(actionbs[:, 1, :], [batch_size , 1, -1])
    statebs = np.reshape(statebs[:, -1, :], [batch_size , 1, -1])

    return obsbs, actionbs, statebs


def handle_action(actionas, actionbs, stateas, statebs, if_train=True):

    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size

    actionas = actionas[:, :, :FLAGS.clip_action_size]
    # print('actionas is', actionas.shape)
    actionas = np.reshape(actionas[:, 1, :], [batch_size, 1, -1])
    actionas = np.tile(actionas, (1, 3, 1))

    actionbs = actionbs[:, :, :FLAGS.clip_action_size]
    actionbs = np.reshape(actionbs[:, 1, :], [batch_size, 1, -1])
    actionbs = np.tile(actionbs, (1, 3, 1))


    stateas = stateas[:, :, :FLAGS.state_size]
    statebs = statebs[:, :, :FLAGS.state_size]

    return actionas, actionbs, stateas, statebs

def handle_atmaml_action(actionas, actionbs, stateas, statebs, if_train=True):

    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size

    actionas = actionas[:, :, :FLAGS.clip_action_size]
    pick = np.reshape(actionas[:, 0, :], [batch_size, 1, -1])
    place = np.reshape(actionas[:, 1, :], [batch_size, 1, -1])
    pick_place = np.concatenate([pick, place], axis=2)
    # place_place = np.concatenate([place, place], axis=2)
    actionas = np.tile(pick_place, (1, 3, 1))
    # actionas = np.tile(pick_place, (1, 2, 1))
    # actionas = np.concatenate([actionas, place_place], axis=1)


    actionbs = actionbs[:, :, :FLAGS.clip_action_size]
    pick = np.reshape(actionbs[:, 0, :], [batch_size, 1, -1])
    place = np.reshape(actionbs[:, 1, :], [batch_size, 1, -1])
    pick_place = np.concatenate([pick, place], axis=2)
    # place_place = np.concatenate([place, place], axis=2)
    actionbs = np.tile(pick_place, (1, 3, 1))
    # actionbs = np.tile(pick_place, (1, 2, 1))
    # actionbs = np.concatenate([actionbs, place_place], axis=1)


    stateas = stateas[:, :, :FLAGS.state_size]
    statebs = statebs[:, :, :FLAGS.state_size]

    return actionas, actionbs, stateas, statebs


def handle_atmaml_action2(actionas, actionbs, stateas, statebs, if_train=True):

    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size

    actionas = actionas[:, :, :FLAGS.clip_action_size]
    place = np.reshape(actionas[:, 1, :], [batch_size, 1, -1])
    actionas = np.tile(place, (1, 3, 1))


    actionbs = actionbs[:, :, :FLAGS.clip_action_size]
    place = np.reshape(actionbs[:, 1, :], [batch_size, 1, -1])
    actionbs = np.tile(place, (1, 3, 1))

    # print('statebs',statebs)

    stateas = stateas[:, :, :FLAGS.state_size]
    statebs = statebs[:, :, :FLAGS.state_size]

    return actionas, actionbs, stateas, statebs


def train(graph, model, saver, sess, save_dir, restore_itr=0):
    """
    Train the model.
    """
    PRINT_INTERVAL = 100
    TEST_PRINT_INTERVAL = 100
    # TEST_PRINT_INTERVAL = 1
    SUMMARY_INTERVAL = 100
    # SAVE_INTERVAL = 10
    SAVE_INTERVAL = 1000
    TOTAL_ITERS = FLAGS.metatrain_iterations
    prelosses, postlosses = [], []
    save_dir = save_dir + '/model'
    train_writer = tf.summary.FileWriter(save_dir, graph)
    # actual training.
    if restore_itr == 0:
        training_range = range(TOTAL_ITERS)
    else:
        training_range = range(restore_itr + 1, TOTAL_ITERS)

    with graph.as_default():
        nums= get_num_params()
        print('total parameters', nums)

    for itr in training_range:

        if FLAGS.cross_domain:
            obsas, obsbs, obscs, obses,  actionas, actionbs, actioncs, actiones, \
            stateas, statebs, statecs, statees, labelabs, labelcs = generate_atmaml_data(if_train=True)
            actionas, actionbs, stateas, statebs = handle_atmaml_action2(actionas, actionbs, stateas, statebs,
                                                                         if_train=True)
            actioncs, actiones, statecs, statees = handle_atmaml_action2(actioncs, actiones, statecs, statees,
                                                                         if_train=True)
            # print('actiones', actiones.shape, actiones[0])
            if FLAGS.cross_adapt_domain:
                # probas = np.ones([FLAGS.meta_batch_size, 1])

                # probas = np.random.randint(0, 100, [FLAGS.meta_batch_size, 1]) / 100.0
                if FLAGS.random_adapt_domain:
                    probas = np.random.randint(0, 100, [FLAGS.meta_batch_size, 1]) / 100.0
                    probes = np.random.randint(0, 100, [FLAGS.meta_batch_size, 1]) / 100.0

                    if FLAGS.random_single_domain:
                        single_probas = np.random.randint(0, 2, [FLAGS.meta_batch_size, 1]) / 1.0
                        single_probes = 1.0 - single_probas
                        for mbs in range(0, FLAGS.meta_batch_size):
                            m = np.random.randint(0, 100)
                            if m < FLAGS.random_single_domain_margin:
                                probas[mbs] = single_probas[mbs]
                                probes[mbs] = single_probes[mbs]

                else:
                    probas = probes = np.ones([FLAGS.meta_batch_size, 1])
                # print('probas', probas, 'probes', probes)
                print('probas', np.squeeze(probas))
                print('probes', np.squeeze(probes))

                feed_dict = {
                    model.obsa: obsas,
                    model.obsb: obsbs,
                    model.obsc: obscs,
                    model.obse: obses,
                    model.statea: stateas,
                    model.stateb: statebs,
                    model.actiona: actionas,
                    model.actionb: actionbs,
                    model.actione: actiones,
                    model.labelab: labelabs,
                    model.prop_a: probas,
                    model.prop_e: probes
                }
            else:
                feed_dict = {
                    model.obsa: obsas,
                    model.obsb: obsbs,
                    model.obsc: obscs,
                    model.obse: obses,
                    model.statea: stateas,
                    model.stateb: statebs,
                    model.actiona: actionas,
                    model.actionb: actionbs,
                    model.actione: actiones,
                    model.labelab: labelabs
                }

        else:
            obsas, obsbs, obscs, actionas, actionbs, actioncs, stateas, statebs, statecs, labelabs, labelcs = generate_atmaml_data(if_train=True)
            actionas, actionbs, stateas, statebs = handle_atmaml_action2(actionas, actionbs, stateas, statebs,
                                                                         if_train=True)
            feed_dict = {
                model.obsa: obsas,
                model.obsb: obsbs,
                model.obsc: obscs,
                model.statea: stateas,
                model.stateb: statebs,
                model.actiona: actionas,
                model.actionb: actionbs,
                model.labelab: labelabs
            }


        # print('obsas, obsbs, obscs, actionas, actionbs, actioncs, stateas, statebs, statecs, labelabs, labelcs',
        #       obsas.shape, obsbs.shape, obscs.shape, actionas.shape, actionbs.shape, actioncs.shape,
        #       stateas.shape, statebs.shape, statecs.shape, labelabs.shape, labelcs.shape)



        # input_tensors = [model.train_op, model.total_mix_loss, model.total_losses2[model.num_updates - 1],
        #                  model.test_act_op, model.total_semantic_loss, model.total_loss1, model.train_summ_op]
        input_tensors = [model.train_op, model.total_mix_loss, model.total_losses2[model.num_updates - 1],
                         model.test_act_op, model.total_semantic_loss, model.total_loss1, model.outputas,
                         model.semantic_outputs, model.post_pre_semantic_loss, model.train_summ_op]

        with graph.as_default():
            # nums= get_num_params()
            # print('total parameters', nums)
            results = sess.run(input_tensors, feed_dict=feed_dict)
            if FLAGS.record_loss:
                train_writer.add_summary(results[-1], itr)

        print('Iteration %d: average preloss is %.2f, post_pre_semantic_loss is %.2f, total_mix_loss is %.2f, average postloss is %.2f,  '
              'average semantic is %.2f' % (itr, np.mean(results[5]), np.mean(results[8]), np.mean(results[1]), np.mean(results[2]), np.mean(results[4])))
        # print('demo state', stateas[-1])
        # print('target state', statebs[-1])
        # print ('results',results[3].shape,results[3][0][0], 'demos',actionbs[0][0])
        # print('demo actionas', actionas[-1])
        print('predict outputa', results[6][-1].shape)
        print('predict pre_semantic_outputa', results[7][0][-1].shape, results[7][0][-1])
        print('predict pre_semantic_outputb', results[7][1][-1].shape, results[7][1][-1])
        print('predict pre_semantic_outputc', results[7][2][-1].shape, results[7][2][-1])
        print('demo actionbs', actionbs[-1])
        print ('predicted actions', results[3][-1])

        if itr != 0 and (itr % TEST_PRINT_INTERVAL == 0):
            if FLAGS.cross_domain:
                obsas, obsbs, obscs, obses, actionas, actionbs, actioncs, actiones, \
                stateas, statebs, statecs, extra_states, labelabs, labelcs = generate_atmaml_data(if_train=False)
            else:
                obsas, obsbs, obscs, actionas, actionbs, actioncs, stateas, statebs, statecs, labelabs, labelcs = generate_atmaml_data(if_train=False)

            actionas, actionbs, stateas, statebs = handle_atmaml_action2(actionas, actionbs, stateas, statebs, if_train=False)
            obsbs, actionbs, statebs = generate_single_atmaml_place_data(obsbs, actionbs, statebs, if_train=False)

            print('actionas, actionbs, stateas, statebs', actionas.shape, actionbs.shape, stateas.shape, statebs.shape)

            if FLAGS.cross_domain:
                if FLAGS.cross_adapt_domain:
                    if not FLAGS.zero_robot_domain:
                        actioncs, actiones, statecs, statees = handle_atmaml_action2(actioncs, actiones, statecs,
                                                                                     statees, if_train=False)
                    probas = np.ones([1, 1])
                    probes = np.random.randint(0, 1, [1, 1]) / 100.0
                    feed_dict = {
                        model.obsa: obsas,
                        model.obsb: obsbs,
                        model.obsc: obscs,
                        model.obse: obses,
                        model.statea: stateas,
                        model.stateb: statebs,
                        model.actiona: actionas,
                        model.actionb: actionbs,
                        model.actione: actiones,
                        model.labelab: labelabs,
                        model.prop_a: probas,
                        model.prop_e: probes
                    }
                else:
                    feed_dict = {
                        model.obsa: obsas,
                        model.obsb: obsbs,
                        model.obsc: obscs,
                        model.obse: obses,
                        model.statea: stateas,
                        model.stateb: statebs,
                        model.actiona: actionas,
                        model.actionb: actionbs,
                        model.actione: actiones,
                        model.labelab: labelabs
                    }


            else:
                feed_dict = {
                    model.obsa: obsas,
                    model.obsb: obsbs,
                    model.statea: stateas,
                    model.stateb: statebs,
                    model.actiona: actionas,
                    model.actionb: actionbs,
                    model.labelab: labelcs
                }



            input_tensors = [model.total_loss1, model.total_losses2[model.num_updates - 1],
                             model.test_act_op]
            with graph.as_default():
                results = sess.run(input_tensors, feed_dict=feed_dict)
                # train_writer.add_summary(results[-1], itr)
                print('test loss', np.mean(results[1]) )
                # print('test demoa', actionas)
                print('test demob', actionbs)
                print('test predict', results[2])


        if itr != 0 and (itr % SAVE_INTERVAL == 0 or itr == training_range[-1]):
            print ('Saving model to: %s' % (save_dir + '_%d' % itr))
            with graph.as_default():
                saver.save(sess, save_dir + '_%d' % itr)


def main():
    tf.set_random_seed(FLAGS.random_seed)
    np.random.seed(FLAGS.random_seed)
    random.seed(FLAGS.random_seed)

    graph = tf.Graph()

    # gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=FLAGS.gpu_memory_fraction)
    # tf_config = tf.ConfigProto(gpu_options=gpu_options)

    tf_config = tf.ConfigProto()
    tf_config.gpu_options.allow_growth = True
    sess = tf.Session(graph=graph, config=tf_config)
    # sess = tf.Session(graph=graph)
    network_config = {
        'num_filters': [FLAGS.num_filters] * FLAGS.num_conv_layers,
        'strides': [[1, 2, 2, 1]] * FLAGS.num_strides + [[1, 1, 1, 1]] * (FLAGS.num_conv_layers - FLAGS.num_strides),
        'filter_size': FLAGS.filter_size,
        'image_width': FLAGS.im_width,
        'image_height': FLAGS.im_height,
        'image_channels': FLAGS.num_channels,
        'n_layers': FLAGS.num_fc_layers,
        'layer_size': FLAGS.layer_size,
        'initialization': FLAGS.init,
    }
    # data_generator = DataGenerator()
    state_idx = range(FLAGS.state_size)
    img_idx = range(len(state_idx), len(state_idx) + FLAGS.im_height * FLAGS.im_width * FLAGS.num_channels)
    # need to compute x_idx and img_idx from data_generator

    model = MIL(FLAGS.action_size, state_idx=state_idx, img_idx=img_idx, network_config=network_config)
    # TODO: figure out how to save summaries and checkpoints
    exp_string = FLAGS.experiment + '.' + FLAGS.init + '_init.' + str(FLAGS.num_conv_layers) + '_conv' + '.' + str(
        FLAGS.num_strides) + '_strides' + '.' + str(FLAGS.num_filters) + '_filters' + \
                 '.' + str(FLAGS.num_fc_layers) + '_fc' + '.' + str(FLAGS.layer_size) + '_dim' + '.bt_dim_' + str(
        FLAGS.bt_dim) + '.mbs_' + str(FLAGS.meta_batch_size) + \
                 '.ubs_' + str(FLAGS.update_batch_size) + '.numstep_' + str(FLAGS.num_updates) + '.updatelr_' + str(
        FLAGS.train_update_lr)

    if FLAGS.clip:
        exp_string += '.clip_' + str(int(FLAGS.clip_max))
    if FLAGS.conv_bt:
        exp_string += '.conv_bt'
    if FLAGS.all_fc_bt:
        exp_string += '.all_fc_bt'
    if FLAGS.fp:
        exp_string += '.fp'
    if FLAGS.learn_final_eept:
        exp_string += '.learn_ee_pos'
    if FLAGS.no_action:
        exp_string += '.no_action'
    if FLAGS.zero_state:
        exp_string += '.zero_state'
    if FLAGS.two_head:
        exp_string += '.two_heads'
    if FLAGS.two_arms:
        exp_string += '.two_arms'
    if FLAGS.temporal_conv_2_head:
        exp_string += '.1d_conv_act_' + str(FLAGS.temporal_num_layers) + '_' + str(FLAGS.temporal_num_filters)
        if FLAGS.temporal_conv_2_head_ee:
            exp_string += '_ee_' + str(FLAGS.temporal_num_layers_ee) + '_' + str(FLAGS.temporal_num_filters_ee)
        exp_string += '_' + str(FLAGS.temporal_filter_size) + 'x1_filters'
    if FLAGS.training_set_size != -1:
        exp_string += '.' + str(FLAGS.training_set_size) + '_trials'

    save_dir = FLAGS.save_dir + '/' + exp_string

    # put here for now
    if FLAGS.train:
        # data_generator.generate_batches(noisy=FLAGS.use_noisy_demos)
        # with graph.as_default():
        #     train_image_tensors = data_generator.make_batch_tensor(network_config, restore_iter=FLAGS.restore_iter)
        #     inputa = train_image_tensors[:, :FLAGS.update_batch_size*FLAGS.T, :]
        #     inputb = train_image_tensors[:, FLAGS.update_batch_size*FLAGS.T:, :]
        #     train_input_tensors = {'inputa': inputa, 'inputb': inputb}

        model.init_network(graph, input_tensors=None, restore_iter=FLAGS.restore_iter)
        # model.init_network(graph, input_tensors=val_input_tensors, restore_iter=FLAGS.restore_iter, prefix='Validation_')
    else:
        model.init_network(graph, prefix='Testing')
    with graph.as_default():
        # Set up saver.
        saver = tf.train.Saver(max_to_keep=10)
        # Initialize variables.
        init_op = tf.global_variables_initializer()
        sess.run(init_op, feed_dict=None)
        # Start queue runners (used for loading videos on the fly)
        tf.train.start_queue_runners(sess=sess)
    if FLAGS.resume:
        model_file = tf.train.latest_checkpoint(save_dir)
        if FLAGS.restore_iter > 0:
            model_file = model_file[:model_file.index('model')] + 'model_' + str(FLAGS.restore_iter)
        if model_file:
            ind1 = model_file.index('model')
            resume_itr = int(model_file[ind1 + 6:])
            print("Restoring model weights from " + model_file)
            with graph.as_default():
                saver.restore(sess, model_file)
    if FLAGS.train:
        train(graph, model, saver, sess, save_dir, restore_itr=FLAGS.restore_iter)


if __name__ == "__main__":
    main()
# coding: utf-8
import numpy as np
import random
import tensorflow as tf
import logging
import imageio
import read_data
from keras.utils import to_categorical

# from data_generator import DataGenerator
from place_pick_mil import MIL
# from evaluation.eval_reach import evaluate_vision_reach
# from evaluation.eval_push import evaluate_push
from tensorflow.python.platform import flags
import os

from functools import reduce
from operator import mul
def get_num_params():
    nums=0
    for variable in tf.trainable_variables():
        shape= variable.get_shape()
        nums+=reduce(mul, [dim.value for dim in shape], 1)
    return nums

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
FLAGS = flags.FLAGS
LOGGER = logging.getLogger(__name__)

## Dataset/method options
flags.DEFINE_string('experiment', 'pick_place', 'sim_vision_reach or sim_push')
flags.DEFINE_string('data_path', './pick_dataset_origin/human_robot_dataset/',
                    'path to the directory where demo files that containing robot states and actions are stored')
flags.DEFINE_string('demo_gif_dir', 'data', 'path to the videos of demonstrations')
flags.DEFINE_string('gif_prefix', 'object', 'prefix of the video directory for each task, e.g. object_0 for task 0')
flags.DEFINE_integer('im_width', 264,
                     'width of the images in the demo videos,  125 for sim_push, and 80 for sim_vision_reach')
flags.DEFINE_integer('im_height', 196,
                     'height of the images in the demo videos, 125 for sim_push, and 64 for sim_vision_reach')
flags.DEFINE_integer('num_channels', 3, 'number of channels of the images in the demo videos')
flags.DEFINE_integer('T', 3, 'time horizon of the demo videos, 50 for reach, 100 for push')
flags.DEFINE_bool('hsv', False, 'convert the image to HSV format')
flags.DEFINE_bool('use_noisy_demos', False, 'use noisy demonstrations or not (for domain shift)')
flags.DEFINE_string('noisy_demo_gif_dir', None, 'path to the videos of noisy demonstrations')
flags.DEFINE_string('noisy_demo_file', None,
                    'path to the directory where noisy demo files that containing robot states and actions are stored')
flags.DEFINE_bool('no_action', True, 'do not include actions in the demonstrations for inner update')
flags.DEFINE_bool('no_state', False, 'do not include states in the demonstrations during training')
flags.DEFINE_bool('no_final_eept', False, 'do not include final ee pos in the demonstrations for inner update')
flags.DEFINE_bool('zero_state', True,
                  'zero-out states (meta-learn state) in the demonstrations for inner update (used in the paper with video-only demos)')
flags.DEFINE_bool('two_arms', False, 'use two-arm structure when state is zeroed-out')
flags.DEFINE_integer('training_set_size', -1, 'size of the training set, 1500 for sim_reach, 693 for sim push, anzero_stated \
                                                -1 for all data except those in validation set')
flags.DEFINE_integer('val_set_size', 150, 'size of the training set, 150 for sim_reach and 76 for sim push')

## Training options
flags.DEFINE_integer('metatrain_iterations', 30000,  'number of metatraining iterations.')  # 30k for pushing, 50k for reaching and placing
flags.DEFINE_integer('meta_batch_size', 16,  'number of tasks sampled per meta-update')  # 15 for reaching, 15 for pushing, 12 for placing
flags.DEFINE_integer('meta_test_batch_size', 1,  'number of tasks sampled per meta-update')  # 15 for reaching, 15 for pushing, 12 for placing
flags.DEFINE_float('meta_lr', 1e-4, 'the base learning rate of the generator')
flags.DEFINE_integer('update_batch_size', 1,
                     'number of examples used for inner gradient update (K for K-shot learning).')
flags.DEFINE_float('train_update_lr', 1e-4,
                   'step size alpha for inner gradient update.')  # 0.001 for reaching, 0.01 for pushing and placing
flags.DEFINE_integer('num_updates', 1, 'number of inner gradient updates during training.')  # 5 for placing
flags.DEFINE_bool('clip', True, 'use gradient clipping for fast gradient')
flags.DEFINE_float('clip_max', 100.0, 'maximum clipping value for fast gradient')
flags.DEFINE_float('clip_min', -100.0, 'minimum clipping value for fast gradient')
# flags.DEFINE_float('clip_max', 20.0, 'maximum clipping value for fast gradient')
# flags.DEFINE_float('clip_min', -20.0, 'minimum clipping value for fast gradient')
flags.DEFINE_bool('fc_bt', True, 'use bias transformation for the first fc layer')
flags.DEFINE_bool('all_fc_bt', False, 'use bias transformation for all fc layers')
flags.DEFINE_bool('conv_bt', False, 'use bias transformation for the first conv layer, N/A for using pretraining')
flags.DEFINE_integer('bt_dim', 10, 'the dimension of bias transformation for FC layers')
flags.DEFINE_string('pretrain_weight_path', 'N/A', 'path to pretrained weights')
flags.DEFINE_bool('train_pretrain_conv1', False, 'whether to finetune the pretrained weights')
flags.DEFINE_bool('two_head', True, 'use two-head architecture')
flags.DEFINE_bool('learn_final_eept', False, 'learn an auxiliary loss for predicting final end-effector pose')
flags.DEFINE_bool('learn_final_eept_whole_traj', False, 'learn an auxiliary loss for predicting final end-effector pose \
                                                         by passing the whole trajectory of eepts (used for video-only models)')
flags.DEFINE_bool('stopgrad_final_eept', True,
                  'stop the gradient when concatenate the predicted final eept with the feature points')
flags.DEFINE_integer('final_eept_min', 6, 'first index of the final eept in the action array')
flags.DEFINE_integer('final_eept_max', 8, 'last index of the final eept in the action array')
flags.DEFINE_float('final_eept_loss_eps', 0.1, 'the coefficient of the auxiliary loss')
flags.DEFINE_float('act_loss_eps', 1.0, 'the coefficient of the action loss')
flags.DEFINE_float('loss_multiplier', 100.0,
                   'the constant multiplied with the loss value, 100 for reach and 50 for push')
flags.DEFINE_bool('use_l1_l2_loss', False, 'use a loss with combination of l1 and l2')
flags.DEFINE_float('l2_eps', 0.01, 'coeffcient of l2 loss')
flags.DEFINE_bool('shuffle_val', False, 'whether to choose the validation set via shuffling or not')

## Model options
flags.DEFINE_integer('random_seed', 0, 'random seed for training')
flags.DEFINE_bool('fp', True, 'use spatial soft-argmax or not')
flags.DEFINE_string('norm', 'layer_norm', 'batch_norm, layer_norm, or None')
flags.DEFINE_bool('dropout', False, 'use dropout for fc layers or not')
flags.DEFINE_float('keep_prob', 0.5, 'keep probability for dropout')
flags.DEFINE_integer('num_filters', 64,
                     'number of filters for conv nets -- 64 for placing, 16 for pushing, 40 for reaching.')
flags.DEFINE_integer('filter_size', 3, 'filter size for conv nets -- 3 for placing, 5 for pushing, 3 for reaching.')
flags.DEFINE_integer('num_conv_layers', 5, 'number of conv layers -- 5 for placing, 4 for pushing, 3 for reaching.')
flags.DEFINE_integer('num_strides', 3,
                     'number of conv layers with strided filters -- 3 for placing, 4 for pushing, 3 for reaching.')
flags.DEFINE_bool('conv', True, 'whether or not to use a convolutional network, only applicable in some cases')
flags.DEFINE_integer('num_fc_layers', 3, 'number of fully-connected layers')
flags.DEFINE_integer('layer_size', 200, 'hidden dimension of fully-connected layers')
flags.DEFINE_bool('temporal_conv_2_head', True,
                  'whether or not to use temporal convolutions for the two-head architecture in video-only setting.')
flags.DEFINE_bool('temporal_conv_2_head_ee', False, 'whether or not to use temporal convolutions for the two-head architecture in video-only setting \
                for predicting the ee pose.')
flags.DEFINE_integer('temporal_filter_size', 10, 'filter size for temporal convolution')
flags.DEFINE_integer('temporal_num_filters', 32, 'number of filters for temporal convolution')
flags.DEFINE_integer('temporal_num_filters_ee', 32, 'number of filters for temporal convolution for ee pose prediction')
flags.DEFINE_integer('temporal_num_layers', 3, 'number of layers for temporal convolution for ee pose prediction')
flags.DEFINE_integer('temporal_num_layers_ee', 3, 'number of layers for temporal convolution for ee pose prediction')
flags.DEFINE_string('init', 'xavier', 'initializer for conv weights. Choose among random, xavier, and he')
flags.DEFINE_bool('max_pool', False, 'Whether or not to use max pooling rather than strided convolutions')
flags.DEFINE_bool('stop_grad', False, 'if True, do not use second derivatives in meta-optimization (for axis_angle)')

## Logging, saving, and testing options
flags.DEFINE_bool('log', True, 'if false, do not log summaries, for debugging code.')
flags.DEFINE_string('save_dir', './atmaml_pick_logs', 'directory for summaries and checkpoints.')
# flags.DEFINE_string('save_dir', './amaml_human_pick_logs', 'directory for summaries and checkpoints.')
# flags.DEFINE_bool('resume', True, 'resume training if there is a model available')
flags.DEFINE_bool('resume', False, 'resume training if there is a model available')
flags.DEFINE_bool('train', True, 'True to train, False to test.')
flags.DEFINE_integer('restore_iter', -1, 'iteration to load model (-1 for latest model)')
# flags.DEFINE_integer('restore_iter', 20000, 'iteration to load model (-1 for latest model)')
flags.DEFINE_integer('train_update_batch_size', -1, 'number of examples used for gradient update during training \
                    (use if you want to test with a different number).')
flags.DEFINE_integer('test_update_batch_size', 1, 'number of demos used during test time')
flags.DEFINE_float('gpu_memory_fraction', 0.8, 'fraction of memory used in gpu')
flags.DEFINE_bool('record_gifs', True, 'record gifs during evaluation')
flags.DEFINE_integer('color_num', 4, '')
flags.DEFINE_integer('object_num', 4, '')
flags.DEFINE_integer('train_task_num', 6, '')
flags.DEFINE_integer('task_num', 8, '')
flags.DEFINE_integer('demo_num', 5, '')
flags.DEFINE_integer('index_range', 20, '')
flags.DEFINE_integer('index_train_range', 20, '')
# flags.DEFINE_string('demo_type', 'robot', 'robot or human')
flags.DEFINE_string('demo_type', 'human', 'robot or human')
flags.DEFINE_string('extra_type', 'robot', 'robot or human') #opposite to demo_type
flags.DEFINE_string('compare_type', 'robot', 'robot or human')
flags.DEFINE_string('target_type', 'robot', '')
# flags.DEFINE_float('weight_xy', 0.999, '')
# flags.DEFINE_float('weight_z', 0.001, '')
# flags.DEFINE_float('weight_rxyz', 0.001, '')
flags.DEFINE_integer('clip_action_size', 2, 'size of embedding')
flags.DEFINE_integer('action_size', 2, 'size of embedding')
flags.DEFINE_integer('state_size', 6, 'size of embedding')
flags.DEFINE_integer('output_data', 6, '')
# flags.DEFINE_float('margin', 1.0, 'margin of loss')
# flags.DEFINE_float('pre_margin', 1.0, 'margin of loss')
flags.DEFINE_float('pre_margin_clip', 1.0, '')
flags.DEFINE_float('pre_margin', 2.0, 'pre_margin of loss')
flags.DEFINE_float('margin', 2.0, 'margin of loss')
flags.DEFINE_float('pre_margin_coefficient', 10.0, 'margin of pre-contrastive conloss')
flags.DEFINE_float('margin_coefficient', 10.0, 'margin of post-contrastive loss')
flags.DEFINE_float('pre_tar_coefficient', 0.0, '')
flags.DEFINE_float('post_tar_coefficient', 0.0, '')
flags.DEFINE_bool('norm_pre_contrastive', True, 'norm for pre_task contrastive ')
flags.DEFINE_bool('norm_post_contrastive', True, 'norm for post_task contrastive ')
flags.DEFINE_bool('all_frame_contrastive', True, '')
flags.DEFINE_bool('cos_contrastive_loss', True, 'True for cos loss, False for kl loss ')
flags.DEFINE_bool('pre_contrastive', False, 'task contrastive for pre-update ')
flags.DEFINE_bool('post_contrastive', False, 'task contrastive for post-update ')
# flags.DEFINE_bool('amaml', False, 'true for amaml')
flags.DEFINE_bool('amaml', True, 'true for amaml')
flags.DEFINE_bool('amaml_extra_domain', True, 'true for extra_domain')
flags.DEFINE_bool('cross_domain', True, 'use human and robot ')
flags.DEFINE_bool('cross_adapt_domain', True, 'adpat human/robot demos')
flags.DEFINE_bool('zero_robot_domain', True, 'adpat human/robot demos')
flags.DEFINE_bool('random_adapt_domain', True, 'randomly adpat human/robot demos')
flags.DEFINE_bool('random_single_domain', True, 'randomly adpat human/robot demos')
flags.DEFINE_integer('random_single_domain_margin', 50, 'random single domain margin')
flags.DEFINE_integer('embed_size', 0, 'size of embedding')
flags.DEFINE_bool('tar_mil', False, 'with target loss')
flags.DEFINE_bool('record_loss', False, 'true for amaml')
flags.DEFINE_bool('contrastive_fc', False, '')
# flags.DEFINE_bool('amaml', False, 'true for amaml')

def generate_data(if_train=True):
    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size

    color_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.color_num
    print('color_list', color_list)

    object_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.object_num
    print('object_list', object_list)
    if if_train:
        task_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.train_task_num
    else:
        task_list = np.random.randint(FLAGS.train_task_num, FLAGS.task_num, size=batch_size)
    print('task_list', task_list)

    demo_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.demo_num
    print('demo_list', demo_list)

    target_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.demo_num
    print('target_list', target_list)

    obsas = []
    obsbs = []
    stateas = []
    statebs = []
    actionas = []
    actionbs = []

    color_num = ['color_blue', 'color_green', 'color_orange', 'color_yellow']
    # color_num = ['color_blue', 'color_green', 'color_orange']
    object_num = ['object_type_animal', 'object_type_car', 'object_type_dinosaur', 'object_type_tool']

    for element in range(0, batch_size):
        demo_path = '%s/%s/%s/%s/task_%d/demo_%d' % (
            FLAGS.data_path, color_num[color_list[element]], object_num[object_list[element]], FLAGS.demo_type,
            task_list[element], demo_list[element])
        target_path = '%s/%s/%s/%s/task_%d/demo_%d' % (
            FLAGS.data_path, color_num[color_list[element]], object_num[object_list[element]], FLAGS.target_type,
            task_list[element], target_list[element])

        print('demo_path', demo_path)
        print('target_path', target_path)

        index = np.random.randint(0, FLAGS.index_train_range)

        # if if_train:
        #     index = np.random.randint(0, FLAGS.index_train_range)
        # else:
        #     index = np.random.randint(FLAGS.index_train_range, FLAGS.index_range)

        if FLAGS.demo_type == 'robot':
            obsa, statea, actiona = read_data.Read_Robot_Data2(demo_path, FLAGS.T, index)
        elif FLAGS.demo_type == 'human':
            obsa, statea, actiona = read_data.Read_Human_Data2(demo_path, FLAGS.T, index)
        obsb, stateb, actionb = read_data.Read_Robot_Data2(target_path, FLAGS.T, index)

        obsas.append(obsa)
        obsbs.append(obsb)
        stateas.append(statea)
        statebs.append(stateb)
        actionas.append(actiona)
        actionbs.append(actionb)


    obsas = np.reshape(obsas,  [batch_size, FLAGS.T, FLAGS.im_width * FLAGS.im_height * FLAGS.num_channels])
    obsbs = np.reshape(obsbs,    [batch_size, FLAGS.T, FLAGS.im_width * FLAGS.im_height * FLAGS.num_channels])

    actionas = np.reshape(actionas, [batch_size, FLAGS.T, FLAGS.output_data])
    actionbs = np.reshape(actionbs, [batch_size, FLAGS.T, FLAGS.output_data])

    stateas = np.zeros([batch_size, FLAGS.T, FLAGS.output_data])
    statebs = np.zeros([batch_size, FLAGS.T, FLAGS.output_data])

    return obsas, obsbs, actionas, actionbs, stateas, statebs


def generate_atmaml_data(if_train=True):
    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size

    color_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.color_num
    # print('color_list', color_list)
    compare_color_list = (color_list + np.random.randint(1, FLAGS.color_num-1)) %  FLAGS.color_num
    # print('compare_color_list', compare_color_list)

    object_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.object_num
    # print('object_list', object_list)
    compare_object_list = (object_list + np.random.randint(1, FLAGS.object_num-1)) % FLAGS.object_num
    # print('compare_object_list', compare_object_list)

    if if_train:
        task_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.train_task_num
        compare_task_list = (task_list + np.random.randint(1, FLAGS.train_task_num)) % FLAGS.train_task_num
        # compare_task_list = (task_list + np.random.randint(1, FLAGS.task_num - 1)) % FLAGS.task_num
    else:
        task_list = np.random.randint(FLAGS.train_task_num, FLAGS.task_num, size=batch_size)
        compare_task_list = np.random.randint(FLAGS.train_task_num, FLAGS.task_num,  size=batch_size)

    # print('task_list', task_list)
    # print('compare_task_list', compare_task_list)


    demo_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.demo_num
    # print('demo_list', demo_list)
    target_list = (np.random.randint(0, 100, size=batch_size) + 1) % FLAGS.demo_num
    # print('target_list', target_list)
    compare_demo_list = (task_list +  np.random.randint(1, FLAGS.demo_num-1)) %  FLAGS.demo_num
    # print('compare_demo_list', compare_demo_list)



    obsas = []
    obsbs = []
    obscs = []
    extra_obses = []

    stateas = []
    statebs = []
    statecs = []
    extra_states = []

    actionas = []
    actionbs = []
    actioncs = []
    extra_actions = []

    labelabs = []
    labelcs = []


    color_num = ['color_blue', 'color_green', 'color_orange', 'color_yellow']
    # color_num = ['color_blue', 'color_green', 'color_orange']
    object_num = ['object_type_animal', 'object_type_car', 'object_type_dinosaur', 'object_type_tool']

    for element in range(0, batch_size):

        demo_path = '%s/%s/%s/%s/task_%d/demo_%d' % (
            FLAGS.data_path, color_num[color_list[element]], object_num[object_list[element]], FLAGS.demo_type,
            task_list[element], demo_list[element])

        target_path = '%s/%s/%s/%s/task_%d/demo_%d' % (
            FLAGS.data_path, color_num[color_list[element]], object_num[object_list[element]], FLAGS.target_type,
            task_list[element], target_list[element])

        compare_path = '%s/%s/%s/%s/task_%d/demo_%d' % (
            FLAGS.data_path, color_num[compare_color_list[element]], object_num[compare_object_list[element]],
            FLAGS.compare_type, compare_task_list[element], compare_demo_list[element])

        if FLAGS.cross_domain:
                extra_path = '%s/%s/%s/%s/task_%d/demo_%d' % (FLAGS.data_path, color_num[color_list[element]], object_num[object_list[element]],
                                                              FLAGS.extra_type, task_list[element], demo_list[element])
        if  element==0:
            if FLAGS.cross_domain:
                print('extra_path', extra_path)
            print('demo_path', demo_path)
            print('target_path', target_path)
            print('compare_path', compare_path)


        index = np.random.randint(0, FLAGS.index_train_range)

        # if if_train:
        #     index = np.random.randint(0, FLAGS.index_train_range)
        # else:
        #     index = np.random.randint(FLAGS.index_train_range, FLAGS.index_range)

        if FLAGS.demo_type == 'robot':
            obsa, statea, actiona = read_data.Read_Robot_Data2(demo_path, FLAGS.T, index)
        else :
            obsa, statea, actiona = read_data.Read_Human_Data2(demo_path, FLAGS.T, index)

        if FLAGS.compare_type == 'robot':
            obsc, statec, actionc = read_data.Read_Robot_Data2(compare_path, FLAGS.T, index)
        else:
            obsc, statec, actionc = read_data.Read_Human_Data2(compare_path, FLAGS.T, index)

        if FLAGS.cross_domain:
            if FLAGS.extra_type == 'robot':
                extra_obs, extra_state, extra_action = read_data.Read_Robot_Data2(extra_path, FLAGS.T, index)
            else:
                extra_obs, extra_state, extra_action = read_data.Read_Human_Data2(extra_path, FLAGS.T, index)

        obsb, stateb, actionb = read_data.Read_Robot_Data2(target_path, FLAGS.T, index)
        if FLAGS.tar_mil:
            labelab  = to_categorical(color_list[element] * FLAGS.index_train_range + index, FLAGS.color_num*FLAGS.index_train_range)
            labelc = to_categorical(compare_color_list[element] * FLAGS.index_train_range + index, FLAGS.color_num*FLAGS.index_train_range)
            labelabs.append(labelab)
            labelcs.append(labelc)
            # print('element', element, 'labela and labelb', labelab, 'labelc', labelc)

        # print('----------------------------------------------------------------')

        obsas.append(obsa)
        obsbs.append(obsb)
        obscs.append(obsc)
        stateas.append(statea)
        statebs.append(stateb)
        statecs.append(statec)
        actionas.append(actiona)
        actionbs.append(actionb)
        actioncs.append(actionc)
        if FLAGS.cross_domain:
            extra_obses.append(extra_obs)
            extra_states.append(extra_state)
            extra_actions.append(extra_action)


    obsas = np.reshape(obsas,  [batch_size, FLAGS.T, FLAGS.im_width * FLAGS.im_height * FLAGS.num_channels])
    obsbs = np.reshape(obsbs,    [batch_size, FLAGS.T, FLAGS.im_width * FLAGS.im_height * FLAGS.num_channels])
    obscs = np.reshape(obscs, [batch_size, FLAGS.T, FLAGS.im_width * FLAGS.im_height * FLAGS.num_channels])


    actionas = np.reshape(actionas, [batch_size, FLAGS.T, FLAGS.output_data])
    actionbs = np.reshape(actionbs, [batch_size, FLAGS.T, FLAGS.output_data])
    actioncs = np.reshape(actioncs, [batch_size, FLAGS.T, FLAGS.output_data])


    stateas = np.zeros([batch_size, FLAGS.T, FLAGS.output_data])
    statebs = np.zeros([batch_size, FLAGS.T, FLAGS.output_data])
    statecs = np.zeros([batch_size, FLAGS.T, FLAGS.output_data])


    if FLAGS.tar_mil:
        labelabs = np.reshape(labelabs, [batch_size, 1, -1])
        labelcs = np.reshape(labelcs, [batch_size, 1, -1])
        # print('labelabs', labelabs.shape, 'labelcs', labelcs)
    else:
        labelabs = np.zeros([batch_size, 1, FLAGS.color_num*FLAGS.index_train_range])
        labelcs = np.zeros( [batch_size, 1, FLAGS.color_num*FLAGS.index_train_range])

    if FLAGS.cross_domain:
        extra_obses = np.reshape(extra_obses, [batch_size, FLAGS.T, FLAGS.im_width * FLAGS.im_height * FLAGS.num_channels])
        extra_actions = np.reshape(extra_actions, [batch_size, FLAGS.T, FLAGS.output_data])
        extra_states = np.zeros([batch_size, FLAGS.T, FLAGS.output_data])

        return obsas, obsbs, obscs, extra_obses, \
               actionas, actionbs, actioncs, extra_actions, \
               stateas, statebs, statecs, extra_states, \
               labelabs, labelcs
    else:
        return obsas, obsbs, obscs, actionas, actionbs, actioncs, stateas, statebs, statecs, labelabs, labelcs

def generate_single_place_data(obsas, obsbs, actionas, actionbs, stateas, statebs, if_train=True):

    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size

    # print('obsas, obsbs, actionas, actionbs, stateas, statebs', obsas.shape, obsbs.shape, actionas.shape, actionbs.shape, stateas.shape, statebs.shape)
    obsbs = np.reshape(obsbs[:, 0, :], [batch_size , 1, -1])
    actionbs = np.reshape(actionbs[:, 1, :], [batch_size , 1, -1])
    statebs = np.reshape(statebs[:, -1, :], [batch_size , 1, -1])
    # print('obsas, obsbs, actionas, actionbs, stateas, statebs of generate_single_place_data', obsas.shape, obsbs.shape, actionas.shape, actionbs.shape,
    # stateas.shape, statebs.shape)
    return obsas, obsbs, actionas, actionbs, stateas, statebs


def generate_single_atmaml_place_data(obsbs, actionbs, statebs, if_train=True):

    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size



    # print('obsas, obsbs, actionas, actionbs, stateas, statebs', obsbs.shape, actionbs.shape, statebs.shape)
    obsbs = np.reshape(obsbs[:, 0, :], [batch_size , 1, -1])
    actionbs = np.reshape(actionbs[:, 1, :], [batch_size , 1, -1])
    statebs = np.reshape(statebs[:, -1, :], [batch_size , 1, -1])

    return obsbs, actionbs, statebs


def handle_action(actionas, actionbs, stateas, statebs, if_train=True):

    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size

    actionas = actionas[:, :, :FLAGS.clip_action_size]
    # print('actionas is', actionas.shape)
    actionas = np.reshape(actionas[:, 1, :], [batch_size, 1, -1])
    actionas = np.tile(actionas, (1, 3, 1))

    actionbs = actionbs[:, :, :FLAGS.clip_action_size]
    actionbs = np.reshape(actionbs[:, 1, :], [batch_size, 1, -1])
    actionbs = np.tile(actionbs, (1, 3, 1))


    stateas = stateas[:, :, :FLAGS.state_size]
    statebs = statebs[:, :, :FLAGS.state_size]

    return actionas, actionbs, stateas, statebs

def handle_atmaml_action(actionas, actionbs, stateas, statebs, if_train=True):

    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size

    actionas = actionas[:, :, :FLAGS.clip_action_size]
    pick = np.reshape(actionas[:, 0, :], [batch_size, 1, -1])
    place = np.reshape(actionas[:, 1, :], [batch_size, 1, -1])
    pick_place = np.concatenate([pick, place], axis=2)
    # place_place = np.concatenate([place, place], axis=2)
    actionas = np.tile(pick_place, (1, 3, 1))
    # actionas = np.tile(pick_place, (1, 2, 1))
    # actionas = np.concatenate([actionas, place_place], axis=1)


    actionbs = actionbs[:, :, :FLAGS.clip_action_size]
    pick = np.reshape(actionbs[:, 0, :], [batch_size, 1, -1])
    place = np.reshape(actionbs[:, 1, :], [batch_size, 1, -1])
    pick_place = np.concatenate([pick, place], axis=2)
    # place_place = np.concatenate([place, place], axis=2)
    actionbs = np.tile(pick_place, (1, 3, 1))
    # actionbs = np.tile(pick_place, (1, 2, 1))
    # actionbs = np.concatenate([actionbs, place_place], axis=1)


    stateas = stateas[:, :, :FLAGS.state_size]
    statebs = statebs[:, :, :FLAGS.state_size]

    return actionas, actionbs, stateas, statebs


def handle_atmaml_action2(actionas, actionbs, stateas, statebs, if_train=True):

    if if_train:
        batch_size = FLAGS.meta_batch_size
    else:
        batch_size = FLAGS.meta_test_batch_size

    actionas = actionas[:, :, :FLAGS.clip_action_size]
    place = np.reshape(actionas[:, 1, :], [batch_size, 1, -1])
    actionas = np.tile(place, (1, 3, 1))


    actionbs = actionbs[:, :, :FLAGS.clip_action_size]
    place = np.reshape(actionbs[:, 1, :], [batch_size, 1, -1])
    actionbs = np.tile(place, (1, 3, 1))

    # print('statebs',statebs)

    stateas = stateas[:, :, :FLAGS.state_size]
    statebs = statebs[:, :, :FLAGS.state_size]

    return actionas, actionbs, stateas, statebs


def train(graph, model, saver, sess, save_dir, restore_itr=0):
    """
    Train the model.
    """
    PRINT_INTERVAL = 100
    TEST_PRINT_INTERVAL = 100
    # TEST_PRINT_INTERVAL = 1
    SUMMARY_INTERVAL = 100
    # SAVE_INTERVAL = 10
    SAVE_INTERVAL = 1000
    TOTAL_ITERS = FLAGS.metatrain_iterations
    prelosses, postlosses = [], []
    save_dir = save_dir + '/model'
    train_writer = tf.summary.FileWriter(save_dir, graph)
    # actual training.
    if restore_itr == 0:
        training_range = range(TOTAL_ITERS)
    else:
        training_range = range(restore_itr + 1, TOTAL_ITERS)

    with graph.as_default():
        nums= get_num_params()
        print('total parameters', nums)

    for itr in training_range:

        if FLAGS.cross_domain:
            obsas, obsbs, obscs, obses,  actionas, actionbs, actioncs, actiones, \
            stateas, statebs, statecs, statees, labelabs, labelcs = generate_atmaml_data(if_train=True)
            actionas, actionbs, stateas, statebs = handle_atmaml_action2(actionas, actionbs, stateas, statebs,
                                                                         if_train=True)
            actioncs, actiones, statecs, statees = handle_atmaml_action2(actioncs, actiones, statecs, statees,
                                                                         if_train=True)
            # print('actiones', actiones.shape, actiones[0])
            if FLAGS.cross_adapt_domain:
                # probas = np.ones([FLAGS.meta_batch_size, 1])

                # probas = np.random.randint(0, 100, [FLAGS.meta_batch_size, 1]) / 100.0
                if FLAGS.random_adapt_domain:
                    probas = np.random.randint(0, 100, [FLAGS.meta_batch_size, 1]) / 100.0
                    probes = np.random.randint(0, 100, [FLAGS.meta_batch_size, 1]) / 100.0

                    if FLAGS.random_single_domain:
                        single_probas = np.random.randint(0, 2, [FLAGS.meta_batch_size, 1]) / 1.0
                        single_probes = 1.0 - single_probas
                        for mbs in range(0, FLAGS.meta_batch_size):
                            m = np.random.randint(0, 100)
                            if m < FLAGS.random_single_domain_margin:
                                probas[mbs] = single_probas[mbs]
                                probes[mbs] = single_probes[mbs]

                else:
                    probas = probes = np.ones([FLAGS.meta_batch_size, 1])
                # print('probas', probas, 'probes', probes)
                print('probas', np.squeeze(probas))
                print('probes', np.squeeze(probes))

                feed_dict = {
                    model.obsa: obsas,
                    model.obsb: obsbs,
                    model.obsc: obscs,
                    model.obse: obses,
                    model.statea: stateas,
                    model.stateb: statebs,
                    model.actiona: actionas,
                    model.actionb: actionbs,
                    model.actione: actiones,
                    model.labelab: labelabs,
                    model.prop_a: probas,
                    model.prop_e: probes
                }
            else:
                feed_dict = {
                    model.obsa: obsas,
                    model.obsb: obsbs,
                    model.obsc: obscs,
                    model.obse: obses,
                    model.statea: stateas,
                    model.stateb: statebs,
                    model.actiona: actionas,
                    model.actionb: actionbs,
                    model.actione: actiones,
                    model.labelab: labelabs
                }

        else:
            obsas, obsbs, obscs, actionas, actionbs, actioncs, stateas, statebs, statecs, labelabs, labelcs = generate_atmaml_data(if_train=True)
            actionas, actionbs, stateas, statebs = handle_atmaml_action2(actionas, actionbs, stateas, statebs,
                                                                         if_train=True)
            feed_dict = {
                model.obsa: obsas,
                model.obsb: obsbs,
                model.obsc: obscs,
                model.statea: stateas,
                model.stateb: statebs,
                model.actiona: actionas,
                model.actionb: actionbs,
                model.labelab: labelabs
            }


        # print('obsas, obsbs, obscs, actionas, actionbs, actioncs, stateas, statebs, statecs, labelabs, labelcs',
        #       obsas.shape, obsbs.shape, obscs.shape, actionas.shape, actionbs.shape, actioncs.shape,
        #       stateas.shape, statebs.shape, statecs.shape, labelabs.shape, labelcs.shape)



        # input_tensors = [model.train_op, model.total_mix_loss, model.total_losses2[model.num_updates - 1],
        #                  model.test_act_op, model.total_semantic_loss, model.total_loss1, model.train_summ_op]
        input_tensors = [model.train_op, model.total_mix_loss, model.total_losses2[model.num_updates - 1],
                         model.test_act_op, model.total_semantic_loss, model.total_loss1, model.outputas,
                         model.semantic_outputs, model.post_pre_semantic_loss, model.train_summ_op]

        with graph.as_default():
            # nums= get_num_params()
            # print('total parameters', nums)
            results = sess.run(input_tensors, feed_dict=feed_dict)
            if FLAGS.record_loss:
                train_writer.add_summary(results[-1], itr)

        print('Iteration %d: average preloss is %.2f, post_pre_semantic_loss is %.2f, total_mix_loss is %.2f, average postloss is %.2f,  '
              'average semantic is %.2f' % (itr, np.mean(results[5]), np.mean(results[8]), np.mean(results[1]), np.mean(results[2]), np.mean(results[4])))
        # print('demo state', stateas[-1])
        # print('target state', statebs[-1])
        # print ('results',results[3].shape,results[3][0][0], 'demos',actionbs[0][0])
        # print('demo actionas', actionas[-1])
        print('predict outputa', results[6][-1].shape)
        print('predict pre_semantic_outputa', results[7][0][-1].shape, results[7][0][-1])
        print('predict pre_semantic_outputb', results[7][1][-1].shape, results[7][1][-1])
        print('predict pre_semantic_outputc', results[7][2][-1].shape, results[7][2][-1])
        print('demo actionbs', actionbs[-1])
        print ('predicted actions', results[3][-1])

        if itr != 0 and (itr % TEST_PRINT_INTERVAL == 0):
            if FLAGS.cross_domain:
                obsas, obsbs, obscs, obses, actionas, actionbs, actioncs, actiones, \
                stateas, statebs, statecs, extra_states, labelabs, labelcs = generate_atmaml_data(if_train=False)
            else:
                obsas, obsbs, obscs, actionas, actionbs, actioncs, stateas, statebs, statecs, labelabs, labelcs = generate_atmaml_data(if_train=False)

            actionas, actionbs, stateas, statebs = handle_atmaml_action2(actionas, actionbs, stateas, statebs, if_train=False)
            obsbs, actionbs, statebs = generate_single_atmaml_place_data(obsbs, actionbs, statebs, if_train=False)

            print('actionas, actionbs, stateas, statebs', actionas.shape, actionbs.shape, stateas.shape, statebs.shape)

            if FLAGS.cross_domain:
                if FLAGS.cross_adapt_domain:
                    if not FLAGS.zero_robot_domain:
                        actioncs, actiones, statecs, statees = handle_atmaml_action2(actioncs, actiones, statecs,
                                                                                     statees, if_train=False)
                    probas = np.ones([1, 1])
                    probes = np.random.randint(0, 1, [1, 1]) / 100.0
                    feed_dict = {
                        model.obsa: obsas,
                        model.obsb: obsbs,
                        model.obsc: obscs,
                        model.obse: obses,
                        model.statea: stateas,
                        model.stateb: statebs,
                        model.actiona: actionas,
                        model.actionb: actionbs,
                        model.actione: actiones,
                        model.labelab: labelabs,
                        model.prop_a: probas,
                        model.prop_e: probes
                    }
                else:
                    feed_dict = {
                        model.obsa: obsas,
                        model.obsb: obsbs,
                        model.obsc: obscs,
                        model.obse: obses,
                        model.statea: stateas,
                        model.stateb: statebs,
                        model.actiona: actionas,
                        model.actionb: actionbs,
                        model.actione: actiones,
                        model.labelab: labelabs
                    }


            else:
                feed_dict = {
                    model.obsa: obsas,
                    model.obsb: obsbs,
                    model.statea: stateas,
                    model.stateb: statebs,
                    model.actiona: actionas,
                    model.actionb: actionbs,
                    model.labelab: labelcs
                }



            input_tensors = [model.total_loss1, model.total_losses2[model.num_updates - 1],
                             model.test_act_op]
            with graph.as_default():
                results = sess.run(input_tensors, feed_dict=feed_dict)
                # train_writer.add_summary(results[-1], itr)
                print('test loss', np.mean(results[1]) )
                # print('test demoa', actionas)
                print('test demob', actionbs)
                print('test predict', results[2])


        if itr != 0 and (itr % SAVE_INTERVAL == 0 or itr == training_range[-1]):
            print ('Saving model to: %s' % (save_dir + '_%d' % itr))
            with graph.as_default():
                saver.save(sess, save_dir + '_%d' % itr)


def main():
    tf.set_random_seed(FLAGS.random_seed)
    np.random.seed(FLAGS.random_seed)
    random.seed(FLAGS.random_seed)

    graph = tf.Graph()

    # gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=FLAGS.gpu_memory_fraction)
    # tf_config = tf.ConfigProto(gpu_options=gpu_options)

    tf_config = tf.ConfigProto()
    tf_config.gpu_options.allow_growth = True
    sess = tf.Session(graph=graph, config=tf_config)
    # sess = tf.Session(graph=graph)
    network_config = {
        'num_filters': [FLAGS.num_filters] * FLAGS.num_conv_layers,
        'strides': [[1, 2, 2, 1]] * FLAGS.num_strides + [[1, 1, 1, 1]] * (FLAGS.num_conv_layers - FLAGS.num_strides),
        'filter_size': FLAGS.filter_size,
        'image_width': FLAGS.im_width,
        'image_height': FLAGS.im_height,
        'image_channels': FLAGS.num_channels,
        'n_layers': FLAGS.num_fc_layers,
        'layer_size': FLAGS.layer_size,
        'initialization': FLAGS.init,
    }
    # data_generator = DataGenerator()
    state_idx = range(FLAGS.state_size)
    img_idx = range(len(state_idx), len(state_idx) + FLAGS.im_height * FLAGS.im_width * FLAGS.num_channels)
    # need to compute x_idx and img_idx from data_generator

    model = MIL(FLAGS.action_size, state_idx=state_idx, img_idx=img_idx, network_config=network_config)
    # TODO: figure out how to save summaries and checkpoints
    exp_string = FLAGS.experiment + '.' + FLAGS.init + '_init.' + str(FLAGS.num_conv_layers) + '_conv' + '.' + str(
        FLAGS.num_strides) + '_strides' + '.' + str(FLAGS.num_filters) + '_filters' + \
                 '.' + str(FLAGS.num_fc_layers) + '_fc' + '.' + str(FLAGS.layer_size) + '_dim' + '.bt_dim_' + str(
        FLAGS.bt_dim) + '.mbs_' + str(FLAGS.meta_batch_size) + \
                 '.ubs_' + str(FLAGS.update_batch_size) + '.numstep_' + str(FLAGS.num_updates) + '.updatelr_' + str(
        FLAGS.train_update_lr)

    if FLAGS.clip:
        exp_string += '.clip_' + str(int(FLAGS.clip_max))
    if FLAGS.conv_bt:
        exp_string += '.conv_bt'
    if FLAGS.all_fc_bt:
        exp_string += '.all_fc_bt'
    if FLAGS.fp:
        exp_string += '.fp'
    if FLAGS.learn_final_eept:
        exp_string += '.learn_ee_pos'
    if FLAGS.no_action:
        exp_string += '.no_action'
    if FLAGS.zero_state:
        exp_string += '.zero_state'
    if FLAGS.two_head:
        exp_string += '.two_heads'
    if FLAGS.two_arms:
        exp_string += '.two_arms'
    if FLAGS.temporal_conv_2_head:
        exp_string += '.1d_conv_act_' + str(FLAGS.temporal_num_layers) + '_' + str(FLAGS.temporal_num_filters)
        if FLAGS.temporal_conv_2_head_ee:
            exp_string += '_ee_' + str(FLAGS.temporal_num_layers_ee) + '_' + str(FLAGS.temporal_num_filters_ee)
        exp_string += '_' + str(FLAGS.temporal_filter_size) + 'x1_filters'
    if FLAGS.training_set_size != -1:
        exp_string += '.' + str(FLAGS.training_set_size) + '_trials'

    save_dir = FLAGS.save_dir + '/' + exp_string

    # put here for now
    if FLAGS.train:
        # data_generator.generate_batches(noisy=FLAGS.use_noisy_demos)
        # with graph.as_default():
        #     train_image_tensors = data_generator.make_batch_tensor(network_config, restore_iter=FLAGS.restore_iter)
        #     inputa = train_image_tensors[:, :FLAGS.update_batch_size*FLAGS.T, :]
        #     inputb = train_image_tensors[:, FLAGS.update_batch_size*FLAGS.T:, :]
        #     train_input_tensors = {'inputa': inputa, 'inputb': inputb}

        model.init_network(graph, input_tensors=None, restore_iter=FLAGS.restore_iter)
        # model.init_network(graph, input_tensors=val_input_tensors, restore_iter=FLAGS.restore_iter, prefix='Validation_')
    else:
        model.init_network(graph, prefix='Testing')
    with graph.as_default():
        # Set up saver.
        saver = tf.train.Saver(max_to_keep=10)
        # Initialize variables.
        init_op = tf.global_variables_initializer()
        sess.run(init_op, feed_dict=None)
        # Start queue runners (used for loading videos on the fly)
        tf.train.start_queue_runners(sess=sess)
    if FLAGS.resume:
        model_file = tf.train.latest_checkpoint(save_dir)
        if FLAGS.restore_iter > 0:
            model_file = model_file[:model_file.index('model')] + 'model_' + str(FLAGS.restore_iter)
        if model_file:
            ind1 = model_file.index('model')
            resume_itr = int(model_file[ind1 + 6:])
            print("Restoring model weights from " + model_file)
            with graph.as_default():
                saver.restore(sess, model_file)
    if FLAGS.train:
        train(graph, model, saver, sess, save_dir, restore_itr=FLAGS.restore_iter)


if __name__ == "__main__":
    main()
# coding: utf-8 # from data_generator import DataGenerator # from evaluation.eval_reach import evaluate_vision_reach # from evaluation.eval_push import evaluate_push ## Dataset/method options ## Training options # 30k for pushing, 50k for reaching and placing # 15 for reaching, 15 for pushing, 12 for placing # 15 for reaching, 15 for pushing, 12 for placing # 0.001 for reaching, 0.01 for pushing and placing # 5 for placing # flags.DEFINE_float('clip_max', 20.0, 'maximum clipping value for fast gradient') # flags.DEFINE_float('clip_min', -20.0, 'minimum clipping value for fast gradient') ## Model options ## Logging, saving, and testing options # flags.DEFINE_string('save_dir', './amaml_human_pick_logs', 'directory for summaries and checkpoints.') # flags.DEFINE_bool('resume', True, 'resume training if there is a model available') # flags.DEFINE_integer('restore_iter', 20000, 'iteration to load model (-1 for latest model)') # flags.DEFINE_string('demo_type', 'robot', 'robot or human') #opposite to demo_type # flags.DEFINE_float('weight_xy', 0.999, '') # flags.DEFINE_float('weight_z', 0.001, '') # flags.DEFINE_float('weight_rxyz', 0.001, '') # flags.DEFINE_float('margin', 1.0, 'margin of loss') # flags.DEFINE_float('pre_margin', 1.0, 'margin of loss') # flags.DEFINE_bool('amaml', False, 'true for amaml') # flags.DEFINE_bool('amaml', False, 'true for amaml') # color_num = ['color_blue', 'color_green', 'color_orange'] # if if_train: #     index = np.random.randint(0, FLAGS.index_train_range) # else: #     index = np.random.randint(FLAGS.index_train_range, FLAGS.index_range) # print('color_list', color_list) # print('compare_color_list', compare_color_list) # print('object_list', object_list) # print('compare_object_list', compare_object_list) # compare_task_list = (task_list + np.random.randint(1, FLAGS.task_num - 1)) % FLAGS.task_num # print('task_list', task_list) # print('compare_task_list', compare_task_list) # print('demo_list', demo_list) # print('target_list', target_list) # print('compare_demo_list', compare_demo_list) # color_num = ['color_blue', 'color_green', 'color_orange'] # if if_train: #     index = np.random.randint(0, FLAGS.index_train_range) # else: #     index = np.random.randint(FLAGS.index_train_range, FLAGS.index_range) # print('element', element, 'labela and labelb', labelab, 'labelc', labelc) # print('----------------------------------------------------------------') # print('labelabs', labelabs.shape, 'labelcs', labelcs) # print('obsas, obsbs, actionas, actionbs, stateas, statebs', obsas.shape, obsbs.shape, actionas.shape, actionbs.shape, stateas.shape, statebs.shape) # print('obsas, obsbs, actionas, actionbs, stateas, statebs of generate_single_place_data', obsas.shape, obsbs.shape, actionas.shape, actionbs.shape, # stateas.shape, statebs.shape) # print('obsas, obsbs, actionas, actionbs, stateas, statebs', obsbs.shape, actionbs.shape, statebs.shape) # print('actionas is', actionas.shape) # place_place = np.concatenate([place, place], axis=2) # actionas = np.tile(pick_place, (1, 2, 1)) # actionas = np.concatenate([actionas, place_place], axis=1) # place_place = np.concatenate([place, place], axis=2) # actionbs = np.tile(pick_place, (1, 2, 1)) # actionbs = np.concatenate([actionbs, place_place], axis=1) # print('statebs',statebs) Train the model. # TEST_PRINT_INTERVAL = 1 # SAVE_INTERVAL = 10 # actual training. # print('actiones', actiones.shape, actiones[0]) # probas = np.ones([FLAGS.meta_batch_size, 1]) # probas = np.random.randint(0, 100, [FLAGS.meta_batch_size, 1]) / 100.0 # print('probas', probas, 'probes', probes) # print('obsas, obsbs, obscs, actionas, actionbs, actioncs, stateas, statebs, statecs, labelabs, labelcs', #       obsas.shape, obsbs.shape, obscs.shape, actionas.shape, actionbs.shape, actioncs.shape, #       stateas.shape, statebs.shape, statecs.shape, labelabs.shape, labelcs.shape) # input_tensors = [model.train_op, model.total_mix_loss, model.total_losses2[model.num_updates - 1], #                  model.test_act_op, model.total_semantic_loss, model.total_loss1, model.train_summ_op] # nums= get_num_params() # print('total parameters', nums) # print('demo state', stateas[-1]) # print('target state', statebs[-1]) # print ('results',results[3].shape,results[3][0][0], 'demos',actionbs[0][0]) # print('demo actionas', actionas[-1]) # train_writer.add_summary(results[-1], itr) # print('test demoa', actionas) # gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=FLAGS.gpu_memory_fraction) # tf_config = tf.ConfigProto(gpu_options=gpu_options) # sess = tf.Session(graph=graph) # data_generator = DataGenerator() # need to compute x_idx and img_idx from data_generator # TODO: figure out how to save summaries and checkpoints # put here for now # data_generator.generate_batches(noisy=FLAGS.use_noisy_demos) # with graph.as_default(): #     train_image_tensors = data_generator.make_batch_tensor(network_config, restore_iter=FLAGS.restore_iter) #     inputa = train_image_tensors[:, :FLAGS.update_batch_size*FLAGS.T, :] #     inputb = train_image_tensors[:, FLAGS.update_batch_size*FLAGS.T:, :] #     train_input_tensors = {'inputa': inputa, 'inputb': inputb} # model.init_network(graph, input_tensors=val_input_tensors, restore_iter=FLAGS.restore_iter, prefix='Validation_') # Set up saver. # Initialize variables. # Start queue runners (used for loading videos on the fly)