"""Baselines sharing SDILNet's exact architecture and initialisation, so
comparisons are apples-to-apples.

- BP: exact backprop (autograd) -- the alignment/performance upper bound.
- DFA is NOT a separate class: it is exactly SDIL with a fixed-random apical
  vectorizer, no residual, no predictor, no perturbation. See dfa_config().
"""

import torch
import torch.nn.functional as F

from .core import SDILNet, SDILConfig


class BPNet(SDILNet):
    """Same tanh MLP, trained by exact backprop via autograd."""

    def bp_step(self, x, y, eta, momentum=0.0):
        for W in self.W:
            W.requires_grad_(True)
        for b in self.b:
            b.requires_grad_(True)
        logits = self.logits(x)
        loss = F.cross_entropy(logits, y)
        grads = torch.autograd.grad(loss, self.W + self.b)
        nW = len(self.W)
        with torch.no_grad():
            for i in range(nW):
                if momentum:
                    self.mW[i].mul_(momentum).add_(grads[i])
                    self.mb[i].mul_(momentum).add_(grads[nW + i])
                    self.W[i] -= eta * self.mW[i]
                    self.b[i] -= eta * self.mb[i]
                else:
                    self.W[i] -= eta * grads[i]
                    self.b[i] -= eta * grads[nW + i]
        for W in self.W:
            W.requires_grad_(False)
        for b in self.b:
            b.requires_grad_(False)
        return loss.item()


def dfa_config(eta=0.05, momentum=0.0):
    """SDIL configured to reproduce Direct Feedback Alignment exactly."""
    return SDILConfig(eta=eta, use_residual=False, learn_A=False, learn_P=False,
                      pert_every=10**9, momentum=momentum)


@torch.no_grad()
def evaluate(net, test_loader):
    correct, total, tot_loss = 0, 0, 0.0
    for x, y in test_loader:
        logits = net.logits(x)
        tot_loss += F.cross_entropy(logits, y, reduction="sum").item()
        correct += (logits.argmax(1) == y).sum().item()
        total += y.shape[0]
    return correct / total, tot_loss / total