# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_synapses.ipynb.

# %% auto 0
__all__ = ['Synapse', 'SimpleDenseSynapse', 'DenseSynapse', 'DenseTensorSynapseWithHiddenLayer',
           'DenseMatrixSynapseWithHiddenLayer', 'ConvSynapse', 'ConvSynapseWithPool', 'AttentionSynapse',
           'SelfAttentionSynapse', 'BinaryMixerSynapse']

# %% ../nbs/02_synapses.ipynb 6
import jax
import jax.numpy as jnp
from typing import *
import treex as tx
from abc import ABC, abstractmethod
from flax import linen as nn
from .lagrangians import *
import functools as ft
from fastcore.meta import delegates
from fastcore.utils import *
from fastcore.basics import *
from string import ascii_letters
from flax.linen.pooling import max_pool, avg_pool
from einops import rearrange

# %% ../nbs/02_synapses.ipynb 8
class Synapse(tx.Module, ABC):
    """The simple interface class for any synapse. Define an alignment function through `__call__` that returns a scalar.

    The energy is simply the negative of this function.
    """

    def energy(self, *gs):
        return -self(*gs)

    @abstractmethod
    def __call__(self, *gs):
        """The alignment function of a synapse"""
        pass

# %% ../nbs/02_synapses.ipynb 11
class SimpleDenseSynapse(Synapse):
    """The simplest of dense synapses that connects two layers (with vectorized activations) together"""
    W: jnp.ndarray = tx.Parameter.node() # treex's preferred way of declaring an attribute as a parameter
    def __call__(self, g1, g2):
        if self.initializing():
            self.W = nn.initializers.normal(0.02)(tx.next_key(), g1.shape + g2.shape)
        return g1 @ self.W @ g2

# %% ../nbs/02_synapses.ipynb 15
class DenseSynapse(Synapse):
    """A dense synapse that aligns the representations of any number of `gs`.

    The one learnable parameter `W` is a tensor with a dimension for each connected layer.
    In the case of 2 layers, this is the traditional learnable matrix synapse.
    In cases `N>2` layers this is a new kind of layer where the learnable parameter is an N dimensional tensor.

    By default, this will flatten all inputs as needed to treat all activations as vectors. 
    
    The number of layers we can align with this synapse is capped at the number of ranks that JAX stores (<255),
    but you'll probably run out of memory first..
    """

    W: jnp.ndarray = tx.Parameter.node()
    stdinit: float = 0.02
    flatten_args: bool = True

    def __init__(self, stdinit: float = 0.02, flatten_args=True):
        self.stdinit = stdinit
        self.flatten_args = flatten_args

    def __call__(self, *gs):
        if self.initializing():
            ndims_total = jnp.sum(jnp.array([len(g.shape) for g in gs]))
            assert (
                ndims_total <= 52
            ), f"We are limited to english ASCII letters. We cannot connect more than 52 dimensions. Got {ndims_total} total dimensions."
            if self.flatten_args:
                gshapes = tuple([g_.size for g_ in gs])
            else:
                gshapes = tuple([g_.shape[-1] for g_ in gs])
            self.W = nn.initializers.normal(self.stdinit)(tx.next_key(), gshapes)
        if self.flatten_args:
            gs = [g_.ravel() for g_ in gs]
            abcs = ascii_letters[: len(gs)]
            einsum_arg = ",".join([abcs, ",".join(abcs)]) + "->"
            return jnp.einsum(einsum_arg, self.W, *gs)
        else:
            # Design the einsum to take letter positions corresponding to the
            Wabcs = ""
            gabcs = []
            i = 0
            for g in gs:
                ndims = len(g.shape)
                Wabcs += ascii_letters[(i - 1) + ndims]
                gabcs.append(ascii_letters[i : i + ndims])
                i = i + ndims
            einsum_arg = ",".join([Wabcs, ",".join(gabcs)]) + "->"
            return jnp.einsum(einsum_arg, self.W, *gs)

