import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import einops
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.modeling_utils import ModelMixin
from typing import Any, Tuple, Optional
from flash_attn import flash_attn_varlen_func
from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa

from .layers import LLamaFeedForward, RMSNorm

# import frasch


def modulate(x, scale):
    return x * (1 + scale)

class TimestepEmbedder(nn.Module):
    """
    Embeds scalar timesteps into vector representations.
    """
    def __init__(self, hidden_size, frequency_embedding_size=256):
        super().__init__()
        self.hidden_size = hidden_size
        self.frequency_embedding_size = frequency_embedding_size
        self.mlp = nn.Sequential(
            nn.Linear(self.frequency_embedding_size, self.hidden_size),
            nn.SiLU(),
            nn.Linear(self.hidden_size, self.hidden_size),
        )

    @staticmethod
    def timestep_embedding(t, dim, max_period=10000):
        """
        Create sinusoidal timestep embeddings.
        :param t: a 1-D Tensor of N indices, one per batch element.
        :param dim: the dimension of the output.
        :param max_period: controls the minimum frequency of the embeddings.
        :return: an (N, D) Tensor of positional embeddings.
        """
        half = dim // 2
        freqs = torch.exp(
            -np.log(max_period) * torch.arange(0, half, dtype=t.dtype) / half
        ).to(t.device)
        args = t[:, :, None] * 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):
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
        t_freq = t_freq.to(self.mlp[0].weight.dtype)
        return self.mlp(t_freq)

class FinalLayer(nn.Module):
    def __init__(self, hidden_size, num_patches, out_channels):
        super().__init__()
        self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
        self.linear = nn.Linear(hidden_size, num_patches * out_channels)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(min(hidden_size, 1024), hidden_size),
        )
        
    def forward(self, x, c):
        scale = self.adaLN_modulation(c)
        x = modulate(self.norm_final(x), scale)
        x = self.linear(x)
        return x

