"""Original reduced MHA/MQA/GQA comparison (Shazeer 2019; Ainslie et al. 2023).

Eight query heads, two features per head, causal attention. Repeated K/V tensors
are a transparent CPU reference, not an optimized grouped-attention kernel.
"""
import math
import torch
import torch.nn as nn


class SharedKVAttention(nn.Module):
    def __init__(self, kv_heads):
        super().__init__()
        self.kv_heads = kv_heads
        self.query = nn.Linear(16, 16, bias=False)
        self.key = nn.Linear(16, kv_heads * 2, bias=False)
        self.value = nn.Linear(16, kv_heads * 2, bias=False)
        self.softmax = nn.Softmax(dim=-1)
        self.project = nn.Linear(16, 16, bias=False)

    def forward(self, x):
        b, t, _ = x.shape
        q = self.query(x).reshape(b, t, 8, 2).transpose(1, 2)
        k = self.key(x).reshape(b, t, self.kv_heads, 2).transpose(1, 2)
        v = self.value(x).reshape(b, t, self.kv_heads, 2).transpose(1, 2)
        k = k.repeat_interleave(8 // self.kv_heads, dim=1)
        v = v.repeat_interleave(8 // self.kv_heads, dim=1)
        scores = q @ k.transpose(-2, -1) / math.sqrt(2)
        mask = torch.ones(t, t, dtype=torch.bool, device=x.device).triu(1)
        weights = self.softmax(scores.masked_fill(mask, float('-inf')))
        mixed = (weights @ v).transpose(1, 2).reshape(b, t, 16)
        return self.project(mixed)


# tensorviz: input=1,4,16
class HeadSharingComparison(nn.Module):
    def __init__(self):
        super().__init__()
        self.mha = SharedKVAttention(8)
        self.gqa = SharedKVAttention(2)
        self.mqa = SharedKVAttention(1)

    def forward(self, x):
        return torch.stack([self.mha(x), self.gqa(x), self.mqa(x)])


def cache_bytes(tokens, kv_heads):
    return 2 * 1 * tokens * kv_heads * 2 * 4


def verify():
    torch.manual_seed(0)
    x = torch.randn(1, 4, 16, dtype=torch.float64)
    for count in [8, 2, 1]:
        model = SharedKVAttention(count).double().eval()
        dense = SharedKVAttention(8).double().eval()
        with torch.no_grad():
            dense.query.weight.copy_(model.query.weight)
            dense.project.weight.copy_(model.project.weight)
            for compact, expanded in [(model.key, dense.key), (model.value, dense.value)]:
                shared = compact.weight.reshape(count, 2, 16)
                expanded.weight.copy_(shared.repeat_interleave(8 // count, 0).reshape(16, 16))
            torch.testing.assert_close(model(x), dense(x))
            changed = x.clone(); changed[:, 3] += 20
            torch.testing.assert_close(model(x)[:, :3], model(changed)[:, :3])
            # Independent last-query calculation with compact K/V cache.
            q = model.query(x[:, -1:]).reshape(1, 1, 8, 2).transpose(1, 2)
            k = model.key(x).reshape(1, 4, count, 2).transpose(1, 2)
            v = model.value(x).reshape(1, 4, count, 2).transpose(1, 2)
            heads = []
            for h in range(8):
                group = h // (8 // count)
                p = torch.softmax(q[:, h] @ k[:, group].transpose(-2, -1) / math.sqrt(2), -1)
                heads.append(p @ v[:, group])
            cached = model.project(torch.stack(heads, dim=2).reshape(1, 1, 16))
            torch.testing.assert_close(cached, model(x)[:, -1:])
    assert cache_bytes(2048, 8) == 262144
    assert cache_bytes(2048, 2) * 4 == cache_bytes(2048, 8)
    assert cache_bytes(2048, 1) * 8 == cache_bytes(2048, 8)
    return ['MHA with duplicated K/V weights matches grouped attention for 8, 2 and 1 KV heads', 'Future token changes cannot affect earlier causal outputs', 'Compact per-group K/V calculation reproduces the final query output', 'Analytical KV bytes scale with the number of KV heads']


def experiment():
    torch.manual_seed(7)
    dense = SharedKVAttention(8).double().eval()
    x = torch.randn(1, 4, 16, dtype=torch.float64)
    cases = []
    for name, count in [('MHA', 8), ('GQA', 2), ('MQA', 1)]:
        model = SharedKVAttention(count).double().eval()
        with torch.no_grad():
            model.query.weight.copy_(dense.query.weight)
            model.project.weight.copy_(dense.project.weight)
            for expanded, compact in [(dense.key, model.key), (dense.value, model.value)]:
                grouped = expanded.weight.reshape(count, 8 // count, 2, 16).mean(1)
                compact.weight.copy_(grouped.reshape(count * 2, 16))
            output = model(x)[0, -1].tolist()
        for context in [128, 2048]:
            cases.append({'id':f'{name.lower()}-{context}', 'label':f'{name} · {count} KV heads · {context} cached tokens',
                'note':'All eight query heads remain distinct. Lines show which K/V head each query consults. K/V weights are group means of the same MHA weights, so outputs need not match. Cache size is an analytical estimate for one layer, batch 1, head width 2 and float32, excluding temporary expansion and framework overhead.',
                'target':name.lower(), 'routing':{'sourceLabel':'Query heads', 'destinationLabel':'Shared key/value heads', 'sources':[f'Q {h}' for h in range(8)], 'destinations':[f'KV {h}' for h in range(count)], 'weights':[[float(g == h // (8 // count)) for g in range(count)] for h in range(8)]},
                'vectors':[{'label':'Last token output · four-token forward', 'values':output}],
                'metrics':[{'label':'Analytical K + V cache · bytes', 'value':cache_bytes(context, count)}, {'label':'KV heads', 'value':count}, {'label':'Query heads per KV head', 'value':8 // count}]})
    return {'kind':'routing', 'title':'Keep queries, share keys and values.', 'description':'Head connections from the recorded configuration; cache bytes calculated as 2 × batch × tokens × KV heads × head width × bytes per value. This CPU reference does not measure decoding speed.', 'controlLabel':'Head and cache configuration', 'cases':cases}