# %% ../nbs/02_synapses.ipynb 18
class DenseTensorSynapseWithHiddenLayer(Synapse):
    """A generalized DenseTensorSynapse that has a hidden lagrangian (non-linearity).

    We can specify a Lagrangian non-linearity for the hidden neuron layer with tau=0 and shape `(nhid,)`.

    The lagrangian can have its own learnable parameters, for example:

    ```
    from hamux.core.lagrangians import *

    syn = DenseTerminusSynapse(20, hidden_lagrangian=LSoftmax(beta_init=0.2)).init(jax.random.PRNGKey(0), tuple(gs))
    ```
    """

    W: jnp.ndarray = tx.Parameter.node()
    hidden_lagrangian: Optional[tx.Module]  # An already initialized lagrangian function

    def __init__(
        self,
        nhid: int,
        num_heads=1,
        stdinit: float = 0.02,
        hidden_lagrangian: tx.Module = LRelu(),
    ):
        self.nhid = nhid
        self.stdinit = stdinit
        self.num_heads = num_heads
        self.hidden_lagrangian = hidden_lagrangian

    def __call__(self, *gs):
        if self.initializing():
            assert (
                len(gs) <= 52
            ), "We are limited to english ASCII letters. We cannot connect more than 50 layers if you include our hidden layer and number of heads"
            key = tx.next_key()
            gshapes = (self.num_heads, self.nhid) + tuple([g_.size for g_ in gs])
            self.W = nn.initializers.normal(self.stdinit)(key, gshapes)

        gs = [g_.ravel() for g_ in gs]
        abcs = ascii_letters[: len(gs)]
        einsum_arg = (
            ",".join(["YZ" + abcs, ",".join(abcs)]) + "->YZ"
        )  # We call the nhid dimension "Z" and number of heads "Y"
        x = jnp.einsum(einsum_arg, self.W, *gs)
        return self.hidden_lagrangian(x).sum()

# %% ../nbs/02_synapses.ipynb 19
class DenseMatrixSynapseWithHiddenLayer(Synapse):
    """A modified DenseSynapse that has a hidden lagrangian (non-linearity).

    We can specify a Lagrangian non-linearity for the hidden neuron layer with tau=0 and shape `(nhid,)`.

    Unlike the DenseTensorSynapseWithHiddenLayer, treat layers as if they are concatenated on the same
    visible layer dimension instead of giving each its own dimension of the tensor space.
    """

    Ws: List[jnp.ndarray] = tx.Parameter.node()
    hidden_lagrangian: Optional[tx.Module]  # An already initialized lagrangian function

    def __init__(
        self, nhid: int, stdinit: float = 0.02, hidden_lagrangian: tx.Module = LRelu(), do_ravel=True, do_norm=False,
    ):
        self.nhid = nhid
        self.stdinit = stdinit
        self.hidden_lagrangian = hidden_lagrangian
        self.do_ravel = do_ravel
        self.do_norm = do_norm

    def __call__(self, *gs):
        if self.initializing():
            def initw(g):
                if self.do_ravel:
                    gsize = g.size
                else:
                    gsize = g.shape[-1]
                return nn.initializers.normal(self.stdinit)(tx.next_key(), (gsize, self.nhid))
            self.Ws = [initw(g_) for g_ in gs]

        if self.do_ravel:
            gs = [g_.ravel() for g_ in gs]
        if self.do_norm:
            Ws = [W / jnp.sqrt((W**2).sum(0,keepdims=True)) for W in self.Ws]
        else:
            Ws = self.Ws
        hid_state = jnp.stack([g @ W for (W, g) in zip(Ws, gs)]).sum(0)
        return self.hidden_lagrangian(hid_state)

