import math
import os
from dataclasses import dataclass
from typing import Optional

from huggingface_hub import hf_hub_download
import lm_eval as evaluator
import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file
from torchtune.modules import RotaryPositionalEmbeddings
from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    PreTrainedModel,
    PretrainedConfig,
)
from transformers.modeling_outputs import CausalLMOutput

try:
    from flashfftconv import FlashFFTConv

    flash_fft_available = True
except ImportError as e:
    print(f"Unable to import FlashFFTConv: {e}. Falling back to PyTorch implementation.")
    flash_fft_available = False

try:
    from flash_attn import flash_attn_func
except ImportError as e:
    print(f"Unable to import Triton-based flash attention: {e}. No alternative currently available.")

os.environ["HF_ALLOW_CODE_EVAL"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

loss_fn = nn.CrossEntropyLoss()


def nearest_power_of_two(n: int, round_up: bool = False) -> int:
    if n <= 1:
        return 1
    return 1 << ((n - 1).bit_length() if round_up else (n).bit_length() - 1)


def find_multiple(n: int, k: int) -> int:
    if n % k == 0:
        return n
    return n + k - (n % k)


def get_hankel(seq_len: int, use_hankel_L: bool = False) -> torch.Tensor:
    entries = torch.arange(1, seq_len + 1, dtype=torch.float64)
    i_plus_j = entries.reshape(-1, 1) + entries.reshape(1, -1)

    if use_hankel_L:
        sgn = (-1.0) ** (i_plus_j - 2.0) + 1.0
        denom = (i_plus_j + 3.0) * (i_plus_j - 1.0) * (i_plus_j + 1.0)
        Z = sgn * (8.0 / denom)
    elif not use_hankel_L:
        Z = 2.0 / (i_plus_j**3 - i_plus_j)
    else:
        raise ValueError("use_hankel_L must be a boolean")

    return Z


def get_spectral_filters(
    seq_len: int,
    K: int,
    use_hankel_L: bool = False,
    device: torch.device = None,
    dtype: torch.dtype = torch.float64,
) -> torch.Tensor:
    Z = get_hankel(seq_len, use_hankel_L).to(device)
    sigma, phi = torch.linalg.eigh(Z)
    sigma_k, phi_k = sigma[-K:], phi[:, -K:]
    phi_k *= sigma_k**0.25
    return phi_k.to(device=device, dtype=dtype)


class BaseConfigForCausalLM(PretrainedConfig):
    """Base PretrainedConfig class to be decorated with dataclass"""

    model_type = "base_model"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)


