"""
Step 3 — Train the network to predict the change in 10 ms.

    input  : s_t = [x, z, vx, vz]            output : s_(t+1) - s_t
    model  : 4 -> 16 ReLU neurons -> 4        loss   : mean squared error, in raw units (m, m/s)
    update : plain SGD, learning rate 0.1, batches of 256, 240 passes over the data

Run:  python 03_train.py                      (the lesson's network)
      python 03_train.py --hidden 2           (change one thing; see 05_compare.py)
      python 03_train.py --hidden 8 8         (two hidden layers of 8)
Out:  results/<tag>/network.pt, results/<tag>/losses.json
"""

import argparse, json, os
import numpy as np
import torch
from torch import nn
from model import make_network

parser = argparse.ArgumentParser()
parser.add_argument("--hidden", type=int, nargs="+", default=[16])
parser.add_argument("--lr", type=float, default=0.1)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--epochs", type=int, default=240)
parser.add_argument("--batch", type=int, default=256)
args = parser.parse_args()
EPOCHS, BATCH = args.epochs, args.batch
tag = f"hidden{'x'.join(map(str, args.hidden))}_lr{args.lr:g}_seed{args.seed}"
tag += (f"_e{EPOCHS}" if EPOCHS != 240 else "") + (f"_b{BATCH}" if BATCH != 256 else "")

# Data: every pair from the training throws.
data = np.load("data/throws.npz")
throws = data["states"][data["split"] == "train"]             # (500, 201, 4)
now, after = throws[:, :-1].reshape(-1, 4), throws[:, 1:].reshape(-1, 4)
inputs = torch.tensor(now, dtype=torch.float32)               # (100000, 4)
targets = torch.tensor(after - now, dtype=torch.float32)      # the change in 10 ms

torch.manual_seed(args.seed)
network = make_network(args.hidden[0] if len(args.hidden) == 1 else args.hidden)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(network.parameters(), lr=args.lr)

losses = []
for epoch in range(EPOCHS):
    order = torch.randperm(len(inputs))                       # a new shuffle every pass
    for i in range(0, len(order), BATCH):
        batch = order[i:i + BATCH]
        loss = loss_fn(network(inputs[batch]), targets[batch])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    with torch.no_grad():
        losses.append(loss_fn(network(inputs), targets).item())
    if (epoch + 1) % 40 == 0:
        print(f"epoch {epoch + 1:3d}  loss {losses[-1]:.5f}")

os.makedirs(f"results/{tag}", exist_ok=True)
torch.save({"hidden": args.hidden[0] if len(args.hidden) == 1 else args.hidden, "weights": network.state_dict()}, f"results/{tag}/network.pt")
json.dump({"args": vars(args), "losses": losses}, open(f"results/{tag}/losses.json", "w"))
print(f"saved results/{tag}/")