import torch
from torch import nn, einsum
from einops.layers.torch import Rearrange, Reduce

class Affine(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.g = nn.Parameter(torch.ones(1, 1, dim))
        self.b = nn.Parameter(torch.zeros(1, 1, dim))
        self.fn = fn

    def forward(self, x):
        return self.fn(x) * self.g + self.b

class LayerScaleResidual(nn.Module): # https://arxiv.org/abs/2103.17239
    def __init__(self, dim, depth, fn):
        super().__init__()
        if depth <= 18:
            init_eps = 0.1
        elif depth > 18 and depth <= 24:
            init_eps = 1e-5
        else:
            init_eps = 1e-6

        scale = torch.zeros(1, 1, dim).fill_(init_eps)
        self.scale = nn.Parameter(scale)
        self.fn = fn

    def forward(self, x):
        return self.fn(x) * self.scale + x

def ResMLP(*, image_size, patch_size, dim, depth, num_classes, expansion_factor = 4):
    assert (image_size % patch_size) == 0, 'image must be divisible by patch size'
    num_patches = (image_size // patch_size) ** 2
    wrapper = lambda i, fn: LayerScaleResidual(dim, i + 1, Affine(dim, fn))

    return nn.Sequential(
        Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
        nn.Linear((patch_size ** 2) * 3, dim),
        *[nn.Sequential(
            wrapper(i, nn.Conv1d(num_patches, num_patches, 1)),
            wrapper(i, nn.Sequential(
                nn.Linear(dim, dim * expansion_factor),
                nn.GELU(),
                nn.Linear(dim * expansion_factor, dim)
            ))
        ) for i in range(depth)],
        Reduce('b n c -> b c', 'mean'),
        nn.Linear(dim, num_classes)
    )
