"""Bridge: keep explicit matrix operations; replace manual derivatives with autograd."""
import argparse
import torch
from common import settings, load_data
from network_numpy import initialize
from network_torch import objective

# region tensor_step

def tensor_step(params, X, y, cfg):
    weights = {name: torch.tensor(value, requires_grad=True)
               for name, value in params.items()}
    X_tensor, y_tensor = torch.from_numpy(X), torch.from_numpy(y)
    Z1 = X_tensor @ weights['W1'].T + weights['b1']
    A1 = torch.relu(Z1)
    logits = A1 @ weights['W2'].T + weights['b2']
    loss = objective(logits, y_tensor, cfg['loss'])
    loss.backward()
    gradients = {name: p.grad.detach().numpy().copy() for name, p in weights.items()}
    with torch.no_grad():
        for p in weights.values():
            p -= cfg['learning_rate'] * p.grad
    after = {name: p.detach().numpy().copy() for name, p in weights.items()}
    return loss.item(), gradients, after

# endregion tensor_step


def main():
    parser = argparse.ArgumentParser(description='Perform one update using raw PyTorch tensors')
    parser.add_argument('--config', required=True)
    args = parser.parse_args()
    cfg = settings(args.config)
    torch.set_num_threads(cfg['threads'])
    X, y = load_data(args.config, cfg, 'train')
    if cfg['batch_size'] > len(X):
        raise ValueError('Batch is larger than the training set')
    loss, gradients, _ = tensor_step(initialize(cfg), X[:cfg['batch_size']], y[:cfg['batch_size']], cfg)
    print('Cost before this update:', loss)
    print('Gradient shapes:', {name: g.shape for name, g in gradients.items()})
    print('The forward math is explicit. Autograd supplied all four gradient arrays.')


if __name__ == '__main__':
    main()