"""
model.py — the learned part: a network that maps a state to its change over the next 10 ms.
One hidden layer of ReLU neurons gives the rule the bends it needs to turn a fall into a bounce.
"""

import torch
from torch import nn


def make_network(hidden=16):
    """4 numbers in (the state) -> hidden layer(s) of ReLU neurons -> 4 numbers out (the change).

    hidden=16 gives one hidden layer of 16 neurons; hidden=[8, 8] gives two hidden layers of 8.
    """
    sizes = [hidden] if isinstance(hidden, int) else list(hidden)
    layers, width = [], 4
    for size in sizes:
        layers += [nn.Linear(width, size), nn.ReLU()]
        width = size
    return nn.Sequential(*layers, nn.Linear(width, 4))


def load_network(path):
    saved = torch.load(path)
    network = make_network(saved["hidden"])
    network.load_state_dict(saved["weights"])
    return network


def predict_next(network, states):
    """Next state = current state + the network's guessed change. Works on one state or many."""
    with torch.no_grad():
        change = network(torch.tensor(states, dtype=torch.float32)).numpy()
    return states + change