import torch
import torch.nn as nn

bce, sigmoid, softmax = nn.BCELoss(), nn.Sigmoid(), nn.Softmax(dim=1)

# TODO: Document how this shit works because wtf
class ASM2VEC(nn.Module):
    def __init__(self, vocab_size, function_size, embedding_size):
        super(ASM2VEC, self).__init__()
        # Dictionary of the token embeddings: v_t
        self.embeddings = nn.Embedding(
            vocab_size,
            embedding_size,
            _weight=(torch.rand(vocab_size, embedding_size) - 0.5) / embedding_size / 2,
        )
        # Dictionary of the function embeddings, \theta_f_s
        self.embeddings_f = nn.Embedding(
            function_size,
            2 * embedding_size,
            _weight=(torch.rand(function_size, 2 * embedding_size) - 0.5)
            / embedding_size
            / 2,
        )
        # Dictionary of outputs: Transposed v'_t
        self.embeddings_r = nn.Embedding(
            vocab_size,
            2 * embedding_size,
            _weight=torch.zeros(vocab_size, 2 * embedding_size),
        )
        # Where the old embeddings are stored once the training step is done
        self.old_embeddings_f = None

    def init_estimation_mode(self, function_size_new):
        device = self.embeddings.weight.device
        embedding_size = self.embeddings.embedding_dim

        embedding_weights = self.embeddings.weight
        self.embeddings = nn.Embedding.from_pretrained(embedding_weights)

        output_weights = self.embeddings_r.weight
        self.embeddings_r = nn.Embedding.from_pretrained(output_weights)

        self.old_embeddings_f = self.embeddings_f
        self.embeddings_f = nn.Embedding(
            function_size_new,
            2 * embedding_size,
            _weight=(
                (torch.rand(function_size_new, 2 * embedding_size) - 0.5)
                / embedding_size
                / 2
            ).to(device),
        )

    def v(self, inp):
        # Retrieve the embeddings for all the context tokens
        e = self.embeddings(inp[:, 1:])
        # Retrieve the embedding for the function, \theta_f_s
        v_f = self.embeddings_f(inp[:, 0])
        # Calculate CT(in_(j-1))
        v_prev = torch.cat([e[:, 0], (e[:, 1] + e[:, 2]) / 2], dim=1)
        # Calculate CT(in_(j+1))
        v_next = torch.cat([e[:, 3], (e[:, 4] + e[:, 5]) / 2], dim=1)
        # delta(in_j, f_s) = 1/3 * (\theta_f_s + CT(in_(j-1)) + CT(in_(j+1)))
        v = ((v_f + v_prev + v_next) / 3).unsqueeze(2)
        return v

    def forward(self, inp, pos, neg):
        device, batch_size = inp.device, inp.shape[0]
        v = self.v(inp)
        # negative sampling loss
        pred = torch.bmm(self.embeddings_r(torch.cat([pos, neg], dim=1)), v).squeeze()
        label = torch.cat(
            [torch.ones(batch_size, 1), torch.zeros(batch_size, neg.shape[1])], dim=1
        ).to(device)
        return bce(sigmoid(pred), label)

    def predict(self, inp, pos):
        device, batch_size = inp.device, inp.shape[0]
        v = self.v(inp)
        probs = torch.bmm(
            self.embeddings_r(
                torch.arange(self.embeddings_r.num_embeddings)
                .repeat(batch_size, 1)
                .to(device)
            ),
            v,
        ).squeeze(dim=2)
        return softmax(probs)