import torch
from torch import nn


class CausalSelfAttention(nn.Module):

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

    def forward(self, x):
        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values 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).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        # efficient attention using Flash Attention CUDA kernels
        y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=self.causal)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

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


class MLP(nn.Module):

    def __init__(self, n_embd, bias=False):
        super().__init__()
        self.c_fc    = nn.Linear(n_embd, 4 * n_embd, bias=bias)
        self.c_proj  = nn.Linear(4 * n_embd, n_embd, bias=bias)
        self.gelu = nn.GELU()

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


class Block(nn.Module):

    def __init__(self, n_embd, n_head, causal=True, bias=False):
        super().__init__()
        self.ln_1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, causal=causal, bias=bias)
        self.ln_2 = nn.LayerNorm(n_embd)
        self.mlp = MLP(n_embd, bias=bias)

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


class Transformer(nn.Module):
    def __init__(
        self,
        dimension,
        causal=True,
        num_layers=2,
        skip=True,
        max_block_size=50,
        n_head=16,
        bias=False,
    ):
        super().__init__()
        assert(dimension in (1024, 2048))
        if dimension == 1024:
            n_head = 8
        elif dimension == 2048:
            n_head = 16
        assert(dimension % n_head == 0)
        self.skip = skip
        self.dimension = dimension
        self.max_block_size = max_block_size
        self.transformer = nn.ModuleDict(dict(
            wpe=nn.Embedding(max_block_size, dimension),
            h=nn.ModuleList([
                Block(dimension, n_head, bias=bias, causal=causal) for _ in range(num_layers)
            ]),
        ))

    def forward(self, x):
        # torch.Size([16, 2048, 50])
        device = x.device
        b, n_emb, t = x.size()
        if n_emb != self.dimension:
            print("robin", x.size())
            print("alfred", self.dimension)
        assert(n_emb == self.dimension)
        x = x.permute(0, 2, 1)
        # (b, t, n_embd)
        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0)  # shape (1, t)
        pos_emb = self.transformer.wpe(pos)  # position embeddings of shape (1, t, n_embd)
        y = x + pos_emb
        for block in self.transformer.h:
            y = block(y)
        if self.skip:
            y = y + x
        y = y.permute(0, 2, 1)
        # (b, n_embd, t)
        return y


class StreamableLSTM(nn.Module):
    """LSTM without worrying about the hidden state, nor the layout of the data.
    Expects input as convolutional layout.
    """
    def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True):
        super().__init__()
        self.skip = skip
        self.lstm = nn.LSTM(dimension, dimension, num_layers)

    def forward(self, x):
        x = x.permute(2, 0, 1)
        y, _ = self.lstm(x)
        if self.skip:
            y = y + x
        y = y.permute(1, 2, 0)
        return y