"""
HybridFourierLM – Model Architecture
=====================================
Self-contained module that registers the custom config + model with
HuggingFace `transformers` so that `AutoConfig` / `AutoModelForCausalLM`
can load checkpoints transparently.
"""

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    GenerationMixin,
    PretrainedConfig,
    PreTrainedModel,
)
from transformers.modeling_outputs import CausalLMOutput


# ──────────────────────────────────────────────────────────────────────
# Config
# ──────────────────────────────────────────────────────────────────────

class HybridFourierConfig(PretrainedConfig):
    model_type = "hybrid_fourier_lm"

    def __init__(
        self,
        vocab_size=50304,
        latent_dim=768,
        num_layers=12,
        num_modes=64,
        layer_types=None,
        time_scale=128.0,
        dropout=0.05,
        pad_token_id=0,
        bos_token_id=1,
        eos_token_id=2,
        tie_word_embeddings=True,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.latent_dim = latent_dim
        self.num_layers = num_layers
        self.num_modes = num_modes
        self.time_scale = time_scale
        self.dropout = dropout

        if layer_types is None:
            layer_types = [
                "softmax" if (i % 4 == 3) else "linear"
                for i in range(num_layers)
            ]
        assert len(layer_types) == num_layers, (
            f"layer_types length ({len(layer_types)}) must equal num_layers ({num_layers})"
        )
        self.layer_types = layer_types

        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )


# ──────────────────────────────────────────────────────────────────────
# Mixer layers
# ──────────────────────────────────────────────────────────────────────

