from functools import partial, reduce
from typing import Callable, Optional

from einops import rearrange
from flash_attn import flash_attn_with_kvcache, flash_attn_func
import torch
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper,
    CheckpointImpl,
    apply_activation_checkpointing,
)
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.attention.flex_attention import flex_attention, create_block_mask

# drastic speedup
flex_attention = torch.compile(flex_attention)


class KVCache(nn.Module):
    def __init__(
        self,
        batch_size,
        seqlen,
        n_heads,
        head_size,
        device="cuda",
        dtype=torch.bfloat16,
    ):
        super().__init__()

        cache_shape = (batch_size, seqlen, n_heads, head_size)
        self.register_buffer(
            "k_cache",
            torch.zeros(cache_shape, device=device, dtype=dtype),
            persistent=False,
        )
        self.register_buffer(
            "v_cache",
            torch.zeros(cache_shape, device=device, dtype=dtype),
            persistent=False,
        )
        self.is_initialized = False

    def __len__(self):
        return self.k_cache.shape[1]


# 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 RotaryEmbedding(nn.Module):
    def __init__(self, dim, base=50_000):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)

    @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,
        qk_norm=False,
    ):
        super().__init__()
        self.dim = dim
        self.dim_heads = dim_heads
        self.qk_norm = qk_norm

        self.num_heads = dim // dim_heads

        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)

        self.kv_cache: Optional[KVCache] = None

    def forward(
        self,
        x,
        rotary_pos_emb=None,
        block_mask=None,
        pos=None,
    ):
        h = self.num_heads
        b, n, _ = x.shape

        # 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:
            freqs = rotary_pos_emb

            q_dtype, k_dtype = q.dtype, k.dtype
            q, k, freqs = [t.to(torch.float32) for t in (q, k, freqs)]

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

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

        # flex attention
        q, k, v = map(lambda t: t.to(torch.bfloat16), (q, k, v))

        if self.kv_cache is not None:
            q, k, v = map(lambda t: rearrange(t, "b h n d -> b n h d", h=h), (q, k, v))
            if pos.ndim == 1:
                pos = pos.unsqueeze(0).repeat_interleave(q.shape[0], dim=0)
            if not self.kv_cache.is_initialized:
                out = flash_attn_with_kvcache(
                    q,
                    self.kv_cache.k_cache,
                    self.kv_cache.v_cache,
                    k=k,
                    v=v,
                    cache_seqlens=pos[:, 0].int(),
                    causal=False,
                )
            else:
                out = flash_attn_func(
                    q,
                    torch.concat([self.kv_cache.k_cache, k], dim=1),
                    torch.concat([self.kv_cache.v_cache, v], dim=1),
                    causal=False,
                )
            out = rearrange(out, " b n h d -> b n (h d)")
            self.kv_cache.is_initialized = True
        else:
            out = flex_attention(q, k, v, block_mask=block_mask)
            out = rearrange(out, " b h n d -> b n (h d)")

        out = out.to(q.dtype)

        # 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,
        layer_ix=-1,
        qk_norm=False,
    ):
        super().__init__()
        self.dim = dim
        self.dim_heads = dim_heads

        self.pre_norm = LayerNorm(dim)

        self.self_attn = Attention(
            dim,
            dim_heads=dim_heads,
            qk_norm=qk_norm,
        )

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

        self.layer_ix = layer_ix

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

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

        return x


def _custom_mask_mod(cumulative_sizes, q_idx, kv_idx):
    q_in_misc = q_idx < cumulative_sizes[0]
    q_in_audio = (q_idx < cumulative_sizes[1]) & (q_idx >= cumulative_sizes[0])

    kv_in_misc = kv_idx < cumulative_sizes[0]
    kv_in_audio = (kv_idx < cumulative_sizes[1]) & (kv_idx >= cumulative_sizes[0])

    same_block = (q_in_misc == kv_in_misc) | (q_in_audio == kv_in_audio)

    return same_block | q_in_audio


