# Path: src/net.py
# class Net(object):
#     # Sugar syntax
#     def __init__(self):
#         self._server = ParameterServer()
#         self._finalized = False
#         self._dagraph = DAG()
#         coated = sugar.__dict__
#         for name, fun in coated.items():
#             if callable(fun): setattr(Net, name, fun)
#  
#     # To be sugar coated
#     def _x(self, name, *args):
#         assert not self._finalized, \
#         'Graph is finalized by setting an optimizer'
#         return self._dagraph.register_for_slot(
#             name, self._server, *args)
# 
#     def optimize(self, ptr, optimizer,
#                  *args, **kwargs):
#         self._finalized = True
#         self._dagraph.set_closure(ptr)
#         self._server.set_optimizer(
#             optimizer, *args, **kwargs)
# 
#     def forward(self, fetches, feed):
#         return self._dagraph.forward(fetches, feed)
# 
#     def train(self, fetches, feed):
#         val = self._dagraph.forward_to_leaf(fetches, feed)
#         self._dagraph.backward()
#         self._server.apply_optimizer()
#         return val
# 
#     def save_checkpoint(self, file_name):
#         self._server.save(file_name)
# 
#     def load_checkpoint(self, file_name, *args, **kwargs):
#         self._server.load(file_name, *args, **kwargs)
# 
# Path: src/utils/misc.py
# def accuracy(predict, truth):
# def read_mnist():
# def xavier(shape):
# def extract(name, dfault, **kwargs):
# def nxshape(volume, shape):
# def parse_for(class_type, *args):
# 
# Path: src/agents/qfunction.py
# class Qfunction(object):
#     def __init__(self, is_target, 
#                  inp_dim, hid_dim, out_dim):
#         self._net = Net()
#         self._inp, fc1_param, fc2_param, self._out = \
#             _MLP2(self._net, inp_dim, hid_dim, out_dim)
#         self._vars = fc1_param + fc2_param
#         if is_target:
#             self._feed = list()
#             self._assign_ops = list()
#             self._build_assign(fc1_param, (inp_dim, hid_dim))
#             self._build_assign(fc2_param, (hid_dim, out_dim))
#         elif not is_target:
#             self._build_loss()
#     
#     def forward(self, inp_feed):
#         out = self._net.forward(
#             [self._out],
#             {self._inp: inp_feed})[0]
#         return out[0, 0]
#     
#     def train(self, x_feed, y_feed):
#         loss, = self._net.train([], {
#             self._inp: x_feed,
#             self._y: y_feed })
#         return loss
# 
#     def _build_loss(self):
#         self._y = self._net.portal((1,))
#         self._loss = self._net.l2(self._out, self._y)
#         self._net.optimize(self._loss, 'adam', 1e-3)
#      
#     def _build_assign(self, fc_weights, shape):
#         pw, pb = (
#             self._net.portal(shape),
#             self._net.portal(shape[1:2]))
#         self._feed += [pw, pb]
#         self._assign_ops += [
#             self._net.assign(fc_weights[0], pw),
#             self._net.assign(fc_weights[1], pb)]
#     
#     def yield_params_values(self):
#         returns = list()
#         for var in self._vars:
#             returns.append(
#                 self._net.forward([var], {})[0])
#         return returns
#     
#     def assign(self, values):
#         feed_dict = dict(zip(self._feed, values))
#         self._net.forward(self._assign_ops, feed_dict)
#     
#     def save(self, file_name):
#         self._net.save_checkpoint(file_name)
#     
#     def load(self, file_name):
#         self._net.load_checkpoint(file_name)


from src.net import Net
from src.utils.misc import uniform, xavier, perm
from .qfunction import Qfunction
import numpy as np
        self._moving_q = Qfunction(False, *mlp_dims)
        self._target_q = Qfunction(True, *mlp_dims)
    
    def _best_act(self, qfunc, observe):
        max_q_out = None
        observe = list(observe)
        for action in self._action_space:
            q_inp = np.array([observe + action])
            q_out = qfunc.forward(q_inp)

            if not max_q_out or max_q_out < q_out: 
                best_act = action
                max_q_out = q_out
                
        return best_act, max_q_out

    def act(self, observe, epsilon):
        rand = uniform()
        if rand <= epsilon: # exploration
            best_act = None
            q_val = None
        elif rand > epsilon: # exploitation
            best_act, q_val = \
                self._best_act(
                    self._moving_q, observe)
        return best_act, q_val
    
    def store_and_learn(self, transition):
        self._exreplay.append(transition)


