"""How a network decides which inputs matter: attention, on a table with a few objects.

Run: python3 site/ml/foundations/examples/pytorch_attention.py

Each scene is a 1 m × 1 m table with 4 objects of different kinds at random places. A request names one kind.
Task "find": where is the requested object? Task "nearest": where is the object nearest to it?
Every network trains with Adam.
"""

import math

import torch
from torch import nn

ITEMS = ["cube", "ball", "cup", "pen", "box", "can", "key", "toy"]
KINDS = len(ITEMS)


def scenes(count, objects, generator):
    """Object tokens (count, objects, 10): kind as one-hot (8 numbers), then x and y in metres.
    Requests (count, 8): the requested kind as one-hot. Answers for both tasks: (count, 2) positions."""
    kind = torch.stack([torch.randperm(KINDS, generator=generator)[:objects] for _ in range(count)])
    position = torch.rand(count, objects, 2, generator=generator)
    asked = torch.randint(0, objects, (count,), generator=generator)
    rows = torch.arange(count)
    tokens = torch.cat([nn.functional.one_hot(kind, KINDS).float(), position], dim=2)
    request = nn.functional.one_hot(kind[rows, asked], KINDS).float()
    distance = (position - position[rows, asked][:, None]).norm(dim=2)
    distance[rows, asked] = float("inf")                           # not itself
    return tokens, request, {"find": position[rows, asked], "nearest": position[rows, distance.argmin(dim=1)]}


def attention(query, keys, values):
    """Scaled dot-product attention, by hand: score every key against the query, softmax, average the values."""
    scores = query @ keys.T / math.sqrt(keys.shape[1])
    weights = scores.softmax(dim=-1)
    return weights @ values, weights, scores


class Lookup(nn.Module):
    """One attention layer: the request asks, the objects answer."""
    def __init__(self, width=16):
        super().__init__()
        self.embed_object = nn.Linear(KINDS + 2, width)
        self.embed_request = nn.Linear(KINDS, width)
        self.attend = nn.MultiheadAttention(width, num_heads=1, batch_first=True)
        self.out = nn.Linear(width, 2)

    def forward(self, tokens, request):
        objects = self.embed_object(tokens)                        # (N, objects, 16)
        query = self.embed_request(request)[:, None]               # (N, 1, 16)
        answer, _ = self.attend(query, objects, objects)           # query, keys, values
        return self.out(answer[:, 0])


class Transformer(nn.Module):
    """The request token and the object tokens attend to each other, `layers` times; the request's token answers."""
    def __init__(self, layers=2, width=32, heads=2):
        super().__init__()
        self.embed_object = nn.Linear(KINDS + 2, width)
        self.embed_request = nn.Linear(KINDS, width)
        layer = nn.TransformerEncoderLayer(width, heads, dim_feedforward=64, dropout=0.0, batch_first=True)
        self.encoder = nn.TransformerEncoder(layer, layers)
        self.out = nn.Linear(width, 2)

    def forward(self, tokens, request):
        sequence = torch.cat([self.embed_request(request)[:, None], self.embed_object(tokens)], dim=1)
        return self.out(self.encoder(sequence)[:, 0])


class Plain(nn.Module):
    """A plain network on all numbers at once: 8 object slots of 11 numbers (10 + 'filled'), plus the request."""
    SLOTS = 8

    def __init__(self, hidden=256):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(self.SLOTS * (KINDS + 3) + KINDS, hidden), nn.ReLU(),
                                 nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 2))

    def forward(self, tokens, request):
        count, objects, _ = tokens.shape
        slots = torch.zeros(count, self.SLOTS, KINDS + 3)
        slots[:, :objects, :KINDS + 2] = tokens
        slots[:, :objects, KINDS + 2] = 1
        return self.net(torch.cat([slots.flatten(1), request], dim=1))


MODELS = {"lookup": Lookup, "transformer-1": lambda: Transformer(1), "transformer-2": lambda: Transformer(2),
          "plain-16": lambda: Plain(16), "plain-256": lambda: Plain(256)}


LEARNING_RATE = {"find/lookup": 0.01, "find/plain-16": 0.003, "find/plain-256": 0.003,     # Adam rates chosen on
                 "nearest/lookup": 0.001, "nearest/transformer-1": 0.01, "nearest/transformer-2": 0.001,  # validation scenes
                 "nearest/plain-256": 0.0003}


def train(model, task, seed=0, lr=None, epochs=None):
    lr = lr or LEARNING_RATE[f"{task}/{model}"]
    epochs = epochs or {"find": 30, "nearest": 40}[task]
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    generator = torch.Generator().manual_seed(seed)
    tokens, request, answers = scenes(20000, 4, generator)
    network = MODELS[model]()
    optimizer = torch.optim.Adam(network.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    for epoch in range(epochs):
        for batch in torch.randperm(len(tokens), generator=generator).split(128):
            loss = loss_fn(network(tokens[batch], request[batch]), answers[task][batch])
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    return network


def miss_cm(network, task, objects, seed=99):
    tokens, request, answers = scenes(2000, objects, torch.Generator().manual_seed(seed))
    with torch.no_grad():
        return 100 * (network(tokens, request) - answers[task]).norm(dim=1).mean().item()


if __name__ == "__main__":
    network = train("lookup", "find")
    tokens, request, answers = scenes(1, 4, torch.Generator().manual_seed(5))
    with torch.no_grad():
        objects = network.embed_object(tokens[0])                  # the 4 object tokens, 16 numbers each
        query = network.embed_request(request)                     # the request token
        wq, wk, wv = network.attend.in_proj_weight.split(16)       # the layer's three projections
        bq, bk, bv = network.attend.in_proj_bias.split(16)
        mixed, weights, scores = attention(query @ wq.T + bq, objects @ wk.T + bk, objects @ wv.T + bv)
        by_hand = network.out(network.attend.out_proj(mixed))[0]
        by_torch = network(tokens, request)[0]
    kinds = [ITEMS[i] for i in tokens[0, :, :KINDS].argmax(dim=1).tolist()]
    print("objects:", kinds, "request:", ITEMS[request[0].argmax().item()])
    print("scores:", [round(v, 3) for v in scores[0].tolist()], " weights:", [round(v, 4) for v in weights[0].tolist()])
    print("answer by hand:", [round(v, 4) for v in by_hand.tolist()], " by the module:", [round(v, 4) for v in by_torch.tolist()],
          " true:", [round(v, 4) for v in answers["find"][0].tolist()])
    shuffled = tokens[:, torch.tensor([2, 0, 3, 1])]
    with torch.no_grad():
        print("same scene, objects in another order:", [round(v, 4) for v in network(shuffled, request)[0].tolist()])

    for model in ("lookup", "plain-16", "plain-256"):
        net = network if model == "lookup" else train(model, "find")
        print(f"find    {model:13s} {sum(p.numel() for p in net.parameters()):6d} parameters: miss {miss_cm(net, 'find', 4):.2f} cm "
              f"with 4 objects, {miss_cm(net, 'find', 6):.2f} cm with 6")
    for model in ("lookup", "transformer-1", "transformer-2"):
        net = train(model, "nearest")
        print(f"nearest {model:13s} {sum(p.numel() for p in net.parameters()):6d} parameters: miss {miss_cm(net, 'nearest', 4):.2f} cm")