@dataclass
class FlashSTUConfig(BaseConfigForCausalLM):
    model_type = "FlashSTU"

    # Define fields with defaults (as before)
    bsz: int = 1
    dim: int = 1024
    r: int = 1024
    num_heads: int = 12
    num_local_heads: Optional[int] = -1
    num_layers: int = 12
    seq_len: int = 4096
    n: int = 8191
    window_size: int = 2048
    vocab_size: int = 200064
    inter_dim: Optional[int] = 3072
    mlp_scale: Optional[float] = 12.0
    weight_tying: Optional[bool] = True
    bias: Optional[bool] = False
    rope_theta: Optional[float] = 10000.0
    softcap: Optional[float] = 50.0
    num_eigh: Optional[int] = 24
    use_hankel_L: Optional[bool] = False
    use_flash_fft: Optional[bool] = True
    use_tensordot: Optional[bool] = True
    use_attn: Optional[bool] = True
    use_alibi: Optional[bool] = False
    torch_dtype: torch.dtype = torch.bfloat16
    device: torch.device = None

    # Explicit __init__ to handle **kwargs for PretrainedConfig compatibility
    def __init__(
        self,
        bsz: int = 1,
        dim: int = 1024,
        r: int = 1024,
        num_heads: int = 12,
        num_local_heads: Optional[int] = -1,
        num_layers: int = 12,
        seq_len: int = 4096,
        n: int = 8191,
        window_size: int = 2048,
        vocab_size: int = 200064,
        inter_dim: Optional[int] = 3072,
        mlp_scale: Optional[float] = 12.0,
        weight_tying: Optional[bool] = True,
        bias: Optional[bool] = False,
        rope_theta: Optional[float] = 10000.0,
        softcap: Optional[float] = 50.0,
        num_eigh: Optional[int] = 24,
        use_hankel_L: Optional[bool] = False,
        use_flash_fft: Optional[bool] = True,
        use_tensordot: Optional[bool] = True,
        use_attn: Optional[bool] = True,
        use_alibi: Optional[bool] = False,
        torch_dtype: torch.dtype = torch.bfloat16,
        device: torch.device = None,
        **kwargs,  # Catch extra arguments like model_type
    ):
        super().__init__(**kwargs)  # Pass kwargs to parent __init__

        # Assign fields from arguments
        self.bsz = bsz
        self.dim = dim
        self.r = r
        self.num_heads = num_heads
        self.num_local_heads = num_local_heads
        self.num_layers = num_layers
        self.seq_len = seq_len
        self.n = n
        self.window_size = window_size
        self.vocab_size = vocab_size
        self.inter_dim = inter_dim
        self.mlp_scale = mlp_scale
        self.weight_tying = weight_tying
        self.bias = bias
        self.rope_theta = rope_theta
        self.softcap = softcap
        self.num_eigh = num_eigh
        self.use_hankel_L = use_hankel_L
        self.use_flash_fft = use_flash_fft
        self.use_tensordot = use_tensordot
        self.use_attn = use_attn
        self.use_alibi = use_alibi
        self.torch_dtype = torch_dtype
        self.device = device

        # Explicitly call __post_init__ if defined and needed
        self.__post_init__()

    def __post_init__(self):
        # Ensure torch_dtype is a torch.dtype object, not a string
        if isinstance(self.torch_dtype, str):
            try:
                self.torch_dtype = getattr(torch, self.torch_dtype)
            except AttributeError:
                raise ValueError(f"Invalid torch_dtype string: {self.torch_dtype}")

        if self.num_local_heads == -1:
            self.num_local_heads = self.num_heads
        if self.inter_dim is None:
            hidden_dim = self.mlp_scale * self.dim
            num_hidden = int(2 * hidden_dim / 3)
            self.inter_dim = find_multiple(num_hidden, 256)
        self.head_dim = self.dim // self.num_heads

    @classmethod
    def from_name(cls, name: str):
        # presets = {
        #     "tiny": dict(dim=128, num_heads=4, num_layers=2, vocab_size=10000),
        #     "small": dict(dim=256, num_heads=8, num_layers=4, vocab_size=20000),
        #     "gpt2-small": dict(dim=768, num_heads=12, num_layers=12, vocab_size=50257),
        #     # add more as needed
        # }
        # if name not in presets:
        #     raise ValueError(f"Unknown model config name: {name}")

        # return cls(**presets[name])
        print("Not yet implemented")
        pass


class MLP(nn.Module):
    def __init__(self, config: FlashSTUConfig) -> None:
        super().__init__()
        self.w1 = nn.Linear(config.dim, config.inter_dim)
        self.w2 = nn.Linear(config.inter_dim, config.dim)
        self.w2.SCALE_INIT = 1

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(F.gelu(self.w1(x), approximate="tanh"))


class SlidingWindowAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.wq = nn.Linear(config.dim, config.dim)
        self.wk = nn.Linear(config.dim, config.dim)
        self.wv = nn.Linear(config.dim, config.dim)
        self.wo = nn.Linear(config.dim, config.dim)
        self.wo.SCALE_INIT = 1

        self.dim = config.dim
        self.head_dim = config.head_dim
        self.num_heads = config.num_heads
        self.num_local_heads = config.num_local_heads
        self.window_size = config.window_size
        self.softcap = config.softcap

        self.alibi_slopes = self._get_alibi_slopes(self.num_heads) if config.use_alibi else None
        self.rotary_emb = RotaryPositionalEmbeddings(
            dim=self.head_dim,
            max_seq_len=config.seq_len,
            base=config.rope_theta,
        )

    def forward(self, x):
        bsz, seq_len, dim = x.shape

        q, k, v = self.wq(x), self.wk(x), self.wv(x)
        q = q.view(bsz, seq_len, self.num_heads, self.head_dim)
        k = k.view(bsz, seq_len, self.num_local_heads, self.head_dim)
        v = v.view(bsz, seq_len, self.num_local_heads, self.head_dim)

        if self.alibi_slopes is None:
            q, k = self.rotary_emb(q), self.rotary_emb(k)

        y = flash_attn_func(
            q=q,
            k=k,
            v=v,
            causal=True,
            window_size=(self.window_size, 0),
            # softcap=self.softcap,
            alibi_slopes=self.alibi_slopes,
        )

        out = y.reshape(bsz, seq_len, -1)
        out = self.wo(out)

        return out

    def _generate_slopes(self, n: int):
        start = 2 ** (-(2 ** -(math.log2(n) - 3)))
        return [start * (start**i) for i in range(n)]

    def _get_alibi_slopes(self, num_heads: int, interpolation_factor: float = 0.25):
        # If n_heads is a power of 2, generate slopes directly
        if math.log2(num_heads).is_integer():
            slopes = self._generate_slopes(num_heads)
        else:
            # Get slopes for the nearest power of two
            n = nearest_power_of_two(num_heads, round_up=False)
            slopes_power_of_two = self._generate_slopes(n)

            # Generate extra slopes
            extra_slopes = self._generate_slopes(2 * n)
            extra_slopes_trunc = extra_slopes[0::2][: num_heads - n]
            slopes = slopes_power_of_two + extra_slopes_trunc
        slopes = torch.tensor(slopes, device=torch.device("cuda"))  # FA ALiBi must be on CUDA
        slopes = slopes * interpolation_factor  # https://arxiv.org/pdf/2310.13017
        return slopes


