"""Improved sublinear-attention anchor selection for AGILLM-4.

V2 used by the live AGILLM-4 DBlock line:
- fixes gathered ALiBi distance so past causal keys receive distance penalty
- suppresses local/anchor duplicate candidates before softmax
- uses hybrid full-span + recent-tail anchors under the same max-anchor budget
- exposes first-token attention sinks as `--sublinear_sinks`
- includes optional pooled landmark K/V summaries behind `--sublinear_pooled_landmarks`
"""
import torch


def select_hybrid_anchors(k_len, stride, max_anchors, sinks=4, recent_anchors=-1, device='cpu'):
    """Full-span + recent-tail landmark positions, plus attention sinks."""
    device = torch.device(device)
    start = stride - 1
    if stride <= 0 or max_anchors <= 0 or start >= k_len:
        anchors = torch.empty(0, device=device, dtype=torch.long)
    else:
        all_anchors = torch.arange(start, k_len, stride, device=device, dtype=torch.long)
        if all_anchors.numel() <= max_anchors:
            anchors = all_anchors
        else:
            if recent_anchors < 0:
                recent_anchors = max_anchors // 2
            recent_budget = min(max(0, int(recent_anchors)), max_anchors)
            span_budget = max(0, max_anchors - recent_budget)
            parts = []
            if span_budget > 0:
                sel = torch.linspace(0, all_anchors.numel() - 1, span_budget, device=device).round().long().unique()
                parts.append(all_anchors[sel])
            if recent_budget > 0:
                parts.append(all_anchors[-recent_budget:])
            anchors = torch.cat(parts).unique() if parts else torch.empty(0, device=device, dtype=torch.long)
    if sinks > 0 and k_len > 0:
        sink_idx = torch.arange(min(int(sinks), k_len), device=device, dtype=torch.long)
        anchors = torch.cat([sink_idx, anchors]).unique() if anchors.numel() else sink_idx
    return anchors


def local_anchor_valid(q_pos, anchors, window, k_len):
    """False where an anchor is already inside that query's local window."""
    anchor_idx = anchors.view(1, -1).expand(q_pos.numel(), -1)
    local_lo = (q_pos - window).clamp_min(0).view(-1, 1)
    local_hi = (q_pos + window).clamp_max(max(0, k_len - 1)).view(-1, 1)
    return (anchor_idx < local_lo) | (anchor_idx > local_hi)


if __name__ == '__main__':
    N, window, stride, maxA, sinks, recent = 32768, 128, 128, 128, 4, 64
    old_all = list(range(stride - 1, N, stride))
    old = sorted(set(range(N - window - 1, N)) | set(old_all[-maxA:]))
    new = select_hybrid_anchors(N, stride, maxA, sinks, recent).tolist()
    print(f'N={N} stride={stride} maxA={maxA} sinks={sinks} recent={recent}')
    print(f'OLD anchor coverage: {min(old)}..{max(old)} first_half={sum(x < N//2 for x in old)}')
    print(f'NEW anchor coverage: {min(new)}..{max(new)} first_half={sum(x < N//2 for x in new)} recent_tail={sum(x >= N-8192 for x in new)}')
    q = torch.tensor([N - 1])
    print(f'duplicate-suppressed anchors for final query: {int(local_anchor_valid(q, torch.tensor(new), window, N).sum())}/{len(new)}')