import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPooling, SequenceClassifierOutput
import os

from .configuration_bbb import BBBConfig

class BBBModelForSequenceClassification(PreTrainedModel):
  config_class = BBBConfig

  def __init__(self, config: BBBConfig):
    super().__init__(config)

    self.config = config

    self.proj_tab = nn.Sequential(
      nn.LayerNorm(config.d_tab),
      nn.Linear(config.d_tab, config.proj_dim),
      nn.ReLU(),
      nn.Dropout(config.dropout)
    )

    self.proj_img = nn.Sequential(
      nn.LayerNorm(config.d_img),
      nn.Linear(config.d_img, config.proj_dim),
      nn.ReLU(),
      nn.Dropout(config.dropout)
    )

    self.proj_txt = nn.Sequential(
      nn.LayerNorm(config.d_txt),
      nn.Linear(config.d_txt, config.proj_dim),
      nn.ReLU(),
      nn.Dropout(config.dropout)
    )

    self.attention_pooling = nn.Sequential(
      nn.Linear(config.proj_dim, config.proj_dim),
      nn.Tanh(),
      nn.Linear(config.proj_dim, 1, bias=False)
    )

    self.classifier = nn.Sequential(
      nn.Linear(config.proj_dim, config.proj_dim), 
      nn.ReLU(),
      nn.Dropout(config.dropout),
      nn.Linear(config.proj_dim, 1)
    )

  def _init_weights(self, module):
    if isinstance(module, nn.Linear):
      module.weight.data.normal_(mean=0.0, std=1.0)
      if module.bias is not None:
        module.bias.data.zero_()
    elif isinstance(module, nn.Embedding):
      module.weight.data.normal_(mean=0.0, std=1.0)
      if module.padding_idx is not None:
        module.weight.data[module.padding_idx].zero_()
    elif isinstance(module, nn.LayerNorm):
      module.bias.data.zero_()
      module.weight.data.fill_(1.0)

  def forward(self,
              tab: torch.Tensor = None,
              img: torch.Tensor = None,
              txt: torch.Tensor = None,
              labels: torch.Tensor = None):

    if tab is None or img is None or txt is None:
      raise ValueError('You have to specify tabular, image, and text features')

    h_tab = self.proj_tab(tab)
    h_img = self.proj_img(img)
    h_txt = self.proj_txt(txt)

    embeddings = torch.stack([h_tab, h_img, h_txt], dim=1)
    scores = self.attention_pooling(embeddings)
    weights = F.softmax(scores, dim=1)
    weighted_embeddings = embeddings * weights
    pooled_embeddings = torch.sum(weighted_embeddings, dim=1)
    logits = self.classifier(pooled_embeddings).squeeze(-1)

    loss = None
    if labels is not None:
      if self.config.task == "regression":
        loss_fct = nn.MSELoss()
        loss = loss_fct(logits.squeeze(-1), labels.squeeze(-1))
      elif self.config.task == "classification":
        loss_fct = nn.BCEWithLogitsLoss()
        loss = loss_fct(logits, labels.float().unsqueeze(-1))

    return SequenceClassifierOutput(
      loss=loss,
      logits=logits,
      hidden_states=pooled_embeddings,
      attentions=weights.squeeze(-1)
    )