class STU(nn.Module):
    def __init__(self, config):
        super().__init__()

        # Set at top-level post- model init
        self.stu_filters = None
        self.stu_filters_fft = None  # TODO: Optimization: Precompute FFT of filters

        self.n = config.n
        self.num_eigh = config.num_eigh
        self.d_in = config.dim
        self.d_out = config.dim
        self.r = config.r
        self.use_hankel_L = config.use_hankel_L
        self.use_tensordot = config.use_tensordot
        self.flash_fft = (
            FlashFFTConv(self.n, dtype=torch.bfloat16) if config.use_flash_fft and flash_fft_available else None
        )

        # TODO: Add dimensionality reduction `r` here.
        if self.use_tensordot:
            self.M_inputs = nn.Parameter(torch.zeros(self.d_in, self.d_out))
            self.M_filters = nn.Parameter(torch.zeros(self.num_eigh, self.d_in))
        else:
            self.M_phi_plus = nn.Parameter(torch.zeros(self.num_eigh, self.d_in, self.d_out))
            if not self.use_hankel_L:
                self.M_phi_minus = nn.Parameter(torch.zeros(self.num_eigh, self.d_in, self.d_out))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, L, D = x.shape

        if self.use_tensordot:
            # Contract inputs and filters over (K, D) dims first, then convolve
            x_proj = x @ self.M_inputs
            phi_proj = self.stu_filters @ self.M_filters
            if self.flash_fft:
                spectral_plus, spectral_minus = self.flash_conv(x_proj, phi_proj, self.flash_fft, self.use_tensordot)
            else:
                spectral_plus, spectral_minus = self.conv(x_proj, phi_proj, self.n, self.use_tensordot)

        else:
            # Convolve inputs and filters first, then contract over (K, D) dims
            if self.flash_fft:
                U_plus, U_minus = self.flash_conv(x, self.stu_filters, self.flash_fft, self.use_tensordot)
            else:
                U_plus, U_minus = self.conv(x, self.stu_filters, self.n, self.use_tensordot)

            B, L, K, D = U_plus.shape
            spectral_plus = U_plus.reshape(B, L, K * D) @ self.M_phi_plus.reshape(K * D, self.d_out)
            if not self.use_hankel_L:
                spectral_minus = U_minus.reshape(B, L, K * D) @ self.M_phi_minus.reshape(K * D, self.d_out)

        out = spectral_plus if self.use_hankel_L else spectral_plus + spectral_minus
        return out

    def conv(
        self, u: torch.Tensor, v: torch.Tensor, n: int, use_tensordot: bool = True
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Performs convolution via FFT with causal alignment using a negative featurization.

        The input tensor u is modulated by an alternating sign tensor (sgn) that multiplies every other
        time step by -1. This "negative featurization" modulates the phase so that in this implementation
        the correct causal output is obtained by simply slicing the first L elements (i.e. [:seq_len]).
        Note: Using a conventional slice [seq_len-1:2*seq_len-1] would yield a flipped alignment, resulting in leakage.

        Args:
            u: Input tensor of shape (bsz, seq_len, d_in).
            v: Kernel tensor; expected shape is (seq_len, d_out) if use_tensordot is True.
            n: FFT length (typically set to 2*seq_len - 1 for linear convolution with implicit right zero-padding).
            use_tensordot: Boolean flag to control kernel reshaping.

        Returns:
            A tuple (U_plus, U_minus) where:
            - U_plus is the primary convolution output.
            - U_minus is the secondary output, corrected by the sign tensor.
        """
        bsz, seq_len, d_in = u.shape

        sgn = torch.full((1, seq_len, 1), 1, device=u.device)
        sgn[:, 1::2] *= -1  # Apply negative featurization: multiply every other element by -1.

        if use_tensordot:
            _, d_out = v.shape
            v = v.view(1, -1, d_out, 1).to(torch.float32).contiguous()
        else:
            _, K = v.shape
            sgn = sgn.unsqueeze(-1)
            v = v.view(1, -1, K, 1, 1).to(torch.float32).contiguous()  # (bsz, seq_len, K, d_in, stack)
            u = u.view(bsz, -1, 1, d_in).expand(bsz, -1, K, d_in)

        # Cast kernel to float32 for FFT
        v_fft = torch.fft.rfft(v.to(torch.float32), n=n, dim=1)

        U = torch.stack([u, u * sgn], dim=-1).to(torch.float32).contiguous()
        # Cast input stack to float32 for FFT
        U_fft = torch.fft.rfft(U.to(torch.float32), n=n, dim=1)

        # Slicing the first seq_len outputs yields the proper causal convolution given the negative modulation.
        # Perform convolution in float32 and cast back
        U_conv = torch.fft.irfft(v_fft * U_fft, n=n, dim=1)[:, :seq_len].to(u.dtype)
        U_plus, U_minus = torch.unbind(U_conv, dim=-1)
        U_minus = U_minus * sgn

        return U_plus.type_as(u), U_minus.type_as(u)

    def flash_conv(
        self,
        u: torch.Tensor,
        v: torch.Tensor,
        flash_fft: FlashFFTConv,
        use_tensordot: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Flash FFT convolution.

        Args:
            u (torch.Tensor): Input tensor of shape `(B, L, d_in)`, where:
                - `B` is the batch size,
                - `L` is the sequence length,
                - `d_in` is the input dimension.
            v (torch.Tensor): Filter tensor of shape `(K, d_in)`, where:
                - `K` is the number of filters,
                - `d_in` is the input dimension.
            flash_fft (FlashFFTConv): An instance of the FlashFFTConv module, used to perform the convolution.
            use_tensordot (bool, optional): If `True`, performs the tensordot approximation (default is `True`).

        Returns:
            tuple[torch.Tensor, torch.Tensor]: A tuple `(U_plus, U_minus)`:
                - `U_plus`: Convolved output tensor with positive eigenvalues.
                - Shape depends on `use_tensordot`:
                    - If `use_tensordot=True`: `(B, L, d_in)`
                    - If `use_tensordot=False`: `(B, L, K, d_in)`
                - `U_minus`: Convolved output tensor with negative eigenvalues.
                - Shape depends on `use_tensordot`:
                    - If `use_tensordot=True`: `(B, L, d_in)`
                    - If `use_tensordot=False`: `(B, L, K, d_in)`

        Raises:
            ValueError: If the input tensor shapes do not conform to the expected dimensions.

        Example:
            >>> u = torch.randn(4, 16, 32)  # (B, L, d_in)
            >>> v = torch.randn(8, 32)      # (K, d_in)
            >>> flash_fft = FlashFFTConv(n=16, dtype=torch.float32)
            >>> U_plus, U_minus = flash_convolve(u, v, flash_fft, use_tensordot=True)
            >>> print(U_plus.shape, U_minus.shape)
            torch.Size([4, 16, 32]) torch.Size([4, 16, 32])

        """
        bsz, seq_len, d_in = u.shape
        _, K = v.shape

        padded_len = nearest_power_of_two(seq_len, round_up=True)
        pad_len = padded_len - seq_len

        sgn = torch.full((1, 1, padded_len), 1, device=u.device)
        sgn[:, :, 1::2] = -1

        if use_tensordot:
            u_padded = F.pad(u.transpose(1, 2), (0, pad_len)).to(torch.bfloat16)
            v_padded = F.pad(v.transpose(0, 1), (0, pad_len)).to(torch.float32)
            u_conv = torch.stack([u_padded, u_padded * sgn], dim=0).reshape(2 * bsz, d_in, padded_len)
        else:
            u_k_padded = F.pad(u.transpose(1, 2), (0, pad_len)).repeat_interleave(K, dim=1)
            v_padded = F.pad(v.transpose(0, 1), (0, pad_len)).to(torch.float32).repeat(d_in, 1)
            u_conv = torch.stack([u_k_padded, u_k_padded * sgn], dim=0).reshape(2 * bsz, K * d_in, padded_len)

        # Ensure inputs to flash_fft are bfloat16 (input) and float32 (kernel)
        U_conv = flash_fft(u_conv.to(torch.bfloat16), v_padded.to(torch.float32))

        # Trim the output back to the original sequence length
        U_conv = U_conv[..., :seq_len]
        u_plus, u_minus = torch.chunk(U_conv, 2, dim=0)

        if use_tensordot:
            u_minus = u_minus * sgn[:, :, :seq_len]
            U_plus, U_minus = u_plus.transpose(1, 2), u_minus.transpose(1, 2)
        else:
            sgn = sgn[:, :, :seq_len].unsqueeze(-1).transpose(1, 2)
            U_plus = u_plus.view(bsz, d_in, K, seq_len).permute(0, 3, 2, 1).contiguous()
            U_minus = u_minus.view(bsz, d_in, K, seq_len).permute(0, 3, 2, 1).contiguous() * sgn

        return U_plus, U_minus


class SlidingWindowAttentionLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.swa_norm = nn.LayerNorm(config.dim)
        self.swa = SlidingWindowAttention(config)
        self.mlp_norm = nn.LayerNorm(config.dim)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.swa(self.swa_norm(x))
        x = x + self.mlp(self.mlp_norm(x))
        return x


class STULayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.stu_norm = nn.LayerNorm(config.dim)
        self.stu = STU(config)
        self.mlp_norm = nn.LayerNorm(config.dim)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.stu(self.stu_norm(x))
        x = x + self.mlp(self.mlp_norm(x))
        return x


class FlashSTU(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.tok_emb = nn.Embedding(config.vocab_size, config.dim)
        self.layers = nn.ModuleList()

        for layer_idx in range(config.num_layers):
            # For more complex %-split arrangements, see https://arxiv.org/pdf/2406.07887
            if layer_idx % 2 == 0:
                self.layers.append(STULayer(config))
            else:
                self.layers.append(SlidingWindowAttentionLayer(config)) if config.use_attn else self.layers.append(
                    STULayer(config)
                )

        self.norm_f = nn.LayerNorm(config.dim)
        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)

        if self.config.weight_tying:
            self.tok_emb.weight = self.lm_head.weight

        self.std = self.config.dim**-0.5

    def init_weights(self, module):
        std = self.std
        if isinstance(module, nn.Linear):
            if hasattr(module, "SCALE_INIT"):
                std *= (2 * self.config.num_layers) ** -0.5
            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=std)

    def forward(self, input_ids: torch.Tensor, labels: torch.Tensor = None, **kwargs) -> CausalLMOutput:
        x = self.tok_emb(input_ids)

        for layer in self.layers:
            x = layer(x)

        x = self.norm_f(x)
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            loss = loss_fn(logits.flatten(0, 1), labels.flatten(0, 1))

        return CausalLMOutput(
            loss=loss,
            logits=logits,
        )

    def setup_filters(
        self,
        spectral_filters: torch.Tensor,
        spectral_filters_fft: torch.Tensor,
    ):
        for layer in self.layers:
            if isinstance(layer, STULayer):
                layer.stu.stu_filters = spectral_filters
                layer.stu.stu_filters_fft = spectral_filters_fft

    def get_num_params(self):
        """
        Return the number of parameters in the model.
        For non-embedding count (default), the position embeddings get subtracted.
        """
        n_params = sum(p.numel() for p in self.parameters())
        return n_params


