# 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()


def train_value():
    policy_model = get_model('policy')

    environ = Reactor(policy_model, 'train_log/epoch-399')
    value_model = get_value_model('value')
    sess = tf.Session()

    target_y = tf.placeholder(config.dtype, shape=[None], name='target_y')
    q_loss = tf.nn.l2_loss(value_model.loss - target_y)
    trian_step = config.optimizer.minimize(q_loss)

    value_pred = value_model.pred

    dataset = load_data('train')
    info, label = dataset.next_batch(config.minibatch_size)
    state = info[:2]
    move = info[2]
    params = info[3:]
    # for loop, choose move and react, then train
    while True:
        # get state transfer
        input_dict = {}
        for var, data in zip(value_model.inputs, state +[move]+params):
            input_dict[var] = data
        pred = sess.run(value_pred, feed_dict=input_dict)
        if np.random.random() > config.randplay:
            choose_move = batch_max_to_onehot(pred)
        else:

# 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()



def train_value():
    policy_model = get_model('policy')

    environ = Reactor(policy_model, 'train_log/epoch-399')
    value_model = get_value_model('value')
    sess = tf.Session()

    target_y = tf.placeholder(config.dtype, shape=[None], name='target_y')
    q_loss = tf.nn.l2_loss(value_model.loss - target_y)
    trian_step = config.optimizer.minimize(q_loss)

    value_pred = value_model.pred

    dataset = load_data('train')
    info, label = dataset.next_batch(config.minibatch_size)
    state = info[:2]
    move = info[2]
    params = info[3:]
    # for loop, choose move and react, then train
    while True:
        # get state transfer
        input_dict = {}
        for var, data in zip(value_model.inputs, state +[move]+params):
            input_dict[var] = data
        pred = sess.run(value_pred, feed_dict=input_dict)
        if np.random.random() > config.randplay:
            choose_move = batch_max_to_onehot(pred)
        else:

            choose_move = batch_max_to_onehot(np.random.random(pred.shape))