import math, torch
import torch.nn as nn


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=256):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]


class MTModel(nn.Module):
    def __init__(self, vocab, d_model=512, nhead=8, layers=6, dim_ff=2048,
                 dropout=0.1, max_len=256):
        super().__init__()
        self.d_model = d_model
        self.embed = nn.Embedding(vocab, d_model, padding_idx=0)
        self.pos = PositionalEncoding(d_model, max_len)
        self.transformer = nn.Transformer(
            d_model=d_model, nhead=nhead,
            num_encoder_layers=layers, num_decoder_layers=layers,
            dim_feedforward=dim_ff, dropout=dropout,
            batch_first=True, norm_first=True)
        self.out = nn.Linear(d_model, vocab)
        self.out.weight = self.embed.weight
        nn.init.normal_(self.embed.weight, std=d_model ** -0.5)
        nn.init.zeros_(self.out.bias)
        with torch.no_grad():
            self.embed.weight[0].fill_(0)

    def forward(self, src, tgt_in):
        sp, tp = src == 0, tgt_in == 0
        causal = nn.Transformer.generate_square_subsequent_mask(tgt_in.size(1), device=src.device)
        s = self.pos(self.embed(src) * math.sqrt(self.d_model))
        t = self.pos(self.embed(tgt_in) * math.sqrt(self.d_model))
        h = self.transformer(s, t, tgt_mask=causal,
                             src_key_padding_mask=sp, tgt_key_padding_mask=tp,
                             memory_key_padding_mask=sp)
        return self.out(h)