def create_base_model_components(model_name_or_path=None, **kwargs):
    """Create config and filters needed for model initialization"""
    if model_name_or_path is not None:
        config = FlashSTUConfig.from_pretrained(model_name_or_path, **kwargs)
    else:
        config = FlashSTUConfig(**kwargs)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    filters = get_spectral_filters(
        seq_len=config.seq_len,
        K=config.num_eigh,
        use_hankel_L=config.use_hankel_L,
        device=device,
        dtype=config.torch_dtype,
    )
    assert filters.dtype == config.torch_dtype, f"filters dtype is {filters.dtype}, expected {config.torch_dtype}"
    return config, filters


class FlashSTUForCausalLM(PreTrainedModel):
    """Thin wrapper to comply with HuggingFace's expected interface"""

    config_class = FlashSTUConfig
    base_model_prefix = "FlashSTU"

    def __init__(self, config):
        super().__init__(config)

        self.flash_stu = FlashSTU(config)
        self.flash_stu.apply(self.flash_stu.init_weights)

        device = (
            config.device
            if config.device is not None
            else torch.device("cuda" if torch.cuda.is_available() else "cpu")
        )
        torch_dtype = config.torch_dtype  # Assumes __post_init__ already converted it to torch.dtype

        spectral_filters = get_spectral_filters(
            seq_len=config.seq_len,
            K=config.num_eigh,
            use_hankel_L=config.use_hankel_L,
            device=device,
            # Note: get_spectral_filters returns float64, cast later
        )
        spectral_filters_fft = torch.fft.rfft(spectral_filters, n=config.n, dim=1)

        # Setup filters in the model, casting to the target dtype
        self.flash_stu.setup_filters(
            spectral_filters.to(dtype=torch_dtype), spectral_filters_fft.to(dtype=torch_dtype)
        )
        # Note: Moving the entire model to device happens later, after loading weights.

    def forward(
        self, input_ids: torch.Tensor, labels: torch.Tensor = None, attention_mask: torch.Tensor = None, **kwargs
    ) -> CausalLMOutput:
        outputs = self.flash_stu(input_ids, labels=labels, **kwargs)
        return outputs

    def generate(
        self,
        input_ids: torch.Tensor,
        max_length: int = 32,
        num_return_sequences: int = 4,
        temperature: float = 0.8,
        top_k: int = 50,
        top_p: float = 0.95,
        repetition_penalty: float = 1.2,
        seed: int = 42,
    ) -> torch.Tensor:
        """Generate text using top-k and nucleus sampling with temperature and repetition penalty.

        Args:
            input_ids: Input token ids of shape (batch_size, seq_len)
            max_length: Maximum length of generated sequence
            num_return_sequences: Number of sequences to generate per input
            temperature: Sampling temperature. Higher = more random, lower = more focused
            top_k: Number of highest probability tokens to keep for top-k sampling
            top_p: Cumulative probability cutoff for nucleus sampling
            repetition_penalty: Penalty factor for repeating tokens. 1.0 = no penalty
            seed: Random seed for reproducibility

        Returns:
            Generated token ids of shape (num_return_sequences, max_length)
        """
        self.eval()  # Set to eval mode
        device = input_ids.device

        # Expand input for multiple sequences
        input_ids = input_ids.repeat(num_return_sequences, 1)
        generated = input_ids

        # Set up generator for reproducible sampling
        sample_rng = torch.Generator(device=device)
        sample_rng.manual_seed(seed)

        # Generate tokens until we reach max_length
        with torch.no_grad():
            while generated.size(1) < max_length:
                # Get logits for next token
                outputs = self.flash_stu(generated)
                next_token_logits = outputs.logits[:, -1, :]

                # Apply repetition penalty
                if repetition_penalty != 1.0:
                    for i in range(generated.shape[0]):
                        for token in generated[i]:
                            if token in next_token_logits[i]:
                                next_token_logits[i, token] /= repetition_penalty

                # Apply temperature
                if temperature != 1.0:
                    next_token_logits = next_token_logits / temperature

                # Get probabilities
                probs = torch.nn.functional.softmax(next_token_logits, dim=-1)

                # Top-k sampling
                if top_k > 0:
                    indices_to_remove = probs < torch.topk(probs, top_k)[0][..., -1, None]
                    probs[indices_to_remove] = 0

                # Nucleus (top-p) sampling
                if top_p < 1.0:
                    sorted_probs, sorted_indices = torch.sort(probs, descending=True)
                    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

                    # Remove tokens with cumulative probability above the threshold
                    sorted_indices_to_remove = cumulative_probs > top_p
                    # Shift the indices to the right to keep also the first token above the threshold
                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                    sorted_indices_to_remove[..., 0] = 0

                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                    probs[indices_to_remove] = 0

                # Renormalize probabilities
                probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-8)

                # Sample next token
                next_token = torch.multinomial(probs, num_samples=1, generator=sample_rng)

                # Append to generated sequence
                generated = torch.cat([generated, next_token], dim=1)

        return generated

    def get_num_params(self):
        return self.flash_stu.get_num_params()

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        # Get config and create model
        config, _ = create_base_model_components(pretrained_model_name_or_path, **kwargs)
        model = cls(config)

        # Download safetensors file from hub
        weights_path = hf_hub_download(
            repo_id=pretrained_model_name_or_path,
            filename="model.safetensors",
            cache_dir=kwargs.get("cache_dir"),
            force_download=kwargs.get("force_download", False),
            proxies=kwargs.get("proxies", None),
            local_files_only=kwargs.get("local_files_only", False),
            use_auth_token=kwargs.get("use_auth_token", None),
            revision=kwargs.get("revision", None),
            subfolder=kwargs.get("subfolder", ""),
        )

        state_dict = load_file(weights_path)

        # Reconstruct weight tying for tok_emb and lm_head
        tok_emb_key = "tok_emb.weight"
        lm_head_key = "lm_head.weight"

        tok_emb_present = tok_emb_key in state_dict
        lm_head_present = lm_head_key in state_dict

        if tok_emb_present and not lm_head_present:
            print(f"Reconstructing weight tying: Linking missing '{lm_head_key}' to existing '{tok_emb_key}'")
            state_dict[lm_head_key] = state_dict[tok_emb_key]
        elif lm_head_present and not tok_emb_present:
            print(f"Reconstructing weight tying: Linking missing '{tok_emb_key}' to existing '{lm_head_key}'")
            state_dict[tok_emb_key] = state_dict[lm_head_key]
        elif not tok_emb_present and not lm_head_present:
            # This case should ideally not happen if the file is valid
            print(
                f"Warning: Neither '{tok_emb_key}' nor '{lm_head_key}' found in state_dict. Weight tying cannot be reconstructed."
            )
        # If both are present, assume they are loaded correctly (or were never tied)

        # Prepend 'flash_stu.' to all keys to match wrapper's state dict
        final_state_dict = {f"flash_stu.{k}": v for k, v in state_dict.items()}
        model.load_state_dict(final_state_dict)

        # Move to GPU if available
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device=device, dtype=torch.bfloat16)
        model.eval()

        # Print parameter count as a sanity check
        num_params = model.get_num_params()
        print(f"\nModel loaded: {pretrained_model_name_or_path}")
        print(f"Parameter count: {num_params / 1e6:.2f}M")

        return model


