# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
from torch import nn
from einops import rearrange
from torch.nn import functional as F
from ..utils.util import cosine_loss

import torch.nn as nn
import torch.nn.functional as F

from diffusers.models.attention import CrossAttention, FeedForward
from diffusers.utils.import_utils import is_xformers_available
from einops import rearrange


class SyncNet(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.audio_encoder = DownEncoder2D(
            in_channels=config["audio_encoder"]["in_channels"],
            block_out_channels=config["audio_encoder"]["block_out_channels"],
            downsample_factors=config["audio_encoder"]["downsample_factors"],
            dropout=config["audio_encoder"]["dropout"],
            attn_blocks=config["audio_encoder"]["attn_blocks"],
        )

        self.visual_encoder = DownEncoder2D(
            in_channels=config["visual_encoder"]["in_channels"],
            block_out_channels=config["visual_encoder"]["block_out_channels"],
            downsample_factors=config["visual_encoder"]["downsample_factors"],
            dropout=config["visual_encoder"]["dropout"],
            attn_blocks=config["visual_encoder"]["attn_blocks"],
        )

        self.eval()

    def forward(self, image_sequences, audio_sequences):
        vision_embeds = self.visual_encoder(image_sequences)  # (b, c, 1, 1)
        audio_embeds = self.audio_encoder(audio_sequences)  # (b, c, 1, 1)

        vision_embeds = vision_embeds.reshape(vision_embeds.shape[0], -1)  # (b, c)
        audio_embeds = audio_embeds.reshape(audio_embeds.shape[0], -1)  # (b, c)

        # Make them unit vectors
        vision_embeds = F.normalize(vision_embeds, p=2, dim=1)
        audio_embeds = F.normalize(audio_embeds, p=2, dim=1)

        return vision_embeds, audio_embeds


class ResnetBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        dropout: float = 0.0,
        norm_num_groups: int = 32,
        eps: float = 1e-6,
        act_fn: str = "silu",
        downsample_factor=2,
    ):
        super().__init__()

        self.norm1 = nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=eps, affine=True)
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)

        self.norm2 = nn.GroupNorm(num_groups=norm_num_groups, num_channels=out_channels, eps=eps, affine=True)
        self.dropout = nn.Dropout(dropout)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)

        if act_fn == "relu":
            self.act_fn = nn.ReLU()
        elif act_fn == "silu":
            self.act_fn = nn.SiLU()

        if in_channels != out_channels:
            self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
        else:
            self.conv_shortcut = None

        if isinstance(downsample_factor, list):
            downsample_factor = tuple(downsample_factor)

        if downsample_factor == 1:
            self.downsample_conv = None
        else:
            self.downsample_conv = nn.Conv2d(
                out_channels, out_channels, kernel_size=3, stride=downsample_factor, padding=0
            )
            self.pad = (0, 1, 0, 1)
            if isinstance(downsample_factor, tuple):
                if downsample_factor[0] == 1:
                    self.pad = (0, 1, 1, 1)  # The padding order is from back to front
                elif downsample_factor[1] == 1:
                    self.pad = (1, 1, 0, 1)

    def forward(self, input_tensor):
        hidden_states = input_tensor

        hidden_states = self.norm1(hidden_states)
        hidden_states = self.act_fn(hidden_states)

        hidden_states = self.conv1(hidden_states)
        hidden_states = self.norm2(hidden_states)
        hidden_states = self.act_fn(hidden_states)

        hidden_states = self.dropout(hidden_states)
        hidden_states = self.conv2(hidden_states)

        if self.conv_shortcut is not None:
            input_tensor = self.conv_shortcut(input_tensor)

        hidden_states += input_tensor

        if self.downsample_conv is not None:
            hidden_states = F.pad(hidden_states, self.pad, mode="constant", value=0)
            hidden_states = self.downsample_conv(hidden_states)

        return hidden_states


class AttentionBlock2D(nn.Module):
    def __init__(self, query_dim, norm_num_groups=32, dropout=0.0):
        super().__init__()
        if not is_xformers_available():
            raise ModuleNotFoundError(
                "You have to install xformers to enable memory efficient attetion", name="xformers"
            )
        # inner_dim = dim_head * heads
        self.norm1 = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=query_dim, eps=1e-6, affine=True)
        self.norm2 = nn.LayerNorm(query_dim)
        self.norm3 = nn.LayerNorm(query_dim)

        self.ff = FeedForward(query_dim, dropout=dropout, activation_fn="geglu")

        self.conv_in = nn.Conv2d(query_dim, query_dim, kernel_size=1, stride=1, padding=0)
        self.conv_out = nn.Conv2d(query_dim, query_dim, kernel_size=1, stride=1, padding=0)

        self.attn = CrossAttention(query_dim=query_dim, heads=8, dim_head=query_dim // 8, dropout=dropout, bias=True)
        self.attn._use_memory_efficient_attention_xformers = True

    def forward(self, hidden_states):
        assert hidden_states.dim() == 4, f"Expected hidden_states to have ndim=4, but got ndim={hidden_states.dim()}."

        batch, channel, height, width = hidden_states.shape
        residual = hidden_states

        hidden_states = self.norm1(hidden_states)
        hidden_states = self.conv_in(hidden_states)
        hidden_states = rearrange(hidden_states, "b c h w -> b (h w) c")

        norm_hidden_states = self.norm2(hidden_states)
        hidden_states = self.attn(norm_hidden_states, attention_mask=None) + hidden_states
        hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states

        hidden_states = rearrange(hidden_states, "b (h w) c -> b c h w", h=height, w=width)
        hidden_states = self.conv_out(hidden_states)

        hidden_states = hidden_states + residual
        return hidden_states


class DownEncoder2D(nn.Module):
    def __init__(
        self,
        in_channels=4 * 16,
        block_out_channels=[64, 128, 256, 256],
        downsample_factors=[2, 2, 2, 2],
        layers_per_block=2,
        norm_num_groups=32,
        attn_blocks=[1, 1, 1, 1],
        dropout: float = 0.0,
        act_fn="silu",
    ):
        super().__init__()
        self.layers_per_block = layers_per_block

        # in
        self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1)

        # down
        self.down_blocks = nn.ModuleList([])

        output_channels = block_out_channels[0]
        for i, block_out_channel in enumerate(block_out_channels):
            input_channels = output_channels
            output_channels = block_out_channel
            # is_final_block = i == len(block_out_channels) - 1

            down_block = ResnetBlock2D(
                in_channels=input_channels,
                out_channels=output_channels,
                downsample_factor=downsample_factors[i],
                norm_num_groups=norm_num_groups,
                dropout=dropout,
                act_fn=act_fn,
            )

            self.down_blocks.append(down_block)

            if attn_blocks[i] == 1:
                attention_block = AttentionBlock2D(query_dim=output_channels, dropout=dropout)
                self.down_blocks.append(attention_block)

        # out
        self.norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
        self.act_fn_out = nn.ReLU()

    def forward(self, hidden_states):
        hidden_states = self.conv_in(hidden_states)

        # down
        for down_block in self.down_blocks:
            hidden_states = down_block(hidden_states)

        # post-process
        hidden_states = self.norm_out(hidden_states)
        hidden_states = self.act_fn_out(hidden_states)

        return hidden_states