# %% ../nbs/02_synapses.ipynb 20
class ConvSynapse(Synapse):
    """A convolutional, binary synapse. Can automatically detect the number of output features from the 2 layers it connects"""

    conv: tx.Conv
    # Delegate arguments to conv EXCEPT the features_out, which we calculate from the output layer

    @delegates(tx.Conv)
    def __init__(self, kernel_size: Union[int, Iterable[int]], **kwargs):
        # assert pool_type in ["max", "avg"]
        self.kernel_size = kernel_size
        conv_kwargs = {"use_bias": False}
        conv_kwargs.update(kwargs)
        self.conv_kwargs = conv_kwargs

    def example_output(self, g1, features_out=1):
        """Test the shape output of the convolutional layer. If unspecified, output features are 1"""
        conv = tx.Conv(features_out, self.kernel_size, **self.conv_kwargs).init(
            jax.random.PRNGKey(0), g1
        )
        return conv(g1)

    def __call__(self, g1, g2):
        """The convolutional operation. g2 is assumed to be the "output" of the convolution"""
        if self.initializing():
            features_out = g2.shape[-1]
            self.conv = tx.Conv(
                features_out, self.kernel_size, **self.conv_kwargs
            ).init(tx.next_key(), g1)
        return jnp.multiply(self.conv(g1), g2).sum()

# %% ../nbs/02_synapses.ipynb 22
class ConvSynapseWithPool(Synapse):
    """A convolutional, binary synapse. Can automatically detect the number of output features from the 2 layers it connects"""

    conv: tx.Conv
    # Delegate arguments to conv EXCEPT the features_out, which we calculate from the output layer

    @delegates(tx.Conv)
    def __init__(
        self,
        kernel_size: Union[int, Iterable[int]],
        pool_window=(5, 5),
        pool_stride=(2, 2),
        pool_type="avg",
        **kwargs,
    ):
        # assert pool_type in ["max", "avg"]
        self.kernel_size = kernel_size
        self.conv_kwargs = kwargs
        self.pool_window = pool_window
        self.pool_stride = pool_stride
        self.pool_type = pool_type

        self.pooler = max_pool if self.pool_type == "max" else avg_pool
        # self.conv = None

    def example_output(self, g1, features_out=1):
        """Test the shape output of the convolutional layer. If unspecified, output features are 1"""
        conv = tx.Conv(features_out, self.kernel_size, **self.conv_kwargs).init(
            jax.random.PRNGKey(0), g1
        )
        output = self.pooler(conv(g1), self.pool_window, strides=self.pool_stride)
        return output

    def __call__(self, g1, g2):
        """The convolutional operation. g2 is assumed to be the "output" of the convolution"""
        if self.initializing():
            features_out = g2.shape[-1]
            self.conv = tx.Conv(
                features_out, self.kernel_size, **self.conv_kwargs
            ).init(tx.next_key(), g1)
        output = self.pooler(self.conv(g1), self.pool_window, strides=self.pool_stride)
        return jnp.multiply(output, g2).sum()