# Path: src/net.py
# class Net(object):
#     # Sugar syntax
#     def __init__(self):
#         self._server = ParameterServer()
#         self._finalized = False
#         self._dagraph = DAG()
#         coated = sugar.__dict__
#         for name, fun in coated.items():
#             if callable(fun): setattr(Net, name, fun)
#  
#     # To be sugar coated
#     def _x(self, name, *args):
#         assert not self._finalized, \
#         'Graph is finalized by setting an optimizer'
#         return self._dagraph.register_for_slot(
#             name, self._server, *args)
# 
#     def optimize(self, ptr, optimizer,
#                  *args, **kwargs):
#         self._finalized = True
#         self._dagraph.set_closure(ptr)
#         self._server.set_optimizer(
#             optimizer, *args, **kwargs)
# 
#     def forward(self, fetches, feed):
#         return self._dagraph.forward(fetches, feed)
# 
#     def train(self, fetches, feed):
#         val = self._dagraph.forward_to_leaf(fetches, feed)
#         self._dagraph.backward()
#         self._server.apply_optimizer()
#         return val
# 
#     def save_checkpoint(self, file_name):
#         self._server.save(file_name)
# 
#     def load_checkpoint(self, file_name, *args, **kwargs):
#         self._server.load(file_name, *args, **kwargs)
# 
# Path: src/utils/misc.py
# def accuracy(predict, truth):
# def read_mnist():
# def xavier(shape):
# def extract(name, dfault, **kwargs):
# def nxshape(volume, shape):
# def parse_for(class_type, *args):
# 
# Path: src/agents/qfunction.py
# class Qfunction(object):
#     def __init__(self, is_target, 
#                  inp_dim, hid_dim, out_dim):
#         self._net = Net()
#         self._inp, fc1_param, fc2_param, self._out = \
#             _MLP2(self._net, inp_dim, hid_dim, out_dim)
#         self._vars = fc1_param + fc2_param
#         if is_target:
#             self._feed = list()
#             self._assign_ops = list()
#             self._build_assign(fc1_param, (inp_dim, hid_dim))
#             self._build_assign(fc2_param, (hid_dim, out_dim))
#         elif not is_target:
#             self._build_loss()
#     
#     def forward(self, inp_feed):
#         out = self._net.forward(
#             [self._out],
#             {self._inp: inp_feed})[0]
#         return out[0, 0]
#     
#     def train(self, x_feed, y_feed):
#         loss, = self._net.train([], {
#             self._inp: x_feed,
#             self._y: y_feed })
#         return loss
# 
#     def _build_loss(self):
#         self._y = self._net.portal((1,))
#         self._loss = self._net.l2(self._out, self._y)
#         self._net.optimize(self._loss, 'adam', 1e-3)
#      
#     def _build_assign(self, fc_weights, shape):
#         pw, pb = (
#             self._net.portal(shape),
#             self._net.portal(shape[1:2]))
#         self._feed += [pw, pb]
#         self._assign_ops += [
#             self._net.assign(fc_weights[0], pw),
#             self._net.assign(fc_weights[1], pb)]
#     
#     def yield_params_values(self):
#         returns = list()
#         for var in self._vars:
#             returns.append(
#                 self._net.forward([var], {})[0])
#         return returns
#     
#     def assign(self, values):
#         feed_dict = dict(zip(self._feed, values))
#         self._net.forward(self._assign_ops, feed_dict)
#     
#     def save(self, file_name):
#         self._net.save_checkpoint(file_name)
#     
#     def load(self, file_name):
#         self._net.load_checkpoint(file_name)


# Path: src/agents/agent.py
from src.net import Net
from src.utils.misc import uniform, xavier, perm
from .qfunction import Qfunction
import numpy as np

        self._moving_q = Qfunction(False, *mlp_dims)
        self._target_q = Qfunction(True, *mlp_dims)
    
    def _best_act(self, qfunc, observe):
        max_q_out = None
        observe = list(observe)
        for action in self._action_space:
            q_inp = np.array([observe + action])
            q_out = qfunc.forward(q_inp)

            if not max_q_out or max_q_out < q_out: 
                best_act = action
                max_q_out = q_out
                
        return best_act, max_q_out

    def act(self, observe, epsilon):
        rand = uniform()
        if rand <= epsilon: # exploration
            best_act = None
            q_val = None
        elif rand > epsilon: # exploitation
            best_act, q_val = \
                self._best_act(
                    self._moving_q, observe)
        return best_act, q_val
    
    def store_and_learn(self, transition):
        self._exreplay.append(transition)


        shuffle = perm(len(self._exreplay))[:self._batch]