class Attention(nn.Module):
    def __init__(
        self,
        dim,
        n_heads,
        n_kv_heads=None,
        qk_norm=False,
        y_dim=0,
        base_seqlen=None,
        proportional_attn=False,
        attention_dropout=0.0,
        max_position_embeddings=384,
    ):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads or n_heads
        self.qk_norm = qk_norm
        self.y_dim = y_dim
        self.base_seqlen = base_seqlen
        self.proportional_attn = proportional_attn
        self.attention_dropout = attention_dropout
        self.max_position_embeddings = max_position_embeddings

        self.head_dim = dim // n_heads

        self.wq = nn.Linear(dim, n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)

        if y_dim > 0:
            self.wk_y = nn.Linear(y_dim, self.n_kv_heads * self.head_dim, bias=False)
            self.wv_y = nn.Linear(y_dim, self.n_kv_heads * self.head_dim, bias=False)
            self.gate = nn.Parameter(torch.zeros(n_heads))

        self.wo = nn.Linear(n_heads * self.head_dim, dim, bias=False)

        if qk_norm:
            self.q_norm = nn.LayerNorm(self.n_heads * self.head_dim)
            self.k_norm = nn.LayerNorm(self.n_kv_heads * self.head_dim)
            if y_dim > 0:
                self.ky_norm = nn.LayerNorm(self.n_kv_heads * self.head_dim, eps=1e-6)
            else:
                self.ky_norm = nn.Identity()
        else:
            self.q_norm = nn.Identity()
            self.k_norm = nn.Identity()
            self.ky_norm = nn.Identity()


    @staticmethod
    def apply_rotary_emb(xq, xk, freqs_cis):
        # xq, xk: [batch_size, seq_len, n_heads, head_dim]
        # freqs_cis: [1, seq_len, 1, head_dim]
        xq_ = xq.float().reshape(*xq.shape[:-1], -1, 2)
        xk_ = xk.float().reshape(*xk.shape[:-1], -1, 2)

        xq_complex = torch.view_as_complex(xq_)
        xk_complex = torch.view_as_complex(xk_)
        
        freqs_cis = freqs_cis.unsqueeze(2)

        # Apply freqs_cis
        xq_out = xq_complex * freqs_cis
        xk_out = xk_complex * freqs_cis

        # Convert back to real numbers
        xq_out = torch.view_as_real(xq_out).flatten(-2)
        xk_out = torch.view_as_real(xk_out).flatten(-2)

        return xq_out.type_as(xq), xk_out.type_as(xk)
    
    # copied from huggingface modeling_llama.py
    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
        def _get_unpad_data(attention_mask):
            seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
            indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
            max_seqlen_in_batch = seqlens_in_batch.max().item()
            cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
            return (
                indices,
                cu_seqlens,
                max_seqlen_in_batch,
            )

        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape

        key_layer = index_first_axis(
            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),
            indices_k,
        )
        value_layer = index_first_axis(
            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),
            indices_k,
        )
        if query_length == kv_seq_len:
            query_layer = index_first_axis(
                query_layer.reshape(batch_size * kv_seq_len, self.n_heads, head_dim),
                indices_k,
            )
            cu_seqlens_q = cu_seqlens_k
            max_seqlen_in_batch_q = max_seqlen_in_batch_k
            indices_q = indices_k
        elif query_length == 1:
            max_seqlen_in_batch_q = 1
            cu_seqlens_q = torch.arange(
                batch_size + 1, dtype=torch.int32, device=query_layer.device
            )  # There is a memcpy here, that is very bad.
            indices_q = cu_seqlens_q[:-1]
            query_layer = query_layer.squeeze(1)
        else:
            # The -q_len: slice assumes left padding.
            attention_mask = attention_mask[:, -query_length:]
            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)

        return (
            query_layer,
            key_layer,
            value_layer,
            indices_q,
            (cu_seqlens_q, cu_seqlens_k),
            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
        )

    def forward(
        self,
        x,
        x_mask,
        freqs_cis,
        y=None,
        y_mask=None,
        init_cache=False,
    ):
        bsz, seqlen, _ = x.size()
        xq = self.wq(x)
        xk = self.wk(x)
        xv = self.wv(x)

        if x_mask is None:
            x_mask = torch.ones(bsz, seqlen, dtype=torch.bool, device=x.device)
        inp_dtype = xq.dtype

        xq = self.q_norm(xq)
        xk = self.k_norm(xk)

        xq = xq.view(bsz, seqlen, self.n_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_kv_heads, self.head_dim)

        if self.n_kv_heads != self.n_heads:
            n_rep = self.n_heads // self.n_kv_heads
            xk = xk.repeat_interleave(n_rep, dim=2)
            xv = xv.repeat_interleave(n_rep, dim=2)

        freqs_cis = freqs_cis.to(xq.device)
        xq, xk = self.apply_rotary_emb(xq, xk, freqs_cis)

        if inp_dtype in [torch.float16, torch.bfloat16]:
            # begin var_len flash attn
            (
                query_states,
                key_states,
                value_states,
                indices_q,
                cu_seq_lens,
                max_seq_lens,
            ) = self._upad_input(xq, xk, xv, x_mask, seqlen)

            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens

            attn_output_unpad = flash_attn_varlen_func(
                query_states.to(inp_dtype),
                key_states.to(inp_dtype),
                value_states.to(inp_dtype),
                cu_seqlens_q=cu_seqlens_q,
                cu_seqlens_k=cu_seqlens_k,
                max_seqlen_q=max_seqlen_in_batch_q,
                max_seqlen_k=max_seqlen_in_batch_k,
                dropout_p=0.0,
                causal=False,
                softmax_scale=None,
                softcap=30,
            )
            output = pad_input(attn_output_unpad, indices_q, bsz, seqlen)
        else:
            output = (
                F.scaled_dot_product_attention(
                    xq.permute(0, 2, 1, 3),
                    xk.permute(0, 2, 1, 3),
                    xv.permute(0, 2, 1, 3),
                    attn_mask=x_mask.bool().view(bsz, 1, 1, seqlen).expand(-1, self.n_heads, seqlen, -1),
                    scale=None,
                )
                .permute(0, 2, 1, 3)
                .to(inp_dtype)
            ) #ok


        if hasattr(self, "wk_y"):
            yk = self.ky_norm(self.wk_y(y)).view(bsz, -1, self.n_kv_heads, self.head_dim)
            yv = self.wv_y(y).view(bsz, -1, self.n_kv_heads, self.head_dim)
            n_rep = self.n_heads // self.n_kv_heads
            # if n_rep >= 1:
            #     yk = yk.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3)
            #     yv = yv.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3)
            if n_rep >= 1:
                yk = einops.repeat(yk, "b l h d -> b l (repeat h) d", repeat=n_rep)
                yv = einops.repeat(yv, "b l h d -> b l (repeat h) d", repeat=n_rep)
            output_y = F.scaled_dot_product_attention(
                xq.permute(0, 2, 1, 3),
                yk.permute(0, 2, 1, 3),
                yv.permute(0, 2, 1, 3),
                y_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_heads, seqlen, -1).to(torch.bool),
            ).permute(0, 2, 1, 3)
            output_y = output_y * self.gate.tanh().view(1, 1, -1, 1)
            output = output + output_y

        output = output.flatten(-2)
        output = self.wo(output)

        return output.to(inp_dtype)

