"""
Other biologically-motivated / non-backprop local learning baselines, sharing
SDILNet's architecture and initialisation for a fair comparison:

- FANet     : Feedback Alignment (Lillicrap 2016) -- backprop through FIXED random
              feedback matrices instead of W^T (sequential, layer-wise). DFA is the
              direct variant (already in baselines.dfa_config); FA is the sequential one.
- PEPITANet : PEPITA (Dellaferrera & Kreiman 2022) -- forward-only. A second forward
              pass on error-modulated input; weights follow the activation change.
- FFNet     : Forward-Forward (Hinton 2022) -- each layer locally maximises a
              "goodness" (sum of squares) on positive (real label) data and minimises
              it on negative (wrong label) data; inference picks max total goodness.

All train with no global backward graph. Autograd, where used (FF's per-layer
local loss), never crosses layer boundaries -> the update stays local.
"""
import math
import torch
import torch.nn.functional as F

from .core import SDILNet, ACTS


# --------------------------------------------------------------------------
# Feedback Alignment (sequential, layer-wise random feedback)
# --------------------------------------------------------------------------
class FANet(SDILNet):
    def __init__(self, *args, b_scale=1.0, **kw):
        super().__init__(*args, **kw)
        g = torch.Generator(device="cpu").manual_seed(4242)
        # fixed random feedback B[l] with the same shape as W[l], for l=1..L-1
        self.B = [None]
        for l in range(1, self.L):
            shape = self.W[l].shape
            self.B.append((torch.randn(*shape, generator=g) * (b_scale / math.sqrt(shape[1]))
                           ).to(self.W[l].device, self.dtype))

    def fa_step(self, x, y, yoh, eta, momentum=0.0):
        with torch.no_grad():
            fwd = self.forward(x)
            h, u = fwd["h"], fwd["u"]
            B = x.shape[0]
            e = torch.softmax(h[-1], 1) - yoh          # grad wrt logits
            loss = F.cross_entropy(h[-1], y).item()
            # output layer (exact local)
            delta = e                                   # grad wrt u[L-1]
            self._upd(self.L - 1, delta, h[self.L - 1], eta, momentum)
            # hidden layers: propagate with fixed random B
            for l in range(self.L - 2, -1, -1):
                delta = (delta @ self.B[l + 1]) * self.act_prime(u[l])   # grad wrt u[l]
                self._upd(l, delta, h[l], eta, momentum)
            return loss

    def _upd(self, l, delta, pre, eta, momentum):
        dW = -(delta.t() @ pre) / delta.shape[0]
        db = -delta.mean(0)
        if momentum:
            self.mW[l].mul_(momentum).add_(dW); self.mb[l].mul_(momentum).add_(db)
            self.W[l] += eta * self.mW[l]; self.b[l] += eta * self.mb[l]
        else:
            self.W[l] += eta * dW; self.b[l] += eta * db


# --------------------------------------------------------------------------
# PEPITA (forward-only, error-modulated second pass)
# --------------------------------------------------------------------------
class PEPITANet(SDILNet):
    def __init__(self, *args, f_scale=0.1, **kw):
        super().__init__(*args, **kw)
        g = torch.Generator(device="cpu").manual_seed(7777)
        n_in = self.sizes[0]
        # projection of output error back onto the input (fixed random)
        self.Fproj = (torch.randn(n_in, self.n_classes, generator=g)
                      * (f_scale / math.sqrt(self.n_classes))).to(self.W[0].device, self.dtype)

    def pepita_step(self, x, y, yoh, eta, momentum=0.0):
        with torch.no_grad():
            std = self.forward(x)
            hs = std["h"]
            e = torch.softmax(hs[-1], 1) - yoh
            loss = F.cross_entropy(hs[-1], y).item()
            x_mod = x + (e @ self.Fproj.t())           # modulate the input by the error
            mod = self.forward(x_mod)
            hm = mod["h"]
            B = x.shape[0]
            for l in range(self.L):
                # change in this layer's output between clean and modulated passes
                dpost = hs[l + 1] - hm[l + 1]
                # PEPITA: first layer uses the CLEAN input as presynaptic; deeper
                # layers use the modulated presynaptic.
                pre = x if l == 0 else hm[l]
                dW = -(dpost.t() @ pre) / B
                db = -dpost.mean(0)
                if momentum:
                    self.mW[l].mul_(momentum).add_(dW); self.mb[l].mul_(momentum).add_(db)
                    self.W[l] += eta * self.mW[l]; self.b[l] += eta * self.mb[l]
                else:
                    self.W[l] += eta * dW; self.b[l] += eta * db
            return loss


