import math
import torch
import torch.nn.functional as F
from einops import rearrange
from torch import nn
from transformers import PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
try:
    from fla.layers.utils import get_layer_cache, update_layer_cache
    from fla.models.utils import Cache, FLAUnsupportedCacheGenerationMixin
    from fla.modules import FusedRMSNormSwishGate, RMSNorm, ShortConvolution
    from fla.ops.gdn2 import chunk_gdn2, fused_recurrent_gdn2
except (ImportError, AttributeError) as e:
    raise ImportError('Linear-50M requires a current CUDA build of flash-linear-attention: pip install "flash-linear-attention[cuda]"') from e

class Linear50MConfig(PretrainedConfig):
    model_type = "linear_50m"
    def __init__(self, vocab_size=32000, hidden_size=384, num_hidden_layers=16, num_heads=6,
                 head_dim=64, intermediate_size=1120, conv_size=4, norm_eps=1e-5,
                 max_position_embeddings=4096, use_cache=True, **kwargs):
        for key, value in dict(bos_token_id=1, eos_token_id=2, pad_token_id=2,
                               tie_word_embeddings=True).items(): kwargs.setdefault(key, value)
        super().__init__(**kwargs); self.vocab_size, self.hidden_size, self.num_hidden_layers = vocab_size, hidden_size, num_hidden_layers
        self.num_heads, self.head_dim, self.intermediate_size = num_heads, head_dim, intermediate_size
        self.conv_size, self.norm_eps, self.max_position_embeddings, self.use_cache = conv_size, norm_eps, max_position_embeddings, use_cache

