# Path: brainforge/learner/backpropagation.py
# class Backpropagation(Learner):
# 
#     def __init__(self, layerstack, cost="mse", optimizer="sgd", name="", **kw):
#         super().__init__(layerstack, cost, name, **kw)
#         self.optimizer = (
#             optimizer if isinstance(optimizer, GradientDescent) else optimizers[optimizer]()
#         )
#         self.optimizer.initialize(nparams=self.layers.num_params)
# 
#     def learn_batch(self, X, Y, w=None, metrics=(), update=True):
#         m = len(X)
#         preds = self.predict(X)
#         delta = self.cost.derivative(preds, Y)
#         if w is not None:
#             delta *= w[:, None]
#         self.backpropagate(delta)
#         if update:
#             self.update(m)
# 
#         train_metrics = {"cost": self.cost(self.output, Y) / m}
#         if metrics:
#             for metric in metrics:
#                 train_metrics[str(metric).lower()] = metric(preds, Y) / m
#         return train_metrics
# 
#     def backpropagate(self, error):
#         for layer in self.layers[-1:0:-1]:
#             error = layer.backpropagate(error)
#         return error
# 
#     def update(self, m):
#         W = self.layers.get_weights(unfold=True)
#         gW = self.get_gradients(unfold=True)
#         oW = self.optimizer.optimize(W, gW, m)
#         self.layers.set_weights(oW, fold=True)
#         self.zero_gradients()
# 
#     def get_weights(self, unfold=True):
#         return self.layers.get_weights(unfold=unfold)
# 
#     def set_weigts(self, ws, fold=True):
#         self.layers.set_weights(ws=ws, fold=fold)
# 
#     def get_gradients(self, unfold=True):
#         grads = [l.gradients for l in self.layers if l.trainable]
#         if unfold:
#             grads = np.concatenate(grads)
#         return grads
# 
#     def set_gradients(self, gradients, 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(gradients[start:end])
#                 start = end
#         else:
#             for w, layer in zip(gradients, trl):
#                 layer.set_weights(w)
# 
#     def zero_gradients(self):
#         for layer in self.layers:
#             if not layer.trainable:
#                 continue
#             layer.nabla_w *= 0
#             layer.nabla_b *= 0
# 
#     @property
#     def num_params(self):
#         return self.layers.num_params
# 
#     @property
#     def output_shape(self):
#         return self.layers.output_shape
# 
#     @property
#     def input_shape(self):
#         return self.layers.input_shape
# 
# Path: brainforge/metrics/costs.py
# class CostFunction:
# class _MeanSquaredError(CostFunction):
# class _CategoricalCrossEntropy(CostFunction):
# class _BinaryCrossEntropy(CostFunction):
# class _Hinge(CostFunction):
#     def __call__(self, outputs, targets):
#     def __str__(self):
#     def derivative(outputs, targets):
#     def __call__(self, outputs, targets):
#     def true_derivative(outputs, targets):
#     def __call__(self, outputs: np.ndarray, targets: np.ndarray):
#     def true_derivative(outputs, targets):
#     def __call__(self, outputs: np.ndarray, targets: np.ndarray):
#     def true_derivative(outputs, targets):
#     def __call__(self, outputs, targets):
#     def derivative(outputs, targets):
# def get(cost_function):


import numpy as np
from .backpropagation import Backpropagation
from ..metrics import mse



class LocalCorrelationAligment(Backpropagation):

    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

    def learn_batch(self, X, Y, w=None, metrics=(), update=True):
        m = len(X)
        Y_correlations = np.corrcoef(Y.reshape((m, -1)))

        for layer in self.trainable_layers:
            h_correlations = np.corrcoef(layer.output.reshape(m, -1))

# Path: brainforge/learner/backpropagation.py
# class Backpropagation(Learner):
# 
#     def __init__(self, layerstack, cost="mse", optimizer="sgd", name="", **kw):
#         super().__init__(layerstack, cost, name, **kw)
#         self.optimizer = (
#             optimizer if isinstance(optimizer, GradientDescent) else optimizers[optimizer]()
#         )
#         self.optimizer.initialize(nparams=self.layers.num_params)
# 
#     def learn_batch(self, X, Y, w=None, metrics=(), update=True):
#         m = len(X)
#         preds = self.predict(X)
#         delta = self.cost.derivative(preds, Y)
#         if w is not None:
#             delta *= w[:, None]
#         self.backpropagate(delta)
#         if update:
#             self.update(m)
# 
#         train_metrics = {"cost": self.cost(self.output, Y) / m}
#         if metrics:
#             for metric in metrics:
#                 train_metrics[str(metric).lower()] = metric(preds, Y) / m
#         return train_metrics
# 
#     def backpropagate(self, error):
#         for layer in self.layers[-1:0:-1]:
#             error = layer.backpropagate(error)
#         return error
# 
#     def update(self, m):
#         W = self.layers.get_weights(unfold=True)
#         gW = self.get_gradients(unfold=True)
#         oW = self.optimizer.optimize(W, gW, m)
#         self.layers.set_weights(oW, fold=True)
#         self.zero_gradients()
# 
#     def get_weights(self, unfold=True):
#         return self.layers.get_weights(unfold=unfold)
# 
#     def set_weigts(self, ws, fold=True):
#         self.layers.set_weights(ws=ws, fold=fold)
# 
#     def get_gradients(self, unfold=True):
#         grads = [l.gradients for l in self.layers if l.trainable]
#         if unfold:
#             grads = np.concatenate(grads)
#         return grads
# 
#     def set_gradients(self, gradients, 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(gradients[start:end])
#                 start = end
#         else:
#             for w, layer in zip(gradients, trl):
#                 layer.set_weights(w)
# 
#     def zero_gradients(self):
#         for layer in self.layers:
#             if not layer.trainable:
#                 continue
#             layer.nabla_w *= 0
#             layer.nabla_b *= 0
# 
#     @property
#     def num_params(self):
#         return self.layers.num_params
# 
#     @property
#     def output_shape(self):
#         return self.layers.output_shape
# 
#     @property
#     def input_shape(self):
#         return self.layers.input_shape
# 
# Path: brainforge/metrics/costs.py
# class CostFunction:
# class _MeanSquaredError(CostFunction):
# class _CategoricalCrossEntropy(CostFunction):
# class _BinaryCrossEntropy(CostFunction):
# class _Hinge(CostFunction):
#     def __call__(self, outputs, targets):
#     def __str__(self):
#     def derivative(outputs, targets):
#     def __call__(self, outputs, targets):
#     def true_derivative(outputs, targets):
#     def __call__(self, outputs: np.ndarray, targets: np.ndarray):
#     def true_derivative(outputs, targets):
#     def __call__(self, outputs: np.ndarray, targets: np.ndarray):
#     def true_derivative(outputs, targets):
#     def __call__(self, outputs, targets):
#     def derivative(outputs, targets):
# def get(cost_function):


# Path: brainforge/learner/local_correlation.py
import numpy as np
from .backpropagation import Backpropagation
from ..metrics import mse




class LocalCorrelationAligment(Backpropagation):

    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

    def learn_batch(self, X, Y, w=None, metrics=(), update=True):
        m = len(X)
        Y_correlations = np.corrcoef(Y.reshape((m, -1)))

        for layer in self.trainable_layers:
            h_correlations = np.corrcoef(layer.output.reshape(m, -1))

            local_error = mse(h_correlations.flat, Y_correlations.flat)