"""
RNN trainable weight strategies
===============================

Functions expect a block adjacency matrix of the form


.. code::

     |From
     |-------
   T |+ + 0 -
   o |0 + + -
     |0 0 + -
     |+ + + -

Where  "+" denotes the exc connections, "-" the inh connections
and "0" no connection.


This module implements strategies for the BlockRnn class.
"""

from .blockstructures import BlockMatrix, BlockVector
from abc import ABC, abstractmethod


class Trainer(ABC):
    """
    Abstract BlockRnn training strategy.
    """
    def get_config(self):
        return {"name": self.__class__.__name__}

    @abstractmethod
    def adjacency_trainable_weights(self, adjacency: BlockMatrix):
        return []

    @abstractmethod
    def input_trainable_weights(self, W, b):
        return []

    def input_non_trainable_weights(self, W, b):
        trainable = self.input_trainable_weights(W, b)
        return [weight for weight in [W, b] if weight not in trainable]

    def adjacency_non_trainable_weights(self, adjacency: BlockMatrix):
        trainable = self.adjacency_trainable_weights(adjacency)
        return [weight for weight in adjacency.flatten() if weight not in trainable and weight != 0]

    def get_trainable_weights(self, W, b, adjacency):
        return self.input_trainable_weights(W, b) + self.adjacency_trainable_weights(adjacency)

    def get_non_trainable_weights(self, W, b, adjacency):
        return self.input_non_trainable_weights(W, b) + self.adjacency_non_trainable_weights(adjacency)


class ExcitatoryForwardTrainer(Trainer):
    """
    Train E -> E forward connections.

    Do not train input.
    """
    def input_trainable_weights(self, W, b):
        return []

    def adjacency_trainable_weights(self, adjacency: BlockMatrix):
        return list(adjacency.get_diag(1)[:-1])
