import torch
import torch.nn as nn
import torch.nn.functional as F
from loguru import logger


class ModelP2V(torch.nn.Module):
    def __init__(self, n_products, size, batch_size, n_negative_samples):
        super().__init__()
        # trainable variables
        logger.info(f"initialize trainable variables: w_i and w_o")
        self.wi = torch.nn.Embedding(n_products, size, sparse=True)
        with torch.no_grad():
            self.wi.weight.uniform_(-0.025, 0.025)
        self.wo = torch.nn.Embedding(n_products, size, sparse=True)
        with torch.no_grad():
            self.wo.weight.uniform_(-0.025, 0.025)

        self.ones = torch.ones(batch_size)
        self.zeros = torch.zeros(batch_size, n_negative_samples)

    def forward(self, center, context, negative_samples):

        # embed products (center, context, negative_samples)
        wi_center = self.wi(center)
        wo_positive_samples = self.wo(context)
        wo_negative_samples = self.wo(negative_samples)

        # logits
        logits_positive_samples = torch.einsum(
            "ij,ij->i", (wi_center, wo_positive_samples)
        )
        logits_negative_samples = torch.einsum(
            "ik,ijk->ij", (wi_center, wo_negative_samples)
        )

        # loss
        loss_positive_samples = F.binary_cross_entropy_with_logits(
            input=logits_positive_samples,
            target=torch.ones_like(logits_positive_samples),
            reduction="sum",
        )
        loss_negative_samples = F.binary_cross_entropy_with_logits(
            input=logits_negative_samples,
            target=torch.zeros_like(logits_negative_samples),
            reduction="sum",
        )

        n_samples = logits_positive_samples.shape[0] * (
            logits_negative_samples.shape[1] + 1
        )
        return (loss_positive_samples + loss_negative_samples) / n_samples
