from transformers import PretrainedConfig


class FelaPdmConfig(PretrainedConfig):
    model_type = "fela-pdm"

    def __init__(
        self,
        variant="cmapss_FD001",
        task="rul",
        in_channels=14,
        patch=1,
        n_embd=64,
        n_layer=4,
        n_head=4,
        fno_modes=32,
        gla_chunk=32,
        ffn_hidden=128,
        dropout=0.0,
        use_gdn=False,
        gdn_every=4,
        n_classes=0,
        rul_head=True,
        seq_len=30,
        default_variant=None,
        variants=None,
        **kwargs,
    ):
        if isinstance(variants, dict):
            name = variant or default_variant
            if name in variants:
                v = variants[name]
                task = v.get("task", task)
                in_channels = v.get("in_channels", in_channels)
                patch = v.get("patch", patch)
                n_embd = v.get("n_embd", n_embd)
                n_layer = v.get("n_layer", n_layer)
                n_head = v.get("n_head", n_head)
                fno_modes = v.get("fno_modes", fno_modes)
                gla_chunk = v.get("gla_chunk", gla_chunk)
                ffn_hidden = v.get("ffn_hidden", ffn_hidden)
                dropout = v.get("dropout", dropout)
                use_gdn = v.get("use_gdn", use_gdn)
                gdn_every = v.get("gdn_every", gdn_every)
                n_classes = v.get("n_classes", n_classes)
                rul_head = v.get("rul_head", rul_head)
                seq_len = v.get("seq_len", seq_len)
        self.variant = variant
        self.task = task
        self.in_channels = in_channels
        self.patch = patch
        self.n_embd = n_embd
        self.n_layer = n_layer
        self.n_head = n_head
        self.fno_modes = fno_modes
        self.gla_chunk = gla_chunk
        self.ffn_hidden = ffn_hidden
        self.dropout = dropout
        self.use_gdn = use_gdn
        self.gdn_every = gdn_every
        self.n_classes = n_classes
        self.rul_head = rul_head
        self.seq_len = seq_len
        if default_variant is not None:
            self.default_variant = default_variant
        if variants is not None:
            self.variants = variants
        super().__init__(**kwargs)