"""Will the T-shirt network work on customers it has not seen? Train, validation and test sets, overfitting,
early stopping and distribution shift, on the shop from Neural network: several inputs.

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

The shop now learns from its sales records instead of its chart: what each customer actually kept.
Most kept their chart size, but 15% kept a neighbouring size (they like it looser or tighter).
"""

import copy

import torch
from torch import nn

from pytorch_choices import size_chart, make_people, scale

KEPT_OTHER = 0.15     # share of customers who kept a neighbouring size


def sales_records(people, generator):
    """The size each customer kept: the chart's size, or for 15% of them one size up or down."""
    chart = torch.tensor([size_chart(h, w) for h, w in people.tolist()])
    other = torch.rand(len(chart), generator=generator) < KEPT_OTHER
    step = torch.where(torch.rand(len(chart), generator=generator) < 0.5, -1, 1)
    return torch.where(other, (chart + step).clamp(0, 2), chart)


def customers(count, generator, heights=(150, 195), weights=(45, 105)):
    """`count` customers in a height and weight range, with their records. The default is the shop's usual range."""
    if (heights, weights) == ((150, 195), (45, 105)):
        people, _ = make_people(count, generator)
    else:
        h = heights[0] + (heights[1] - heights[0]) * torch.rand(count, generator=generator)
        w = weights[0] + (weights[1] - weights[0]) * torch.rand(count, generator=generator)
        people = torch.stack([h, w], dim=1)
    return people, sales_records(people, generator)


def split(seed):
    """40 customers to train on, 40 for validation, and a fixed test set of 1,000 used once at the end."""
    generator = torch.Generator().manual_seed(seed)
    train_set = customers(40, generator)
    validation_set = customers(40, generator)
    test_set = customers(1000, torch.Generator().manual_seed(1000))
    return train_set, validation_set, test_set


def make_network(hidden=64):
    """Far more parameters (4,547 with 64) than the 40 customers it learns from."""
    return nn.Sequential(nn.Linear(2, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 3))


def accuracy(network, data):
    people, sizes = data
    with torch.no_grad():
        return (network(scale(people)).argmax(dim=1) == sizes).float().mean().item()


def train(seed=0, steps=3000, lr=0.05, hidden=64, observe=None):
    """Plain SGD on all 40 training customers at once. Early stopping keeps the network from the first step
    with the highest validation accuracy."""
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    train_set, validation_set, _ = split(seed)
    network = make_network(hidden)
    optimizer = torch.optim.SGD(network.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    inputs, answers = scale(train_set[0]), train_set[1]
    best_accuracy, best_step, best = -1.0, 0, None
    for step in range(steps + 1):
        validation_accuracy = accuracy(network, validation_set)
        if validation_accuracy > best_accuracy:              # early stopping: remember the best network so far
            best_accuracy, best_step, best = validation_accuracy, step, copy.deepcopy(network)
        if observe:
            observe(step, network)
        if step == steps:
            break
        loss = loss_fn(network(inputs), answers)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    return network, best, best_step


SHIFTS = {"taller and heavier": ((195, 215), (105, 130)), "smaller": ((125, 150), (25, 45))}


def report(seed):
    train_set, validation_set, test_set = split(seed)
    final, best, best_step = train(seed)
    chart = lambda data: (torch.tensor([size_chart(h, w) for h, w in data[0].tolist()]) == data[1]).float().mean().item()
    print(f"seed {seed}: the chart itself matches {chart(test_set):.3f} of the 1,000 test records")
    print(f"  after 3000 steps: training accuracy {accuracy(final, train_set):.3f}, validation {accuracy(final, validation_set):.3f}")
    print(f"  early stopping: highest validation accuracy first at step {best_step}; training {accuracy(best, train_set):.3f}, "
          f"validation {accuracy(best, validation_set):.3f}")
    print(f"  test, used once: early-stopped {accuracy(best, test_set):.3f}, after 3000 steps {accuracy(final, test_set):.3f}")
    for name, (heights, weights) in SHIFTS.items():
        shifted = customers(500, torch.Generator().manual_seed(2000), heights, weights)
        print(f"  {name} customers ({heights[0]}-{heights[1]} cm, {weights[0]}-{weights[1]} kg): "
              f"early-stopped {accuracy(best, shifted):.3f}, chart {chart(shifted):.3f}")


if __name__ == "__main__":
    report(6)        # the seed the overfitting lab shows
    report(0)        # the seed the shift lab opens on: the best test score, and the smaller customers mostly wrong