# Path: util/dataset.py
# def load_data(_type):
#     '''
#     return dataset which yeild minibatch data
#     '''
#     data = Dataset('../data', _type)
#     return data
# 
# Path: util/tools.py
# def batch_max_to_onehot(tensor):
#     shape = tensor.shape
#     index = np.argmax(tensor.reshape(shape[0], -1), axis=1)
#     oht = np.zeros([shape[0], np.prod(shape[1:])])
#     oht[np.arange(shape[0]), index] = 1
#     oht = oht.reshape(shape)
#     return oht


import tensorflow as tf
import numpy as np
from util.dataset import load_data
from util.tools import batch_max_to_onehot
from config import config
from value_dataset import Reactor
from value_model import get_value_model
from policy_model import get_model
        from IPython import embed; embed()
            next_dict[var] = data[np.logic_not(if_terminated)]
        next_pred = sess.run(value_pred, feed_dict=next_dict)
        next_max = next_pred.reshape(next_pred.shape[0]).max(axis=1)
        reward[if_terminated] += config.gamma * next_max

        # train value_nets
        input_dict[value_model.label] = choose_move
        input_dict[target_y] = reward
        sess.run(train_step, feed_dict=input_dict)

        # refresh data
        state = next_state
        move = next_move
        params = next_params
        if any(if_terminated):
            info, label = dataset.next_batch(if_terminated.sum())
            for i in range(len(state)):
                state[i][if_terminated] = info[i]
            move[if_terminated] = info[2]
            for i in range(len(params)):
                params[i][if_terminated] = info[3+i]

        # save step
        if (i+1)%config.check_point == 0:
            save_path = saver.save(sess, "%s/value-epoch-%d" %(config.save_path, i))
            print("Model saved in file: %s" % save_path)

def train_policy(load_path):
    policy_model = get_model('policy')
    value_model = get_value_model('value')

# Path: util/dataset.py
# def load_data(_type):
#     '''
#     return dataset which yeild minibatch data
#     '''
#     data = Dataset('../data', _type)
#     return data
# 
# Path: util/tools.py
# def batch_max_to_onehot(tensor):
#     shape = tensor.shape
#     index = np.argmax(tensor.reshape(shape[0], -1), axis=1)
#     oht = np.zeros([shape[0], np.prod(shape[1:])])
#     oht[np.arange(shape[0]), index] = 1
#     oht = oht.reshape(shape)
#     return oht


# Path: rl_train/trainer.py
import tensorflow as tf
import numpy as np
from util.dataset import load_data
from util.tools import batch_max_to_onehot
from config import config
from value_dataset import Reactor
from value_model import get_value_model
from policy_model import get_model
        from IPython import embed; embed()

            next_dict[var] = data[np.logic_not(if_terminated)]
        next_pred = sess.run(value_pred, feed_dict=next_dict)
        next_max = next_pred.reshape(next_pred.shape[0]).max(axis=1)
        reward[if_terminated] += config.gamma * next_max

        # train value_nets
        input_dict[value_model.label] = choose_move
        input_dict[target_y] = reward
        sess.run(train_step, feed_dict=input_dict)

        # refresh data
        state = next_state
        move = next_move
        params = next_params
        if any(if_terminated):
            info, label = dataset.next_batch(if_terminated.sum())
            for i in range(len(state)):
                state[i][if_terminated] = info[i]
            move[if_terminated] = info[2]
            for i in range(len(params)):
                params[i][if_terminated] = info[3+i]

        # save step
        if (i+1)%config.check_point == 0:
            save_path = saver.save(sess, "%s/value-epoch-%d" %(config.save_path, i))
            print("Model saved in file: %s" % save_path)

def train_policy(load_path):
    policy_model = get_model('policy')
    value_model = get_value_model('value')

    train_data = load_data(value_model, 'train_log/value.ckpt')