class LinearFourierMixer(nn.Module):
    def __init__(self, channels, num_modes=64, num_heads=12, time_scale=128, dropout=0.05):
        super().__init__()
        assert channels % num_heads == 0, (
            f"channels ({channels}) must be perfectly divisible by num_heads ({num_heads})"
        )
        self.channels = channels
        self.num_modes = num_modes
        self.num_heads = num_heads
        self.head_dim = channels // num_heads
        self.time_scale = time_scale

        freq_bands = torch.exp(torch.linspace(math.log(0.0001), math.log(num_modes), num_modes))
        self.num_modes = freq_bands.shape[0]
        self.register_buffer("frequencies", freq_bands)

        self.q_proj = nn.Linear(channels, self.num_heads * self.num_modes)
        self.k_proj = nn.Linear(channels, self.num_heads * self.num_modes)
        self.v_proj = nn.Linear(channels, channels)
        self.proj_v2 = nn.Linear(channels, channels)
        self.out_proj = nn.Linear(channels, channels)
        self.activation = nn.SiLU()
        self.norm_in = nn.LayerNorm(channels)
        self.norm_out = nn.LayerNorm(channels)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, attention_mask=None):
        B, seq_len, C = x.shape
        norm_x = self.norm_in(x)

        Q = F.elu(self.q_proj(norm_x)).view(B, seq_len, self.num_heads, self.num_modes) + 1.0
        K = F.elu(self.k_proj(norm_x)).view(B, seq_len, self.num_heads, self.num_modes) + 1.0

        v1 = self.v_proj(norm_x)
        v2 = self.activation(self.proj_v2(norm_x))

        t = (torch.arange(seq_len, device=x.device, dtype=x.dtype) / self.time_scale).view(-1, 1)
        omega_t = 2 * math.pi * t * self.frequencies.unsqueeze(0)
        U = torch.cos(omega_t).unsqueeze(0).unsqueeze(2)
        V = torch.sin(omega_t).unsqueeze(0).unsqueeze(2)

        Q_cos = Q * U
        Q_sin = Q * V
        K_cos = K * U
        K_sin = K * V

        Q_rot = torch.cat([Q_cos, Q_sin], dim=-1)
        K_rot = torch.cat([K_cos, K_sin], dim=-1)

        # Linear-attention normalization in fp32 for numerical stability
        orig_dtype = Q_rot.dtype
        Q_rot = Q_rot.float()
        K_rot = K_rot.float()

        if seq_len > 512:
            v1_heads = v1.view(B, seq_len, self.num_heads, self.head_dim)
            out_chunks = []
            chunk_size = 256 if seq_len > 1024 else 512
            for i_start in range(0, seq_len, chunk_size):
                i_end = min(i_start + chunk_size, seq_len)
                Q_chunk = Q_rot[:, i_start:i_end, :, :]       # [B, C, H, 2M]
                K_past = K_rot[:, :i_end, :, :]               # [B, j_max, H, 2M]
                v1_past = v1_heads[:, :i_end, :, :].float()   # [B, j_max, H, D]

                A_chunk = torch.einsum('b i h m, b j h m -> b h i j', Q_chunk, K_past)
                A_chunk = A_chunk / math.sqrt(self.num_modes * 2)

                i_abs = torch.arange(i_start, i_end, device=x.device).view(-1, 1)
                j_abs = torch.arange(i_end, device=x.device).view(1, -1)
                causal_mask = (j_abs <= i_abs).to(dtype=A_chunk.dtype)
                A_chunk = A_chunk * causal_mask.unsqueeze(0).unsqueeze(0)

                if attention_mask is not None:
                    pad_mask = attention_mask[:, None, None, :i_end].to(dtype=A_chunk.dtype)
                    A_chunk = A_chunk * pad_mask

                row_denom = torch.abs(A_chunk.sum(dim=-1, keepdim=True)) + 1.0
                A_chunk = A_chunk / row_denom

                v1_chunk = torch.einsum('b h i j, b j h d -> b i h d', A_chunk, v1_past)
                out_chunks.append(v1_chunk.to(orig_dtype))
            v1_token_mixed = torch.cat(out_chunks, dim=1).reshape(B, seq_len, C)
            v1_token_mixed = self.dropout(v1_token_mixed)
            if attention_mask is not None:
                v1_token_mixed = torch.nan_to_num(v1_token_mixed, nan=0.0, posinf=0.0, neginf=0.0)
                v1_token_mixed = v1_token_mixed * attention_mask.unsqueeze(-1).to(dtype=v1_token_mixed.dtype)
        else:
            A = torch.einsum('b i h m, b j h m -> b h i j', Q_rot, K_rot)
            A = A / math.sqrt(self.num_modes * 2)

            causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device, dtype=A.dtype))
            A = A * causal_mask.unsqueeze(0).unsqueeze(0)

            if attention_mask is not None:
                pad_mask = attention_mask[:, None, None, :].to(dtype=A.dtype)
                A = A * pad_mask

            row_denom = torch.abs(A.sum(dim=-1, keepdim=True)) + 1.0
            A = A / row_denom
            A = A.to(orig_dtype)

            v1_heads = v1.view(B, seq_len, self.num_heads, self.head_dim)
            v1_token_mixed = torch.einsum('b h i j, b j h d -> b i h d', A, v1_heads)
            v1_token_mixed = v1_token_mixed.reshape(B, seq_len, C)
            v1_token_mixed = self.dropout(v1_token_mixed)
            if attention_mask is not None:
                v1_token_mixed = torch.nan_to_num(v1_token_mixed, nan=0.0, posinf=0.0, neginf=0.0)
                v1_token_mixed = v1_token_mixed * attention_mask.unsqueeze(-1).to(dtype=v1_token_mixed.dtype)

        v3 = v1_token_mixed * v2
        return self.norm_out(self.out_proj(v3)) + x


