from functools import partial, reduce
from typing import Callable

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange


try:
    from flash_attn_interface import flash_attn_func
except ImportError:
    print("Failed to import flash_attn_3.")
    try:
        from flash_attn import flash_attn_func
    except ImportError:
        print("Failed to import flash_attn. You won't be able to use the Diffusion model.")
        flash_attn_func = None
        flash_attn_varlen_func = None

# from flash_attn_interface import flash_attn_func
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper,
    CheckpointImpl,
    apply_activation_checkpointing,
)


# Custom Layer Normalization with bias control
class LayerNorm(nn.Module):
    def __init__(self, dim, bias=False):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(dim))
        if bias:
            self.beta = nn.Parameter(torch.zeros(dim))
        else:
            self.register_buffer("beta", torch.zeros(dim))

    def forward(self, x):
        return F.layer_norm(x, x.shape[-1:], weight=self.gamma, bias=self.beta)


class ScaledSinusoidalEmbedding(nn.Module):
    def __init__(self, dim, theta=10000):
        super().__init__()
        assert (dim % 2) == 0, "dimension must be divisible by 2"
        self.scale = nn.Parameter(torch.ones(1) * dim**-0.5)

        half_dim = dim // 2
        freq_seq = torch.arange(half_dim, dtype=torch.float32) / half_dim
        inv_freq = theta**-freq_seq
        self.register_buffer("inv_freq", inv_freq, persistent=True)

    @torch.autocast("cuda", enabled=False)
    def forward(self, x, pos=None, seq_start_pos=None):
        seq_len, device = x.shape[1], x.device

        if pos is None:
            pos = torch.arange(seq_len, device=device, dtype=torch.float32)

        if seq_start_pos is not None:
            pos = pos - seq_start_pos[..., None]

        emb = torch.outer(pos, self.inv_freq)
        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
        return emb * self.scale


class RotaryEmbedding(nn.Module):
    def __init__(self, dim, base=10000):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)

    def forward_from_seq_len(self, seq_len):
        device = self.inv_freq.device
        t = torch.arange(seq_len, device=device)
        return self.forward(t)

    @torch.autocast("cuda", enabled=False)
    def forward(self, t):
        t = t.to(torch.float32)
        freqs = torch.outer(t, self.inv_freq)
        freqs = torch.cat((freqs, freqs), dim=-1)
        return freqs


def rotate_half(x):
    x = rearrange(x, "... (j d) -> ... j d", j=2)
    x1, x2 = x.unbind(dim=-2)
    return torch.cat((-x2, x1), dim=-1)


@torch.autocast("cuda", enabled=False)
def apply_rotary_pos_emb(t, freqs):
    out_dtype = t.dtype

    # cast to float32 if necessary for numerical stability
    dtype = reduce(torch.promote_types, (t.dtype, freqs.dtype, torch.float32))
    rot_dim, seq_len = freqs.shape[-1], t.shape[-2]
    freqs, t = freqs.to(dtype), t.to(dtype)
    freqs = freqs[-seq_len:, :]

    if t.ndim == 4 and freqs.ndim == 3:
        freqs = rearrange(freqs, "b n d -> b 1 n d")

    # partial rotary embeddings, Wang et al. GPT-J
    t, t_unrotated = t[..., :rot_dim], t[..., rot_dim:]
    t = t * freqs.cos() + rotate_half(t) * freqs.sin()

    t, t_unrotated = t.to(out_dtype), t_unrotated.to(out_dtype)

    return torch.cat((t, t_unrotated), dim=-1)


class Attention(nn.Module):
    def __init__(
        self,
        dim,
        dim_heads=64,
        causal=False,
        is_cross=False,
        qk_norm=False,
    ):
        super().__init__()
        self.dim = dim
        self.dim_heads = dim_heads
        self.causal = causal
        self.is_cross = is_cross
        self.qk_norm = qk_norm

        self.num_heads = dim // dim_heads

        if self.is_cross:
            self.to_q = nn.Linear(dim, dim, bias=False)
            self.to_kv = nn.Linear(dim, dim * 2, bias=False)
        else:
            self.to_qkv = nn.Linear(dim, dim * 3, bias=False)

        if self.qk_norm:
            self.q_ln = nn.LayerNorm(dim_heads)
            self.k_ln = nn.LayerNorm(dim_heads)

        self.to_out = nn.Linear(dim, dim, bias=False)

    def forward(
        self,
        x,
        context=None,
        rotary_pos_emb=None,
        causal=None,
    ):
        h, has_context = self.num_heads, context is not None
        assert has_context == self.is_cross

        if has_context:
            # Use separate linear projections for q and k/v
            q = self.to_q(x)
            k, v = self.to_kv(context).chunk(2, dim=-1)

        else:
            # Use fused linear projection
            q, k, v = self.to_qkv(x).chunk(3, dim=-1)

        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))

        # Normalize q and k for cosine sim attention
        if self.qk_norm:
            q = self.q_ln(q)
            k = self.k_ln(k)

        if rotary_pos_emb is not None and not has_context:
            freqs = rotary_pos_emb

            q_dtype = q.dtype
            k_dtype = k.dtype

            q = q.to(torch.float32)
            k = k.to(torch.float32)
            freqs = freqs.to(torch.float32)

            q = apply_rotary_pos_emb(q, freqs)
            k = apply_rotary_pos_emb(k, freqs)

            q = q.to(q_dtype)
            k = k.to(k_dtype)

        n, device = q.shape[-2], q.device

        causal = self.causal if causal is None else causal

        if n == 1 and causal:
            causal = False

        # TODO: is it an issue to always cast to bfloat16 here??
        fa_dtype_in = q.dtype
        q, k, v = map(
            lambda t: rearrange(t, "b h n d -> b n h d").to(torch.bfloat16),
            (q, k, v),
        )

        out = flash_attn_func(q, k, v, causal=causal)
        if isinstance(out, tuple):  # not sure why this is a tuple with fa3
            out = out[0]
        out = rearrange(out.to(fa_dtype_in), "b n h d -> b h n d")

        # merge heads
        out = rearrange(out, " b h n d -> b n (h d)")

        # Communicate between heads
        out = self.to_out(out)
        return out


