# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
import math
import pdb

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

from agents.networks.shared.general import LinearLayer, ResidualLayer, LayerNorm


class SelfAttentionLayer(nn.Module):
    """Basic Transformer encoder. Adds no positional encodings. Performs no special pooling of the output."""
    def __init__(self, vec_size, num_attention_heads, attention_head_size, obs_embed_size, is_last_layer=False):
        super().__init__()
        self.is_last_layer = is_last_layer
        self.vec_size = vec_size
        self.obs_embed_size = obs_embed_size
        self.num_attention_heads = num_attention_heads
        self.attention_head_size = attention_head_size

        self.query = LinearLayer(obs_embed_size, vec_size)
        self.key = LinearLayer(obs_embed_size, vec_size)
        self.value = LinearLayer(obs_embed_size, vec_size)

        self.dot_product_scale = 1.0 / math.sqrt(attention_head_size)

    def split_heads_apart(self, x):
        new_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_shape)  # Changes the tensor shape without moving or copying data.
        return x.permute(0, 2, 1, 3)  # Moves data to maintain the meaning of the reordered list of dimensions.

    def concatenate_heads(self, x):
        x = x.permute(0, 2, 1, 3).contiguous()
        new_shape = x.size()[:-2] + (self.vec_size,)
        return x.view(*new_shape)

    def forward(self, tens):  # [B, N, C]
        # B = batch size
        # C = num vector columns
        # H = num heads
        # S = head size
        # V = vec_size = H * S

        # Project the input to get the QKV vectors.
        queries = self.query(tens)  # [B, C, V]
        keys = self.key(tens)       # [B, C, V]
        values = self.value(tens)   # [B, C, V]

        # Split the heads apart to operate on the head vectors.
        split_queries = self.split_heads_apart(queries)  # [B, H, C, S]
        if self.is_last_layer:
            split_queries = split_queries[:, :, 0, :].unsqueeze(2)
        split_keys = self.split_heads_apart(keys)        # [B, H, C, S]
        split_values = self.split_heads_apart(values)    # [B, H, C, S]

        # Take the dot product between each query and key to get the raw attention scores.
        transposed_keys = split_keys.transpose(-1, -2)   # [B, H, S, C]
        attention_scores = torch.matmul(split_queries, transposed_keys)  # [B, H, C, C]
        attention_scores = attention_scores * self.dot_product_scale     # [B, H, C, C]

        # Normalize the attention scores to probabilities.
        attention_probs = nn.Softmax(dim=-1)(attention_scores)  # [B, H, C, C]

        # Calculate the weighted value vectors.
        weighted_values = torch.matmul(attention_probs, split_values)  # [B, H, C, S]

        # Concatenate the heads for output.
        output = self.concatenate_heads(weighted_values)  # [B, C, V]
        return output, attention_probs.detach()


class TransformerLayer(nn.Module):
    def __init__(
            self,
            vec_size,
            num_attention_heads,
            attention_head_size,
            hidden_layer_size,
            obs_embed_size,
            is_last_layer=False
    ):
        super(TransformerLayer, self).__init__()
        self.is_last_layer = is_last_layer
        self.attention = SelfAttentionLayer(
            vec_size, num_attention_heads, attention_head_size, obs_embed_size, is_last_layer
        )
        self.attention_residual = ResidualLayer(vec_size, obs_embed_size)
        self.attention_layer_norm = LayerNorm(obs_embed_size)

        self.feedforward = LinearLayer(vec_size, hidden_layer_size)
        self.feedforward_residual = ResidualLayer(hidden_layer_size, vec_size)
        self.feedforward_layer_norm = LayerNorm(vec_size)

    def forward(self, input):
        # Attention phase.
        att_output, att_probs = self.attention(input)
        if self.is_last_layer:
            att_output = self.attention_residual(att_output, input[:, 0, :])
        else:
            att_output = self.attention_residual(att_output, input)
        att_output = self.attention_layer_norm(att_output)

        # Feedforward phase.
        output = self.feedforward(att_output)
        output = F.relu(output)
        output = self.feedforward_residual(output, att_output)
        output = self.feedforward_layer_norm(output)
        return output, att_probs


