import inspect
import math
from dataclasses import dataclass
from functools import partial
from typing import (
    Callable,
    Dict,
    Literal,
    Optional,
    Protocol,
    Sequence,
    Tuple,
    Type,
    Union,
)

import torch
from torch import Tensor, nn
from torch.nn.functional import pad
from transformers import PretrainedConfig, PreTrainedModel

from eir.models.layers.attention_layers import Transformer
from eir.models.layers.lcl_layers import _find_lcl_padding_needed
from eir.utils.logging import get_logger

logger = get_logger(name=__name__, tqdm_compatible=True)


@dataclass
class SequenceModelConfig:
    """
    :param model_init_config:
          Configuration / arguments used to initialise model.

    :param model_type:
         Which type of image model to use.

    :param embedding_dim:
        Which dimension to use for the embeddings. If ``None``, will automatically set
        this value based on the number of tokens and attention heads.

    :param position:
        Whether to encode the token position or use learnable position embeddings.

    :param position_dropout:
        Dropout for the positional encoding / embedding.

    :param window_size:
        If set to more than 0, will apply a sliding window of feature
        extraction over the input, meaning the model (e.g. transformer) will only
        see a part of the input at a time. Can be Useful to avoid the O(n²)
        complexity of transformers, as it becomes O(window_size² * n_windows) instead.

    :param pool:
        Whether and how to pool (max / avg) the final feature maps before being
        passed to the final fusion module / predictor. Meaning we pool over the
        sequence (i.e. time) dimension, so the resulting dimensions is embedding_dim
        instead of sequence_length * embedding_dim. If using windowed / conv
        transformers, this becomes embedding_dim * number_of_chunks.

    :param pretrained_model:
          Specify whether the model type is assumed to be pretrained and from the
          Pytorch Image Models repository.

    :param freeze_pretrained_model:
          Whether to freeze the pretrained model weights.

    """

    model_init_config: Union["BasicTransformerFeatureExtractorModelConfig", Dict]

    model_type: Literal["sequence-default"] | str = "sequence-default"

    embedding_dim: int = 64

    position: Literal["encode", "embed"] = "encode"
    position_dropout: float = 0.10
    window_size: int = 0
    pool: Union[Literal["avg"], Literal["max"], None] = None

    pretrained_model: bool = False
    freeze_pretrained_model: bool = False


class TransformerWrapperModel(nn.Module):
    def __init__(
        self,
        feature_extractor: Union["TransformerFeatureExtractor", nn.Module],
        model_config: SequenceModelConfig,
        embedding_dim: int,
        num_tokens: int,
        max_length: int,
        external_feature_extractor: bool,
        device: str,
        embeddings: Optional[nn.Embedding] = None,
        pre_computed_num_out_features: Optional[int] = None,
    ) -> None:
        super().__init__()
        self.model_config = model_config
        self.embedding_dim = embedding_dim
        self.num_tokens = num_tokens
        self.max_length = max_length
        self.external_feature_extractor = external_feature_extractor
        self.pre_computed_num_out_features = pre_computed_num_out_features

        pos_repr_class = get_positional_representation_class(
            position_model_config=self.model_config.position
        )
        self.pos_representation = pos_repr_class(
            embedding_dim=self.embedding_dim,
            dropout=self.model_config.position_dropout,
            max_length=self.max_length,
        )

        if embeddings:
            self.embedding = embeddings
        else:
            self.embedding = nn.Embedding(
                num_embeddings=self.num_tokens, embedding_dim=self.embedding_dim
            )

        self.feature_extractor = feature_extractor

        if not embeddings:
            self.init_embedding_weights()

        (
            self.dynamic_extras,
            self.extract_features,
        ) = _get_transformer_wrapper_feature_extractor(
            feature_extractor=self.feature_extractor,
            window_size=self.model_config.window_size,
            max_length=self.max_length,
            embedding_dim=self.embedding_dim,
            pool=self.model_config.pool,
            device=device,
            external_feature_extractor=self.external_feature_extractor,
        )

    @property
    def num_out_features(self) -> int:
        if self.pre_computed_num_out_features:
            return self.pre_computed_num_out_features

        padding = self.dynamic_extras.get("padding", 0)
        length_with_padding = self.max_length + padding

        if self.model_config.pool in ("avg", "max"):
            num_chunks = 1
            if self.model_config.window_size:
                num_chunks = length_with_padding // self.model_config.window_size

            return self.embedding_dim * num_chunks

        return length_with_padding * self.embedding_dim

    def init_embedding_weights(self) -> None:
        init_range = 0.1
        self.embedding.weight.data.uniform_(-init_range, init_range)

    def forward(self, input: Tensor) -> Tensor:
        out = input * math.sqrt(self.embedding_dim)
        out = self.pos_representation(out)
        out = self.extract_features(out)

        return out