class NotGDN2(nn.Module):
    def __init__(self, c, layer_idx):
        super().__init__(); d, h, hd = c.hidden_size, c.num_heads, c.head_dim; self.h, self.hd, self.layer_idx = h, hd, layer_idx
        self.qkv = nn.Linear(d, 3*d, bias=False); self.conv = ShortConvolution(3*d, c.conv_size, bias=False, activation="silu")
        self.f = nn.Sequential(nn.Linear(d, hd, False), nn.Linear(hd, d, False)); self.bw = nn.Linear(d, 2*d, bias=False)
        self.gate = nn.Sequential(nn.Linear(d, hd, False), nn.Linear(hd, d, True)); self.recover = nn.Sequential(nn.Linear(d, hd, False), nn.Linear(hd, d, True))
        self.A_log = nn.Parameter(torch.empty(h).uniform_(1, 16).log()); dt = torch.exp(torch.empty(d).uniform_(math.log(.001), math.log(.1)))
        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt))); self.A_log._no_weight_decay = self.dt_bias._no_weight_decay = True
        self.norm = FusedRMSNormSwishGate(hd, eps=c.norm_eps); self.out = nn.Linear(d, d, bias=False)
        self.apply(self._init)
        for weight, parts in ((self.qkv.weight, 3), (self.bw.weight, 2)):
            for part in weight.chunk(parts, 0): nn.init.xavier_uniform_(part, gain=2**-2.5)
        nn.init.zeros_(self.recover[-1].weight); nn.init.zeros_(self.recover[-1].bias)
    @staticmethod
    def _init(m):
        if getattr(m, "_is_hf_initialized", False): return
        if isinstance(m, nn.Linear):
            nn.init.xavier_uniform_(m.weight, gain=2**-2.5)
            if m.bias is not None: nn.init.zeros_(m.bias)
        m._is_hf_initialized = True
    def forward(self, x, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
        length = x.shape[1]; last = get_layer_cache(self, past_key_values); conv = None if last is None else last["conv_state"]; cu = kwargs.get("cu_seqlens")
        qkv, conv = self.conv(self.qkv(x), cache=conv, output_final_state=use_cache, cu_seqlens=cu)
        q, k, v = (rearrange(z, "b t (h d) -> b t h d", h=self.h) for z in qkv.chunk(3, -1))
        g = rearrange(-self.A_log.float().exp().repeat_interleave(self.hd) *
                      F.softplus(self.f(x).float() + self.dt_bias.float()),
                      "b t (h d) -> b t h d", h=self.h)
        b, w = (rearrange(z.sigmoid(), "b t (h d) -> b t h d", h=self.h) for z in self.bw(x).chunk(2, -1))
        state = None if last is None else last["recurrent_state"]
        op = fused_recurrent_gdn2 if use_cache and not self.training and not torch.is_grad_enabled() and x.shape[1] <= 64 else chunk_gdn2
        both, state = op(*(torch.cat(z, 2) for z in ((q, q), (k, k), (v, v),
                          (g, torch.zeros_like(g)), (b, torch.zeros_like(b)), (w, w))),
                         scale=self.hd**-.5, initial_state=state, output_final_state=use_cache,
                         use_qk_l2norm_in_kernel=True, use_gate_in_kernel=False, cu_seqlens=cu)
        a, archive = both.chunk(2, 2); protected = archive-a
        support = protected*((a.float().square().mean(-1, True)+1e-6)/(protected.float().square().mean(-1, True)+1e-6)).sqrt().to(protected)
        theta = rearrange(self.recover(x), "b t (h d) -> b t h d", h=self.h); o = a + theta.sin()*support
        if use_cache: update_layer_cache(self, past_key_values, recurrent_state=state, conv_state=conv, offset=length)
        z = rearrange(self.gate(x), "b t (h d) -> b t h d", h=self.h); return self.out(rearrange(self.norm(o, z), "b t h d -> b t (h d)"))

class SwiGLU(nn.Module):
    def __init__(self, c): super().__init__(); self.up, self.down = nn.Linear(c.hidden_size, 2*c.intermediate_size, False), nn.Linear(c.intermediate_size, c.hidden_size, False)
    def forward(self, x): a, b = self.up(x).chunk(2, -1); return self.down(F.silu(a)*b)

class Block(nn.Module):
    def __init__(self, c, i): super().__init__(); self.n1, self.mix, self.n2, self.mlp = RMSNorm(c.hidden_size), NotGDN2(c, i), RMSNorm(c.hidden_size), SwiGLU(c)
    def forward(self, x, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
        x = x + self.mix(self.n1(x), attention_mask, past_key_values, use_cache, **kwargs); return x + self.mlp(self.n2(x))

class Linear50MPreTrainedModel(PreTrainedModel):
    config_class, base_model_prefix, _supports_cache_class, _no_split_modules = Linear50MConfig, "model", True, ["Block"]
    def _init_weights(self, m):
        if getattr(m, "_is_hf_initialized", False): return
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, std=.02)
            if getattr(m, "bias", None) is not None: nn.init.zeros_(m.bias)

class Linear50MModel(Linear50MPreTrainedModel):
    def __init__(self, config):
        super().__init__(config); self.embed = nn.Embedding(config.vocab_size, config.hidden_size); self.blocks = nn.ModuleList(Block(config, i) for i in range(config.num_hidden_layers)); self.norm = RMSNorm(config.hidden_size); self.post_init()
    def get_input_embeddings(self): return self.embed
    def set_input_embeddings(self, value): self.embed = value
    def forward(self, input_ids=None, attention_mask=None, inputs_embeds=None, past_key_values=None, use_cache=None,
                output_hidden_states=None, return_dict=None, **kwargs):
        if (input_ids is None) == (inputs_embeds is None): raise ValueError("Pass exactly one of input_ids or inputs_embeds")
        x = self.embed(input_ids) if inputs_embeds is None else inputs_embeds
        if x.device.type != "cuda": raise RuntimeError("Linear-50M currently requires CUDA/FLA inference")
        if attention_mask is not None:
            if past_key_values is None and not bool(attention_mask.all()): raise ValueError("Variable-length padded batches are not supported yet")
            attention_mask = None
        output_hidden_states = self.config.output_hidden_states if output_hidden_states is None else output_hidden_states
        return_dict = self.config.return_dict if return_dict is None else return_dict
        use_cache = getattr(self.config, "use_cache", True) if use_cache is None and not self.training else bool(use_cache)
        if use_cache and not isinstance(past_key_values, Cache): past_key_values = Cache.from_legacy_cache(past_key_values)
        hidden = () if output_hidden_states else None
        for block in self.blocks:
            if output_hidden_states: hidden += (x,)
            x = block(x, attention_mask, past_key_values, use_cache, **kwargs)
        x = self.norm(x)
        if output_hidden_states: hidden += (x,)
        if not return_dict: return tuple(v for v in (x, past_key_values, hidden) if v is not None)
        return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values, hidden_states=hidden)

class Linear50MForCausalLM(Linear50MPreTrainedModel, FLAUnsupportedCacheGenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embed.weight"}
    def __init__(self, config): super().__init__(config); self.model = Linear50MModel(config); self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False); self.post_init()
    def get_input_embeddings(self): return self.model.embed
    def set_input_embeddings(self, value): self.model.embed = value
    def get_output_embeddings(self): return self.lm_head
    def set_output_embeddings(self, value): self.lm_head = value
    def get_decoder(self): return self.model
    def set_decoder(self, value): self.model = value
    def forward(self, input_ids=None, attention_mask=None, inputs_embeds=None, past_key_values=None, labels=None,
                use_cache=None, output_hidden_states=None, return_dict=None, logits_to_keep=0, **kwargs):
        return_dict = self.config.return_dict if return_dict is None else return_dict
        out = self.model(input_ids, attention_mask, inputs_embeds, past_key_values, use_cache, output_hidden_states, return_dict, **kwargs); h = out[0]
        logits = self.lm_head(h if labels is not None or not logits_to_keep else h[:, -logits_to_keep:]); loss = None
        if labels is not None: loss = F.cross_entropy(logits[:, :-1].float().flatten(0, 1), labels[:, 1:].flatten(), ignore_index=-100)
        if not return_dict: return ((loss,) if loss is not None else ()) + (logits,) + out[1:]
        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=out.past_key_values, hidden_states=out.hidden_states)