# coding=utf-8
"""Gated Cross-Layer Attention (GCLA).

Implements Section VII of the Wiola paper. Each decoder layer performs:

  * GQA self-attention with SRPE on the local sequence, and
  * cross-attention to compressed summaries of up to ``Lambda`` preceding
    layers (supplied by the model as ``context_summaries``),

blended by a scalar gate ``beta = sigmoid(phi)`` and modulated by a sigmoid
output gate ``G``:

    O   = (1 - beta) * O_self + beta * O_ctx
    A   = (G * concat(O)) W_O

The context tensor is provided *per query position* as a causal cumulative mean
of prior-layer outputs (see :class:`WiolaModel`), so a cached incremental decode
reproduces a full forward pass exactly.
"""

import math
from typing import Optional, Tuple

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

from .srpe import SpiralRotaryEmbedding, apply_srpe


def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """Expand [B, H_kv, T, d] -> [B, H_kv*n_rep, T, d] (GQA)."""
    if n_rep == 1:
        return x
    b, kv, t, d = x.shape
    x = x[:, :, None, :, :].expand(b, kv, n_rep, t, d)
    return x.reshape(b, kv * n_rep, t, d)


class GatedCrossLayerAttention(nn.Module):
    def __init__(self, config, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.hidden_size = 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
        self.n_rep = self.num_heads // self.num_kv_heads
        self.lookback = config.gcla_lookback

        q_dim = self.num_heads * self.head_dim
        kv_dim = self.num_kv_heads * self.head_dim

        self.q_proj = nn.Linear(self.hidden_size, q_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, kv_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, kv_dim, bias=False)
        self.o_proj = nn.Linear(q_dim, self.hidden_size, bias=False)

        # Cross-layer context projections.
        self.k_ctx_proj = nn.Linear(self.hidden_size, kv_dim, bias=False)
        self.v_ctx_proj = nn.Linear(self.hidden_size, kv_dim, bias=False)

        # Sigmoid output gate.
        self.gate_proj = nn.Linear(self.hidden_size, q_dim, bias=False)

        # Scalar blend gate beta = sigmoid(phi).
        self.blend_logit = nn.Parameter(torch.tensor(float(config.gcla_gate_init)))

        self.srpe = SpiralRotaryEmbedding(
            head_dim=self.head_dim,
            max_position_embeddings=config.max_position_embeddings,
            theta=config.srpe_theta,
            spiral_divisor=config.srpe_spiral_divisor,
            radial_amplitude=config.srpe_radial_amplitude,
            radial_frequency=config.srpe_radial_frequency,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,  # [B, T, d]
        position_ids: torch.Tensor,  # [B, T]
        attention_mask: Optional[torch.Tensor],  # [B, 1, T, T_k] additive
        context_summaries: Optional[torch.Tensor] = None,  # [B, T, Lambda, d]
        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
    ):
        bsz, q_len, _ = hidden_states.shape

        q = self.q_proj(hidden_states)
        k = self.k_proj(hidden_states)
        v = self.v_proj(hidden_states)

        q = q.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

        # SRPE rotation on q/k (per head).
        cos, sin = self.srpe(hidden_states, position_ids)
        q, k = apply_srpe(q, k, cos, sin)

        # Concatenate cached KV (incremental decoding).
        # Concatenate cached KV (incremental decoding).
        if (
            past_key_value is not None
            and past_key_value[0] is not None  # ← guard against placeholder None
            and past_key_value[1] is not None
        ):
            past_k, past_v = past_key_value
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)
        present = (k, v) if use_cache else None

        # GQA expansion.
        k_rep = repeat_kv(k, self.n_rep)
        v_rep = repeat_kv(v, self.n_rep)

        scale = 1.0 / math.sqrt(self.head_dim)
        scores = torch.matmul(q, k_rep.transpose(-1, -2)) * scale  # [B,H,T,T_k]
        if attention_mask is not None:
            scores = scores + attention_mask
        attn = F.softmax(scores, dim=-1, dtype=torch.float32).to(q.dtype)
        o_self = torch.matmul(attn, v_rep)  # [B,H,T,dh]

        # Cross-layer context attention.
        if context_summaries is not None and context_summaries.shape[2] > 0:
            lam = context_summaries.shape[2]
            k_ctx = self.k_ctx_proj(context_summaries)  # [B,T,Lam,kv_dim]
            v_ctx = self.v_ctx_proj(context_summaries)
            k_ctx = k_ctx.view(bsz, q_len, lam, self.num_kv_heads, self.head_dim)
            v_ctx = v_ctx.view(bsz, q_len, lam, self.num_kv_heads, self.head_dim)
            # Expand kv heads to full head count.
            k_ctx = k_ctx.repeat_interleave(self.n_rep, dim=3)  # [B,T,Lam,H,dh]
            v_ctx = v_ctx.repeat_interleave(self.n_rep, dim=3)
            # scores[b,h,t,l] = q[b,h,t,:] . k_ctx[b,t,l,h,:]
            ctx_scores = torch.einsum("bhtd,btlhd->bhtl", q, k_ctx) * scale
            ctx_attn = F.softmax(ctx_scores, dim=-1, dtype=torch.float32).to(q.dtype)
            o_ctx = torch.einsum("bhtl,btlhd->bhtd", ctx_attn, v_ctx)
            beta = torch.sigmoid(self.blend_logit)
            o = (1.0 - beta) * o_self + beta * o_ctx
        else:
            o = o_self

        # Merge heads.
        o = o.transpose(1, 2).contiguous().view(bsz, q_len, self.num_heads * self.head_dim)

        # Sigmoid output gate.
        g = torch.sigmoid(self.gate_proj(hidden_states))
        o = self.o_proj(g * o)
        return o, present