"""How a network reads an image: a ball on a 16 × 16 screen, and three networks that report where it is.

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

Training images never show the ball in the bottom-right corner (centre x > 9 and y > 9). The test asks for
positions there too. All three networks train with Adam.
"""

import torch
from torch import nn

SIZE, RADIUS = 16, 2.0
ROWS, COLS = torch.meshgrid(torch.arange(SIZE) + 0.5, torch.arange(SIZE) + 0.5, indexing="ij")   # pixel centres


def draw(centres):
    """One image per (x, y) centre: brightness 1 inside the ball, 0 outside, a soft edge one pixel wide."""
    distance = ((COLS - centres[:, 0, None, None]) ** 2 + (ROWS - centres[:, 1, None, None]) ** 2).sqrt()
    return (RADIUS + 0.5 - distance).clamp(0, 1).unsqueeze(1)          # shape (N, 1, 16, 16): N images, 1 channel


def held_out(centres):
    return (centres[:, 0] > 9) & (centres[:, 1] > 9)                   # the bottom-right corner


def images(count, seed, corner=False):
    """`count` images with the ball anywhere except the corner, or (corner=True) only in the corner."""
    generator, kept = torch.Generator().manual_seed(seed), []
    while sum(len(k) for k in kept) < count:
        centres = 2 + 12 * torch.rand(4 * count, 2, generator=generator)
        kept.append(centres[held_out(centres) == corner])
    centres = torch.cat(kept)[:count]
    return draw(centres), centres


class SpatialSoftmax(nn.Module):
    """Softmax over the 256 pixels of a one-channel map, then the average pixel position under those weights."""
    def forward(self, maps):
        weights = maps.flatten(1).softmax(dim=1).view(-1, SIZE, SIZE)
        return torch.stack([(weights * COLS).sum((1, 2)), (weights * ROWS).sum((1, 2))], dim=1)


def make_network(kind):
    if kind == "dense":            # every pixel has its own weights
        return nn.Sequential(nn.Flatten(), nn.Linear(SIZE * SIZE, 3), nn.ReLU(), nn.Linear(3, 2))
    if kind == "cnn":              # shared 3 × 3 filters, pooling, then a dense layer for the answer
        return nn.Sequential(nn.Conv2d(1, 4, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
                             nn.Conv2d(4, 4, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
                             nn.Flatten(), nn.Linear(4 * 4 * 4, 8), nn.ReLU(), nn.Linear(8, 2))
    return nn.Sequential(          # "heatmap": shared filters only, and a fixed formula reads the map
        nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.Conv2d(8, 8, 3, padding=1), nn.ReLU(),
        nn.Conv2d(8, 1, 3, padding=1), SpatialSoftmax())


LEARNING_RATE = {"dense": 0.01, "cnn": 0.01, "heatmap": 0.03}     # Adam rates chosen on validation images


def train(kind, seed=0, lr=None, epochs=60):
    lr = lr or LEARNING_RATE[kind]
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    pictures, centres = images(2000, seed)
    network = make_network(kind)
    optimizer = torch.optim.Adam(network.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    order = torch.Generator().manual_seed(seed)
    for epoch in range(epochs):
        for batch in torch.randperm(len(pictures), generator=order).split(64):
            loss = loss_fn(network(pictures[batch]), centres[batch])     # squared miss in pixels
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    return network


def miss(network, data):
    """Average distance, in pixels, between the predicted and the true centre."""
    pictures, centres = data
    with torch.no_grad():
        return (network(pictures) - centres).norm(dim=1).mean().item()


if __name__ == "__main__":
    pictures, centres = images(1, seed=5)
    print("one image: tensor shape", tuple(pictures.shape), "ball centre", [round(v, 2) for v in centres[0].tolist()])

    # One learned 3 × 3 filter over a 5 × 5 patch around the ball, by hand and by nn.Conv2d.
    heatmap = train("heatmap")
    conv = nn.Conv2d(1, 1, 3)                                         # no padding: a 5 × 5 patch gives 3 × 3
    with torch.no_grad():
        conv.weight.copy_(heatmap[0].weight[:1])                       # the first layer's filter 0
        conv.bias.copy_(heatmap[0].bias[:1])
        patch = pictures[0, 0, 1:6, 9:14]                              # rows 1-5, columns 9-13
        by_torch = conv(patch[None, None])[0, 0]
    by_hand = [[(patch[r:r + 3, c:c + 3] * conv.weight[0, 0]).sum().item() + conv.bias.item() for c in range(3)] for r in range(3)]
    print("patch:", [[round(v, 3) for v in row] for row in patch.tolist()])
    print("filter:", [[round(v, 3) for v in row] for row in conv.weight[0, 0].tolist()], "bias", round(conv.bias.item(), 3))
    print("by hand: ", [[round(v, 4) for v in row] for row in by_hand])
    print("by torch:", [[round(v, 4) for v in row] for row in by_torch.tolist()])

    for kind in ("dense", "cnn", "heatmap"):
        print(f"{kind:8s}: {sum(p.numel() for p in make_network(kind).parameters()):5d} parameters")
    print("a dense layer from 256 pixels to 256 outputs:", sum(p.numel() for p in nn.Linear(256, 256).parameters()),
          "parameters; a 3 × 3 convolution from 1 to 8 channels:", sum(p.numel() for p in nn.Conv2d(1, 8, 3).parameters()))

    test, corner = images(500, seed=77), images(500, seed=78, corner=True)
    for kind in ("dense", "cnn", "heatmap"):
        network = train(kind)
        print(f"{kind:8s} seed 0: miss {miss(network, test):.3f} px on usual positions, {miss(network, corner):.3f} px in the unseen corner")