# --------------------------------------------------------------------------
# Forward-Forward (Hinton 2022)
# --------------------------------------------------------------------------
class FFNet:
    """Fully-connected Forward-Forward net. Label is overlaid on the input;
    each layer is trained by a local goodness objective; classification sums
    goodness across layers over the candidate labels."""

    def __init__(self, sizes, act="tanh", device="cpu", seed=0, threshold=2.0,
                 n_classes=10, dtype=torch.float32, overlay_val=10.0):
        self.sizes = list(sizes)
        self.L = len(sizes) - 1
        self.n_classes = n_classes
        self.device = device
        self.dtype = dtype
        self.act, _ = ACTS[act]
        self.thr = threshold
        self.overlay_val = overlay_val
        g = torch.Generator(device="cpu").manual_seed(seed)
        self.W, self.b, self.mW, self.mb = [], [], [], []
        for i in range(self.L):
            w = (torch.randn(sizes[i + 1], sizes[i], generator=g) / math.sqrt(sizes[i])
                 ).to(device, dtype)
            w.requires_grad_(True)
            self.W.append(w)
            bb = torch.zeros(sizes[i + 1], device=device, dtype=dtype, requires_grad=True)
            self.b.append(bb)

    def __init_overlay_scale__(self):
        pass

    def _overlay(self, x, labels):
        """Overlay one-hot label onto the first n_classes input features.
        The label value must be strong enough to actually shift the goodness
        (10 of 784 pixels is otherwise swamped by the shared image)."""
        xo = x.clone()
        xo[:, :self.n_classes] = 0.0
        xo[torch.arange(x.shape[0]), labels] = self.overlay_val
        return xo

    @staticmethod
    def _norm(h):
        return h / (h.norm(dim=1, keepdim=True) + 1e-8)

    def _layer_forward(self, l, h_in):
        return self.act(h_in @ self.W[l].t() + self.b[l])

    def train_step(self, x, y, eta):
        # positive = correct label; negative = a random wrong label
        neg = (y + torch.randint(1, self.n_classes, y.shape, device=y.device)) % self.n_classes
        xpos = self._overlay(x, y)
        xneg = self._overlay(x, neg)
        hpos, hneg = xpos, xneg
        total = 0.0
        for l in range(self.L):
            hp = self._layer_forward(l, hpos.detach())
            hn = self._layer_forward(l, hneg.detach())
            gp = hp.pow(2).mean(1)                      # goodness (positive)
            gn = hn.pow(2).mean(1)                      # goodness (negative)
            # push positive goodness above threshold, negative below
            loss = (F.softplus(-(gp - self.thr)) + F.softplus(gn - self.thr)).mean()
            gW, gb = torch.autograd.grad(loss, [self.W[l], self.b[l]])
            with torch.no_grad():
                self.W[l] -= eta * gW
                self.b[l] -= eta * gb
            total += loss.item()
            hpos, hneg = self._norm(hp).detach(), self._norm(hn).detach()
        return total / self.L

    @torch.no_grad()
    def predict(self, x):
        scores = torch.zeros(x.shape[0], self.n_classes, device=x.device)
        for c in range(self.n_classes):
            h = self._overlay(x, torch.full((x.shape[0],), c, device=x.device, dtype=torch.long))
            good = torch.zeros(x.shape[0], device=x.device)
            for l in range(self.L):
                h = self._layer_forward(l, h)
                if l > 0:                               # skip first layer for scoring (Hinton)
                    good = good + h.pow(2).mean(1)
                h = self._norm(h)
            scores[:, c] = good
        return scores.argmax(1)

    @torch.no_grad()
    def evaluate(self, test_loader):
        correct, total = 0, 0
        for x, y in test_loader:
            pred = self.predict(x)
            correct += (pred == y).sum().item(); total += y.shape[0]
        return correct / total, 0.0


