import json
import math

import torch
import torch.nn as nn

# Hyperparameters used to train pico_model.pth, they must match the checkpoint
D_MODEL = 512
NHEAD = 8
DIM_FF = 4 * D_MODEL
NUM_ENCODER_LAYERS = 3
NUM_DECODER_LAYERS = 3
MAX_LEN = 122  # 120 characters + start and end symbols


class Embedding(nn.Module):
    def __init__(self, vocab_size, d_model, dropout=0.1, max_len=MAX_LEN):
        super().__init__()
        self.d_model = d_model
        self.input_embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
        # Fixed sinusoidal positional encoding, not learned
        self.pos_embedding = nn.Embedding.from_pretrained(
            self.pos_encoding(max_len, d_model), freeze=True)
        self.dropout = nn.Dropout(dropout)

    def forward(self, X):
        pos_mat = torch.arange(X.size(dim=-1), device=X.device)
        X = self.input_embedding(X) * math.sqrt(self.d_model)
        X += self.pos_embedding(pos_mat)
        return self.dropout(X)

    def pos_encoding(self, max_len, d_model):
        dividend = torch.arange(max_len).unsqueeze(0).T
        divisor = torch.pow(10000.0, torch.arange(0, d_model, 2) / d_model)
        angles = dividend / divisor
        pe = torch.zeros((max_len, d_model))
        pe[:, 0::2] = torch.sin(angles)
        pe[:, 1::2] = torch.cos(angles)
        return pe


class Translator(nn.Module):
    def __init__(self, vocab_size, d_model=D_MODEL, nhead=NHEAD,
                 num_encoder_layers=NUM_ENCODER_LAYERS,
                 num_decoder_layers=NUM_DECODER_LAYERS,
                 dim_feedforward=DIM_FF, dropout=0.1, max_len=MAX_LEN):
        super().__init__()
        self.embedding = Embedding(vocab_size, d_model, max_len=max_len)
        self.transformer = nn.Transformer(
            d_model, nhead, num_encoder_layers, num_decoder_layers,
            dim_feedforward, dropout, batch_first=True)
        self.fc = nn.Linear(d_model, vocab_size)
        self.fc.weight = self.embedding.input_embedding.weight


def load(model_path, vocab_path):
    with open(vocab_path, encoding='utf-8') as f:
        token2idx = json.load(f)
    idx2token = {idx: tok for tok, idx in token2idx.items()}
    model = Translator(vocab_size=len(token2idx))
    model.load_state_dict(torch.load(model_path, map_location='cpu'))
    model.eval()
    return model, token2idx, idx2token


@torch.no_grad()
def translate(model, token2idx, idx2token, sentence):
    """Greedy decoding, one character at a time"""
    chars = list(sentence.strip())[:MAX_LEN - 2]
    if not chars:
        return ''
    src = torch.LongTensor(
        [[token2idx['<s>']] + [token2idx.get(c, token2idx['<unk>'])
                               for c in chars] + [token2idx['</s>']]])
    memory = model.transformer.encoder(model.embedding(src))
    tgt = torch.LongTensor([[token2idx['<s>']]])
    max_len = min(len(chars) + 20, MAX_LEN)
    for _ in range(max_len - 1):
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(
            tgt.size(dim=1))
        out = model.transformer.decoder(
            model.embedding(tgt), memory, tgt_mask=tgt_mask)
        next_char = torch.argmax(model.fc(out[:, -1, :]), dim=-1).unsqueeze(-1)
        tgt = torch.cat((tgt, next_char), dim=1)
        if next_char.item() == token2idx['</s>']:
            break
    tokens = [idx2token[i] for i in tgt[0, 1:].tolist()]
    if tokens and tokens[-1] == '</s>':
        tokens = tokens[:-1]
    return ''.join(tokens)