def _create_block_mask(attn_block_sizes, b, num_heads, device):
    cumulative_sizes = torch.cumsum(
        torch.tensor(attn_block_sizes, dtype=torch.long, device=device), dim=0
    )
    seq_len = cumulative_sizes[-1]

    def mask_mod(b, h, q, k):
        return _custom_mask_mod(cumulative_sizes, q, k)

    block_mask = create_block_mask(
        mask_mod, B=b, H=num_heads, Q_LEN=seq_len, KV_LEN=seq_len, device=device, _compile=True
    )
    return block_mask


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

        self.dim = dim
        self.depth = depth
        self.layers = nn.ModuleList([])
        assert dim % dim_heads == 0
        self.dim_heads = dim_heads
        self.num_heads = dim // dim_heads

        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)

        assert attn_block_sizes is not None
        self.attn_block_sizes = attn_block_sizes
        self.block_mask = None

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

    @property
    def device(self):
        return next(self.parameters()).device

    def setup_caches(self, batch_size: int, seqlen: int):
        for b in self.layers:
            b.self_attn.kv_cache = KVCache(
                batch_size,
                seqlen,
                self.num_heads,
                self.dim_heads,
                dtype=self.layers[0].self_attn.to_out.weight.dtype,
                device=self.device,
            )

    def forward(
        self,
        x,
        prepend_embeds=None,
    ):
        # Direct implementation for compile compatibility
        # Don't call _core_forward to avoid conditional logic with stop_layer_idx
        b, _, 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=1)
        else:
            prepend_length = 0

        assert x.shape[1] == sum(self.attn_block_sizes), f"{x.shape[1]} != {self.attn_block_sizes}"

        if self.layers[0].self_attn.kv_cache is not None:
            assert len(self.attn_block_sizes) == 2
            if not self.layers[0].self_attn.kv_cache.is_initialized:
                x_pre = x[:, : self.attn_block_sizes[0]]
                pos = torch.arange(self.attn_block_sizes[0], device=device)
                rotary_pos_emb = self.rotary_pos_emb(pos)
                for layer in self.layers:
                    x_pre = layer(
                        x_pre,
                        rotary_pos_emb=rotary_pos_emb,
                        pos=pos,
                    )
            x = x[:, self.attn_block_sizes[0] :]
            pos = self.attn_block_sizes[0] + torch.arange(self.attn_block_sizes[1], device=device)
            prepen_cut = 1
        else:
            # refresh block mask if needed
            if self.block_mask is None or self.block_mask.shape[0] != b:
                print("refreshing block mask...")
                self.block_mask = _create_block_mask(self.attn_block_sizes, b, self.num_heads, x.device)

            pos = torch.arange(x.shape[1], device=device)
            prepen_cut = prepend_length

        rotary_pos_emb = self.rotary_pos_emb(pos)

        for layer in self.layers:
            x = layer(
                x,
                rotary_pos_emb=rotary_pos_emb,
                block_mask=self.block_mask,
                pos=pos,
            )

        x = self.project_out(x)
        x = x[:, prepen_cut:]

        return x

    def forward_intermediate(
        self,
        x,
        prepend_embeds=None,
        intermediate_layer_idx=None,
    ):
        if intermediate_layer_idx is None:
            raise ValueError("intermediate_layer_idx must be provided")
        return self._core_forward(
            x,
            prepend_embeds=prepend_embeds,
            stop_layer_idx=intermediate_layer_idx,
        )

    def _core_forward(self, x, prepend_embeds=None, stop_layer_idx=None):
        b, _, 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=1)
        else:
            prepend_length = 0

        assert x.shape[1] == sum(self.attn_block_sizes), f"{x.shape[1]} != {self.attn_block_sizes}"

        if self.layers[0].self_attn.kv_cache is not None:
            assert len(self.attn_block_sizes) == 2
            if not self.layers[0].self_attn.kv_cache.is_initialized:
                x_pre = x[:, : self.attn_block_sizes[0]]
                pos = torch.arange(self.attn_block_sizes[0], device=device)
                rotary_pos_emb = self.rotary_pos_emb(pos)
                for layer in self.layers:
                    x_pre = layer(
                        x_pre,
                        rotary_pos_emb=rotary_pos_emb,
                        pos=pos,
                    )
            x = x[:, self.attn_block_sizes[0] :]
            pos = self.attn_block_sizes[0] + torch.arange(self.attn_block_sizes[1], device=device)
            prepen_cut = 1
        else:
            # refresh block mask if needed
            if self.block_mask is None or self.block_mask.shape[0] != b:
                print("refreshing block mask...")
                self.block_mask = _create_block_mask(self.attn_block_sizes, b, self.num_heads, x.device)

            pos = torch.arange(x.shape[1], device=device)
            prepen_cut = prepend_length

        rotary_pos_emb = self.rotary_pos_emb(pos)

        for i, layer in enumerate(self.layers):
            x = layer(
                x,
                rotary_pos_emb=rotary_pos_emb,
                block_mask=self.block_mask,
                pos=pos,
            )
            if stop_layer_idx is not None and i == stop_layer_idx:
                return x[:, prepen_cut:]

        if stop_layer_idx is not None:
            raise ValueError(f"intermediate_layer_idx {stop_layer_idx} out of range")

        x = self.project_out(x)
        x = x[:, prepen_cut:]

        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)