# Create initial config and filters for registration
config, filters = create_base_model_components()

# Register models
AutoConfig.register("FlashSTU", FlashSTUConfig)
AutoModel.register(FlashSTUConfig, FlashSTU)
AutoModelForCausalLM.register(FlashSTUConfig, FlashSTUForCausalLM)

print("Registered FlashSTU model and configuration.")


def run_model_diagnostics(model, tokenizer, device):
    """Run detailed diagnostics to analyze model behavior."""
    print("\nRunning model diagnostics...")

    # Test cases of varying difficulty and length
    test_cases = [
        # Simple completion
        "2 + 2 =",
        # Medium difficulty
        "The capital of France is Paris. The capital of Germany is",
        # Complex reasoning
        "If a train travels 120 kilometers in 2 hours, its average speed is",
        # Pattern completion
        "1, 2, 3, 4,",
        # Long context
        "The following is a detailed explanation of photosynthesis: Plants use sunlight to",
    ]

    with torch.no_grad():
        for prompt in test_cases:
            print(f"\nAnalyzing prompt: {prompt}")

            # Tokenize
            tokens = tokenizer(prompt, return_tensors="pt")
            input_ids = tokens["input_ids"].to(device)

            outputs = model.flash_stu(input_ids, labels=input_ids)

            labels = input_ids.clone()
            shift_logits = outputs.logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()

            loss_fct = nn.CrossEntropyLoss(reduction="none")
            token_losses = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)).view(
                shift_labels.size()
            )

            # Print token-by-token analysis
            input_tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
            print("\nToken-by-token loss:")
            for i, (token, loss) in enumerate(zip(input_tokens[1:], token_losses[0])):
                print(f"{token}: {loss.item():.3f}")

            print(f"Average loss: {token_losses.mean().item():.3f}")

            # Generate with different temperatures
            temps = [0.5, 0.7, 1.0]
            print("\nGeneration temperature comparison:")
            for temp in temps:
                gen_ids = model.generate(
                    input_ids,
                    max_length=25,
                    num_return_sequences=1,
                    temperature=temp,
                    top_p=0.9,
                    repetition_penalty=1.5,
                    seed=42,
                )
                gen_text = tokenizer.decode(gen_ids[0], skip_special_tokens=True)
                print(f"\nTemp {temp}: {gen_text}")