def get_embedding_dim_for_sequence_model(
    embedding_dim: Union[None, int], num_tokens: int, num_heads: int
) -> int:
    if embedding_dim is None:
        auto_emb_dim = math.ceil((int(num_tokens**0.25) / num_heads)) * num_heads
        logger.info(
            "Setting up automatic embedding dimension of %d based on %d "
            "tokens and %d attention heads.",
            auto_emb_dim,
            num_tokens,
            num_heads,
        )
        embedding_dim = auto_emb_dim

    return embedding_dim


class SimpleFeatureExtractorForwardProtocol(Protocol):
    def __call__(
        self,
        input: torch.Tensor,
        feature_extractor: Union[nn.Module, "TransformerFeatureExtractor"],
    ) -> torch.Tensor: ...


class HFFeatureExtractorForwardProtocol(Protocol):
    def __call__(
        self,
        input: torch.Tensor,
        feature_extractor: Union[nn.Module, "TransformerFeatureExtractor"],
        key: str = "last_hidden_state",
    ) -> torch.Tensor: ...


def _get_transformer_wrapper_feature_extractor(
    feature_extractor: Union["TransformerFeatureExtractor", nn.Module],
    external_feature_extractor: bool,
    window_size: int,
    embedding_dim: int,
    max_length: int,
    device: str,
    pool: Union[Literal["avg"], Literal["max"], None] = None,
) -> Tuple[Dict[str, int], Callable[[torch.Tensor], torch.Tensor]]:
    dynamic_extras = {"padding": 0}

    feature_extractor_forward = _get_feature_extractor_forward(
        is_hf_model=external_feature_extractor,
        feature_extractor=feature_extractor,
        input_length=window_size if window_size else max_length,
        embedding_size=embedding_dim,
        device=device,
        pool=pool,
    )
    if not window_size:
        extractor = partial(
            feature_extractor_forward,
            feature_extractor=feature_extractor,
        )
    else:
        num_chunks = int(math.ceil(max_length / window_size))
        logger.debug(
            "Setting num chunks to %d as window size of %d and maximum sequence length "
            "of %d were passed in.",
            num_chunks,
            window_size,
            max_length,
        )

        padding = _find_lcl_padding_needed(
            input_size=max_length,
            kernel_size=window_size,
            num_chunks=num_chunks,
        )
        dynamic_extras["padding"] = padding

        extractor = partial(
            _conv_transformer_forward,
            feature_extractor=feature_extractor,
            feature_extractor_forward_callable=feature_extractor_forward,
            max_length=max_length,
            window_size=window_size,
            padding=padding,
        )

    return dynamic_extras, extractor


def _get_feature_extractor_forward(
    is_hf_model: bool,
    feature_extractor: Union[nn.Module, PreTrainedModel],
    input_length: int,
    embedding_size: int,
    device: str,
    pool: Union[Literal["avg"], Literal["max"], None] = None,
) -> SimpleFeatureExtractorForwardProtocol | HFFeatureExtractorForwardProtocol:
    if is_hf_model:
        return get_hf_transformer_forward(
            feature_extractor_=feature_extractor,
            input_length=input_length,
            embedding_dim=embedding_size,
            device=device,
            pool=pool,
        )
    return _get_simple_transformer_forward(pool=pool)


