# Path: brainforge/model/layerstack.py
# class LayerStack(Model):
# 
#     def __init__(self, input_shape, layers=()):
#         super().__init__(input_shape)
#         self.layers = []
#         self.architecture = []
#         self.learning = False
#         self._iterme = None
# 
#         self._add_input_layer(input_shape)
#         for layer in layers:
#             self.add(layer)
# 
#     def _add_input_layer(self, input_shape):
#         if isinstance(input_shape, int) or isinstance(input_shape, np.int64):
#             input_shape = (input_shape,)
#         from ..layers.core import InputLayer
#         inl = InputLayer(input_shape)
#         self.layers.append(inl)
#         self.architecture.append(str(inl))
# 
#     def add(self, layer):
#         layer.connect(self)
#         self.layers.append(layer)
#         self.architecture.append(str(layer))
# 
#     def pop(self):
#         self.layers.pop()
#         self.architecture.pop()
# 
#     def feedforward(self, X):
#         for layer in self.layers:
#             X = layer.feedforward(X)
#         return X
# 
#     def get_states(self, unfold=True):
#         hs = [layer.output for layer in self.layers]
#         return np.concatenate(hs) if unfold else hs
# 
#     def get_weights(self, unfold=True):
#         ws = [layer.get_weights(unfold=unfold) for
#               layer in self.layers if layer.trainable]
#         return np.concatenate(ws) if unfold else ws
# 
#     def set_weights(self, ws, fold=True):
#         trl = (l for l in self.layers if l.trainable)
#         if fold:
#             start = 0
#             for layer in trl:
#                 end = start + layer.num_params
#                 layer.set_weights(ws[start:end])
#                 start = end
#         else:
#             for w, layer in zip(ws, trl):
#                 layer.set_weights(w)
# 
#     def describe(self):
#         return "Architecture: " + "->".join(self.architecture),
# 
#     def reset(self):
#         for layer in (l for l in self.layers if l.trainable):
#             layer.reset()
# 
#     @property
#     def outshape(self):
#         return self.layers[-1].outshape
# 
#     @property
#     def num_params(self):
#         return sum(layer.num_params for layer in self.layers if layer.trainable)
# 
#     @property
#     def output(self):
#         return self.layers[-1].output
# 
#     @property
#     def trainable_layers(self):
#         return [layer for layer in self.layers if layer.trainable]
# 
#     def __iter__(self):
#         return self
# 
#     def __next__(self) -> LayerBase:
#         if self._iterme is None:
#             self._iterme = iter(self.layers)
#         try:
#             # noinspection PyTypeChecker
#             return next(self._iterme)
#         except StopIteration:
#             self._iterme = None
#             raise
# 
#     def __getitem__(self, item):
#         return self.layers.__getitem__(item)
# 
#     predict = feedforward
#     output_shape = outshape
# 
# Path: brainforge/learner/feedback_alignment.py
# class DirectFeedbackAlignment(Backpropagation):
# 
#     def __init__(self, layerstack, cost, optimizer, name="", **kw):
#         super().__init__(layerstack, cost, optimizer, name, **kw)
#         self.backwards_weights = np.concatenate(
#             [white(self.outshape[0], np.prod(layer.outshape))
#              for layer in self.trainable_layers[:-1]]
#         )
# 
#     def backpropagate(self, error):
#         m = len(error)
#         self.layers[-1].backpropagate(error)
#         all_deltas = error @ self.backwards_weights  # [m x net_out] [net_out x [layer_outs]] = [m x [layer_outs]]
#         start = 0
#         for layer in self.trainable_layers[1:-1]:
#             num_deltas = np.prod(layer.outshape)
#             end = start + num_deltas
#             delta = all_deltas[:, num_deltas]
#             layer.backpropagate(delta.reshape((m,) + layer.outshape))
#             start = end
# 
# Path: brainforge/layers/core.py
# class Dense(FFBase):
# 
#     def connect(self, brain):
#         inshape = brain.outshape
#         if len(inshape) != 1:
#             err = "Dense only accepts input shapes with 1 dimension!\n"
#             err += "Maybe you should consider placing <Flatten> before <Dense>?"
#             raise RuntimeError(err)
#         self.weights = white(inshape[0], self.neurons)
#         self.biases = zX(self.neurons)
#         super().connect(brain)
#         if self.compiled:
#             from .. import llatomic
#             print("Compiling DenseLayer...")
#             self.op = llatomic.DenseOp()
#         else:
#             self.op = atomic.DenseOp()
# 
#     def feedforward(self, X):
#         self.inputs = X
#         self.output = self.activation.forward(self.op.forward(
#             X, self.weights, self.biases
#         ))
#         return self.output
# 
#     def backpropagate(self, delta):
#         delta *= self.activation.backward(self.output)
#         nabla_w, nabla_b, dX = self.op.backward(self.inputs, delta, self.weights)
#         self.nabla_w += nabla_w
#         self.nabla_b += nabla_b
#         return dX
# 
#     def __str__(self):
#         return "Dense-{}-{}".format(self.neurons, str(self.activation)[:5])
# 
# class Flatten(Reshape):
#     pass


from brainforge.model import LayerStack
from brainforge.learner import DirectFeedbackAlignment
from brainforge.layers import Dense, Flatten
from verres.data import inmemory



def build_layerstack(input_shape, output_dim):
    return LayerStack(input_shape, [
        Flatten(),
        Dense(64, activation="tanh", compiled=True),
        Dense(output_dim, activation="softmax")])


def build_dfa_net(input_shape, output_dim):

