"""
-----------------------------------------------------------------------------
Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.

NVIDIA CORPORATION and its licensors retain all intellectual property
and proprietary rights in and to this software, related documentation
and any modifications thereto. Any use, reproduction, disclosure or
distribution of this software and related documentation without an express
license agreement from NVIDIA CORPORATION is strictly prohibited.
-----------------------------------------------------------------------------
"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

from vae.modules.attention import CrossAttention, SelfAttention


class FeedForward(nn.Module):
    def __init__(self, dim, mult=4):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(dim, dim * mult), nn.GELU(), nn.Linear(dim * mult, dim))

    def forward(self, x):
        return self.net(x)


# Adapted from https://github.com/facebookresearch/DiT/blob/main/models.py#L27
class TimestepEmbedder(nn.Module):
    """
    Embeds scalar timesteps into vector representations.
    """

    def __init__(self, hidden_size, frequency_embedding_size=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )
        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def timestep_embedding(t, dim, max_period=10000):
        """
        Create sinusoidal timestep embeddings.

        Args:
            t: a 1-D Tensor of N indices, one per batch element.
                These may be fractional.
            dim: the dimension of the output.
            max_period: controls the minimum frequency of the embeddings.

        Returns:
            an (N, D) Tensor of positional embeddings.
        """
        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
        half = dim // 2
        freqs = torch.exp(-np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
            device=t.device
        )
        args = t[:, None].float() * freqs[None]
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2:
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
        return embedding

    def forward(self, t):
        dtype = next(self.mlp.parameters()).dtype  # need to determine on the fly...
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
        t_freq = t_freq.to(dtype=dtype)
        t_emb = self.mlp(t_freq)
        return t_emb


class DiTLayer(nn.Module):
    def __init__(self, dim, num_heads, qknorm=False, gradient_checkpointing=True, qknorm_type="LayerNorm"):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.gradient_checkpointing = gradient_checkpointing

        self.norm1 = nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False)
        self.attn1 = SelfAttention(dim, num_heads, qknorm=qknorm, qknorm_type=qknorm_type)
        self.norm2 = nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False)
        self.attn2 = CrossAttention(dim, num_heads, context_dim=dim, qknorm=qknorm, qknorm_type=qknorm_type)
        self.norm3 = nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False)
        self.ff = FeedForward(dim)
        self.adaln_linear = nn.Linear(dim, dim * 6, bias=True)

    def forward(self, x, c, t_emb):
        if self.training and self.gradient_checkpointing:
            return checkpoint(self._forward, x, c, t_emb, use_reentrant=False)
        else:
            return self._forward(x, c, t_emb)

    def _forward(self, x, c, t_emb):
        # x: [B, N, C], hidden states
        # c: [B, M, C], condition (assume normed and projected to C)
        # t_emb: [B, C], timestep embedding of adaln
        # return: [B, N, C], updated hidden states

        B, N, C = x.shape
        t_adaln = self.adaln_linear(F.silu(t_emb)).view(B, 6, -1)  # [B, 6, C]
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = t_adaln.chunk(6, dim=1)

        h = self.norm1(x)
        h = h * (1 + scale_msa) + shift_msa
        x = x + gate_msa * self.attn1(h)

        h = self.norm2(x)
        x = x + self.attn2(h, c)

        h = self.norm3(x)
        h = h * (1 + scale_mlp) + shift_mlp
        x = x + gate_mlp * self.ff(h)

        return x


class DiT(nn.Module):
    def __init__(
        self,
        hidden_dim=1024,
        num_heads=16,
        latent_size=2048,
        latent_dim=8,
        num_layers=24,
        qknorm=False,
        gradient_checkpointing=True,
        qknorm_type="LayerNorm",
        use_pos_embed=False,
        use_parts=False,
        part_embed_mode="part2_only",
    ):
        super().__init__()

        # project in
        self.proj_in = nn.Linear(latent_dim, hidden_dim)

        # positional encoding (just use a learnable positional encoding)
        self.use_pos_embed = use_pos_embed
        if self.use_pos_embed:
            self.pos_embed = nn.Parameter(torch.randn(1, latent_size, hidden_dim) / hidden_dim**0.5)

        # part encoding (a must to distinguish parts!)
        self.use_parts = use_parts
        self.part_embed_mode = part_embed_mode
        if self.use_parts:
            if self.part_embed_mode == "element":
                self.part_embed = nn.Parameter(torch.randn(latent_size, hidden_dim) / hidden_dim**0.5)
            elif self.part_embed_mode == "part":
                self.part_embed = nn.Parameter(torch.randn(2, hidden_dim))
            elif self.part_embed_mode == "part2_only":
                # we only add this to the second part to distinguish from the first part
                self.part_embed = nn.Parameter(torch.randn(1, hidden_dim) / hidden_dim**0.5)

        # timestep encoding
        self.timestep_embed = TimestepEmbedder(hidden_dim)

        # transformer layers
        self.layers = nn.ModuleList(
            [DiTLayer(hidden_dim, num_heads, qknorm, gradient_checkpointing, qknorm_type) for _ in range(num_layers)]
        )

        # project out
        self.norm_out = nn.LayerNorm(hidden_dim, eps=1e-6, elementwise_affine=False)
        self.proj_out = nn.Linear(hidden_dim, latent_dim)

        # init
        self.init_weight()

    def init_weight(self):
        # Initialize transformer layers
        def _basic_init(module):
            if isinstance(module, nn.Linear):
                torch.nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)

        self.apply(_basic_init)

        # Initialize timestep embedding MLP:
        nn.init.normal_(self.timestep_embed.mlp[0].weight, std=0.02)
        nn.init.normal_(self.timestep_embed.mlp[2].weight, std=0.02)

        # Zero-out adaLN modulation layers in DiT blocks:
        for layer in self.layers:
            nn.init.constant_(layer.adaln_linear.weight, 0)
            nn.init.constant_(layer.adaln_linear.bias, 0)

        # Zero-out output layers:
        nn.init.constant_(self.proj_out.weight, 0)
        nn.init.constant_(self.proj_out.bias, 0)

    def forward(self, x, c, t):
        # x: [B, N, C], hidden states
        # c: [B, M, C], condition (assume normed and projected to C)
        # t: [B,], timestep
        # return: [B, N, C], updated hidden states

        B, N, C = x.shape

        # project in
        x = self.proj_in(x)

        # positional encoding
        if self.use_pos_embed:
            x = x + self.pos_embed

        # part encoding
        if self.use_parts:
            if self.part_embed_mode == "element":
                x += self.part_embed
            elif self.part_embed_mode == "part":
                x[:, : x.shape[1] // 2, :] += self.part_embed[0]
                x[:, x.shape[1] // 2 :, :] += self.part_embed[1]
            elif self.part_embed_mode == "part2_only":
                x[:, x.shape[1] // 2 :, :] += self.part_embed[0]

        # timestep encoding
        t_emb = self.timestep_embed(t)  # [B, C]

        # transformer layers
        for layer in self.layers:
            x = layer(x, c, t_emb)

        # project out
        x = self.norm_out(x)
        x = self.proj_out(x)

        return x