# Path: brainforge/reinforcement/experience.py
# def replay_memory_factory(limit, mode, time, downsample=0):
#     if time > 1:
#         return TimeExperience(time, limit, downsample)
#     else:
#         return Experience(limit, mode, downsample)
# 
# 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)


import numpy as np
from .experience import replay_memory_factory
from .abstract_agent import AgentBase
        from pickle import loads, dumps
        R = np.array(self.rewards[1:] + [reward])
        ix = tuple(self.actions)
        Y = Q
        Y[range(len(Y)), ix] = -(R + Y.max(axis=1) * self.cfg.gamma)
        Y[-1, ix[-1]] = -reward
        self.xp.remember(X, Y)
        self.cfg.epsilon *= self.cfg.epsilon_decay
        self.reset()

    def accumulate_multiple(self, states, rewards):
        N = len(states)
        qs = self.net.predict(states)
        self.inputs = np.stack(self.inputs)
        Q = np.stack(self.predictions[N:] + list(qs))
        self.rewards = np.array(self.rewards[N:] + list(rewards))
        ix = tuple(self.actions)
        Y = Q
        Y[range(len(Y)), ix] = -(self.rewards + Y.max(axis=1) * self.cfg.gamma)
        Y[-N:, ix[-1]] = -rewards
        self.xp.remember(self.inputs, Y)
        self.reset()


class DDQN(DQN):

    type = "DDQN"

    def __init__(self, network, num_actions, agentconfig, **kw):
        super().__init__(network, num_actions, agentconfig, **kw)
        self.double = [network, loads(dumps(network))]

# Path: brainforge/reinforcement/experience.py
# def replay_memory_factory(limit, mode, time, downsample=0):
#     if time > 1:
#         return TimeExperience(time, limit, downsample)
#     else:
#         return Experience(limit, mode, downsample)
# 
# 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/reinforcement/qlearning.py
import numpy as np
from .experience import replay_memory_factory
from .abstract_agent import AgentBase
        from pickle import loads, dumps

        R = np.array(self.rewards[1:] + [reward])
        ix = tuple(self.actions)
        Y = Q
        Y[range(len(Y)), ix] = -(R + Y.max(axis=1) * self.cfg.gamma)
        Y[-1, ix[-1]] = -reward
        self.xp.remember(X, Y)
        self.cfg.epsilon *= self.cfg.epsilon_decay
        self.reset()

    def accumulate_multiple(self, states, rewards):
        N = len(states)
        qs = self.net.predict(states)
        self.inputs = np.stack(self.inputs)
        Q = np.stack(self.predictions[N:] + list(qs))
        self.rewards = np.array(self.rewards[N:] + list(rewards))
        ix = tuple(self.actions)
        Y = Q
        Y[range(len(Y)), ix] = -(self.rewards + Y.max(axis=1) * self.cfg.gamma)
        Y[-N:, ix[-1]] = -rewards
        self.xp.remember(self.inputs, Y)
        self.reset()


class DDQN(DQN):

    type = "DDQN"

    def __init__(self, network, num_actions, agentconfig, **kw):
        super().__init__(network, num_actions, agentconfig, **kw)
        self.double = [network, loads(dumps(network))]

        self.doublexp = [replay_memory_factory(agentconfig.xpsize, "drop", agentconfig.time) for _ in range(2)]