import torch
import torch.nn as nn
from DecoderLayer import DecoderLayer
import math

class Decoder(nn.Module):
    def __init__(self, output_dim, embed_dim, num_layers, num_heads, expand_dim, dropout, device, max_length = 30):
        super().__init__()
        self.tok_embedding = nn.Embedding(output_dim, embed_dim)
        #self.pos_embedding = nn.Embedding(max_length, embed_dim)
        self.pos_embedding = nn.Embedding.from_pretrained(self.get_positional_encoding(max_length, embed_dim))
        self.layers = nn.ModuleList([DecoderLayer(embed_dim, num_heads, expand_dim, dropout) for _ in range(num_layers)])
        self.fc_out = nn.Linear(embed_dim, output_dim)
        self.dropout = nn.Dropout(dropout)
        self.scale = torch.sqrt(torch.FloatTensor([embed_dim])).to(device)
        self.device = device
        
    def forward(self, trg, enc_src, trg_mask, src_mask):
        #trg = [batch size, trg len]
        #enc_src = [batch size, src len, embed dim]
        #trg_mask = [batch size, 1, trg len, trg len]
        #src_mask = [batch size, 1, 1, src len]
                
        batch_size = trg.shape[0]
        trg_len = trg.shape[1]
        
        pos = torch.arange(0, trg_len).unsqueeze(0).repeat(batch_size, 1).to(self.device)
                            
        #pos = [batch size, trg len]
            
        trg = self.dropout((self.tok_embedding(trg) * self.scale) + self.pos_embedding(pos))
                
        #trg = [batch size, trg len, embed dim]
        
        for layer in self.layers:
            trg = layer(trg, enc_src, trg_mask, src_mask)
        
        #trg = [batch size, trg len, embed dim]
        
        output = self.fc_out(trg)
        
        #output = [batch size, trg len, output dim]
            
        return output
    
    def get_positional_encoding(self, max_seq_len, embed_dim):
        pos_enc = torch.zeros(max_seq_len, embed_dim)
        position = torch.arange(0, max_seq_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, embed_dim, 2) * (-math.log(10000.0) / embed_dim))

        pos_enc[:, 0::2] = torch.sin(position * div_term)
        pos_enc[:, 1::2] = torch.cos(position * div_term)
        return pos_enc
[((9, 29, 9, 64), 'torch.nn.Embedding', 'nn.Embedding', ({(9, 42, 9, 52): 'output_dim', (9, 54, 9, 63): 'embed_dim'}, {}), '(output_dim, embed_dim)', True, 'import torch.nn as nn\n'), ((13, 22, 13, 54), 'torch.nn.Linear', 'nn.Linear', ({(13, 32, 13, 41): 'embed_dim', (13, 43, 13, 53): 'output_dim'}, {}), '(embed_dim, output_dim)', True, 'import torch.nn as nn\n'), ((14, 23, 14, 42), 'torch.nn.Dropout', 'nn.Dropout', ({(14, 34, 14, 41): 'dropout'}, {}), '(dropout)', True, 'import torch.nn as nn\n'), ((47, 18, 47, 53), 'torch.zeros', 'torch.zeros', ({(47, 30, 47, 41): 'max_seq_len', (47, 43, 47, 52): 'embed_dim'}, {}), '(max_seq_len, embed_dim)', False, 'import torch\n'), ((51, 27, 51, 57), 'torch.sin', 'torch.sin', ({(51, 37, 51, 56): 'position * div_term'}, {}), '(position * div_term)', False, 'import torch\n'), ((52, 27, 52, 57), 'torch.cos', 'torch.cos', ({(52, 37, 52, 56): 'position * div_term'}, {}), '(position * div_term)', False, 'import torch\n'), ((12, 37, 12, 92), 'DecoderLayer.DecoderLayer', 'DecoderLayer', ({(12, 50, 12, 59): 'embed_dim', (12, 61, 12, 70): 'num_heads', (12, 72, 12, 82): 'expand_dim', (12, 84, 12, 91): 'dropout'}, {}), '(embed_dim, num_heads, expand_dim, dropout)', False, 'from DecoderLayer import DecoderLayer\n'), ((48, 19, 48, 47), 'torch.arange', 'torch.arange', ({(48, 32, 48, 33): '0', (48, 35, 48, 46): 'max_seq_len'}, {}), '(0, max_seq_len)', False, 'import torch\n'), ((49, 29, 49, 58), 'torch.arange', 'torch.arange', ({(49, 42, 49, 43): '0', (49, 45, 49, 54): 'embed_dim', (49, 56, 49, 57): '2'}, {}), '(0, embed_dim, 2)', False, 'import torch\n'), ((15, 32, 15, 62), 'torch.FloatTensor', 'torch.FloatTensor', ({(15, 50, 15, 61): '[embed_dim]'}, {}), '([embed_dim])', False, 'import torch\n'), ((49, 63, 49, 80), 'math.log', 'math.log', ({(49, 72, 49, 79): '10000.0'}, {}), '(10000.0)', False, 'import math\n'), ((27, 14, 27, 38), 'torch.arange', 'torch.arange', ({(27, 27, 27, 28): '0', (27, 30, 27, 37): 'trg_len'}, {}), '(0, trg_len)', False, 'import torch\n')]