from copy import deepcopy
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn
from torch import Tensor
from torch.nn.modules import Module

from pytorch_gleam.modeling.layers.hopfield import Hopfield


class HopfieldEncoderLayer(Module):
    """Module with underlying Hopfield association to be used as an encoder in transformer-like architectures."""

    def __init__(
        self,
        hopfield_association: Hopfield,
        dim_feedforward: int = 2048,
        dropout: float = 0.1,
        activation: str = r"relu",
    ):
        """Initialise a new instance of a Hopfield association-based encoder module.

        :param hopfield_association: instance of Hopfield association module
        :param dim_feedforward: depth of the linear projections applied internally
        :param activation: activation to be applied on the result of the internal linear projections
        :param dropout: dropout probability to be applied internally
        """
        super(HopfieldEncoderLayer, self).__init__()
        self.hopfield_association = deepcopy(hopfield_association)

        self.linear_residual = nn.Linear(self.hopfield_association.state_pattern_dim, dim_feedforward)
        self.dropout_residual = nn.Dropout(dropout)
        self.linear_output = nn.Linear(dim_feedforward, self.hopfield_association.state_pattern_dim)

        self.norm_residual = nn.LayerNorm(self.hopfield_association.state_pattern_dim)
        self.norm_output = nn.LayerNorm(self.hopfield_association.state_pattern_dim)
        self.dropout_hopfield_association = nn.Dropout(dropout)
        self.dropout_output = nn.Dropout(dropout)

        self.activation_residual = getattr(torch, activation, None)
        assert self.activation_residual is not None, r"invalid activation function supplied."

    def forward(
        self,
        src: Tensor,
        src_mask: Optional[Tensor] = None,
        src_key_padding_mask: Optional[Tensor] = None,
    ) -> Tensor:
        """Apply Hopfield encoding on specified data.

        :param src: data to be processed by Hopfield encoder module
        :param src_mask: mask to be applied on association matrix
        :param src_key_padding_mask: mask to be applied on stored patterns
        :return: Hopfield-encoded input data
        """
        data_associated = self.hopfield_association(
            input=src,
            stored_pattern_padding_mask=src_key_padding_mask,
            association_mask=src_mask,
        )
        src = src + self.dropout_hopfield_association(input=data_associated)
        src = self.norm_residual(input=src)

        result_residual_inner = self.activation_residual(input=self.linear_residual(input=src))
        data_associated = self.linear_output(input=self.dropout_residual(input=result_residual_inner))
        src = src + self.dropout_output(input=data_associated)

        return self.norm_output(input=src)

    def get_association_matrix(self, input: Union[Tensor, Tuple[Tensor, Tensor, Tensor]]) -> Tensor:
        """Fetch Hopfield association matrix gathered by passing through the specified data.

        :param input: data to be passed through the Hopfield association
        :return: association matrix as computed by the Hopfield core module
        """
        return self.hopfield_association.get_association_matrix(input=input)

    @property
    def batch_first(self) -> int:
        return self.hopfield_association.batch_first

    @property
    def input_size(self) -> int:
        return self.hopfield_association.input_size

    @property
    def output_size(self) -> int:
        return self.linear_output.out_features


