# Path: brainforge/reinforcement/abstract_agent.py
# class AgentBase(abc.ABC):
# 
#     type = ""
# 
#     def __init__(self, network, agentconfig: AgentConfig, **kw):
#         if agentconfig is None:
#             agentconfig = AgentConfig(**kw)
#         self.net = network
#         # self.shadow_net = network.layers.get_weights()
#         self.xp = agentconfig.replay_memory
#         self.cfg = agentconfig  # type: AgentConfig
#         self.shadow_net = None
# 
#     @abc.abstractmethod
#     def reset(self):
#         raise NotImplementedError
# 
#     @abc.abstractmethod
#     def sample(self, state, reward):
#         raise NotImplementedError
# 
#     @abc.abstractmethod
#     def accumulate(self, state, reward):
#         raise NotImplementedError
# 
#     def learn_batch(self, batch_size=None, callbacks=None):
#         X, Y = self.xp.replay(batch_size or self.cfg.bsize)
#         N = len(X)
#         if N < self.cfg.batch_size:
#             return 0.
#         cost = self.net.fit(X, Y, batch_size=32, verbose=0, epochs=1, callbacks=callbacks)
#         return np.mean(cost.history["loss"])
#         # return np.mean(costs)
# 
#     def push_weights(self):
#         W = self.net.layers.get_weights(unfold=True)
#         D = np.linalg.norm(self.shadow_net - W)
#         self.shadow_net *= (1. - self.cfg.tau)
#         self.shadow_net += self.cfg.tau * self.net.layers.get_weights(unfold=True)
#         return D / len(W)
# 
#     def pull_weights(self):
#         self.net.layers.set_weights(self.shadow_net, fold=True)
# 
# Path: brainforge/util/rl_util.py
# def discount_rewards(rwd, gamma=0.99):
#     """
#     Compute the discounted reward backwards in time
#     """
#     discounted_r = np.zeros_like(rwd)
#     running_add = rwd[-1]
#     for t, r in enumerate(rwd[::-1]):
#         running_add += gamma * r
#         discounted_r[t] = running_add
#     discounted_r[0] = rwd[-1]
#     return discounted_r[::-1]


import numpy as np
from .abstract_agent import AgentBase
from ..util.rl_util import discount_rewards

    """Policy Gradient"""

    type = "PolicyGradient"

    def __init__(self, network, nactions, agentconfig=None, **kw):
        super().__init__(network, agentconfig, **kw)
        self.actions = np.arange(nactions)
        self.action_labels = np.eye(nactions)
        self.X = []
        self.Y = []
        self.rewards = []
        self.grad = np.zeros((network.num_params,))

    def reset(self):
        self.X = []
        self.Y = []
        self.rewards = []

    def sample(self, state, reward):
        self.X.append(state)
        self.rewards.append(reward)
        probabilities = self.net.predict(state[None, ...])[0]
        action = np.random.choice(self.actions, p=probabilities)
        self.Y.append(self.action_labels[action])
        return action

    def accumulate(self, state, reward):
        R = np.array(self.rewards[1:] + [reward])
        if self.cfg.gamma > 0.:

# Path: brainforge/reinforcement/abstract_agent.py
# class AgentBase(abc.ABC):
# 
#     type = ""
# 
#     def __init__(self, network, agentconfig: AgentConfig, **kw):
#         if agentconfig is None:
#             agentconfig = AgentConfig(**kw)
#         self.net = network
#         # self.shadow_net = network.layers.get_weights()
#         self.xp = agentconfig.replay_memory
#         self.cfg = agentconfig  # type: AgentConfig
#         self.shadow_net = None
# 
#     @abc.abstractmethod
#     def reset(self):
#         raise NotImplementedError
# 
#     @abc.abstractmethod
#     def sample(self, state, reward):
#         raise NotImplementedError
# 
#     @abc.abstractmethod
#     def accumulate(self, state, reward):
#         raise NotImplementedError
# 
#     def learn_batch(self, batch_size=None, callbacks=None):
#         X, Y = self.xp.replay(batch_size or self.cfg.bsize)
#         N = len(X)
#         if N < self.cfg.batch_size:
#             return 0.
#         cost = self.net.fit(X, Y, batch_size=32, verbose=0, epochs=1, callbacks=callbacks)
#         return np.mean(cost.history["loss"])
#         # return np.mean(costs)
# 
#     def push_weights(self):
#         W = self.net.layers.get_weights(unfold=True)
#         D = np.linalg.norm(self.shadow_net - W)
#         self.shadow_net *= (1. - self.cfg.tau)
#         self.shadow_net += self.cfg.tau * self.net.layers.get_weights(unfold=True)
#         return D / len(W)
# 
#     def pull_weights(self):
#         self.net.layers.set_weights(self.shadow_net, fold=True)
# 
# Path: brainforge/util/rl_util.py
# def discount_rewards(rwd, gamma=0.99):
#     """
#     Compute the discounted reward backwards in time
#     """
#     discounted_r = np.zeros_like(rwd)
#     running_add = rwd[-1]
#     for t, r in enumerate(rwd[::-1]):
#         running_add += gamma * r
#         discounted_r[t] = running_add
#     discounted_r[0] = rwd[-1]
#     return discounted_r[::-1]


# Path: brainforge/reinforcement/policygradient.py
import numpy as np
from .abstract_agent import AgentBase
from ..util.rl_util import discount_rewards


    """Policy Gradient"""

    type = "PolicyGradient"

    def __init__(self, network, nactions, agentconfig=None, **kw):
        super().__init__(network, agentconfig, **kw)
        self.actions = np.arange(nactions)
        self.action_labels = np.eye(nactions)
        self.X = []
        self.Y = []
        self.rewards = []
        self.grad = np.zeros((network.num_params,))

    def reset(self):
        self.X = []
        self.Y = []
        self.rewards = []

    def sample(self, state, reward):
        self.X.append(state)
        self.rewards.append(reward)
        probabilities = self.net.predict(state[None, ...])[0]
        action = np.random.choice(self.actions, p=probabilities)
        self.Y.append(self.action_labels[action])
        return action

    def accumulate(self, state, reward):
        R = np.array(self.rewards[1:] + [reward])
        if self.cfg.gamma > 0.:

            R = discount_rewards(R, self.cfg.gamma)