"""
Convolutional Transformer

Largely inspired by https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/cvt.py

Original paper:
@article{wu2021cvt,
  title={Cvt: Introducing convolutions to vision transformers},
  author={Wu, Haiping and Xiao, Bin and Codella, Noel and Liu, Mengchen and Dai, Xiyang and Yuan, Lu and Zhang, Lei},
  journal={arXiv preprint arXiv:2103.15808},
  year={2021}
}
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

from einops import rearrange, repeat
from einops.layers.torch import Rearrange

# Helper methods

def group_dict_by_key(cond, d):
    return_val = [dict(), dict()]
    for key in d.keys():
        match = bool(cond(key))
        ind = int(not match)
        return_val[ind][key] = d[key]
    return (*return_val,)

def group_by_key_prefix_and_remove_prefix(prefix, d):
    kwargs_with_prefix, kwargs = group_dict_by_key(lambda x: x.startswith(prefix), d)
    kwargs_without_prefix = dict(map(lambda x: (x[0][len(prefix):], x[1]), tuple(kwargs_with_prefix.items())))
    return kwargs_without_prefix, kwargs

# classes

class LayerNorm(nn.Module):
    """Layer normalization, but done in channel dimension #1"""
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
        self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))

    def forward(self, x):
        std = torch.var(x, dim=1, unbiased=False, keepdim=True).sqrt()
        mean = torch.mean(x, dim=1, keepdim=True)
        return (x - mean) / (std + self.eps) * self.g + self.b

class PreNorm(nn.Module):
    """Pre-Normalization layer"""
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = LayerNorm(dim)
        self.fn = fn

    def forward(self, x, **kwargs):
        x = self.norm(x)
        return self.fn(x, **kwargs)

class FeedForward(nn.Module):
    """Convolutional projection in the transformer."""
    def __init__(self, dim, mult=4, dropout=0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(dim, dim * mult, 1),  # 1x1 convolution
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Conv2d(dim * mult, dim, 1),  # 1x1 convolution
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)

class DepthWiseConv2d(nn.Module):
    """Depthwise convolutional layer"""
    def __init__(self, dim_in, dim_out, kernel_size, padding, stride, bias=True):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(dim_in, dim_in, kernel_size=kernel_size, padding=padding, groups=dim_in, stride=stride, bias=bias),
            nn.BatchNorm2d(dim_in),
            nn.Conv2d(dim_in, dim_out, kernel_size=1, bias=bias)
        )

    def forward(self, x):
        return self.net(x)

class Attention(nn.Module):
    """Custom Attention layer"""
    def __init__(self, dim, proj_kernel, kv_proj_stride, heads=8, dim_head=64, dropout=0.):
        super().__init__()
        inner_dim = dim_head * heads
        padding = proj_kernel // 2
        self.heads = heads
        self.scale = dim_head ** -0.5

        self.attend = nn.Softmax(dim=-1)

        self.to_q = DepthWiseConv2d(dim, inner_dim, proj_kernel, padding=padding, stride=1, bias=False)
        self.to_kv = DepthWiseConv2d(dim, inner_dim * 2, proj_kernel, padding=padding, stride=kv_proj_stride, bias=False)

        self.to_out = nn.Sequential(
            nn.Conv2d(inner_dim, dim, 1),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        shape = x.shape
        b, n, _, y, h = *shape, self.heads
        q, k, v = (self.to_q(x), *self.to_kv(x).chunk(2, dim=1))
        q, k, v = map(lambda t: rearrange(t, 'b (h d) x y -> (b h) (x y) d', h=h), (q, k, v))

        dots = torch.einsum('b i d, b j d -> b i j', q, k) * self.scale

        attn = self.attend(dots)

        out = torch.einsum('b i j, b j d -> b i d', attn, v)
        out = rearrange(out, '(b h) (x y) d -> b (h d) x y', h=h, y=y)
        return self.to_out(out)

class Transformer(nn.Module):
    """Custom Transformer layer."""
    def __init__(self, dim, proj_kernel, kv_proj_stride, depth, heads, dim_head=64, mlp_mult=4, dropout=0.):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                PreNorm(dim, Attention(dim, proj_kernel=proj_kernel, kv_proj_stride=kv_proj_stride, heads=heads, dim_head=dim_head, dropout=dropout)),
                PreNorm(dim, FeedForward(dim, mlp_mult, dropout=dropout))
            ]))

    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x
            x = ff(x) + x
        return x