class HopfieldDecoderLayer(Module):
    def __init__(
        self,
        hopfield_association_self: Hopfield,
        hopfield_association_cross: Hopfield,
        dim_feedforward: int = 2048,
        dropout: float = 0.1,
        activation: str = r"relu",
    ):
        """Initialise a new instance of a Hopfield association-based encoder module.

        :param hopfield_association_self: instance of Hopfield self-association module
        :param hopfield_association_cross: instance of Hopfield cross-association module
        :param dim_feedforward: depth of the linear projections applied internally
        :param dropout: dropout probability to be applied internally
        :param activation: activation to be applied on the result of the internal linear projections
        """
        super(HopfieldDecoderLayer, self).__init__()
        self.hopfield_association_self = deepcopy(hopfield_association_self)
        self.hopfield_association_cross = deepcopy(hopfield_association_cross)

        self.linear_residual = nn.Linear(self.hopfield_association_self.state_pattern_dim, dim_feedforward)
        self.dropout_residual = nn.Dropout(dropout)
        self.linear_output = nn.Linear(dim_feedforward, self.hopfield_association_self.state_pattern_dim)

        self.norm_residual_self = nn.LayerNorm(self.hopfield_association_self.state_pattern_dim)
        self.norm_residual_cross = nn.LayerNorm(self.hopfield_association_self.state_pattern_dim)
        self.norm_output = nn.LayerNorm(self.hopfield_association_self.state_pattern_dim)
        self.dropout_hopfield_association_self = nn.Dropout(dropout)
        self.dropout_hopfield_association_cross = nn.Dropout(dropout)
        self.dropout_output = nn.Dropout(dropout)

        self.activation_residual = getattr(torch, activation, None)
        assert self.activation_residual is not None, r"invalid activation function supplied."

    def forward(
        self,
        tgt: Tensor,
        memory: Tensor,
        tgt_mask: Optional[Tensor] = None,
        memory_mask: Optional[Tensor] = None,
        tgt_key_padding_mask: Optional[Tensor] = None,
        memory_key_padding_mask: Optional[Tensor] = None,
    ) -> Tensor:
        """Apply Hopfield decoding on specified data.

        :param tgt: data to be processed by Hopfield decoder module (self-association)
        :param memory: data to be processed by Hopfield encoder module (cross-association)
        :param tgt_mask: mask to be applied on self-association matrix
        :param memory_mask: mask to be applied on cross-association matrix
        :param tgt_key_padding_mask: mask to be applied on stored patterns
        :param memory_key_padding_mask: mask to be applied on state patterns as well as pattern projection
        :return: Hopfield-decoded input
        """
        data_associated = self.hopfield_association_self(
            input=tgt,
            stored_pattern_padding_mask=tgt_key_padding_mask,
            association_mask=tgt_mask,
        )
        tgt = tgt + self.dropout_hopfield_association_self(input=data_associated)
        tgt = self.norm_residual_self(input=tgt)

        data_associated = self.hopfield_association_cross(
            input=(memory, tgt, memory),
            stored_pattern_padding_mask=memory_key_padding_mask,
            association_mask=memory_mask,
        )
        tgt = tgt + self.dropout_hopfield_association_cross(input=data_associated)
        tgt = self.norm_residual_cross(input=tgt)

        result_residual_inner = self.activation_residual(input=self.linear_residual(input=tgt))
        data_associated = self.linear_output(input=self.dropout_residual(input=result_residual_inner))
        tgt = tgt + self.dropout_output(input=data_associated)
        return self.norm_output(input=tgt)

    def get_association_matrix_self(self, input: Union[Tensor, Tuple[Tensor, Tensor, Tensor]]) -> Tensor:
        """Fetch Hopfield self-association matrix gathered by passing through the specified data.

        :param input: data to be passed through the Hopfield association
        :return: association matrix as computed by the Hopfield core module
        """
        return self.hopfield_association_self.get_association_matrix(input=input)

    def get_association_matrix_cross(self, input: Union[Tensor, Tuple[Tensor, Tensor, Tensor]]) -> Tensor:
        """Fetch Hopfield cross-association matrix gathered by passing through the specified data.

        :param input: data to be passed through the Hopfield association
        :return: association matrix as computed by the Hopfield core module
        """
        return self.hopfield_association_cross.get_association_matrix(input=input)

    @property
    def batch_first(self) -> int:
        return self.hopfield_association_self.batch_first

    @property
    def input_size(self) -> int:
        return self.hopfield_association_self.input_size

    @property
    def output_size(self) -> int:
        return self.linear_output_self.out_features
