"""A network that chooses: T-shirt size S, M or L from height and weight.

Run: python3 site/ml/foundations/examples/pytorch_choices.py
Every person comes with the right size, so this is the PyTorch basics loop again,
with two inputs instead of one and three choices instead of one number.
"""

import torch
from torch import nn

SIZES = ["S", "M", "L"]


def size_chart(height, weight):
    """The shop's chart: the right answers we train on (made up for this lesson)."""
    if height > 180 or weight > 85:
        return 2                                   # L
    if height < 165 and weight < 60:
        return 0                                   # S
    return 1                                       # M


def make_people(count, generator):
    heights = 150 + 45 * torch.rand(count, generator=generator)    # 150-195 cm
    weights = 45 + 60 * torch.rand(count, generator=generator)     # 45-105 kg
    sizes = torch.tensor([size_chart(h, w) for h, w in zip(heights.tolist(), weights.tolist())])
    return torch.stack([heights, weights], dim=1), sizes


def scale(people):
    """Centre both inputs near 0 and give them a similar spread, so one learning rate suits both."""
    return (people - torch.tensor([170.0, 75.0])) / torch.tensor([10.0, 15.0])


def make_network():
    return nn.Sequential(nn.Linear(2, 4), nn.ReLU(), nn.Linear(4, 3))


def train(seed=0, steps=3000, lr=0.5, observe=None):
    torch.set_num_threads(1)                                       # same numbers on every machine
    torch.manual_seed(seed)
    generator = torch.Generator().manual_seed(seed)
    people, answers = make_people(80, generator)
    inputs = scale(people)
    network = make_network()
    optimizer = torch.optim.SGD(network.parameters(), lr=lr)
    for step in range(steps):
        scores = network(inputs)                                   # 80 × 3: one score per size
        log_probabilities = torch.log_softmax(scores, dim=1)       # softmax, then log
        right = log_probabilities.gather(1, answers.unsqueeze(1))  # log p of the right size
        loss = -right.mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if observe:
            observe(step + 1, network, loss.item())
    return network, people, answers


if __name__ == "__main__":
    def report(step, network, loss):
        if step in (1, 100, 500, 3000):
            print(f"step {step:4d}: loss {loss:.3f}")

    network, people, answers = train(observe=report)
    with torch.no_grad():
        # nn.CrossEntropyLoss is the usual one-call name for the same loss.
        scores = network(scale(people))
        print("same loss from nn.CrossEntropyLoss:", round(nn.CrossEntropyLoss()(scores, answers).item(), 3))
        picked = scores.argmax(dim=1)
        print(f"training people sized correctly: {(picked == answers).sum().item()} of {len(answers)}")
        new_people, new_answers = make_people(200, torch.Generator().manual_seed(99))
        new_picked = network(scale(new_people)).argmax(dim=1)
        print(f"200 new people sized correctly: {(new_picked == new_answers).sum().item()}")
        person = torch.tensor([[178.0, 70.0]])
        probabilities = torch.softmax(network(scale(person)), dim=1)[0]
        print("178 cm, 70 kg:", ", ".join(f"{s} {p:.2f}" for s, p in zip(SIZES, probabilities.tolist())))