# Path: brainforge/model/layerstack.py
# class LayerStack(Model):
# 
#     def __init__(self, input_shape, layers=()):
#         super().__init__(input_shape)
#         self.layers = []
#         self.architecture = []
#         self.learning = False
#         self._iterme = None
# 
#         self._add_input_layer(input_shape)
#         for layer in layers:
#             self.add(layer)
# 
#     def _add_input_layer(self, input_shape):
#         if isinstance(input_shape, int) or isinstance(input_shape, np.int64):
#             input_shape = (input_shape,)
#         from ..layers.core import InputLayer
#         inl = InputLayer(input_shape)
#         self.layers.append(inl)
#         self.architecture.append(str(inl))
# 
#     def add(self, layer):
#         layer.connect(self)
#         self.layers.append(layer)
#         self.architecture.append(str(layer))
# 
#     def pop(self):
#         self.layers.pop()
#         self.architecture.pop()
# 
#     def feedforward(self, X):
#         for layer in self.layers:
#             X = layer.feedforward(X)
#         return X
# 
#     def get_states(self, unfold=True):
#         hs = [layer.output for layer in self.layers]
#         return np.concatenate(hs) if unfold else hs
# 
#     def get_weights(self, unfold=True):
#         ws = [layer.get_weights(unfold=unfold) for
#               layer in self.layers if layer.trainable]
#         return np.concatenate(ws) if unfold else ws
# 
#     def set_weights(self, ws, fold=True):
#         trl = (l for l in self.layers if l.trainable)
#         if fold:
#             start = 0
#             for layer in trl:
#                 end = start + layer.num_params
#                 layer.set_weights(ws[start:end])
#                 start = end
#         else:
#             for w, layer in zip(ws, trl):
#                 layer.set_weights(w)
# 
#     def describe(self):
#         return "Architecture: " + "->".join(self.architecture),
# 
#     def reset(self):
#         for layer in (l for l in self.layers if l.trainable):
#             layer.reset()
# 
#     @property
#     def outshape(self):
#         return self.layers[-1].outshape
# 
#     @property
#     def num_params(self):
#         return sum(layer.num_params for layer in self.layers if layer.trainable)
# 
#     @property
#     def output(self):
#         return self.layers[-1].output
# 
#     @property
#     def trainable_layers(self):
#         return [layer for layer in self.layers if layer.trainable]
# 
#     def __iter__(self):
#         return self
# 
#     def __next__(self) -> LayerBase:
#         if self._iterme is None:
#             self._iterme = iter(self.layers)
#         try:
#             # noinspection PyTypeChecker
#             return next(self._iterme)
#         except StopIteration:
#             self._iterme = None
#             raise
# 
#     def __getitem__(self, item):
#         return self.layers.__getitem__(item)
# 
#     predict = feedforward
#     output_shape = outshape
# 
# Path: brainforge/learner/feedback_alignment.py
# class DirectFeedbackAlignment(Backpropagation):
# 
#     def __init__(self, layerstack, cost, optimizer, name="", **kw):
#         super().__init__(layerstack, cost, optimizer, name, **kw)
#         self.backwards_weights = np.concatenate(
#             [white(self.outshape[0], np.prod(layer.outshape))
#              for layer in self.trainable_layers[:-1]]
#         )
# 
#     def backpropagate(self, error):
#         m = len(error)
#         self.layers[-1].backpropagate(error)
#         all_deltas = error @ self.backwards_weights  # [m x net_out] [net_out x [layer_outs]] = [m x [layer_outs]]
#         start = 0
#         for layer in self.trainable_layers[1:-1]:
#             num_deltas = np.prod(layer.outshape)
#             end = start + num_deltas
#             delta = all_deltas[:, num_deltas]
#             layer.backpropagate(delta.reshape((m,) + layer.outshape))
#             start = end
# 
# Path: brainforge/layers/core.py
# class Dense(FFBase):
# 
#     def connect(self, brain):
#         inshape = brain.outshape
#         if len(inshape) != 1:
#             err = "Dense only accepts input shapes with 1 dimension!\n"
#             err += "Maybe you should consider placing <Flatten> before <Dense>?"
#             raise RuntimeError(err)
#         self.weights = white(inshape[0], self.neurons)
#         self.biases = zX(self.neurons)
#         super().connect(brain)
#         if self.compiled:
#             from .. import llatomic
#             print("Compiling DenseLayer...")
#             self.op = llatomic.DenseOp()
#         else:
#             self.op = atomic.DenseOp()
# 
#     def feedforward(self, X):
#         self.inputs = X
#         self.output = self.activation.forward(self.op.forward(
#             X, self.weights, self.biases
#         ))
#         return self.output
# 
#     def backpropagate(self, delta):
#         delta *= self.activation.backward(self.output)
#         nabla_w, nabla_b, dX = self.op.backward(self.inputs, delta, self.weights)
#         self.nabla_w += nabla_w
#         self.nabla_b += nabla_b
#         return dX
# 
#     def __str__(self):
#         return "Dense-{}-{}".format(self.neurons, str(self.activation)[:5])
# 
# class Flatten(Reshape):
#     pass


# Path: xperiments/xp_feedbackalignment.py
from brainforge.model import LayerStack
from brainforge.learner import DirectFeedbackAlignment
from brainforge.layers import Dense, Flatten
from verres.data import inmemory




def build_layerstack(input_shape, output_dim):
    return LayerStack(input_shape, [
        Flatten(),
        Dense(64, activation="tanh", compiled=True),
        Dense(output_dim, activation="softmax")])


def build_dfa_net(input_shape, output_dim):

    return DirectFeedbackAlignment(build_layerstack(input_shape, output_dim),