"""Original reduced additive-attention decoder step, Bahdanau et al. (2014).

Bidirectional GRU annotations, a learned additive alignment score, context sum,
and one GRUCell decoder update. No trained translation or original maxout head.
"""
import torch
import torch.nn as nn


class AdditiveAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.annotation_projection = nn.Linear(8, 8, bias=False)
        self.state_projection = nn.Linear(8, 8, bias=False)
        self.tanh = nn.Tanh()
        self.score = nn.Linear(8, 1, bias=False)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, annotations, previous_state):
        features = self.tanh(
            self.annotation_projection(annotations)
            + self.state_projection(previous_state).unsqueeze(1)
        )
        weights = self.softmax(self.score(features))
        return (weights * annotations).sum(1)


class AdditiveDecoderStep(nn.Module):
    def __init__(self):
        super().__init__()
        self.source_embedding = nn.Embedding(16, 4)
        self.encoder = nn.GRU(4, 4, batch_first=True, bidirectional=True)
        self.alignment = AdditiveAttention()
        self.target_embedding = nn.Embedding(16, 4)
        self.decoder = nn.GRUCell(12, 8)
        self.output_head = nn.Linear(8, 16)

    def forward(self, source_tokens, previous_state, previous_token):
        annotations, _ = self.encoder(self.source_embedding(source_tokens))
        context = self.alignment(annotations, previous_state)
        embedded = self.target_embedding(previous_token)
        next_state = self.decoder(
            torch.cat([embedded, context], dim=-1), previous_state
        )
        return self.output_head(next_state)


def alignment_weights(attention, annotations, state):
    energy = attention.score(
        torch.tanh(
            attention.annotation_projection(annotations)
            + attention.state_projection(state).unsqueeze(1)
        )
    )
    return energy.softmax(1)


def verify():
    torch.manual_seed(21)
    model = AdditiveDecoderStep().double().eval()
    tokens = torch.tensor([[2, 3, 4, 5, 6]])
    state = torch.randn(1, 8, dtype=torch.float64)
    with torch.no_grad():
        annotations, _ = model.encoder(model.source_embedding(tokens))
        weights = alignment_weights(model.alignment, annotations, state)
        torch.testing.assert_close(
            weights.sum(1), torch.ones(1, 1, dtype=torch.float64)
        )
        independent = sum(weights[:, i] * annotations[:, i] for i in range(5))
        torch.testing.assert_close(model.alignment(annotations, state), independent)
        changed = tokens.clone()
        changed[0, -1] = 7
        altered, _ = model.encoder(model.source_embedding(changed))
        assert not torch.allclose(annotations[:, 0], altered[:, 0])
        assert not torch.allclose(
            weights, alignment_weights(model.alignment, annotations, state + 3)
        )
        model.alignment.score.weight.zero_()
        torch.testing.assert_close(
            model.alignment(annotations, state), annotations.mean(1)
        )
    return [
        "Alignment weights are positive and sum to one over source positions",
        "The context equals an independent weighted sum of bidirectional annotations",
        "A changed final source token can affect the first annotation through the backward GRU",
        "Changing decoder state changes alignment; zero scores give the mean source annotation",
    ]


def experiment():
    torch.manual_seed(21)
    model = AdditiveDecoderStep().double().eval()
    tokens = torch.tensor([[2, 3, 4, 5, 6]])
    cases = []
    with torch.no_grad():
        annotations, _ = model.encoder(model.source_embedding(tokens))
        for identity, label, state in [
            (
                "zero",
                "Zero previous decoder state",
                torch.zeros(1, 8, dtype=torch.float64),
            ),
            (
                "positive",
                "Positive previous decoder state",
                torch.ones(1, 8, dtype=torch.float64) * 3,
            ),
            (
                "negative",
                "Negative previous decoder state",
                torch.ones(1, 8, dtype=torch.float64) * -3,
            ),
        ]:
            weights = alignment_weights(model.alignment, annotations, state)
            context = model.alignment(annotations, state)
            cases.append(
                {
                    "id": identity,
                    "label": label,
                    "target": "alignment.softmax",
                    "note": "The same five synthetic source IDs produce fixed bidirectional annotations. Only the previous decoder state changes here. The learned additive score determines the weighting used by this step. These untrained alignments do not map translated words or demonstrate translation quality.",
                    "matrices": [
                        {
                            "label": "Source annotations · five positions × eight features",
                            "values": annotations[0].tolist(),
                        },
                        {
                            "label": "Alignment · one decoder step × five source positions",
                            "values": [weights[0, :, 0].tolist()],
                        },
                    ],
                    "vectors": [
                        {
                            "label": "Previous decoder state",
                            "values": state[0].tolist(),
                        },
                        {"label": "Weighted context", "values": context[0].tolist()},
                    ],
                    "metrics": [
                        {
                            "label": "Alignment weight sum",
                            "value": weights.sum().item(),
                        },
                        {
                            "label": "Most weighted source position · zero-based",
                            "value": weights.argmax(1).item(),
                        },
                    ],
                }
            )
    return {
        "kind": "matrices",
        "title": "Ask the source a different question at each step.",
        "description": "Recorded additive alignment for a fixed source sequence under three decoder states. In translation, the evolving decoder state makes this context change at every generated token. This lesson captures one such step.",
        "controlLabel": "Previous decoder state",
        "cases": cases,
    }