"""Original reduced Evoformer coupling subsystem, not protein structure prediction.

Implements supplementary Algorithms 7, 10, 11 and a pair transition: MSA row
attention with pair bias, outer-product mean, outgoing triangle multiplication.
Omitted: MSA column attention/transition, incoming triangles, triangle attention,
templates, recycling, extra-MSA stack and the entire structure module.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


class MSARowAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.msa_norm = nn.LayerNorm(8)
        self.pair_norm = nn.LayerNorm(4)
        self.query = nn.Linear(8, 8, bias=False)
        self.key = nn.Linear(8, 8, bias=False)
        self.value = nn.Linear(8, 8, bias=False)
        self.pair_bias = nn.Linear(4, 2, bias=False)
        self.gate = nn.Linear(8, 8)
        self.softmax = nn.Softmax(-1)
        self.project = nn.Linear(8, 8)

    def forward(self, msa, pair):
        s, r, _ = msa.shape
        m = self.msa_norm(msa)
        q = self.query(m).reshape(s, r, 2, 4).transpose(1, 2)
        k = self.key(m).reshape(s, r, 2, 4).transpose(1, 2)
        v = self.value(m).reshape(s, r, 2, 4).transpose(1, 2)
        bias = self.pair_bias(self.pair_norm(pair)).permute(2, 0, 1)
        weights = self.softmax(q @ k.transpose(-1, -2) / 2 + bias[None])
        mixed = (weights @ v).transpose(1, 2).reshape(s, r, 8)
        return self.project(self.gate(m).sigmoid() * mixed)


class OuterProductMean(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = nn.LayerNorm(8)
        self.left = nn.Linear(8, 3)
        self.right = nn.Linear(8, 3)
        self.project = nn.Linear(9, 4)

    def forward(self, msa):
        m = self.norm(msa)
        a = self.left(m)
        b = self.right(m)
        pairs = torch.einsum("sic,sjd->ijcd", a, b) / msa.shape[0]
        return self.project(pairs.flatten(2))


class OutgoingTriangle(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = nn.LayerNorm(4)
        self.left = nn.Linear(4, 4)
        self.right = nn.Linear(4, 4)
        self.left_gate = nn.Linear(4, 4)
        self.right_gate = nn.Linear(4, 4)
        self.output_gate = nn.Linear(4, 4)
        self.sum_norm = nn.LayerNorm(4)
        self.project = nn.Linear(4, 4)

    def forward(self, pair):
        z = self.norm(pair)
        a = self.left_gate(z).sigmoid() * self.left(z)
        b = self.right_gate(z).sigmoid() * self.right(z)
        triangles = torch.einsum("ikc,jkc->ijc", a, b)
        return self.output_gate(z).sigmoid() * self.project(self.sum_norm(triangles))


class PairTransition(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = nn.LayerNorm(4)
        self.up = nn.Linear(4, 16)
        self.relu = nn.ReLU()
        self.down = nn.Linear(16, 4)

    def forward(self, pair):
        return self.down(self.relu(self.up(self.norm(pair))))


class EvoformerCoupling(nn.Module):
    def __init__(self):
        super().__init__()
        self.msa_row = MSARowAttention()
        self.outer_product = OuterProductMean()
        self.triangle = OutgoingTriangle()
        self.pair_transition = PairTransition()

    def forward(self, msa, pair):
        msa = msa + self.msa_row(msa, pair)
        pair = pair + self.outer_product(msa)
        pair = pair + self.triangle(pair)
        return pair + self.pair_transition(pair)


def setup():
    torch.manual_seed(61)
    return (
        EvoformerCoupling().double(),
        torch.randn(3, 4, 8, dtype=torch.float64),
        torch.randn(4, 4, 4, dtype=torch.float64),
    )


def verify():
    model, msa, pair = setup()
    base = model(msa, pair)
    torch.testing.assert_close(model(msa[[2, 0, 1]], pair), base)
    perm = torch.tensor([2, 0, 3, 1])
    torch.testing.assert_close(
        model(msa[:, perm], pair[perm][:, perm]), base[perm][:, perm]
    )
    op = model.outer_product
    m = op.norm(msa)
    a = op.left(m)
    b = op.right(m)
    manual = torch.stack(
        [
            torch.stack(
                [
                    torch.stack([torch.outer(a[s, i], b[s, j]) for s in range(3)])
                    .mean(0)
                    .flatten()
                    for j in range(4)
                ]
            )
            for i in range(4)
        ]
    )
    torch.testing.assert_close(op(msa), op.project(manual))
    tri = model.triangle
    z = tri.norm(pair)
    left = tri.left_gate(z).sigmoid() * tri.left(z)
    right = tri.right_gate(z).sigmoid() * tri.right(z)
    manual = torch.stack(
        [
            torch.stack(
                [
                    sum(
                        (left[i, k] * right[j, k] for k in range(4)),
                        torch.zeros(4, dtype=pair.dtype),
                    )
                    for j in range(4)
                ]
            )
            for i in range(4)
        ]
    )
    torch.testing.assert_close(
        tri(pair), tri.output_gate(z).sigmoid() * tri.project(tri.sum_norm(manual))
    )
    changed = msa.clone()
    changed[1, 2] += torch.arange(8, dtype=msa.dtype)
    assert not torch.allclose(model(changed, pair), base)
    msa.requires_grad_()
    pair.requires_grad_()
    model(msa, pair).square().sum().backward()
    assert msa.grad.norm() > 0 and pair.grad.norm() > 0
    assert model.msa_row.pair_bias.weight.grad.norm() > 0
    return [
        "Reordering aligned sequences preserves pair outputs; relabeling residues permutes both pair axes consistently",
        "Outer-product mean matches an explicit per-sequence, per-residue-pair calculation",
        "Outgoing triangle multiplication matches an explicit sum over shared third residues k",
        "An MSA edit changes pair features; gradients flow through both streams and the pair-to-attention bias",
    ]


def experiment():
    model, msa, pair = setup()
    cases = []
    with torch.no_grad():
        for identity, label, stage, edit in [
            ("msa", "After pair-biased MSA attention", 0, False),
            ("outer", "After outer-product mean", 1, False),
            ("triangle", "After outgoing triangle update", 2, False),
            ("final", "After pair transition", 3, False),
            ("edit", "Edit one aligned residue · final state", 3, True),
        ]:
            m = msa.clone()
            if edit:
                m[1, 2] += torch.arange(8, dtype=msa.dtype)
            updated = m + model.msa_row(m, pair)
            outer = pair + model.outer_product(updated)
            triangles = outer + model.triangle(outer)
            final = triangles + model.pair_transition(triangles)
            selected = [pair, outer, triangles, final][stage]
            cases.append(
                dict(
                    id=identity,
                    label=label,
                    target=["msa_row", "outer_product", "triangle", "pair_transition"][
                        stage
                    ],
                    note="Three aligned sequences describe four residue positions. MSA features and pair features are synthetic continuous vectors, not amino-acid identities or distances. "
                    + [
                        "Pair features bias attention along each aligned sequence; the pair tensor has not yet been updated.",
                        "The mean outer product over aligned sequences transfers MSA information into each residue pair.",
                        "The outgoing update for pair (i,j) sums products from (i,k) and (j,k) over every third residue k.",
                        "A pointwise feed-forward transition changes the pair channels after the triangle update.",
                    ][stage]
                    + (
                        " One MSA vector at sequence 1, residue 2 was changed before this chain."
                        if edit
                        else ""
                    ),
                    matrices=[
                        dict(
                            label="Updated first MSA sequence · residues × features",
                            values=updated[0].tolist(),
                        ),
                        dict(
                            label="Current pair feature · channel 0",
                            values=selected[:, :, 0].tolist(),
                        ),
                        dict(
                            label="Change from initial pair · channel 0",
                            values=(selected - pair)[:, :, 0].tolist(),
                        ),
                    ],
                    vectors=[
                        dict(
                            label="Pair (0, 1) feature vector",
                            values=selected[0, 1].tolist(),
                        )
                    ],
                    metrics=[
                        dict(label="Aligned sequences", value=3),
                        dict(label="Residue positions", value=4),
                        dict(
                            label="Pair change · L2 norm",
                            value=(selected - pair).norm().item(),
                        ),
                    ],
                )
            )
    return dict(
        kind="matrices",
        title="Exchange sequence and pair information.",
        description="Recorded Evoformer coupling on synthetic MSA and pair tensors. Pair features are learned representation channels, not a distance map or predicted structure.",
        controlLabel="Evoformer stage",
        cases=cases,
    )