"""
Other biologically-motivated / non-backprop local learning baselines. FA and
PEPITA reuse SDILNet's tensor layout; FF and EP retain their native model
classes because their state spaces and objectives are fundamentally different:

- 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.  The
implementations follow the authors' published algorithms and reference code;
their deliberately different architectures are exposed rather than hidden
behind an allegedly apples-to-apples SDILNet wrapper.
"""
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, feedback_seed=None, **kw):
        model_seed = kw.get("seed", 0)
        super().__init__(*args, **kw)
        g = torch.Generator(device="cpu").manual_seed(
            model_seed + 4242 if feedback_seed is None else feedback_seed)
        # 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_hidden_gradients(self, fwd, yoh):
        """Return FA estimates of dL/dh for each hidden block output.

        The returned tensors include the residual-branch scale used by the
        corresponding local weight update.  Identity skips transport ``dh``
        exactly; only the nonlinear branch is transported through fixed B.
        Keeping this calculation separate makes the signal directly
        measurable without using autograd for learning.
        """
        h, u = fwd["h"], fwd["u"]
        e = torch.softmax(h[-1], 1) - yoh
        dh = e @ self.B[self.L - 1]
        estimates = [None] * (self.L - 1)
        for l in range(self.L - 2, -1, -1):
            branch_scale = self.res_alpha if self.residual and l >= 1 else 1.0
            estimates[l] = branch_scale * dh
            du = estimates[l] * self.act_prime(u[l])
            if l > 0:
                feedback = du @ self.B[l]
                dh = dh + feedback if self.residual and l >= 1 else feedback
        return estimates

    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"]
            e = torch.softmax(h[-1], 1) - yoh          # grad wrt logits
            loss = F.cross_entropy(h[-1], y).item()
            hidden_grad = self.fa_hidden_gradients(fwd, yoh)
            # output layer (exact local)
            self._upd(self.L - 1, e, h[self.L - 1], eta, momentum)

            for l in range(self.L - 1):
                du = hidden_grad[l] * self.act_prime(u[l])
                self._upd(l, du, 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):
    """Original two-forward-pass PEPITA/ERIN rule.

    Dellaferrera & Kreiman initialise both feedforward weights and the input
    error projection with He-uniform samples, then multiply the latter by 0.05.
    Their fully-connected implementation has ReLU hidden units, a softmax
    output, no biases, and optionally shares one dropout mask across the two
    presentations.  We keep logits internally but form exactly the same
    softmax errors.
    """

    def __init__(self, *args, f_scale=0.05, original_init=True,
                 use_bias=False, keep_prob=1.0, **kw):
        model_seed = kw.get("seed", 0)
        super().__init__(*args, **kw)
        g = torch.Generator(device="cpu").manual_seed(model_seed + 7777)
        n_in = self.sizes[0]
        if original_init:
            for l, w in enumerate(self.W):
                limit = math.sqrt(6.0 / self.sizes[l])
                w.copy_((torch.rand(w.shape, generator=g) * 2.0 * limit - limit)
                        .to(w.device, w.dtype))
        limit = math.sqrt(6.0 / n_in)
        self.Fproj = ((torch.rand(n_in, self.n_classes, generator=g) * 2.0 * limit - limit)
                      * f_scale).to(self.W[0].device, self.dtype)
        self.use_bias = use_bias
        self.keep_prob = keep_prob

    def _pepita_forward(self, x, masks):
        h = [x]
        for l in range(self.L):
            cur = h[-1] @ self.W[l].t()
            if self.use_bias:
                cur = cur + self.b[l]
            if l < self.L - 1:
                cur = self.act(cur)
                if masks is not None:
                    cur = cur * masks[l]
            h.append(cur)
        return h

    def pepita_step(self, x, y, yoh, eta, momentum=0.0):
        with torch.no_grad():
            masks = None
            if self.keep_prob < 1.0:
                masks = [(torch.rand(x.shape[0], self.sizes[l + 1], device=x.device)
                          < self.keep_prob).to(x.dtype) / self.keep_prob
                         for l in range(self.L - 1)]
            hs = self._pepita_forward(x, masks)
            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
            hm = self._pepita_forward(x_mod, masks)
            mod_error = torch.softmax(hm[-1], 1) - yoh
            B = x.shape[0]
            for l in range(self.L):
                pre = hm[l]                            # modulated presynaptic state
                if l == self.L - 1:
                    update = -(mod_error.t() @ pre) / B
                    update_b = -mod_error.mean(0)
                else:
                    diff = hs[l + 1] - hm[l + 1]
                    update = -(diff.t() @ pre) / B
                    update_b = -diff.mean(0)
                if momentum:
                    self.mW[l].mul_(momentum).add_(update)
                    self.W[l] += eta * self.mW[l]
                    if self.use_bias:
                        self.mb[l].mul_(momentum).add_(update_b)
                        self.b[l] += eta * self.mb[l]
                else:
                    self.W[l] += eta * update
                    if self.use_bias:
                        self.b[l] += eta * update_b
            return loss

    def forward(self, x, **_):
        h = self._pepita_forward(x, masks=None)
        return {"h": h, "u": []}


