"""Standard (layer-wise) feedback alignment: Linear / Conv2d layers whose backward uses a fixed random
weight B in place of W^T (Lillicrap et al. 2016).  Nonlinearity derivatives are the true ones, taken at
the forward activations, as in the original algorithm.  Use `apply_fa(module)` to convert all Linear/Conv2d
layers of a block in place."""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class _FALinearFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, W, b, B):
        ctx.save_for_backward(x, B)
        ctx.has_bias = b is not None
        y = x.matmul(W.t())
        return y + b if b is not None else y

    @staticmethod
    def backward(ctx, gy):
        x, B = ctx.saved_tensors
        gx = gy.matmul(B)                                   # random feedback instead of W
        gW = gy.reshape(-1, gy.shape[-1]).t().matmul(x.reshape(-1, x.shape[-1]))
        gb = gy.reshape(-1, gy.shape[-1]).sum(0) if ctx.has_bias else None
        return gx, gW, gb, None


class _FAConvFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, W, b, B, stride, padding):
        ctx.save_for_backward(x, W, B)
        ctx.stride, ctx.padding = stride, padding
        ctx.has_bias = b is not None
        return F.conv2d(x, W, b, stride, padding)

    @staticmethod
    def backward(ctx, gy):
        x, W, B = ctx.saved_tensors
        gx = torch.nn.grad.conv2d_input(x.shape, B, gy, ctx.stride, ctx.padding)     # random feedback kernels
        gW = torch.nn.grad.conv2d_weight(x, W.shape, gy, ctx.stride, ctx.padding)
        gb = gy.sum((0, 2, 3)) if ctx.has_bias else None
        return gx, gW, gb, None, None, None


class FALinear(nn.Linear):
    def __init__(self, *a, **k):
        super().__init__(*a, **k)
        self.register_buffer("B", torch.randn_like(self.weight) / math.sqrt(self.in_features))

    def forward(self, x):
        return _FALinearFn.apply(x, self.weight, self.bias, self.B)


class FAConv2d(nn.Conv2d):
    def __init__(self, *a, **k):
        super().__init__(*a, **k)
        fan_in = self.in_channels * self.kernel_size[0] * self.kernel_size[1]
        self.register_buffer("B", torch.randn_like(self.weight) / math.sqrt(fan_in))

    def forward(self, x):
        return _FAConvFn.apply(x, self.weight, self.bias, self.B, self.stride, self.padding)


def apply_fa(module):
    """Replace every nn.Linear / nn.Conv2d inside `module` (recursively) by its FA variant, keeping weights."""
    for name, child in list(module.named_children()):
        if type(child) is nn.Linear:
            new = FALinear(child.in_features, child.out_features, bias=child.bias is not None)
            new.weight = child.weight
            if child.bias is not None:
                new.bias = child.bias
            new.B = new.B.to(child.weight.device)
            setattr(module, name, new)
        elif type(child) is nn.Conv2d:
            new = FAConv2d(child.in_channels, child.out_channels, child.kernel_size, child.stride, child.padding,
                           bias=child.bias is not None)
            new.weight = child.weight
            if child.bias is not None:
                new.bias = child.bias
            new.B = new.B.to(child.weight.device)
            setattr(module, name, new)
        else:
            apply_fa(child)
    return module