"""Catch: a network learns to catch a falling ball from +1 / -1 alone.

Run: python3 site/ml/foundations/examples/reinforcement_catch.py
Nobody tells the network which move is right. After each drop, the moves it
made become its training labels: pushed up after a catch, down after a miss.
"""

import torch
from torch import nn

ROWS, COLS = 8, 7          # the screen: 8 rows of 7 pixels
MOVES = [-1, 0, 1]         # left, stay, right


def screen(ball_row, ball_col, paddle):
    """What the network sees: one bright pixel for the ball, one for the paddle."""
    pixels = torch.zeros(ROWS, COLS)
    pixels[ball_row, ball_col] = 1.0
    pixels[ROWS - 1, paddle] = 1.0
    return pixels.flatten()


def play(model, generator, greedy=False):
    """Drop one ball. Returns the screens seen, the moves made and the reward."""
    ball_col = int(torch.randint(COLS, (1,), generator=generator))
    paddle = int(torch.randint(COLS, (1,), generator=generator))
    return play_from(model, ball_col, paddle, generator, greedy)


def play_from(model, ball_col, paddle, generator=None, greedy=False):
    screens, moves = [], []
    for ball_row in range(ROWS - 1):              # 7 moves while the ball falls
        pixels = screen(ball_row, ball_col, paddle)
        with torch.no_grad():
            probabilities = torch.softmax(model(pixels), dim=0)
        if greedy:
            move = int(probabilities.argmax())
        else:
            move = int(torch.multinomial(probabilities, 1, generator=generator))
        screens.append(pixels)
        moves.append(move)
        paddle = min(COLS - 1, max(0, paddle + MOVES[move]))
    reward = 1.0 if paddle == ball_col else -1.0  # the only feedback
    return screens, moves, reward


def make_model():
    return nn.Sequential(nn.Linear(ROWS * COLS, 8), nn.ReLU(), nn.Linear(8, 3))


def catch_rate(model):
    """Frozen test: every ball column × every paddle start, best move each time."""
    caught = sum(play_from(model, b, p, greedy=True)[2] > 0
                 for b in range(COLS) for p in range(COLS))
    return caught / (COLS * COLS)


def train(seed=0, updates=1000, drops=128, lr=2.0, observe=None):
    torch.set_num_threads(1)                      # same numbers on every machine
    torch.manual_seed(seed)
    generator = torch.Generator().manual_seed(seed)
    model = make_model()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    for update in range(updates):
        # 1. Play a batch of drops with the current network.
        screens, moves, rewards = [], [], []
        for _ in range(drops):
            s, m, r = play(model, generator)
            screens += s
            moves += m
            rewards += [r] * len(m)               # every move in a drop shares its outcome
        screens = torch.stack(screens)
        moves = torch.tensor(moves).unsqueeze(1)
        rewards = torch.tensor(rewards)

        # 2. The moves it made are the labels; the reward says which way to push.
        log_probabilities = torch.log_softmax(model(screens), dim=1)
        taken = log_probabilities.gather(1, moves).squeeze(1)
        loss = -(rewards * taken).mean()

        # 3. The familiar PyTorch step.
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if observe:
            observe(update + 1, model, rewards)
    return model


if __name__ == "__main__":
    torch.manual_seed(0)
    print(f"Untrained network catches {catch_rate(make_model()):.0%} of the 49 test drops")
    def report(update, model, rewards):
        if update in (1, 100, 250, 500, 1000):
            practice = (rewards > 0).float().mean()   # share of this batch's drops caught
            print(f"update {update:4d}: practice catches {practice:.0%}, test {catch_rate(model):.0%}")
    train(observe=report)