class GLU(nn.Module):
    def __init__(
        self,
        dim_in,
        dim_out,
        activation: Callable,
    ):
        super().__init__()
        self.act = activation
        self.proj = nn.Linear(dim_in, dim_out * 2)

    def forward(self, x):
        x = self.proj(x)
        x, gate = x.chunk(2, dim=-1)
        return x * self.act(gate)


class FeedForward(nn.Module):
    def __init__(
        self,
        dim,
        dim_out=None,
        mult=4,
    ):
        super().__init__()
        inner_dim = int(dim * mult)

        # Default to SwiGLU

        activation = nn.SiLU()

        dim_out = dim if dim_out is None else dim_out

        linear_in = GLU(dim, inner_dim, activation)

        linear_out = nn.Linear(inner_dim, dim_out, bias=True)

        self.ff = nn.Sequential(
            linear_in,
            linear_out,
        )

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


class TransformerBlock(nn.Module):
    def __init__(
        self,
        dim,
        dim_heads=64,
        causal=False,
        layer_ix=-1,
        qk_norm=False,
    ):
        super().__init__()
        self.dim = dim
        self.dim_heads = dim_heads
        self.causal = causal

        self.pre_norm = LayerNorm(dim)

        self.self_attn = Attention(
            dim,
            dim_heads=dim_heads,
            causal=causal,
            is_cross=False,
            qk_norm=qk_norm,
        )

        self.cross_attend_norm = LayerNorm(dim)
        self.cross_attn = Attention(
            dim,
            dim_heads=dim_heads,
            causal=causal,
            is_cross=True,
            qk_norm=qk_norm,
        )

        self.ff_norm = LayerNorm(dim)
        self.ff = FeedForward(dim)

        self.layer_ix = layer_ix

    def forward(
        self,
        x,
        context=None,
        rotary_pos_emb=None,
    ):
        x = x + self.self_attn(
            self.pre_norm(x),
            rotary_pos_emb=rotary_pos_emb,
        )

        if context is not None:
            x = x + self.cross_attn(
                self.cross_attend_norm(x),
                context=context,
            )

        x = x + self.ff(self.ff_norm(x))

        return x


class ContinuousTransformer(nn.Module):
    def __init__(
        self,
        dim,
        depth,
        dim_heads=64,
        dim_in=None,
        dim_out=None,
        causal=False,
        qk_norm=False,
    ):
        super().__init__()

        self.dim = dim
        self.depth = depth
        self.causal = causal
        self.layers = nn.ModuleList([])

        self.project_in = nn.Linear(dim_in, dim, bias=False)
        self.project_out = nn.Linear(dim, dim_out, bias=False)

        self.rotary_pos_emb = RotaryEmbedding(dim_heads)

        for i in range(depth):
            self.layers.append(
                TransformerBlock(
                    dim,
                    dim_heads=dim_heads,
                    causal=causal,
                    layer_ix=i,
                    qk_norm=qk_norm,
                )
            )

    def forward(
        self,
        x,
        context=None,
        prepend_embeds=None,
    ):
        batch, seq, device = *x.shape[:2], x.device

        x = self.project_in(x)

        if prepend_embeds is not None:
            prepend_length, prepend_dim = prepend_embeds.shape[1:]

            assert prepend_dim == x.shape[-1], "prepend dimension must match sequence dimension"

            x = torch.cat((prepend_embeds, x), dim=-2)

        # Attention layers

        rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1])

        # Iterate over the transformer layers
        for layer in self.layers:
            x = layer(
                x,
                context=context,
                rotary_pos_emb=rotary_pos_emb,
            )

        x = self.project_out(x)

        return x


non_reentrant_wrapper = partial(
    checkpoint_wrapper,
    checkpoint_impl=CheckpointImpl.NO_REENTRANT,
)


def check_fn(submodule):
    return isinstance(submodule, TransformerBlock)


def apply_fsdp_checkpointing(model):
    apply_activation_checkpointing(model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=check_fn)