class SoftmaxFourierMixer(nn.Module):
    def __init__(self, channels, num_modes=64, num_heads=12, time_scale=128.0, dropout=0.05):
        super().__init__()
        assert channels % num_heads == 0, (
            f"channels ({channels}) must be perfectly divisible by num_heads ({num_heads})"
        )
        self.channels = channels
        self.num_modes = num_modes
        self.num_heads = num_heads
        self.head_dim = channels // num_heads
        self.time_scale = time_scale

        freq_bands = torch.exp(torch.linspace(math.log(0.0001), math.log(num_modes), num_modes))
        self.num_modes = freq_bands.shape[0]
        self.register_buffer("frequencies", freq_bands)

        self.q_proj = nn.Linear(channels, self.num_heads * self.num_modes)
        self.k_proj = nn.Linear(channels, self.num_heads * self.num_modes)
        self.v_proj = nn.Linear(channels, channels)
        self.proj_v2 = nn.Linear(channels, channels)
        self.out_proj = nn.Linear(channels, channels)
        self.activation = nn.SiLU()
        self.norm_in = nn.LayerNorm(channels)
        self.norm_out = nn.LayerNorm(channels)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, attention_mask=None):
        B, seq_len, C = x.shape
        norm_x = self.norm_in(x)

        Q = self.q_proj(norm_x).view(B, seq_len, self.num_heads, self.num_modes)
        K = self.k_proj(norm_x).view(B, seq_len, self.num_heads, self.num_modes)

        v1 = self.v_proj(norm_x)
        v2 = self.activation(self.proj_v2(norm_x))

        t = (torch.arange(seq_len, device=x.device, dtype=x.dtype) / self.time_scale).view(-1, 1)
        omega_t = 2 * math.pi * t * self.frequencies.unsqueeze(0)
        U = torch.cos(omega_t).unsqueeze(0).unsqueeze(2)
        V = torch.sin(omega_t).unsqueeze(0).unsqueeze(2)

        Q_cos = Q * U
        Q_sin = Q * V
        K_cos = K * U
        K_sin = K * V

        Q_rot = torch.cat([Q_cos, Q_sin], dim=-1)
        K_rot = torch.cat([K_cos, K_sin], dim=-1)

        v1_heads = v1.view(B, seq_len, self.num_heads, self.head_dim)

        # Actual softmax attention via SDPA (B, H, L, 2M)
        Q_b = Q_rot.transpose(1, 2)
        K_b = K_rot.transpose(1, 2)
        V_b = v1_heads.transpose(1, 2)

        # PyTorch's MPS backend has a known bug/crash in C++ kernel
        # (`-[__NSPlaceholderDictionary initWithObjects:forKeys:count:]`)
        # when calling F.scaled_dot_product_attention on certain shapes or when
        # attn_mask and is_causal are combined. We use explicit math attention
        # on MPS (or fallback if SDPA fails) to guarantee stability across all PyTorch versions.
        if x.device.type == "mps" or seq_len > 512:
            scale = 1.0 / math.sqrt(Q_b.size(-1))
            if seq_len > 256:
                out_chunks = []
                chunk_size = 256 if seq_len > 1024 else 512
                for i_start in range(0, seq_len, chunk_size):
                    i_end = min(i_start + chunk_size, seq_len)
                    Q_chunk = Q_b[:, :, i_start:i_end, :]         # [B, H, C, 2M]
                    K_past = K_b[:, :, :i_end, :]                 # [B, H, 2M, j_max]
                    V_past = V_b[:, :, :i_end, :]                 # [B, H, j_max, D]

                    scores_chunk = torch.matmul(Q_chunk, K_past.transpose(-2, -1)) * scale

                    i_abs = torch.arange(i_start, i_end, device=x.device).view(-1, 1)
                    j_abs = torch.arange(i_end, device=x.device).view(1, -1)
                    causal_mask = (j_abs <= i_abs)
                    scores_chunk = scores_chunk.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf"))

                    if attention_mask is not None:
                        pad_mask = attention_mask[:, None, None, :i_end].to(dtype=torch.bool)
                        scores_chunk = scores_chunk.masked_fill(~pad_mask, float("-inf"))

                    attn_weights = F.softmax(scores_chunk, dim=-1)
                    out_chunk = torch.matmul(attn_weights, V_past)  # [B, H, C, D]
                    out_chunks.append(out_chunk)
                v1_token_mixed = torch.cat(out_chunks, dim=2)
            else:
                scores = torch.matmul(Q_b, K_b.transpose(-2, -1)) * scale
                causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool))
                scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf"))
                if attention_mask is not None:
                    pad_mask = attention_mask[:, None, None, :].to(dtype=torch.bool)
                    scores = scores.masked_fill(~pad_mask, float("-inf"))
                attn_weights = F.softmax(scores, dim=-1)
                v1_token_mixed = torch.matmul(attn_weights, V_b)
        else:
            attn_mask = None
            if attention_mask is not None:
                attn_mask = attention_mask[:, None, None, :].to(dtype=Q_b.dtype)
                attn_mask = (1.0 - attn_mask) * torch.finfo(Q_b.dtype).min

            try:
                v1_token_mixed = F.scaled_dot_product_attention(
                    Q_b, K_b, V_b,
                    attn_mask=attn_mask,
                    is_causal=True,
                )
            except Exception:
                scale = 1.0 / math.sqrt(Q_b.size(-1))
                scores = torch.matmul(Q_b, K_b.transpose(-2, -1)) * scale
                causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool))
                scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf"))
                if attention_mask is not None:
                    pad_mask = attention_mask[:, None, None, :].to(dtype=torch.bool)
                    scores = scores.masked_fill(~pad_mask, float("-inf"))
                attn_weights = F.softmax(scores, dim=-1)
                v1_token_mixed = torch.matmul(attn_weights, V_b)

        v1_token_mixed = v1_token_mixed.transpose(1, 2).reshape(B, seq_len, C)
        v1_token_mixed = self.dropout(v1_token_mixed)
        if attention_mask is not None:
            v1_token_mixed = torch.nan_to_num(v1_token_mixed, nan=0.0, posinf=0.0, neginf=0.0)
            v1_token_mixed = v1_token_mixed * attention_mask.unsqueeze(-1).to(dtype=v1_token_mixed.dtype)

        v3 = v1_token_mixed * v2
        return self.norm_out(self.out_proj(v3)) + x


