"""Compare sequential and multiprocessing PyTorch sweeps.
Run as a script so spawn can import the worker function:
    python parallel-multiprocessing.py --epochs 10 --workers 5
Each worker uses one CPU thread and the same seed, data split and initialization.
"""

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision.datasets import MNIST
import numpy as np

def make_model(hidden_sizes=(128,), activation=nn.ReLU):
    torch.manual_seed(42)
    layers = []
    n_inputs = 784
    for size in hidden_sizes:
        layers.extend([nn.Linear(n_inputs, size), activation()])
        n_inputs = size
    layers.append(nn.Linear(n_inputs, 10))
    return nn.Sequential(*layers)

def run_epoch(model, X, y, batch_size=32, optimizer=None):
    training = optimizer is not None
    model.train(training)
    device = next(model.parameters()).device
    loader = DataLoader(
        TensorDataset(torch.as_tensor(X, dtype=torch.float32),
                      torch.as_tensor(y, dtype=torch.long)),
        batch_size=batch_size, shuffle=training,
    )
    loss_fn = nn.CrossEntropyLoss()
    total_loss, correct = 0.0, 0
    with torch.enable_grad() if training else torch.no_grad():
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            if training:
                optimizer.zero_grad()
            logits = model(xb)
            loss = loss_fn(logits, yb)
            if training:
                loss.backward()
                optimizer.step()
            total_loss += loss.item() * len(yb)
            correct += (logits.argmax(dim=1) == yb).sum().item()
    return total_loss / len(y), correct / len(y)


def train_torch(model, X, y, epochs=5, lr=0.1, batch_size=32,
                val_fraction=0.2, optimizer=None, include_initial=False):
    split = int(len(X) * (1 - val_fraction))
    X_fit, y_fit = X[:split], y[:split]
    X_val, y_val = X[split:], y[split:]
    if optimizer is None:
        optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    history = {key: [] for key in ("loss", "accuracy", "val_loss", "val_accuracy")}
    for epoch in range(0 if include_initial else 1, epochs + 1):
        loss, accuracy = run_epoch(
            model, X_fit, y_fit, batch_size, optimizer if epoch else None)
        val_loss, val_accuracy = run_epoch(model, X_val, y_val, batch_size)
        for key, value in zip(history, (loss, accuracy, val_loss, val_accuracy)):
            history[key].append(value)
        print(f"Epoch {epoch}/{epochs} — loss: {loss:.4f}, accuracy: {accuracy:.4f}, "
              f"val_loss: {val_loss:.4f}, val_accuracy: {val_accuracy:.4f}")
    return history

def train_one_config(args):
    lr, X_train, y_train, epochs = args
    torch.set_num_threads(1)
    model = make_model()
    history = train_torch(model, X_train, y_train, epochs=epochs, lr=lr,
                          include_initial=True)
    return lr, history


if __name__ == "__main__":
    import argparse
    import multiprocessing as mp
    from time import perf_counter

    parser = argparse.ArgumentParser()
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--samples", type=int, default=60000)
    parser.add_argument("--workers", type=int, default=5)
    args = parser.parse_args()
    if args.epochs < 1 or not 2 <= args.samples <= 60000 or args.workers < 1:
        parser.error("Require positive epochs/workers and 2–60000 samples")
    data = MNIST("./data", train=True, download=True)
    X_train = data.data[:args.samples].reshape(-1, 784).float().numpy() / 255.0
    y_train = data.targets[:args.samples].numpy()
    configs = [(lr, X_train, y_train, args.epochs) for lr in [0.001, 0.01, 0.1, 1.0, 10.0]]
    start = perf_counter()
    sequential = dict(map(train_one_config, configs))
    sequential_time = perf_counter() - start
    start = perf_counter()
    with mp.get_context("spawn").Pool(min(args.workers, len(configs))) as pool:
        parallel = dict(pool.map(train_one_config, configs))
    parallel_time = perf_counter() - start
    print(f"Sequential: {sequential_time:.2f}s; parallel including worker startup: {parallel_time:.2f}s")
    for lr in sequential:
        print(lr, sequential[lr]["val_accuracy"][-1], parallel[lr]["val_accuracy"][-1])