# --------------------------------------------------------------------------
# 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="relu", device="cpu", seed=0, threshold=2.0,
                 n_classes=10, dtype=torch.float32, overlay_val=None,
                 score_from_layer=1):
        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
        self.score_from_layer = min(score_from_layer, max(0, self.L - 1))
        g = torch.Generator(device="cpu").manual_seed(seed)
        self.W, self.b, self.optimizers = [], [], []
        for i in range(self.L):
            # Match torch.nn.Linear initialization used by the public reference.
            bound = 1.0 / math.sqrt(sizes[i])
            w = (torch.rand(sizes[i + 1], sizes[i], generator=g) * 2.0 * bound - bound).to(
                device, dtype).requires_grad_(True)
            self.W.append(w)
            bb = (torch.rand(sizes[i + 1], generator=g) * 2.0 * bound - bound).to(
                device, dtype).requires_grad_(True)
            self.b.append(bb)
            self.optimizers.append(torch.optim.Adam([w, bb], lr=0.03))

    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
        value = x.max().detach() if self.overlay_val is None else self.overlay_val
        xo[torch.arange(x.shape[0], device=x.device), labels] = value
        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(self._norm(h_in) @ self.W[l].t() + self.b[l])

    def _inputs_to_layer(self, l, x):
        h = x
        with torch.no_grad():
            for k in range(l):
                h = self._layer_forward(k, h)
        return h.detach()

    def train_layer(self, l, x, y, eta=0.03, negative_labels=None):
        """One minibatch update of one greedy FF layer (local autograd only)."""
        if negative_labels is None:
            negative_labels = ((y + torch.randint(
                1, self.n_classes, y.shape, device=y.device)) % self.n_classes)
        hpos = self._inputs_to_layer(l, self._overlay(x, y))
        hneg = self._inputs_to_layer(l, self._overlay(x, negative_labels))
        hp = self._layer_forward(l, hpos)
        hn = self._layer_forward(l, hneg)
        gp = hp.pow(2).mean(1)
        gn = hn.pow(2).mean(1)
        loss = (F.softplus(-gp + self.thr) + F.softplus(gn - self.thr)).mean()
        opt = self.optimizers[l]
        opt.param_groups[0]["lr"] = eta
        opt.zero_grad()
        loss.backward()
        opt.step()
        return loss.item()

    def train_step(self, x, y, eta):
        """Compatibility path; canonical runner trains greedily layer-by-layer."""
        neg = (y + torch.randint(1, self.n_classes, y.shape, device=y.device)) % self.n_classes
        total = 0.0
        for l in range(self.L):
            total += self.train_layer(l, x, y, eta, neg)
        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 >= self.score_from_layer:           # paper excludes first layer
                    good = good + h.pow(2).mean(1)
            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=4, dtype=torch.float32,
                 random_beta_sign=True):
        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
        self.random_beta_sign = random_beta_sign
        g = torch.Generator(device="cpu").manual_seed(seed)
        self.W = []
        for i in range(self.L):
            limit = math.sqrt(6.0 / (sizes[i] + sizes[i + 1]))
            self.W.append((torch.rand(sizes[i + 1], sizes[i], generator=g)
                           * 2.0 * limit - limit).to(device, dtype))
        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):
        # Theano's clip derivative used by the authors is active at the bounds.
        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:
            # The reference implementation starts persistent particles at zero.
            s = [torch.zeros(x.shape[0], n, device=x.device, dtype=x.dtype)
                 for n in self.sizes[1:]]
        for _ in range(T):
            new = []
            for k in range(self.L):
                below = rx if k == 0 else self.rho(s[k - 1])
                # -dE/ds for E = ||rho(s)||^2/2 - b*rho(s)
                #                 - rho(s_below) W^T rho(s).
                drive = -self.rho(s[k]) + below @ self.W[k].t() + self.b[k]
                if k < self.L - 1:                       # top-down from layer above
                    drive = drive + self.rho(s[k + 1]) @ self.W[k + 1]
                if k == self.L - 1 and beta:             # nudge output toward target
                    # Original cost is ||s_out-y||^2, hence the factor two.
                    drive = drive + 2.0 * beta * (y - s[k])
                ds = self.rhop(s[k]) * drive
                new.append((s[k] + self.dt * ds).clamp(0, 1))
            s = new
        return s

    def train_step(self, x, y, yoh, eta, free_state=None, return_free_state=False):
        with torch.no_grad():
            s0 = self._settle(x, beta=0.0, s=free_state, T=self.T_free)   # free phase
            # The authors keep one free-phase particle per training example
            # and reuse it on the next epoch.  Save the free state, not the
            # subsequently nudged state.
            persistent_state = [t.detach().clone() for t in s0]
            beta = self.beta
            if self.random_beta_sign and torch.randint(0, 2, ()).item() == 0:
                beta = -beta
            sb = self._settle(x, yoh, beta=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) / (beta * B)
                db = (self.rho(sb[k]) - self.rho(s0[k])).mean(0) / beta
                layer_eta = eta[k] if isinstance(eta, (list, tuple)) else eta
                self.W[k] += layer_eta * dW
                self.b[k] += layer_eta * db
            # free-phase output as prediction proxy for loss logging
            loss = F.mse_loss(s0[-1], yoh).item()
            return (loss, persistent_state) if return_free_state else loss

    @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