# --------------------------------------------------------------------------
# Equilibrium Propagation (Scellier & Bengio 2017), real-valued, layered MLP
# --------------------------------------------------------------------------
class EPNet:
    """Prototypical EP: bidirectional layered net, hard-sigmoid rho. Free phase
    settles to equilibrium; a weakly-nudged phase perturbs the output toward the
    target; the local contrastive update approximates the loss gradient.
    Updates are local (products of adjacent-layer rho's) with no backprop."""

    def __init__(self, sizes, device="cpu", seed=0, beta=0.5, dt=0.5,
                 T_free=20, T_nudge=8, dtype=torch.float32):
        self.sizes = list(sizes)
        self.L = len(sizes) - 1          # number of weight layers
        self.device = device
        self.dtype = dtype
        self.beta, self.dt, self.T_free, self.T_nudge = beta, dt, T_free, T_nudge
        g = torch.Generator(device="cpu").manual_seed(seed)
        self.W = [(torch.randn(sizes[i + 1], sizes[i], generator=g) / math.sqrt(sizes[i])
                   ).to(device, dtype) for i in range(self.L)]
        self.b = [torch.zeros(sizes[i + 1], device=device, dtype=dtype) for i in range(self.L)]

    @staticmethod
    def rho(s):
        return s.clamp(0, 1)

    @staticmethod
    def rhop(s):
        return ((s > 0) & (s < 1)).to(s.dtype)

    def _settle(self, x, y=None, beta=0.0, s=None, T=20):
        rx = self.rho(x)
        if s is None:
            # init in the active region (rhop=0 at the clamp boundaries would
            # otherwise freeze the dynamics); a feedforward warm start settles fast.
            s = []
            below = rx
            for i in range(self.L):
                below = (below @ self.W[i].t() + self.b[i]).clamp(0, 1)
                s.append(below)
        for _ in range(T):
            new = []
            for k in range(self.L):
                below = rx if k == 0 else self.rho(s[k - 1])
                pre = below @ self.W[k].t() + self.b[k]
                if k < self.L - 1:                       # top-down from layer above
                    pre = pre + self.rho(s[k + 1]) @ self.W[k + 1]
                if k == self.L - 1 and beta:             # nudge output toward target
                    pre = pre + beta * (y - s[k])
                ds = self.rhop(s[k]) * pre - s[k]
                new.append((s[k] + self.dt * ds).clamp(0, 1))
            s = new
        return s

    def train_step(self, x, y, yoh, eta):
        with torch.no_grad():
            s0 = self._settle(x, beta=0.0, T=self.T_free)                 # free phase
            sb = self._settle(x, yoh, beta=self.beta, s=[t.clone() for t in s0], T=self.T_nudge)
            B = x.shape[0]
            for k in range(self.L):
                below0 = self.rho(x) if k == 0 else self.rho(s0[k - 1])
                belowb = self.rho(x) if k == 0 else self.rho(sb[k - 1])
                dW = (self.rho(sb[k]).t() @ belowb - self.rho(s0[k]).t() @ below0) / (self.beta * B)
                db = (self.rho(sb[k]) - self.rho(s0[k])).mean(0) / self.beta
                self.W[k] += eta * dW
                self.b[k] += eta * db
            # free-phase output as prediction proxy for loss logging
            return F.mse_loss(s0[-1], yoh).item()

    @torch.no_grad()
    def predict(self, x):
        s = self._settle(x, beta=0.0, T=self.T_free)
        return s[-1].argmax(1)

    @torch.no_grad()
    def evaluate(self, test_loader):
        correct, total = 0, 0
        for x, y in test_loader:
            correct += (self.predict(x) == y).sum().item(); total += y.shape[0]
        return correct / total, 0.0