class TransformerBlock(nn.Module):
    """
    Corresponds to the Transformer block in the JAX code.
    """
    def __init__(
        self,
        dim,
        n_heads,
        n_kv_heads,
        multiple_of,
        ffn_dim_multiplier,
        norm_eps,
        qk_norm,
        y_dim,
        max_position_embeddings,
    ):
        super().__init__()
        self.attention = Attention(dim, n_heads, n_kv_heads, qk_norm, y_dim=y_dim, max_position_embeddings=max_position_embeddings)
        self.feed_forward = LLamaFeedForward(
            dim=dim,
            hidden_dim=4 * dim,
            multiple_of=multiple_of,
            ffn_dim_multiplier=ffn_dim_multiplier,
        )
        self.attention_norm1 = RMSNorm(dim, eps=norm_eps)
        self.attention_norm2 = RMSNorm(dim, eps=norm_eps)
        self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
        self.ffn_norm2 = RMSNorm(dim, eps=norm_eps)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(min(dim, 1024), 4 * dim),
        )
        self.attention_y_norm = RMSNorm(y_dim, eps=norm_eps)

    def forward(
        self,
        x,
        x_mask,
        freqs_cis,
        y,
        y_mask,
        adaln_input=None,
    ):
        if adaln_input is not None:
            scales_gates = self.adaLN_modulation(adaln_input)
            # TODO: Duong - check the dimension of chunking
            # scale_msa, gate_msa, scale_mlp, gate_mlp = scales_gates.chunk(4, dim=-1)
            scale_msa, gate_msa, scale_mlp, gate_mlp = scales_gates.chunk(4, dim=-1)
            x = x + torch.tanh(gate_msa) * self.attention_norm2(
                self.attention(
                    modulate(self.attention_norm1(x), scale_msa), # ok
                    x_mask,
                    freqs_cis,
                    self.attention_y_norm(y), # ok
                    y_mask,
                )
            )
            x = x + torch.tanh(gate_mlp) * self.ffn_norm2(
                self.feed_forward(
                    modulate(self.ffn_norm1(x), scale_mlp),
                )
            )
        else:
            x = x + self.attention_norm2(
                self.attention(
                    self.attention_norm1(x),
                    x_mask,
                    freqs_cis,
                    self.attention_y_norm(y),
                    y_mask,
                )
            )
            x = x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x)))
        return x


