import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput

from .configuration_pdm import FelaPdmConfig
from .modeling import FELAPDM, PDMConfig


class FelaPdmModel(PreTrainedModel):
    config_class = FelaPdmConfig
    base_model_prefix = "model"
    main_input_name = "x"

    def __init__(self, config):
        super().__init__(config)
        cfg = PDMConfig(
            in_channels=config.in_channels,
            patch=config.patch,
            n_embd=config.n_embd,
            n_layer=config.n_layer,
            n_head=config.n_head,
            fno_modes=config.fno_modes,
            gla_chunk=config.gla_chunk,
            ffn_hidden=config.ffn_hidden,
            dropout=config.dropout,
            use_gdn=config.use_gdn,
            gdn_every=config.gdn_every,
            n_classes=config.n_classes,
            rul_head=config.rul_head,
            seq_len=config.seq_len,
        )
        self.model = FELAPDM(cfg)
        self.task = config.task
        self.post_init()

    def forward(self, x=None, input_values=None, task=None, **kwargs):
        if x is None:
            x = input_values
        out = self.model(x, task=task or self.task)
        return CausalLMOutput(logits=out)