import torch
from transformers import PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput

from .configuration_felatab import FelaTabConfig as _HFConfig
from .modeling import FelaTab
from .modeling import FelaTabConfig as _CoreConfig

_FIELDS = set(_CoreConfig.__dataclass_fields__.keys())


class FelaTabModel(PreTrainedModel):
    config_class = _HFConfig
    base_model_prefix = "model"

    def __init__(self, config):
        super().__init__(config)
        cfg = _CoreConfig(
            **{k: getattr(config, k) for k in _FIELDS if hasattr(config, k)}
        )
        self.model = FelaTab(cfg)
        self.post_init()

    def forward(self, X=None, y=None, n_support=None, task_type_id=None, **kwargs):
        if task_type_id is None:
            task_type_id = torch.tensor(0, device=X.device)
        logits = self.model(X, y, n_support, task_type_id)
        return SequenceClassifierOutput(logits=logits)