import torch
import torch.nn as nn

from .attentions import MultiHeadAttention


class VAEMemoryBank(nn.Module):
    def __init__(
        self,
        bank_size=1000,
        n_hidden_dims=512,
        n_attn_heads=2,
        init_values=None,
        output_channels=192,
    ):
        super().__init__()

        self.bank_size = bank_size
        self.n_hidden_dims = n_hidden_dims
        self.n_attn_heads = n_attn_heads

        self.encoder = MultiHeadAttention(
            channels=n_hidden_dims,
            out_channels=n_hidden_dims,
            n_heads=n_attn_heads,
        )

        self.memory_bank = nn.Parameter(torch.randn(n_hidden_dims, bank_size))
        self.proj = nn.Conv1d(n_hidden_dims, output_channels, 1)
        if init_values is not None:
            with torch.no_grad():
                self.memory_bank.copy_(init_values)

    def forward(self, z: torch.Tensor):
        b, _, _ = z.shape
        ret = self.encoder(
            z, self.memory_bank.unsqueeze(0).repeat(b, 1, 1), attn_mask=None
        )
        ret = self.proj(ret)
        return ret