from transformers import PretrainedConfig


class FelaGridConfig(PretrainedConfig):
    model_type = "fela_grid_renewable"

    def __init__(
        self,
        track="solar",
        Fin=20,
        L=6,
        D=64,
        modes=3,
        nblk=4,
        nq=99,
        arch="dual",
        tracks=None,
        **kwargs,
    ):
        if isinstance(tracks, dict) and track in tracks:
            dims = tracks[track].get("dims", {})
            Fin = dims.get("Fin", Fin)
            L = dims.get("L", L)
            D = dims.get("D", D)
            modes = dims.get("modes", modes)
            nblk = dims.get("nblk", nblk)
            nq = dims.get("nq", nq)
            arch = dims.get("arch", arch)
        self.track = track
        self.Fin = Fin
        self.L = L
        self.D = D
        self.modes = modes
        self.nblk = nblk
        self.nq = nq
        self.arch = arch
        if tracks is not None:
            self.tracks = tracks
        super().__init__(**kwargs)