def _get_simple_transformer_forward(
    pool: Union[Literal["avg"], Literal["max"], None] = None,
):
    pooling_func = _get_sequence_pooling_func(pool=pool)

    def _simple_transformer_forward(
        input: torch.Tensor,
        feature_extractor: "TransformerFeatureExtractor",
    ) -> torch.Tensor:
        tensor_out = feature_extractor(input)
        tensor_pooled = pooling_func(input=tensor_out)
        final_out = tensor_pooled.flatten(1)

        return final_out

    return _simple_transformer_forward


def get_hf_transformer_forward(
    feature_extractor_: PreTrainedModel,
    input_length: int,
    embedding_dim: int,
    device: str,
    pool: Union[Literal["avg"], Literal["max"], None] = None,
):
    forward_argnames = inspect.getfullargspec(feature_extractor_.forward)[0]

    bound_kwargs = _build_transformer_forward_kwargs(
        forward_argnames=forward_argnames,
        config=feature_extractor_.config,
        input_length=input_length,
        embedding_dim=embedding_dim,
        device=device,
    )

    pooling_func = _get_sequence_pooling_func(pool=pool)

    def _hf_transformer_forward(
        input: torch.Tensor,
        feature_extractor: nn.Module,
        key: str = "last_hidden_state",
    ) -> torch.Tensor:
        hf_transformer_out = feature_extractor(inputs_embeds=input, **bound_kwargs)
        tensor_out = getattr(hf_transformer_out, key)
        tensor_pooled = pooling_func(input=tensor_out)
        final_out = tensor_pooled.flatten(1)
        return final_out

    return _hf_transformer_forward


def _get_sequence_pooling_func(
    pool: Union[Literal["avg"], Literal["max"], None]
) -> Callable:
    def _identity(input: torch.Tensor) -> torch.Tensor:
        return input

    def _max(input: torch.Tensor) -> torch.Tensor:
        return input.max(dim=1)[0]

    def _avg(input: torch.Tensor) -> torch.Tensor:
        return input.mean(dim=1)

    if pool is None:
        return _identity

    elif pool == "max":
        return _max

    elif pool == "avg":
        return _avg

    raise ValueError()


def _build_transformer_forward_kwargs(
    forward_argnames: Sequence[str],
    config: PretrainedConfig,
    input_length: int,
    embedding_dim: int,
    device: str,
) -> Dict:
    """
    TODO: Deprecate.
    """
    kwargs = {}

    if "decoder_inputs_embeds" in forward_argnames:
        kwargs["decoder_inputs_embeds"] = torch.randn(1, input_length, embedding_dim)

    return kwargs


def _conv_transformer_forward(
    input: torch.Tensor,
    feature_extractor: Union[nn.Module, "TransformerFeatureExtractor"],
    feature_extractor_forward_callable: (
        SimpleFeatureExtractorForwardProtocol | HFFeatureExtractorForwardProtocol
    ),
    max_length: int,
    window_size: int,
    padding: int,
) -> torch.Tensor:
    out = pad(input=input, pad=[0, 0, padding, 0])
    total_length = max_length + padding

    aggregated_out = None
    for lower_index in range(0, total_length, window_size):
        upper_index = lower_index + window_size

        cur_input = out[:, lower_index:upper_index, :]
        cur_out = feature_extractor_forward_callable(cur_input, feature_extractor)

        if aggregated_out is None:
            aggregated_out = cur_out
        else:
            aggregated_out = torch.cat((aggregated_out, cur_out), dim=1)

    assert aggregated_out is not None
    return aggregated_out


@dataclass
class BasicTransformerFeatureExtractorModelConfig:
    """
    :param num_heads:
        The number of heads in the multi-head attention models

    :param num_layers:
        The number of encoder blocks in the transformer model.

    :param dim_feedforward:
        The dimension of the feedforward layers in the transformer model.

    :param dropout:
         Dropout value to use in the encoder layers.
    """

    num_heads: int = 8
    num_layers: int = 2
    dim_feedforward: Union[int, Literal["auto"]] = "auto"
    dropout: float = 0.10


