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, 'src.layers.attention.MultiHeadedAttention', 'MultiHeadedAttention', False, 'from src.layers.attention import MultiHeadedAttention\n'), (53, 'src.layers.layers.PositionwiseFeedForward', 'PositionwiseFeedForward', False, 'from src.layers.layers import Encoder, EncoderLayer, Decoder, DecoderLayer, PositionwiseFeedForward\n'), (54, 'src.layers.preprocessing.PositionalEncoding', 'PositionalEncoding', False, 'from src.layers.preprocessing import Embeddings, PositionalEncoding\n'), (43, 'torch.nn.Linear', 'nn.Linear', True, 'import torch.nn as nn\n'), (58, 'src.layers.preprocessing.Embeddings', 'Embeddings', False, 'from src.layers.preprocessing import Embeddings, PositionalEncoding\n'), (59, 'src.layers.preprocessing.Embeddings', 'Embeddings', False, 'from src.layers.preprocessing import Embeddings, PositionalEncoding\n'), (67, 'torch.nn.init.xavier_uniform', 'nn.init.xavier_uniform', True, 'import torch.nn as nn\n')]