from typing import Optional, Tuple
import torch, torch.nn as nn, torch.nn.functional as F

from transformers import (
    PretrainedConfig,
    PreTrainedModel,
    GenerationMixin,
    AutoConfig,
    AutoModelForCausalLM,
)
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions

# 1) --------------------------------------------------------------------------
# CONFIG
# -----------------------------------------------------------------------------


class SuperLinearConfig(PretrainedConfig):
    """
    Configuration for the SuperLinear MoE time–series foundation model.
    Only *model_type* must be unique inside transformers; the rest mirrors
    the __init__ arguments of your original Config object.
    """

    model_type = "super_linear"      

    def __init__(
        self,
        seq_len=512,
        pred_len=96,
        inf_pred_len=96,
        max_horizon=96,
        moe_n_experts=8,
        top_k_experts=5,
        moe =1,
        freq_experts='mean_naive_1/6_1/7_1/8_1/12_1/14_1/16_1/21_1/24_1/28_1/30_1/32_1/36_1/42_1/48_1/52_1/56_1/60_1/72_1/84_1/96_1/120_1/144_1/168_1/180_1/224_1/252_1/288_1/336_1/365_1/504_1/672_1/1008_1/1440_1/2016_1/3600',
        auto_regressive= 1,
        con= 0,
        d_model= 128,
        dropout= 0.0,
        fft_len= 10000,
        freeze_experts= 1,
        ker_len= 50,
        layer_type= "RLinear",
        linear_checkpoints_dir= "checkpoints5",
        linear_checkpoints_path= "/cs/azencot_fsas/MoE/",
        load_linear = 1,
        manual_moe = 0,
        misc_moe = 1,
        mlp_gating = 1,
        model_type= "super_linear",
        moe_temp = 1,
        noisy_gating_std = 0.1,
        noisy_gating_std_decay = 1,
        torch_dtype = "float32",
        transformers_version = "4.40.1",
        use_fft = 1,
        **kwargs,                          # any extra CLI args
    ):
        self.seq_len         = seq_len
        self.moe             = moe
        self.pred_len        = pred_len
        self.inf_pred_len    = inf_pred_len
        self.max_horizon     = max_horizon
        self.auto_regressive = auto_regressive
        self.moe_n_experts   = moe_n_experts
        self.top_k_experts   = top_k_experts
        self.freq_experts    = freq_experts
        self.freeze_experts  = freeze_experts
        self.layer_type      = layer_type
        self.linear_checkpoints_path  = linear_checkpoints_path
        self.linear_checkpoints_dir   = linear_checkpoints_dir
        self.load_linear              = load_linear 
        self.manual_moe              = manual_moe
        self.misc_moe                = misc_moe
        self.noisy_gating_std        = noisy_gating_std
        self.noisy_gating_std_decay  = noisy_gating_std_decay
        self.ker_len                 = ker_len
        self.con                     = con
        self.d_model                 = d_model
        self.mlp_gating              = mlp_gating
        self.moe_temp                = moe_temp
        self.use_fft                 = use_fft
        self.fft_len                 = fft_len
        self.dropout                 = dropout
        super().__init__(**kwargs)