"""
Differential Transformer (PRODUCTION - causal-safe)
====================================================

Microsoft 2410.05258 (ICLR 2025)
Cancels noise via dual-softmax attention map subtraction.

Core formula:
  DiffAttn(X) = (softmax(Q1 K1^T / sqrt(d)) - λ × softmax(Q2 K2^T / sqrt(d))) V
  λ = exp(λ_q1 · λ_k1) - exp(λ_q2 · λ_k2) + λ_init

Implementation notes:
  - GroupNorm (sequence-spanning) → LayerNorm(head_dim) (per-token, causal-safe)
  - Causal triu mask explicit (was relying on SDPA, but split + softmax direct now)
  - forward signature: layer_idx kwarg + (Tensor, past_kv) tuple return for modeling compat

Hyperparameter (configuration_aether_v2_7way.py):
  - diff_lambda_init = 0.8
"""
from __future__ import annotations
import math
from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F


class DifferentialAttention(nn.Module):
    """λ-gated dual-softmax differential attention with causal mask + LayerNorm."""

    def __init__(self, config, layer_idx: int = 0):
        super().__init__()
        self.layer_idx = layer_idx
        self.h = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        # Differential splits head_dim in half
        self.diff_head_dim = self.head_dim // 2

        # Layer-depth lambda init (paper Eq. 3): λ_init = 0.8 - 0.6 * exp(-0.3 * (depth - 1))
        # Use config base + depth decay if provided.
        base = getattr(config, "diff_lambda_init", 0.8)
        decay = getattr(config, "diff_lambda_layer_decay", 0.6)
        # Simple per-layer init
        self.lambda_init = base - decay * math.exp(-0.3 * max(layer_idx - 1, 0))

        # Q/K/V/O projections
        self.q_proj = nn.Linear(self.h, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.h, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.h, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.h, bias=False)

        # Lambda parameters per-head
        self.lambda_q1 = nn.Parameter(torch.zeros(self.num_heads, self.diff_head_dim).normal_(0, 0.1))
        self.lambda_k1 = nn.Parameter(torch.zeros(self.num_heads, self.diff_head_dim).normal_(0, 0.1))
        self.lambda_q2 = nn.Parameter(torch.zeros(self.num_heads, self.diff_head_dim).normal_(0, 0.1))
        self.lambda_k2 = nn.Parameter(torch.zeros(self.num_heads, self.diff_head_dim).normal_(0, 0.1))

        # ★ FIX 5/7: LayerNorm(head_dim) — per-token norm, causal-safe.
        # Was nn.GroupNorm(num_groups=num_heads, num_channels=num_heads*head_dim) which
        # normalized over sequence dim and leaked future info.
        self.subln = nn.LayerNorm(self.head_dim, eps=getattr(config, "rms_norm_eps", 1e-6))

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value=None,
        use_cache: bool = False,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[object]]:
        B, S, _ = hidden_states.shape

        q = self.q_proj(hidden_states).view(B, S, self.num_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim)
        v = self.v_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim)

        # AetherCache: cache full k,v (pre-split) in HF layout [B, kv_h, S, D]
        if past_key_value is not None:
            kc, vc = past_key_value.update(k.transpose(1, 2), v.transpose(1, 2), getattr(self, "cache_idx", self.layer_idx))
            k = kc.transpose(1, 2)
            v = vc.transpose(1, 2)
        S_kv = k.shape[1]

        # Split Q, K halves
        q1, q2 = q.split(self.diff_head_dim, dim=-1)  # [B, S, H, d/2]
        k1, k2 = k.split(self.diff_head_dim, dim=-1)  # [B, S_kv, kv_h, d/2]

        # GQA repeat
        repeat = self.num_heads // self.num_kv_heads
        k1 = k1.repeat_interleave(repeat, dim=2)
        k2 = k2.repeat_interleave(repeat, dim=2)
        v = v.repeat_interleave(repeat, dim=2)

        # Transpose to [B, H, S, d]
        q1 = q1.transpose(1, 2)
        q2 = q2.transpose(1, 2)
        k1 = k1.transpose(1, 2)
        k2 = k2.transpose(1, 2)
        v = v.transpose(1, 2)

        scale = 1.0 / math.sqrt(self.diff_head_dim)
        scores1 = torch.matmul(q1, k1.transpose(-2, -1)) * scale
        scores2 = torch.matmul(q2, k2.transpose(-2, -1)) * scale

        # ★ FIX 5/7: causal mask (triu(1) blocks future)
        # AetherCache: query i is at absolute pos (S_kv - S + i) -> shift the diagonal
        causal = torch.ones(S, S_kv, device=q1.device, dtype=torch.bool).triu(1 + S_kv - S)
        scores1 = scores1.masked_fill(causal, float("-inf"))
        scores2 = scores2.masked_fill(causal, float("-inf"))

        attn1 = F.softmax(scores1, dim=-1)
        attn2 = F.softmax(scores2, dim=-1)

        # Lambda per-head (paper Eq. 3)
        lam_q1k1 = (self.lambda_q1 * self.lambda_k1).sum(dim=-1)  # [H]
        lam_q2k2 = (self.lambda_q2 * self.lambda_k2).sum(dim=-1)
        lam = torch.exp(lam_q1k1) - torch.exp(lam_q2k2) + self.lambda_init  # [H]
        lam = lam.view(1, -1, 1, 1)

        diff_attn = attn1 - lam * attn2

        out = torch.matmul(diff_attn, v)  # [B, H, S, head_dim]
        # ★ FIX 5/7: per-token LayerNorm (was sequence-spanning GroupNorm = leak)
        out = self.subln(out)
        out = out * (1 - self.lambda_init)

        out = out.transpose(1, 2).reshape(B, S, -1)
        return self.o_proj(out), past_key_value


__all__ = ["DifferentialAttention"]