class TransformerFeatureExtractor(nn.Module):
    def __init__(
        self,
        model_config: BasicTransformerFeatureExtractorModelConfig,
        embedding_dim: int,
        num_tokens: int,
        max_length: int,
    ) -> None:
        super().__init__()

        self.model_config = model_config
        self.embedding_dim = embedding_dim
        self.num_tokens = num_tokens
        self.max_length = max_length

        dim_feed_forward = parse_dim_feedforward(
            dim_feedforward=self.model_config.dim_feedforward,
            embedding_dim=self.embedding_dim,
        )

        self.transformer_encoder = Transformer(
            d_model=self.embedding_dim,
            nhead=self.model_config.num_heads,
            num_layers=self.model_config.num_layers,
            dim_feedforward=dim_feed_forward,
            dropout=self.model_config.dropout,
            norm_first=True,
        )

    @property
    def num_out_features(self) -> int:
        return self.max_length * self.embedding_dim

    def forward(self, input: Tensor) -> Tensor:
        out = self.transformer_encoder(input)
        return out


def parse_dim_feedforward(
    dim_feedforward: Union[int, Literal["auto"]], embedding_dim: int
) -> int:
    if dim_feedforward == "auto":
        dim_feedforward = embedding_dim * 4
        logger.info(
            "Setting dim_feedfoward to %d based on embedding_dim=%d and 'auto' "
            "option.",
            dim_feedforward,
            embedding_dim,
        )

    return dim_feedforward


class SequenceOutputTransformerFeatureExtractor(TransformerFeatureExtractor):
    def __init__(
        self,
        model_config: BasicTransformerFeatureExtractorModelConfig,
        embedding_dim: int,
        num_tokens: int,
        max_length: int,
    ):
        super().__init__(
            model_config=model_config,
            embedding_dim=embedding_dim,
            num_tokens=num_tokens,
            max_length=max_length,
        )

        mask = torch.triu(
            torch.ones(self.max_length, self.max_length) * float("-inf"), diagonal=1
        )
        self.register_buffer("mask", mask)

    def forward(self, input: Tensor) -> Tensor:
        out = self.transformer_encoder(input, mask=self.mask)
        return out


def get_positional_representation_class(
    position_model_config: Literal["encode", "embed"]
) -> Union[Type["PositionalEncoding"], Type["PositionalEmbedding"]]:
    if position_model_config == "encode":
        return PositionalEncoding
    elif position_model_config == "embed":
        return PositionalEmbedding
    raise ValueError(
        "Unknown value for positional representation. "
        "Expected 'encode' or 'embed' but got '%s'.",
        position_model_config,
    )


class PositionalEncoding(nn.Module):
    def __init__(
        self,
        embedding_dim: int,
        max_length: int,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.max_length = max_length

        position = torch.arange(max_length).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, embedding_dim, 2) * (-math.log(10000.0) / embedding_dim)
        )
        pe: torch.Tensor = torch.zeros(1, max_length, embedding_dim)
        pe[0, :, 0::2] = torch.sin(position * div_term)
        pe[0, :, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe)

        self.pe: Tensor

    def forward(self, x: Tensor) -> Tensor:
        x = x + self.pe[:, : self.max_length, :]
        return self.dropout(x)


class PositionalEmbedding(nn.Module):
    def __init__(
        self,
        embedding_dim: int,
        max_length: int,
        dropout: float = 0.1,
        zero_init: bool = False,
    ) -> None:
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.max_length = max_length

        if zero_init:
            self.embedding = torch.nn.Parameter(
                data=torch.zeros(1, max_length, embedding_dim),
                requires_grad=True,
            )
        else:
            self.embedding = torch.nn.Parameter(
                data=torch.randn(1, max_length, embedding_dim),
                requires_grad=True,
            )

    def forward(self, x: Tensor) -> Tensor:
        x = x + self.embedding[:, : self.max_length, :]
        return self.dropout(x)
