"""A complete manual MLP: NumPy does array operations, not differentiation."""
import numpy as np

PARAMS = ('W1', 'b1', 'W2', 'b2')

# region initialize

def initialize(cfg):
    D, H, K = cfg['input_size'], cfg['hidden_size'], cfg['classes']
    rng = np.random.default_rng(cfg['seed'])
    return {
        'W1': (rng.standard_normal((H, D)) * np.sqrt(cfg['hidden_variance'] / D)).astype(cfg['dtype']),
        'b1': np.zeros(H, dtype=cfg['dtype']),
        'W2': (rng.standard_normal((K, H)) * np.sqrt(cfg['output_variance'] / H)).astype(cfg['dtype']),
        'b2': np.zeros(K, dtype=cfg['dtype']),
    }

# endregion initialize

# region scalar

def one_neuron(x, weights, bias):
    total = float(bias)
    for i in range(len(x)):
        total += float(x[i]) * float(weights[i])
    return total

# endregion scalar

# region forward

def forward(params, X):
    Z1 = X @ params['W1'].T + params['b1']
    A1 = np.maximum(Z1, 0)
    logits = A1 @ params['W2'].T + params['b2']
    shifted = logits - logits.max(axis=1, keepdims=True)
    logP = shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True))
    P = np.exp(logP)
    return {'X': X, 'Z1': Z1, 'A1': A1, 'logits': logits, 'P': P, 'logP': logP}

# endregion forward

# region loss

def loss_signal(cache, y, kind):
    P = cache['P']
    B, K = P.shape
    Y = np.eye(K, dtype=P.dtype)[y]
    if kind == 'ce':
        loss = -cache['logP'][np.arange(B), y].mean()
        dZ2 = (P - Y) / B
    elif kind == 'mse':
        error = P - Y
        loss = (error ** 2).mean()
        dP = 2 * error / (B * K)
        dZ2 = P * (dP - (dP * P).sum(axis=1, keepdims=True))
    else:
        raise ValueError('Loss must be ce or mse')
    return float(loss), dZ2

# endregion loss

# region backward

def backward(params, cache, dZ2):
    dW2 = dZ2.T @ cache['A1']
    db2 = dZ2.sum(axis=0)
    dA1 = dZ2 @ params['W2']
    dZ1 = dA1 * (cache['Z1'] > 0)
    dW1 = dZ1.T @ cache['X']
    db1 = dZ1.sum(axis=0)
    return {'W1': dW1, 'b1': db1, 'W2': dW2, 'b2': db2}

# endregion backward

# region update

def step(params, grads, learning_rate):
    for name in PARAMS:
        params[name] -= learning_rate * grads[name]

# endregion update


def scores(params, X, y, kind):
    cache = forward(params, X)
    loss, _ = loss_signal(cache, y, kind)
    accuracy = float((cache['logits'].argmax(axis=1) == y).mean())
    return loss, accuracy