from xmlrpc.server import MultiPathXMLRPCServer
import torch.nn as nn
import torch.nn.functional as F
import copy

from src.layers.layers import Encoder, EncoderLayer, Decoder, DecoderLayer, PositionwiseFeedForward
from src.layers.preprocessing import Embeddings, PositionalEncoding
from src.layers.attention import MultiHeadedAttention

### Generic EncoderDecoder structure:

class EncoderDecoder(nn.Module):
    """
    A standard Encoder-Decoder architecture. Base for this and many 
    other models.
    """
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator
        
    def forward(self, src, tgt, src_mask, tgt_mask):
        "Take in and process masked src and target sequences."
        encoded_src = self.encode(src, src_mask)
        return self.decode(encoded_src, src_mask, tgt, tgt_mask)
    
    def encode(self, src, src_mask):
        embedded_src = self.src_embed(src)
        return self.encoder(embedded_src, src_mask)
    
    def decode(self, memory, src_mask, tgt, tgt_mask):
        embedded_tgt = self.tgt_embed(tgt)
        return self.decoder(embedded_tgt, memory, src_mask, tgt_mask)


class Generator(nn.Module):
    "Define standard linear + softmax generation step."
    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        return F.log_softmax(self.proj(x), dim=-1)


def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1, alpha=0.5):
    "Helper: Construct a model from hyperparameters."
    c = copy.deepcopy
    attn = MultiHeadedAttention(h, d_model, alpha=alpha)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    position = PositionalEncoding(d_model, dropout)
    model = EncoderDecoder(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
        Generator(d_model, tgt_vocab)
    )

    # This was important from their code.
    # Initialize parameters with Glorot / fan_avg.
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform(p)
    return model


if __name__ == '__main__':
    # Small example model
    tmp_model = make_model(10, 10, 2)
    print(tmp_model)

[((52, 11, 52, 56), 'src.layers.attention.MultiHeadedAttention', 'MultiHeadedAttention', (), '', False, 'from src.layers.attention import MultiHeadedAttention\n'), ((53, 9, 53, 56), 'src.layers.layers.PositionwiseFeedForward', 'PositionwiseFeedForward', ({(53, 33, 53, 40): 'd_model', (53, 42, 53, 46): 'd_ff', (53, 48, 53, 55): 'dropout'}, {}), '(d_model, d_ff, dropout)', False, 'from src.layers.layers import Encoder, EncoderLayer, Decoder, DecoderLayer, PositionwiseFeedForward\n'), ((54, 15, 54, 51), 'src.layers.preprocessing.PositionalEncoding', 'PositionalEncoding', ({(54, 34, 54, 41): 'd_model', (54, 43, 54, 50): 'dropout'}, {}), '(d_model, dropout)', False, 'from src.layers.preprocessing import Embeddings, PositionalEncoding\n'), ((43, 20, 43, 45), 'torch.nn.Linear', 'nn.Linear', ({(43, 30, 43, 37): 'd_model', (43, 39, 43, 44): 'vocab'}, {}), '(d_model, vocab)', True, 'import torch.nn as nn\n'), ((58, 22, 58, 52), 'src.layers.preprocessing.Embeddings', 'Embeddings', ({(58, 33, 58, 40): 'd_model', (58, 42, 58, 51): 'src_vocab'}, {}), '(d_model, src_vocab)', False, 'from src.layers.preprocessing import Embeddings, PositionalEncoding\n'), ((59, 22, 59, 52), 'src.layers.preprocessing.Embeddings', 'Embeddings', ({(59, 33, 59, 40): 'd_model', (59, 42, 59, 51): 'tgt_vocab'}, {}), '(d_model, tgt_vocab)', False, 'from src.layers.preprocessing import Embeddings, PositionalEncoding\n'), ((67, 12, 67, 37), 'torch.nn.init.xavier_uniform', 'nn.init.xavier_uniform', ({(67, 35, 67, 36): 'p'}, {}), '(p)', True, 'import torch.nn as nn\n')]