class TransformerXLInputLayer(nn.Module):
    """Reorders canonical Transformers so that LayerNorm is done first"""
    def __init__(
            self,
            vec_size,
            num_attention_heads,
            attention_head_size,
            hidden_layer_size,
            obs_embed_size,
            is_last_layer=False
    ):
        super().__init__()
        self.is_last_layer = is_last_layer
        self.attention_layer_norm = LayerNorm(obs_embed_size)
        self.attention = SelfAttentionLayer(
            vec_size, num_attention_heads, attention_head_size, obs_embed_size, is_last_layer
        )
        self.attention_residual = ResidualLayer(vec_size, obs_embed_size)

        self.feedforward_layer_norm = LayerNorm(obs_embed_size)
        self.feedforward = LinearLayer(obs_embed_size, hidden_layer_size)
        self.feedforward_residual = ResidualLayer(hidden_layer_size, obs_embed_size)

    def forward(self, input):
        # Attention phase.
        input = self.attention_layer_norm(input)
        att_output, att_probs = self.attention(input)
        att_output = F.relu(att_output)
        if self.is_last_layer:
            att_output = self.attention_residual(att_output, input[:, 0, :])
        else:
            att_output = self.attention_residual(att_output, input)

        # Feedforward phase.
        output = self.feedforward_layer_norm(att_output)
        output = self.feedforward(output)
        output = F.relu(output)
        output = self.feedforward_residual(output, att_output)
        return output, att_probs


class GatedTransformerXLInputLayer(nn.Module):
    """Replaces residual layer with GRUCell."""
    def __init__(
            self,
            vec_size,
            num_attention_heads,
            attention_head_size,
            hidden_layer_size,
            obs_embed_size,
            is_last_layer=False
    ):
        super().__init__()
        self.is_last_layer = is_last_layer
        self.attention_layer_norm = LayerNorm(vec_size)
        self.attention = SelfAttentionLayer(
            vec_size, num_attention_heads, attention_head_size, obs_embed_size, is_last_layer
        )
        self.attention_gated = nn.GRUCell(vec_size, vec_size)

        self.feedforward_layer_norm = LayerNorm(vec_size)
        self.feedforward = LinearLayer(vec_size, hidden_layer_size)
        self.feedforward_gated = nn.GRUCell(hidden_layer_size, vec_size)

    def forward(self, input):
        # Attention phase.
        input = self.attention_layer_norm(input)
        att_output, att_probs = self.attention(input)
        att_output = F.relu(att_output)
        # TODO: should double check the input, if is_last_layer should do input[:, 0, :]
        att_output = self.attention_gated(torch.squeeze(att_output), torch.squeeze(input))
        att_output = torch.unsqueeze(att_output, 0)

        # Feedforward phase.
        output = self.feedforward_layer_norm(att_output)
        output = self.feedforward(output)
        output = F.relu(output)
        output = self.feedforward_gated(torch.squeeze(output), torch.squeeze(att_output))
        output = torch.unsqueeze(output, 0)
        return output, att_probs


class Transformer(nn.Module):
    def __init__(
            self,
            num_attention_heads,
            attention_head_size,
            num_layers,
            hidden_layer_size,
            obs_embed_size,
            layer_type='canonical'
    ):
        super().__init__()
        if layer_type == 'canonical':
            Layer = TransformerLayer
        elif layer_type == 'TrXL-I':
            Layer = TransformerXLInputLayer
        elif layer_type == 'GTrXL-I':
            Layer = GatedTransformerXLInputLayer
        else:
            print('Transformers layer type not recognized.')
            exit()
        vec_size = num_attention_heads * attention_head_size
        self.layers = nn.ModuleList()
        for i in range(num_layers-1):
            self.layers.append(Layer(
                vec_size, num_attention_heads, attention_head_size, hidden_layer_size, obs_embed_size
            )
        )

        self.layers.append(
            Layer(
                vec_size,
                num_attention_heads,
                attention_head_size,
                hidden_layer_size,
                obs_embed_size,
                is_last_layer=True
            )
        )

    def forward(self, tens):
        for layer in self.layers:
            tens, att_probs = layer(tens)
        return tens, att_probs