def validate_model_generation():
    print("\nRunning generation validation test...")

    try:
        from transformers import AutoTokenizer

        # Load model and tokenizer
        # model_id = "Hazan-Lab/Flash_STU_550M"
        model_id = "Hazan-Lab/FlashSTU-340M-0428"
        model = FlashSTUForCausalLM.from_pretrained(model_id)
        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)

        # Move to GPU if available
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device=device, dtype=torch.bfloat16)
        model.eval()

        # Print parameter count as a sanity check
        num_params = model.get_num_params()
        print(f"\nModel loaded: {model_id}")
        print(f"Parameter count: {num_params / 1e6:.2f}M")

        # Run additional diagnostics
        run_model_diagnostics(model, tokenizer, device)

    except Exception as e:
        print(f"\nError during validation: {str(e)}")
        raise


# Run evaluation tasks
tasks = [
    # "mmlu",
    "hellaswag",
    # "piqa",
    # "siqa",
    # "boolq",
    # "winogrande",
    # "commonsense_qa",
    # "openbookqa",
    # "arc",
    # "arc_easy",
    # "arc_challenge",
    # "triviaqa",
    # "nq_open",
    # "humaneval",
    # "mbpp",
    # "gms8k",
    # "hendrycks_math",
    # "mathqa",
    # "minerva_math",
    # "score",
    # "asdiv",
    # "agieval",
    # "bigbench",
]

