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


class SelfAttention(nn.Module):

    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        # key, query, value projections for all heads
        self.key = nn.Linear(config.n_embd, config.n_embd)
        self.query = nn.Linear(config.n_embd, config.n_embd)
        self.value = nn.Linear(config.n_embd, config.n_embd)
        # regularization
        self.attn_drop = nn.Dropout(config.attn_pdrop)
        self.resid_drop = nn.Dropout(config.resid_pdrop)
        # output projection
        self.proj = nn.Linear(config.n_embd, config.n_embd)
        self.n_head = config.n_head

    def forward(self, x, mask):
        B, T, C = x.size()

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        ## [ B x n_heads x T x head_dim ]
        k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = self.value(x).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)
        ## [ B x n_heads x T x T ]
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = att.masked_fill(mask[:,:,:T,:T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        self._attn_map = att.clone()
        att = self.attn_drop(att)
        ## [ B x n_heads x T x head_size ]
        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        ## [ B x T x embedding_dim ]
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

        # output projection
        y = self.resid_drop(self.proj(y))
        return y


class CausalSelfAttention(SelfAttention):
    def __init__(self, config):
        super().__init__(config)
        self.register_buffer("mask", torch.tril(torch.ones(config.block_size, config.block_size))
                                     .view(1, 1, config.block_size, config.block_size))
        transition_dim = config.transition_dim
        if config.mask_values:
            ## mask previous value estimates
            joined_dim = config.observation_dim + config.action_dim + 2
            self.mask.squeeze()[:,joined_dim-1::joined_dim] = 0
        if hasattr(config, 'single_return_token') and config.single_return_token:
            self.mask.squeeze()[:, (1+transition_dim)::transition_dim] = 0

    def forward(self, x):
        return super().forward(x, self.mask)

class FullSelfAttention(SelfAttention):
    def __init__(self, config):
        super().__init__(config)
        self.register_buffer("mask", torch.ones(config.block_size, config.block_size).view(1, 1, config.block_size, config.block_size))
        if config.mask_values:
            ## mask previous value estimates
            joined_dim = config.observation_dim + config.action_dim + 2
            self.mask.squeeze()[:,joined_dim-1::joined_dim] = 0

    def forward(self, x):
        return super().forward(x, self.mask)
