"""The same model and objectives, now using PyTorch's operation-level backward rules."""
import torch
from torch import nn
from torch.nn import functional as F

# region model

def build_model(cfg):
    D, H, K = cfg['input_size'], cfg['hidden_size'], cfg['classes']
    model = nn.Sequential(
        nn.Linear(D, H),
        nn.ReLU(),
        nn.Linear(H, K),
    )
    return model.to(device=cfg['device'], dtype=getattr(torch, cfg['dtype']))

# endregion model

# region copy

def copy_parameters(model, params):
    with torch.no_grad():
        model[0].weight.copy_(torch.from_numpy(params['W1']))
        model[0].bias.copy_(torch.from_numpy(params['b1']))
        model[2].weight.copy_(torch.from_numpy(params['W2']))
        model[2].bias.copy_(torch.from_numpy(params['b2']))

# endregion copy

# region objective

def objective(logits, y, kind):
    if kind == 'ce':
        return F.cross_entropy(logits, y, reduction='mean')
    if kind == 'mse':
        P = logits.softmax(dim=1)
        Y = F.one_hot(y, num_classes=logits.shape[1]).to(dtype=P.dtype)
        return F.mse_loss(P, Y, reduction='mean')
    raise ValueError('Loss must be ce or mse')

# endregion objective


def named_parameters(model):
    return {'W1': model[0].weight, 'b1': model[0].bias,
            'W2': model[2].weight, 'b2': model[2].bias}

# region evaluate

def scores(model, X, y, kind):
    model.eval()
    with torch.no_grad():
        logits = model(X)
        loss = objective(logits, y, kind).item()
        accuracy = (logits.argmax(dim=1) == y).double().mean().item()
    return loss, accuracy

# endregion evaluate