import torch
import torch.nn as nn
from einops import rearrange
from diffusers.models.attention_processor import Attention


class CrossAttention(nn.Module):
    """
    CrossAttention module implements per-pixel temporal attention to fuse the conditional attention module with the base module.

    Args:
        input_channels (int): Number of input channels.
        attention_head_dim (int): Dimension of attention head.
        norm_num_groups (int): Number of groups for GroupNorm normalization (default is 32).

    Attributes:
        attention (Attention): Attention module for computing attention scores.
        norm (torch.nn.GroupNorm): Group normalization layer.
        proj_in (nn.Linear): Linear layer for projecting input data.
        proj_out (nn.Linear): Linear layer for projecting output data.
        dropout (nn.Dropout): Dropout layer for regularization.

    Methods:
        forward(hidden_state, encoder_hidden_states, num_frames, num_conditional_frames):
            Forward pass of the CrossAttention module.

    """

    def __init__(self, input_channels, attention_head_dim, norm_num_groups=32):
        super().__init__()
        self.attention = Attention(
            query_dim=input_channels, cross_attention_dim=input_channels, heads=input_channels//attention_head_dim, dim_head=attention_head_dim, bias=False, upcast_attention=False)
        self.norm = torch.nn.GroupNorm(
            num_groups=norm_num_groups, num_channels=input_channels, eps=1e-6, affine=True)
        self.proj_in = nn.Linear(input_channels, input_channels)
        self.proj_out = nn.Linear(input_channels, input_channels)
        self.dropout = nn.Dropout(p=0.25)

    def forward(self, hidden_state, encoder_hidden_states, num_frames, num_conditional_frames):
        """
        The input hidden state is normalized, then projected using a linear layer.
        Multi-head cross attention is computed between the hidden state (latent of noisy video) and encoder hidden states (CLIP image encoder).
        The output is projected using a linear layer.
        We apply dropout to the newly generated frames (without the control frames).

        Args:
            hidden_state (torch.Tensor): Input hidden state tensor.
            encoder_hidden_states (torch.Tensor): Encoder hidden states tensor.
            num_frames (int): Number of frames.
            num_conditional_frames (int): Number of conditional frames.

        Returns:
            output (torch.Tensor): Output tensor after processing with attention mechanism.

        """
        h, w = hidden_state.shape[2], hidden_state.shape[3]
        hidden_state_norm = rearrange(
            hidden_state, "(B F) C H W -> B C F H W", F=num_frames)
        hidden_state_norm = self.norm(hidden_state_norm)
        hidden_state_norm = rearrange(
            hidden_state_norm, "B C F H W -> (B H W) F C")

        hidden_state_norm = self.proj_in(hidden_state_norm)

        attn = self.attention(hidden_state_norm,
                              encoder_hidden_states=encoder_hidden_states,
                              attention_mask=None,
                              )
        # proj_out

        residual = self.proj_out(attn)  # (B H W) F C
        hidden_state = rearrange(
            hidden_state, "(B F) ... -> B F ...", F=num_frames)
        hidden_state = torch.cat([hidden_state[:, :num_conditional_frames], self.dropout(
            hidden_state[:, num_conditional_frames:])], dim=1)
        hidden_state = rearrange(hidden_state, "B F ... -> (B F) ... ")

        residual = rearrange(
            residual, "(B H W) F C  -> (B F) C H W", H=h, W=w)
        output = hidden_state + residual
        return output


class ConditionalModel(nn.Module):
    """ 
    ConditionalModel module performs the fusion of the conditional attention module to be base model.

    Args:
        input_channels (int): Number of input channels.
        conditional_model (str): Type of conditional model to use. Currently only "cross_attention" is implemented.
        attention_head_dim (int): Dimension of attention head (default is 64).

    Attributes:
        temporal_transformer (CrossAttention): CrossAttention module for temporal transformation.
        conditional_model (str): Type of conditional model used.

    Methods:
        forward(sample, conditioning, num_frames=None, num_conditional_frames=None):
            Forward pass of the ConditionalModel module.

    """

    def __init__(self, input_channels, conditional_model: str, attention_head_dim=64):
        super().__init__()

        if conditional_model == "cross_attention":
            self.temporal_transformer = CrossAttention(
                input_channels=input_channels, attention_head_dim=attention_head_dim)
        else:
            raise NotImplementedError(
                f"mode {conditional_model} not implemented")

        nn.init.zeros_(self.temporal_transformer.proj_out.weight)
        nn.init.zeros_(self.temporal_transformer.proj_out.bias)
        self.conditional_model = conditional_model

    def forward(self, sample, conditioning, num_frames=None, num_conditional_frames=None):
        """
        Forward pass of the ConditionalModel module.

        Args:
            sample (torch.Tensor): Input sample tensor.
            conditioning (torch.Tensor): Conditioning tensor containing the enconding of the conditional frames.
            num_frames (int): Number of frames in the sample.
            num_conditional_frames (int): Number of conditional frames.

        Returns:
            sample (torch.Tensor): Transformed sample tensor.

        """
        sample = rearrange(sample, "(B F) ... -> B F ...", F=num_frames)
        batch_size = sample.shape[0]
        conditioning = rearrange(
            conditioning, "(B F) ... -> B F ...", B=batch_size)

        assert conditioning.ndim == 5
        assert sample.ndim == 5

        conditioning = rearrange(conditioning, "B F C H W -> (B H W) F C")

        sample = rearrange(sample, "B F C H W -> (B F) C H W")

        sample = self.temporal_transformer(
            sample, encoder_hidden_states=conditioning, num_frames=num_frames, num_conditional_frames=num_conditional_frames)

        return sample


if __name__ == "__main__":
    model = CrossAttention(input_channels=320, attention_head_dim=32)