tasks_fewshot = {
    "hellaswag": 0,
    # "mmlu": 5,
    # "piqa": 0,
    # "siqa": 0,
    # "boolq": 0,
    # "winogrande": -1,
    # "commonsense_qa": 7,
    # "openbookqa": -1,
    # "arc": -1,
    # "arc_easy": -1,
    # "arc_challenge": -1,
    # "triviaqa": 5,
    # "nq_open": 5,
    # "humaneval": -1,
    # "mbpp": 3,
    # "gms8k": -1,
    # "hendrycks_math": 4,
    # "mathqa": -1,
    # "minerva_math": -1,
    # "score": -1,
    # "asdiv": -1,
    # "agieval": -1,
    # "bigbench": -1,
}

all_results = {}

# First validate generation works
validate_model_generation()

print("\nStarting evaluation tasks...")
for task in tasks:
    print(f"\nEvaluating task: {task}")
    eval_kwargs = dict(
        model="hf",
        model_args=(
            # "pretrained=Hazan-Lab/Flash_STU_550M,"
            "pretrained=Hazan-Lab/FlashSTU-340M-0428,"
            "trust_remote_code=True,"
            "dtype=bfloat16,"
            "cache_dir=/scratch/gpfs/mn4560/hazan-lab/tensorized_filters/tensorized_filters/eval/cache"
        ),
        tasks=[task],
        batch_size="auto",
        device="cuda:0",
    )
    few_shot_value = tasks_fewshot.get(task, -1)
    if few_shot_value != -1:
        eval_kwargs["num_fewshot"] = few_shot_value
    results = evaluator.simple_evaluate(**eval_kwargs)
    task_result = results["results"].get(task, {})
    all_results[task] = task_result
    print(f"Results for {task}:")
    print(task_result)
    print("\n" + "=" * 50 + "\n")

print("All Evaluation Results:")
for task, result in all_results.items():
    print(f"{task}: {result}")