"""
Step 6 — Batch size, on an equal time budget.

Each run trains the lesson's network (4 -> 16 -> 4, SGD, lr 0.1) for the same 20 seconds of computer time, with
batches of 1, 32, 256, or all 100,000 pairs. It records the loss on all training pairs as it goes, and how many
updates it managed. A fair race: small batches make many cheap, noisy steps; big batches make few, accurate ones.

Run:  python 06_batch_size.py      (about 1.5 minutes; the four runs go in parallel, one CPU core each)
Out:  results/batch<B>_t20/network.pt, losses.json (+ 04_evaluate.py metrics)
"""

import json, os, subprocess, sys, time
import numpy as np
import torch
from torch import nn
from model import make_network

SECONDS, LR = 20.0, 0.1


def race(batch):
    torch.set_num_threads(1)
    torch.manual_seed(0)
    data = np.load("data/throws.npz")
    throws = data["states"][data["split"] == "train"]
    now, after = throws[:, :-1].reshape(-1, 4), throws[:, 1:].reshape(-1, 4)
    inputs = torch.tensor(now, dtype=torch.float32)
    targets = torch.tensor(after - now, dtype=torch.float32)
    network, loss_fn = make_network(16), nn.MSELoss()
    optimizer = torch.optim.SGD(network.parameters(), lr=LR)

    record, updates, trained, next_check = [], 0, 0.0, 0.0     # trained = training time only (checks not counted)
    while trained < SECONDS:
        order = torch.randperm(len(inputs))
        for i in range(0, len(order), batch):
            if trained >= next_check:                           # measure the loss every 0.25 s of training
                with torch.no_grad():
                    record.append([round(trained, 3), updates, loss_fn(network(inputs), targets).item()])
                next_check += 0.25
            start = time.perf_counter()
            rows = order[i:i + batch]
            loss = loss_fn(network(inputs[rows]), targets[rows])
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            trained += time.perf_counter() - start
            updates += 1
            if trained >= SECONDS:
                break
    with torch.no_grad():
        record.append([round(trained, 3), updates, loss_fn(network(inputs), targets).item()])
    tag = f"batch{batch}_t20"
    os.makedirs(f"results/{tag}", exist_ok=True)
    torch.save({"hidden": 16, "weights": network.state_dict()}, f"results/{tag}/network.pt")
    json.dump({"batch": batch, "seconds": SECONDS, "updates": updates, "record": record}, open(f"results/{tag}/losses.json", "w"))
    print(f"batch {batch:6d}: {updates:7d} updates in {SECONDS:.0f} s, final loss {record[-1][2]:.5f}")
    return tag


if __name__ == "__main__":
    if len(sys.argv) > 1:                                       # one race (used by the parallel launcher below)
        race(int(sys.argv[1]))
    else:
        env = {**os.environ, "OMP_NUM_THREADS": "1"}
        procs = [subprocess.Popen(["python3", "06_batch_size.py", str(b)], env=env) for b in (1, 32, 256, 100000)]
        for proc in procs:
            proc.wait()
        for b in (1, 32, 256, 100000):
            subprocess.run(["python3", "04_evaluate.py", "--tag", f"batch{b}_t20"], check=True, env=env)