"""
FTT (Feature Tokenizer Transformer) for tabular UK Biobank data.

A DistilBERT-style transformer adapted to tabular inputs: each row is
treated as a "sentence" whose "tokens" are its features. Continuous
features are embedded via a per-feature learned weight vector scaled by
the (min-max normalized) value; categorical features are embedded via a
per-value lookup table. A separate embedding table encodes which feature
slot ("variable") each token corresponds to, functioning like a position
embedding.

This file is loaded automatically via `trust_remote_code=True` through the
`auto_map` entries in config.json.
"""

from typing import Optional, Tuple, Union

from transformers.models.distilbert.modeling_distilbert import (
    Embeddings,
    DistilBertModel,
    Transformer,
    DistilBertForMaskedLM,
    DistilBertForSequenceClassification,
)
from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask_for_sdpa
from transformers import PretrainedConfig
from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput, SequenceClassifierOutput
from transformers.utils import logging
from transformers.activations import get_activation

import torch
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
import torch.nn.functional as F


logger = logging.get_logger(__name__)


class FTTConfig(PretrainedConfig):
    """Config for the FTT tabular transformer (DistilBERT-style hyperparameters)."""

    model_type = "ftt"

    def __init__(
        self,
        vocab_size: int = 512,
        var_weights_size: int = 512,
        var_bias_size: int = 512,
        max_position_embeddings: int = 256,
        sinusoidal_pos_embds: bool = False,
        n_layers: int = 8,
        n_heads: int = 16,
        dim: int = 128,
        hidden_dim: int = 1024,
        dropout: float = 0.1,
        attention_dropout: float = 0.1,
        activation: str = "gelu",
        initializer_range: float = 0.02,
        qa_dropout: float = 0.1,
        seq_classif_dropout: float = 0.2,
        tie_weights_: bool = True,
        pad_token_id: int = 3,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.var_weights_size = var_weights_size
        self.var_bias_size = var_bias_size
        self.max_position_embeddings = max_position_embeddings
        self.sinusoidal_pos_embds = sinusoidal_pos_embds
        self.n_layers = n_layers
        self.num_hidden_layers = n_layers
        self.n_heads = n_heads
        self.dim = dim
        self.hidden_dim = hidden_dim
        self.dropout = dropout
        self.attention_dropout = attention_dropout
        self.activation = activation
        self.initializer_range = initializer_range
        self.qa_dropout = qa_dropout
        self.seq_classif_dropout = seq_classif_dropout
        self.tie_weights_ = tie_weights_
        super().__init__(pad_token_id=pad_token_id, **kwargs)


class TableEmbeddings(Embeddings):
    def __init__(self, config: PretrainedConfig):
        super().__init__(config)
        self.variable_weights_embeddings = nn.Embedding(config.var_weights_size, config.dim, padding_idx=config.pad_token_id)
        self.variable_bias_embeddings = nn.Embedding(config.max_position_embeddings, config.dim, padding_idx=config.pad_token_id)
        if config.sinusoidal_pos_embds:
            raise NotImplementedError("We don't use sinusoidal position embeddings in gene embedding")

        self.LayerNorm = nn.LayerNorm(config.dim, eps=1e-12)
        self.dropout = nn.Dropout(config.dropout)
        self.config = config

    def forward(
        self,
        num_input_val: Optional[torch.Tensor] = None,
        num_input_ids: Optional[torch.Tensor] = None,
        cat_input_ids: Optional[torch.Tensor] = None,
        num_variable_ids: Optional[torch.Tensor] = None,
        cat_variable_ids: Optional[torch.Tensor] = None,
        input_embeds: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if num_input_ids is None and cat_input_ids is None:
            raise ValueError("any num_input_ids or cat_input_ids is required")

        if num_input_ids is not None:
            num_embed = self.variable_weights_embeddings(num_input_ids)
            num_embed = torch.einsum("bl,bld->bld", num_input_val, num_embed)
        if cat_input_ids is not None:
            cat_embed = self.variable_weights_embeddings(cat_input_ids)

        if num_input_val is None and cat_input_ids is not None:
            weights_embeddings = cat_embed
            bias_embeddings = self.variable_bias_embeddings(cat_variable_ids)
        elif num_input_val is not None and cat_input_ids is None:
            weights_embeddings = num_embed
            bias_embeddings = self.variable_bias_embeddings(num_variable_ids)
        else:
            weights_embeddings = torch.cat([num_embed, cat_embed], dim=1)
            bias_embeddings = self.variable_bias_embeddings(torch.cat([num_variable_ids, cat_variable_ids], dim=1))

        embeddings = weights_embeddings + bias_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)

        return embeddings


class FTTModel(DistilBertModel):
    config_class = FTTConfig

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

        self.embeddings = TableEmbeddings(config)
        self.transformer = Transformer(config)
        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
        self.post_init()

    def forward(
        self,
        num_input_val: Optional[torch.Tensor] = None,
        num_input_ids: Optional[torch.Tensor] = None,
        cat_input_ids: Optional[torch.Tensor] = None,
        num_variable_ids: Optional[torch.Tensor] = None,
        cat_variable_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if inputs_embeds is not None:
            raise ValueError("inputs_embeds is not supported")
        elif num_input_val is None and cat_input_ids is None:
            raise ValueError("any num_input_ids or cat_input_ids is required")

        if num_input_val is not None and cat_input_ids is not None:
            input_shape = torch.tensor(num_input_val.size())
            input_shape[1] = input_shape[1] + cat_input_ids.size(1)

        embeddings = self.embeddings(
            num_input_val,
            num_input_ids,
            cat_input_ids,
            num_variable_ids,
            cat_variable_ids,
        )

        device = embeddings.device
        head_mask_is_none = head_mask is None
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        if self._use_flash_attention_2:
            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
        else:
            if attention_mask is None:
                attention_mask = torch.ones(input_shape, device=device)

            if self._use_sdpa and head_mask_is_none and not output_attentions:
                attention_mask = _prepare_4d_attention_mask_for_sdpa(
                    attention_mask, embeddings.dtype, tgt_len=input_shape[1]
                )

        return self.transformer(
            x=embeddings,
            attn_mask=attention_mask,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )


class IgnoreMaskingMSELoss(nn.Module):
    def __init__(self, ignore_index=-100):
        super(IgnoreMaskingMSELoss, self).__init__()
        self.ignore_index = ignore_index

    def forward(self, prediction, target):
        mask = target != self.ignore_index
        filtered_pred = prediction[mask]
        filtered_target = target[mask]
        return F.mse_loss(filtered_pred, filtered_target)


class FTTForMaskedLM(DistilBertForMaskedLM):
    config_class = FTTConfig
    _tied_weights_keys = ["vocab_projector.weight"]

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

        self.activation = get_activation(config.activation)

        self.distilbert = FTTModel(config)
        self.vocab_transform = nn.Linear(config.dim, config.dim)
        self.vocab_layer_norm = nn.LayerNorm(config.dim, eps=1e-12)
        self.vocab_projector = nn.Linear(config.dim, config.vocab_size)
        self.num_projector = nn.Linear(config.dim, 1)

        self.post_init()

        self.mlm_loss_fct = nn.CrossEntropyLoss()
        self.mlm_loss_mse = IgnoreMaskingMSELoss()

    def forward(
        self,
        num_input_val: Optional[torch.Tensor] = None,
        num_input_ids: Optional[torch.Tensor] = None,
        cat_input_ids: Optional[torch.Tensor] = None,
        num_variable_ids: Optional[torch.Tensor] = None,
        cat_variable_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[MaskedLMOutput, Tuple[torch.Tensor, ...]]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        dlbrt_output = self.distilbert(
            num_input_val=num_input_val,
            num_input_ids=num_input_ids,
            cat_input_ids=cat_input_ids,
            num_variable_ids=num_variable_ids,
            cat_variable_ids=cat_variable_ids,
            attention_mask=attention_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = dlbrt_output[0]
        prediction_logits = self.vocab_transform(hidden_states)
        prediction_logits = self.activation(prediction_logits)
        prediction_logits = self.vocab_layer_norm(prediction_logits)

        continuous_prediction_logits = self.num_projector(prediction_logits[:, : num_variable_ids.size(1), :])
        continuous_prediction_logits = continuous_prediction_logits.squeeze(-1)

        categorical_prediction_logits = self.vocab_projector(prediction_logits[:, num_variable_ids.size(1) :, :])

        mlm_loss = None
        if labels is not None:
            continuous_labels = labels[:, : num_variable_ids.size(1)]
            continuous_mlm_loss = self.mlm_loss_mse(continuous_prediction_logits, continuous_labels)
            categorical_labels = labels[:, num_variable_ids.size(1) :].type(torch.LongTensor).to(device=self.distilbert.device)
            categorical_mlm_loss = self.mlm_loss_fct(
                categorical_prediction_logits.view(-1, categorical_prediction_logits.size(-1)),
                categorical_labels.reshape(-1),
            )
            mlm_loss = continuous_mlm_loss + categorical_mlm_loss
        else:
            return {
                "continuous_prediction_logits": continuous_prediction_logits,
                "categorical_prediction_logits": categorical_prediction_logits,
            }

        return MaskedLMOutput(
            loss=mlm_loss,
            logits=prediction_logits,
            hidden_states=dlbrt_output.hidden_states,
            attentions=dlbrt_output.attentions,
        )


class FTTForSequenceClassification(DistilBertForSequenceClassification):
    config_class = FTTConfig

    def __init__(self, config: PretrainedConfig):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config

        self.distilbert = FTTModel(config)
        self.pre_classifier = nn.Linear(config.dim, config.dim)
        self.classifier = nn.Linear(config.dim, config.num_labels)
        self.dropout = nn.Dropout(config.seq_classif_dropout)
        self.sigmoid = nn.Sigmoid()

        self.post_init()

    def forward(
        self,
        num_input_val: Optional[torch.Tensor] = None,
        num_input_ids: Optional[torch.Tensor] = None,
        cat_input_ids: Optional[torch.Tensor] = None,
        num_variable_ids: Optional[torch.Tensor] = None,
        cat_variable_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[SequenceClassifierOutput, Tuple[torch.Tensor, ...]]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        distilbert_output = self.distilbert(
            num_input_val=num_input_val,
            num_input_ids=num_input_ids,
            cat_input_ids=cat_input_ids,
            num_variable_ids=num_variable_ids,
            cat_variable_ids=cat_variable_ids,
            attention_mask=attention_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_state = distilbert_output[0]
        pooled_output = hidden_state[:, 0]
        pooled_output = self.pre_classifier(pooled_output)
        pooled_output = nn.ReLU()(pooled_output)
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        loss = None
        if labels is not None:
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                labels = labels.float()
                loss_fct = MSELoss()
                if self.num_labels == 1:
                    loss = loss_fct(logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                labels = labels.float()
                loss_fct = BCEWithLogitsLoss()
                loss = loss_fct(logits.squeeze(), labels)

        if not return_dict:
            output = (logits,) + distilbert_output[1:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=distilbert_output.hidden_states,
            attentions=distilbert_output.attentions,
        )