import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_ablation import AblationConfig


class RMSNorm(nn.Module):
    def __init__(self, dim, eps):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        dt = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x * self.weight.float()).to(dt)


def rope_cache(T, dim, base, device):
    inv = 1.0 / (base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
    f = torch.outer(torch.arange(T, device=device, dtype=torch.float32), inv)
    return f.cos(), f.sin()  # (T, dim/2)


def apply_rope(x, cos, sin):
    # x: (B, T, H, D)
    x1, x2 = x.float().chunk(2, dim=-1)
    c, s = cos[None, :, None, :], sin[None, :, None, :]
    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1).to(x.dtype)


class DSAIndexer(nn.Module):
    """DeepSeek Sparse Attention "lightning indexer":
    I[t, s] = sum_j w[t, j] * ReLU(q[t, j] . k[s]).  Used to pick the top-k keys for every query."""

    def __init__(self, cfg):
        super().__init__()
        self.h, self.d, self.base = cfg.dsa_index_heads, cfg.dsa_index_dim, cfg.rope_base
        self.wq = nn.Linear(cfg.hidden_size, self.h * self.d, bias=False)
        self.wk = nn.Linear(cfg.hidden_size, self.d, bias=False)
        self.ww = nn.Linear(cfg.hidden_size, self.h, bias=False)

    def forward(self, x):
        B, T, _ = x.shape
        cos, sin = rope_cache(T, self.d, self.base, x.device)
        q = apply_rope(self.wq(x).view(B, T, self.h, self.d), cos, sin)
        k = apply_rope(self.wk(x).view(B, T, 1, self.d), cos, sin)[:, :, 0]
        w = self.ww(x) * (self.d ** -0.5)
        s = torch.einsum("bthd,bsd->bhts", q, k).relu()
        return torch.einsum("bhts,bth->bts", s, w).float()  # (B, T, T)


class Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.h = cfg.num_attention_heads
        self.d = cfg.hidden_size // self.h
        self.topk = cfg.dsa_topk
        self.qkv = nn.Linear(cfg.hidden_size, 3 * cfg.hidden_size, bias=False)
        self.o = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False)
        self.o.is_resid = True

    def forward(self, x, cos, sin, indexer=None, dense_warmup=False):
        B, T, C = x.shape
        q, k, v = self.qkv(x).view(B, T, 3, self.h, self.d).unbind(2)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        q, k, v = (t.transpose(1, 2) for t in (q, k, v))  # (B, H, T, D)
        aux = None
        if indexer is None or (not self.training and T <= self.topk):
            # standard causal attention (DSA with T <= topk selects every causal token, i.e. is identical)
            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:
            idx_scores = indexer(x.detach())  # indexer is trained only through its own KL loss
            causal = torch.ones(T, T, dtype=torch.bool, device=x.device).tril()
            dense = dense_warmup or T <= self.topk
            if dense:
                sel = causal.expand(B, T, T)
                y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
            else:
                top = idx_scores.masked_fill(~causal, float("-inf")).topk(self.topk, dim=-1).indices
                sel = torch.zeros(B, T, T, dtype=torch.bool, device=x.device).scatter_(-1, top, True) & causal
                y = F.scaled_dot_product_attention(q, k, v, attn_mask=sel[:, None])
            if self.training:
                # indexer distillation: KL(main attention (head-avg, detached) || indexer) over selected tokens
                with torch.no_grad():
                    s = (q @ k.transpose(-1, -2)).float() * (self.d ** -0.5)
                    p = s.masked_fill(~sel[:, None], float("-inf")).softmax(-1).mean(1)
                logq = idx_scores.masked_fill(~sel, float("-inf")).log_softmax(-1).masked_fill(~sel, 0.0)
                aux = (p * (p.clamp_min(1e-9).log() - logq)).sum(-1).mean()
        return self.o(y.transpose(1, 2).reshape(B, T, C)), aux


class MLP(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.gate = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
        self.up = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
        self.down = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)
        self.down.is_resid = True

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))


class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1 = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.attn = Attention(cfg)
        self.norm2 = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.mlp = MLP(cfg)

    def forward(self, x, cos, sin, indexer=None, dense_warmup=False):
        a, aux = self.attn(self.norm1(x), cos, sin, indexer, dense_warmup)
        x = x + a
        return x + self.mlp(self.norm2(x)), aux


class AblationPreTrainedModel(PreTrainedModel):
    config_class = AblationConfig
    base_model_prefix = "ablation"

    def _init_weights(self, m):
        std = self.config.initializer_range
        if isinstance(m, nn.Linear):
            if getattr(m, "is_resid", False):
                std = std / math.sqrt(2 * self.config.num_hidden_layers)
            nn.init.normal_(m.weight, 0.0, std)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, 0.0, std)
        elif isinstance(m, RMSNorm):
            nn.init.ones_(m.weight)


class AblationForCausalLM(AblationPreTrainedModel, GenerationMixin):
    def __init__(self, config):
        super().__init__(config)
        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([Block(config) for _ in range(config.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        # indexers are created last so the shared weights get identical init for both attention types
        use_dsa = config.attention_type == "dsa"
        self.indexers = nn.ModuleList([DSAIndexer(config) for _ in range(config.num_hidden_layers)] if use_dsa else [])
        self.dsa_warmup = False  # set by the trainer: dense attention while the indexer warms up
        self.aux_loss = None
        self.post_init()

    def get_input_embeddings(self):
        return self.embed

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

    def get_output_embeddings(self):
        return None

    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
        x = self.embed(input_ids)
        hd = self.config.hidden_size // self.config.num_attention_heads
        cos, sin = rope_cache(x.shape[1], hd, self.config.rope_base, x.device)
        aux_total = None
        for i, layer in enumerate(self.layers):
            idx = self.indexers[i] if len(self.indexers) else None
            x, aux = layer(x, cos, sin, idx, self.dsa_warmup)
            if aux is not None:
                aux_total = aux if aux_total is None else aux_total + aux
        self.aux_loss = aux_total
        logits = F.linear(self.norm(x), self.embed.weight)  # tied embeddings
        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[:, :-1].float().reshape(-1, logits.size(-1)), labels[:, 1:].reshape(-1), ignore_index=-100
            )
        return CausalLMOutputWithPast(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        # no KV cache: recompute the full sequence every step
        return {"input_ids": input_ids, "use_cache": False}