# %% ../nbs/02_synapses.ipynb 24
class AttentionSynapse(Synapse):
    """A generalized synapse of quadratic order, whose update rule looks very similar to the Attention operation of Transformers.

    We can specify any Lagrangian non-linearity for the hidden neuron layer (which operates with tau=0), but we default to the Softmax Lagrangian.

    To replicate similar configuration to the famous BERT-base models and "Attention is all you need" paper:

    ```
    zspace = 64
    syn = AttentionSynapse(zspace_dim=zspace, num_heads=12, hidden_lagrangian=LSoftmax(beta=1/jnp.sqrt(zspace)))
    ```

    Connecting two layers of shapes: ((Nq, Dq), (Nk, Dk)) and layernorm lagrangians
    """

    Wk: jnp.ndarray = tx.Parameter.node()
    Wq: jnp.ndarray = tx.Parameter.node()
    hidden_lagrangian: Optional[tx.Module]  # An already initialized lagrangian function
    qk_norm: Optional[tx.Module]

    def __init__(
        self,
        num_heads: int = 1,
        zspace_dim: int = 64,
        stdinit: float = 0.02,
        hidden_lagrangian: tx.Module = LSoftmax(),
        do_qk_norm: bool = False,
    ):
        self.zspace_dim = zspace_dim
        self.num_heads = num_heads
        self.stdinit = stdinit
        self.hidden_lagrangian = hidden_lagrangian
        self.do_qk_norm = do_qk_norm

    def __call__(self, gq, gk):
        """Align the queries in gq with the keys in gk"""
        if self.initializing():
            self.Wq = nn.initializers.normal(self.stdinit)(
                tx.next_key(), (gq.shape[-1], self.num_heads, self.zspace_dim)
            )
            self.Wk = nn.initializers.normal(self.stdinit)(
                tx.next_key(), (gk.shape[-1], self.num_heads, self.zspace_dim)
            )
            if self.do_qk_norm:
                self.qk_norm = tx.LayerNorm().init(
                    tx.next_key(), jnp.ones(self.zspace_dim)
                )

        if len(gq.shape) == 1:
            gq = gq[..., None, :]
        if len(gk.shape) == 1:
            gk = gk[..., None, :]
        if len(gq.shape) > 2:
            gq = rearrange(gq, "... d -> (...) d")
        if len(gk.shape) > 2:
            gk = rearrange(gk, "... d -> (...) d")
        Q = jnp.einsum("qd,dhz->qhz", gq, self.Wq)
        K = jnp.einsum("kd,dhz->khz", gk, self.Wk)
        # QK = jnp.einsum("qhz,khz->hqk", Q, K)
        if self.do_qk_norm:
            Q = self.qk_norm(Q)
            K = self.qk_norm(K)
        QK = jnp.einsum("qhz,khz->hqk", Q, K)

        # print("MAX QK: ", QK.max())
        # print("MIN QK: ", QK.min())
        # , QK.min(), gk.max(), gq.max())
        attn = self.hidden_lagrangian(QK)  # h,q
        return attn.sum()


class SelfAttentionSynapse(AttentionSynapse):
    """A special case of the AttentionSynapse where both inputs are of the same layer"""

    def __call__(self, g):
        return super().__call__(g, g)

    
class BinaryMixerSynapse(Synapse):
    """A generalized binary synapse of quadratic order. This synapse is very similar to the Attention synapse but uses a single
    weight matrix instead of a query and key matrix.

    We can specify any Lagrangian non-linearity for the hidden neuron layer (which operates with tau=0), but we default to the Softmax Lagrangian.

    """

    W: jnp.ndarray = tx.Parameter.node()
    hidden_lagrangian: Optional[tx.Module]  # An already initialized lagrangian function

    def __init__(
        self,
        num_heads: int = 1,
        zspace_dim: int = 64,
        stdinit: float = 0.02,
        hidden_lagrangian: tx.Module = LSoftmax(),
    ):
        self.num_heads = num_heads
        self.zspace_dim=zspace_dim
        self.stdinit = stdinit
        self.hidden_lagrangian = hidden_lagrangian

    def __call__(self, ga, gb):
        """Align the activations ga with gb"""
        if self.initializing():
            self.W =  nn.initializers.normal(self.stdinit)(
                tx.next_key(), (self.num_heads, self.zspace_dim, ga.shape[-1], gb.shape[-1])
            )

        if len(ga.shape) == 1:
            ga = ga[...,None,:]
        if len(gb.shape) == 1:
            gb = gb[...,None,:]
        if len(ga.shape) > 2:
            ga = rearrange(ga, "... d -> (...) d")
        if len(gb.shape) > 2:
            gb = rearrange(gb, "... d -> (...) d")

        AB = jnp.einsum("hzab,...a,...b->hz", self.W, ga, gb)
        attn = self.hidden_lagrangian(AB)  # h
        return attn.sum()
