"""Original TensorViz implementation of paired RoPE rotations (Su et al., 2021).

One unmasked width-8 attention head isolates position encoding. Standard RoPE
rotates Q and K; V is unchanged. No trained model or context extrapolation claim.
"""
import math
import torch
import torch.nn as nn


class RotaryPositions(nn.Module):
    def __init__(self, offset=0):
        super().__init__()
        self.offset = offset

    def forward(self, x):
        positions = torch.arange(x.shape[1], device=x.device, dtype=x.dtype) + self.offset
        frequencies = torch.pow(10000., -torch.arange(0, 8, 2, device=x.device, dtype=x.dtype) / 8)
        angles = positions.unsqueeze(-1) * frequencies
        pairs = x.reshape(x.shape[0], x.shape[1], 4, 2)
        first, second = pairs[..., 0], pairs[..., 1]
        rotated_first = first * angles.cos() - second * angles.sin()
        rotated_second = first * angles.sin() + second * angles.cos()
        return torch.stack([rotated_first, rotated_second], dim=-1).flatten(-2)


# tensorviz: input=1,6,8
class RotaryAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.query = nn.Linear(8, 8)
        self.key = nn.Linear(8, 8)
        self.value = nn.Linear(8, 8)
        self.query_rope = RotaryPositions()
        self.key_rope = RotaryPositions()
        self.softmax = nn.Softmax(dim=-1)
        self.project = nn.Linear(8, 8)

    def forward(self, x):
        q = self.query_rope(self.query(x))
        k = self.key_rope(self.key(x))
        v = self.value(x)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(8)
        weights = self.softmax(scores)
        mixed = torch.matmul(weights, v)
        return self.project(mixed)


def rotate_at(vector, position):
    return RotaryPositions(position)(vector.reshape(1, 1, 8)).reshape(8)


def verify():
    q = torch.tensor([1., .5, .2, -.3, .7, -.1, .2, .8], dtype=torch.float64)
    k = torch.tensor([.2, 1., -.4, .6, .1, .3, -.5, .4], dtype=torch.float64)
    for m, n in [(0, 0), (1, 3), (3, 1), (11, 13)]:
        rq, rk = rotate_at(q, m), rotate_at(k, n)
        torch.testing.assert_close(rq.norm(), q.norm(), atol=1e-12, rtol=0)
        torch.testing.assert_close(rk.norm(), k.norm(), atol=1e-12, rtol=0)
        torch.testing.assert_close(rq @ rk, q @ rotate_at(k, n-m), atol=1e-12, rtol=0)
    torch.testing.assert_close(rotate_at(q, 1) @ rotate_at(k, 3), rotate_at(q, 11) @ rotate_at(k, 13), atol=1e-12, rtol=0)
    assert not torch.allclose(rotate_at(q, 1) @ rotate_at(k, 3), rotate_at(q, 1) @ rotate_at(k, 5))
    torch.manual_seed(0)
    model = RotaryAttention().double().eval()
    x = torch.randn(1, 6, 8, dtype=torch.float64)
    with torch.no_grad():
        for projection in [model.query, model.key]:
            projection.weight.zero_()
            projection.bias.zero_()
        expected = model.project(model.value(x).mean(1, keepdim=True).expand_as(x))
        torch.testing.assert_close(model(x), expected)
    return ["Every paired rotation preserves the vector norm", "Shifting both positions equally preserves the Q/K dot product for fixed content", "Changing relative position changes the score in this example", "Zero Q/K yields a uniform mixture of unrotated values"]


def experiment():
    q = torch.tensor([1., .5, .2, -.3, .7, -.1, .2, .8], dtype=torch.float64)
    k = torch.tensor([.2, 1., -.4, .6, .1, .3, -.5, .4], dtype=torch.float64)
    cases = []
    for identity, label, m, n, note in [
        ("same-position", "Both positions 0", 0, 0, "At position zero, all rotation angles are zero. The query and key vectors remain unchanged."),
        ("gap-two", "Query 1 · key 3", 1, 3, "The key is two positions after the query. Each feature pair rotates at its own frequency."),
        ("shift-both", "Shift both by 10: query 11 · key 13", 11, 13, "The absolute angles change, but the relative offset remains two. For these fixed query/key contents, the score matches query 1 / key 3."),
        ("gap-four", "Query 1 · key 5", 1, 5, "Changing the relative offset to four changes the score. RoPE supplies position dependence; this does not prove extrapolation quality."),
        ("reverse", "Query 3 · key 1", 3, 1, "Reversing the relative direction can change the score even when the distance has the same magnitude."),
    ]:
        rq, rk = rotate_at(q, m), rotate_at(k, n)
        cases.append({"id": identity, "label": label, "note": note, "target": "query_rope", "rotation": {
            "queryPosition": m, "keyPosition": n, "query": q[:2].tolist(), "key": k[:2].tolist(), "rotatedQuery": rq[:2].tolist(), "rotatedKey": rk[:2].tolist(),
        }, "vectors": [{"label": "Rotated Q · all 8 features", "values": rq.tolist()}, {"label": "Rotated K · all 8 features", "values": rk.tolist()}], "metrics": [
            {"label": "Key position − query position", "value": n-m},
            {"label": "Scaled Q/K score · all 8 features", "value": (rq @ rk / math.sqrt(8)).item()},
            {"label": "Query norm", "value": rq.norm().item()},
        ]})
    return {"kind": "rotary", "title": "Move position, keep content fixed.", "description": "Recorded float64 rotations of the same synthetic Q/K vectors. The diagram shows their first feature pair (one radian per position); the score uses all four pairs.", "controlLabel": "Token positions", "cases": cases}