# ──────────────────────────────────────────────────────────────────────
# Transformer block
# ──────────────────────────────────────────────────────────────────────

class HybridSpectralBlock(nn.Module):
    def __init__(self, latent_dim, num_modes=64, is_softmax=False,
                 time_scale=128.0, dropout=0.05, num_heads=None):
        super().__init__()
        self.is_softmax = is_softmax
        num_heads = num_heads if num_heads is not None else max(1, latent_dim // 64)

        if is_softmax:
            self.mixer = SoftmaxFourierMixer(latent_dim, num_modes, num_heads, time_scale, dropout)
        else:
            self.mixer = LinearFourierMixer(latent_dim, num_modes, num_heads, time_scale, dropout)

        self.ffn = nn.Sequential(
            nn.LayerNorm(latent_dim),
            nn.Linear(latent_dim, 4 * latent_dim),
            nn.GELU(),
            nn.Linear(4 * latent_dim, latent_dim),
            nn.Dropout(dropout),
        )
        self.gradient_checkpointing = False

    def forward(self, x, attention_mask=None):
        z = self.mixer(x, attention_mask=attention_mask)
        return z + self.ffn(z)


# ──────────────────────────────────────────────────────────────────────
# Full model
# ──────────────────────────────────────────────────────────────────────

class HybridFourierPreTrainedModel(PreTrainedModel):
    config_class = HybridFourierConfig
    base_model_prefix = "hybrid_fourier"
    supports_gradient_checkpointing = True
    _no_split_modules = ["HybridSpectralBlock"]
    _tied_weights_keys = {"lm_head.weight": "embedding.weight"}

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            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=0.02)
        elif isinstance(module, nn.LayerNorm):
            torch.nn.init.zeros_(module.bias)
            torch.nn.init.ones_(module.weight)


class HybridFourierLM(HybridFourierPreTrainedModel, GenerationMixin):
    def __init__(self, config):
        super().__init__(config)
        self.config = config

        self.embedding = nn.Embedding(config.vocab_size, config.latent_dim,
                                       padding_idx=config.pad_token_id)

        blocks = []
        for layer_type in config.layer_types:
            blocks.append(HybridSpectralBlock(
                config.latent_dim,
                config.num_modes,
                is_softmax=(layer_type == "softmax"),
                time_scale=config.time_scale,
                dropout=config.dropout,
            ))
        self.mixers = nn.ModuleList(blocks)

        self.ln_f = nn.LayerNorm(config.latent_dim)
        self.lm_head = nn.Linear(config.latent_dim, config.vocab_size, bias=False)

        self.post_init()

    def get_input_embeddings(self):
        return self.embedding

    def set_input_embeddings(self, value):
        self.embedding = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embedding):
        self.lm_head = new_embedding

    def forward(self, input_ids=None, attention_mask=None, labels=None,
                past_key_values=None, use_cache=None, inputs_embeds=None, **kwargs):
        if inputs_embeds is None:
            z = self.embedding(input_ids)
        else:
            z = inputs_embeds

        for mixer in self.mixers:
            if getattr(self, "gradient_checkpointing", False) and self.training:
                z = checkpoint(mixer, z, attention_mask, use_reentrant=False)
            else:
                z = mixer(z, attention_mask=attention_mask)

        z = self.ln_f(z)
        logits = self.lm_head(z)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
            loss = loss_fct(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
            )

        return CausalLMOutput(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
                                       attention_mask=None, **kwargs):
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "use_cache": False,
        }

    def _reorder_cache(self, past_key_values, beam_idx):
        return past_key_values


# ── Register with AutoClasses ────────────────────────────────────────
AutoConfig.register("hybrid_fourier_lm", HybridFourierConfig)
AutoModelForCausalLM.register(HybridFourierConfig, HybridFourierLM)