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


class LayerNorm(nn.Module):
    """LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False"""

    def __init__(self, ndim):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(ndim))

    def forward(self, input):
        return F.layer_norm(input, self.weight.shape, self.weight, None, 1e-5)


def create_sin_embedding(
    positions: torch.Tensor,
    dim: int,
    max_period: float = 10_000,
    dtype: torch.dtype = torch.float32
) -> torch.Tensor:
    """Create sinusoidal positional embedding, with shape `[B, T, C]`"""
    assert dim % 2 == 0
    half_dim = dim // 2
    positions = positions.to(dtype)
    adim = torch.arange(half_dim, device=positions.device, dtype=dtype).view(1, 1, -1)
    max_period_tensor = torch.full([], max_period, device=positions.device, dtype=dtype)  # avoid sync point
    phase = positions / (max_period_tensor ** (adim / (half_dim - 1)))
    return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)


class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        assert n_embd % n_head == 0
        # main
        self.n_head = n_head
        self.n_embd = n_embd
        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=False)
        # output projection
        self.c_proj = nn.Linear(n_embd, n_embd, bias=False)

    def forward(self, x):
        B, T, C = x.size()  # b_size, sequ_len, emb_dim (n_embd)

        # calculate q, k, v for all heads in batch and move head forward to be the batch dim
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head)
        v = v.view(B, T, self.n_head, C // self.n_head)
        q = q.view(B, T, self.n_head, C // self.n_head)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        k = k.transpose(1, 2)  # (B, nh, T, hs)
        v = v.transpose(1, 2)  # (B, nh, T, hs)
        q = q.transpose(1, 2)  # (B, nh, T, hs)
        # with torch.backends.cuda.sdp_kernel(
        #     enable_flash=True, enable_math=False, enable_mem_efficient=False
        # ):
        # TODO: for compile to work we have to remove this context manager
        y = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, dropout_p=0, is_causal=False
        )
        y = y.transpose(1, 2)

        # re-assemble head outputs side by side
        y = y.contiguous().view(B, T, C)

        # output projection
        y = self.c_proj(y)
        return y


class MLP(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        embd_inner = 4 * n_embd
        self.embd_inner = embd_inner
        self.c_fc = nn.Linear(n_embd, self.embd_inner, bias=False)
        self.c_proj = nn.Linear(self.embd_inner, n_embd, bias=False)
        self.activation = nn.GELU()

    def forward(self, x):
        x = self.c_fc(x)
        x = self.activation(x)
        x = self.c_proj(x)
        return x


class Block(nn.Module):
    def __init__(self, n_head, n_embd):
        super().__init__()
        self.ln_1 = LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head)
        self.ln_2 = LayerNorm(n_embd)
        self.mlp = MLP(n_embd)

    def forward(self, x):
        x = x.swapaxes(-1, -2)
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        x = x.swapaxes(-1, -2)
        return x