class CvT(nn.Module):
    """Convolutional Transformer module.

    Adapted for self-supervised training

    Attributes
    ----------
    s{i}_emb_dim: int
        Embedding dimention at stage i

    s{i}_emb_kernel: int
        Convolutional kernel size at stage i

    s{i}_emb_stride: int
        Convolutional stride at stage i

    s{i}_kv_proj_stride: int
        Convolutional stride in the convolutional projection layers at stage i

    s{i}_heads: int
        Number of attention heads at stage i

    s{i}_depth: int
        Transformer depth at stage i

    s{i}_mlp_mult: int
        MLP ratio at stage i

    dropout: float
        Dropout ratio
    """
    # sample rate and embedding sizes are required model attributes for the HEAR API
    sample_rate = 16000
    embedding_size = 2048
    scene_embedding_size = embedding_size
    timestamp_embedding_size = embedding_size
    def __init__(
        self,
        *,
        s1_emb_dim=64,
        s1_emb_kernel=7,
        s1_emb_stride=4,
        s1_proj_kernel=3,
        s1_kv_proj_stride=2,
        s1_heads=1,
        s1_depth=1,
        s1_mlp_mult=4,
        s2_emb_dim=192,
        s2_emb_kernel=3,
        s2_emb_stride=2,
        s2_proj_kernel=3,
        s2_kv_proj_stride=2,
        s2_heads=3,
        s2_depth=2,
        s2_mlp_mult=4,
        s3_emb_dim=384,
        s3_emb_kernel=3,
        s3_emb_stride=2,
        s3_proj_kernel=3,
        s3_kv_proj_stride=2,
        s3_heads=6,
        s3_depth=10,
        s3_mlp_mult=4,
        dropout=0.,
        pool='mean'
    ):
        super().__init__()
        kwargs = dict(locals())

        dim = 1
        layers = []

        for prefix in ('s1', 's2', 's3'):
            config, kwargs = group_by_key_prefix_and_remove_prefix(f'{prefix}_', kwargs)

            layers.append(nn.Sequential(
                nn.Conv2d(
                    dim,
                    config['emb_dim'],
                    kernel_size=config['emb_kernel'],
                    padding=(config['emb_kernel'] // 2),
                    stride=config['emb_stride']
                ),
                LayerNorm(config['emb_dim']),
                Transformer(
                    dim=config['emb_dim'],
                    proj_kernel=config['proj_kernel'],
                    kv_proj_stride=config['kv_proj_stride'],
                    depth=config['depth'],
                    heads=config['heads'],
                    mlp_mult=config['mlp_mult'],
                    dropout=dropout
                )
            ))

            dim = config['emb_dim']

        self.pool = pool
        assert self.pool in ['mean', 'max', 'mean+max']

        if self.pool == 'mean':
            self.pool_layers = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten())
        elif self.pool == 'max':
            self.pool_layers = nn.Sequential(nn.AdaptiveMaxPool2d(1), nn.Flatten()),
        else:
            self.pool_layers = nn.Sequential(nn.Identity())

        self.layers = nn.Sequential(
            *layers,
            *self.pool_layers,
        )

    def forward(self, x):
        x = self.layers(x)

        if self.pool == 'mean+max':
            x = x.permute(0, 3, 2, 1)  # (batch, time, mel, ch)
            B, T, D, C = x.shape
            x = x.reshape((B, T, C * D))  # (batch, time, mel*ch)
            x = x.mean(1) + x.amax(1)
        return x