class NextDiT(ModelMixin, ConfigMixin):
    """
    Diffusion model with a Transformer backbone for joint image-video training.
    """
    @register_to_config
    def __init__(
        self,
        input_size=(1, 32, 32),
        patch_size=(1, 2, 2),
        in_channels=16,
        hidden_size=4096,
        depth=32,
        num_heads=32,
        num_kv_heads=None,
        multiple_of=256,
        ffn_dim_multiplier=None,
        norm_eps=1e-5,
        pred_sigma=False,
        caption_channels=4096,
        qk_norm=False,
        norm_type="rms",
        model_max_length=120,
        rotary_max_length=384,
        rotary_max_length_t=None
    ):
        super().__init__()
        self.input_size = input_size
        self.patch_size = patch_size
        self.in_channels = in_channels
        self.hidden_size = hidden_size
        self.depth = depth
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads or num_heads
        self.multiple_of = multiple_of
        self.ffn_dim_multiplier = ffn_dim_multiplier
        self.norm_eps = norm_eps
        self.pred_sigma = pred_sigma
        self.caption_channels = caption_channels
        self.qk_norm = qk_norm
        self.norm_type = norm_type
        self.model_max_length = model_max_length
        self.rotary_max_length = rotary_max_length
        self.rotary_max_length_t = rotary_max_length_t
        self.out_channels = in_channels * 2 if pred_sigma else in_channels

        self.x_embedder = nn.Linear(np.prod(self.patch_size) * in_channels, hidden_size)

        self.t_embedder = TimestepEmbedder(min(hidden_size, 1024))
        self.y_embedder = nn.Sequential(
            nn.LayerNorm(caption_channels, eps=1e-6),
            nn.Linear(caption_channels, min(hidden_size, 1024)),
        )

        self.layers = nn.ModuleList([
            TransformerBlock(
                dim=hidden_size,
                n_heads=num_heads,
                n_kv_heads=self.num_kv_heads,
                multiple_of=multiple_of,
                ffn_dim_multiplier=ffn_dim_multiplier,
                norm_eps=norm_eps,
                qk_norm=qk_norm,
                y_dim=caption_channels,
                max_position_embeddings=rotary_max_length,
            )
            for _ in range(depth)
        ])

        self.final_layer = FinalLayer(
            hidden_size=hidden_size,
            num_patches=np.prod(patch_size),
            out_channels=self.out_channels,
        )

        assert (hidden_size // num_heads) % 6 == 0, "3d rope needs head dim to be divisible by 6"

        self.freqs_cis = self.precompute_freqs_cis(
            hidden_size // num_heads,
            self.rotary_max_length,
            end_t=self.rotary_max_length_t
        )
    
    def to(self, *args, **kwargs):
        self = super().to(*args, **kwargs)
        # self.freqs_cis = self.freqs_cis.to(*args, **kwargs)
        return self

    @staticmethod
    def precompute_freqs_cis(
        dim: int,
        end: int,
        end_t: int = None,
        theta: float = 10000.0,
        scale_factor: float = 1.0,
        scale_watershed: float = 1.0,
        timestep: float = 1.0,
    ):
        if timestep < scale_watershed:
            linear_factor = scale_factor
            ntk_factor = 1.0
        else:
            linear_factor = 1.0
            ntk_factor = scale_factor

        theta = theta * ntk_factor
        freqs = 1.0 / (theta ** (torch.arange(0, dim, 6)[: (dim // 6)] / dim)) / linear_factor

        timestep = torch.arange(end, dtype=torch.float32)
        freqs = torch.outer(timestep, freqs).float()
        freqs_cis = torch.exp(1j * freqs)

        if end_t is not None:
            freqs_t = 1.0 / (theta ** (torch.arange(0, dim, 6)[: (dim // 6)] / dim)) / linear_factor
            timestep_t = torch.arange(end_t, dtype=torch.float32)
            freqs_t = torch.outer(timestep_t, freqs_t).float()
            freqs_cis_t = torch.exp(1j * freqs_t)
            freqs_cis_t = freqs_cis_t.view(end_t, 1, 1, dim // 6).repeat(1, end, end, 1)
        else:
            end_t = end
            freqs_cis_t = freqs_cis.view(end_t, 1, 1, dim // 6).repeat(1, end, end, 1)
            
        freqs_cis_h = freqs_cis.view(1, end, 1, dim // 6).repeat(end_t, 1, end, 1)
        freqs_cis_w = freqs_cis.view(1, 1, end, dim // 6).repeat(end_t, end, 1, 1)
        freqs_cis = torch.cat([freqs_cis_t, freqs_cis_h, freqs_cis_w], dim=-1).view(end_t, end, end, -1)
        return freqs_cis

    def forward(
        self, 
        samples, 
        timesteps, 
        encoder_hidden_states,
        encoder_attention_mask,
        scale_factor: float = 1.0, # scale_factor for rotary embedding
        scale_watershed: float = 1.0, # scale_watershed for rotary embedding
    ):
        if samples.ndim == 4: # B C H W
            samples = samples[:, None, ...] # B F C H W
        
        precomputed_freqs_cis = None
        if scale_factor != 1 or scale_watershed != 1:
            precomputed_freqs_cis = self.precompute_freqs_cis(
                self.hidden_size // self.num_heads,
                self.rotary_max_length,
                end_t=self.rotary_max_length_t,
                scale_factor=scale_factor,
                scale_watershed=scale_watershed,
                timestep=torch.max(timesteps.cpu()).item()
            )
            
        if len(timesteps.shape) == 5:
            t, *_ = self.patchify(timesteps, precomputed_freqs_cis)
            timesteps = t.mean(dim=-1)
        elif len(timesteps.shape) == 1:
            timesteps = timesteps[:, None, None, None, None].expand_as(samples)
            t, *_ = self.patchify(timesteps, precomputed_freqs_cis)
            timesteps = t.mean(dim=-1)
        samples, T, H, W, freqs_cis = self.patchify(samples, precomputed_freqs_cis)
        samples = self.x_embedder(samples)
        t = self.t_embedder(timesteps)

        encoder_attention_mask_float = encoder_attention_mask[..., None].float()
        encoder_hidden_states_pool = (encoder_hidden_states * encoder_attention_mask_float).sum(dim=1) / (encoder_attention_mask_float.sum(dim=1) + 1e-8)
        encoder_hidden_states_pool = encoder_hidden_states_pool.to(samples.dtype)
        y = self.y_embedder(encoder_hidden_states_pool)
        y = y.unsqueeze(1).expand(-1, samples.size(1), -1)

        adaln_input = t + y
                                
        for block in self.layers:
            samples = block(samples, None, freqs_cis, encoder_hidden_states, encoder_attention_mask, adaln_input)

        samples = self.final_layer(samples, adaln_input)
        samples = self.unpatchify(samples, T, H, W)

        return samples

    def patchify(self, x, precompute_freqs_cis=None):
        # pytorch is C, H, W
        B, T, C, H, W = x.size()
        pT, pH, pW = self.patch_size
        x = x.view(B, T // pT, pT, C, H // pH, pH, W // pW, pW)
        x = x.permute(0, 1, 4, 6, 2, 5, 7, 3)
        x = x.reshape(B, -1, pT * pH * pW * C)
        if precompute_freqs_cis is None:
            freqs_cis = self.freqs_cis[: T // pT, :H // pH, :W // pW].reshape(-1, * self.freqs_cis.shape[3:])[None].to(x.device)
        else:
            freqs_cis = precompute_freqs_cis[: T // pT, :H // pH, :W // pW].reshape(-1, * precompute_freqs_cis.shape[3:])[None].to(x.device)
        return x, T // pT, H // pH, W // pW, freqs_cis

    def unpatchify(self, x, T, H, W):
        B = x.size(0)
        C = self.out_channels
        pT, pH, pW = self.patch_size
        x = x.view(B, T, H, W, pT, pH, pW, C)
        x = x.permute(0, 1, 4, 7, 2, 5, 3, 6)
        x = x.reshape(B, T * pT, C, H * pH, W * pW)
        return x