import torch
import torch.nn as nn

from .layers import PatchTransformerEncoder, PixelWiseDotProduct


class mViT(nn.Module):
    def __init__(
        self,
        in_channels,
        n_query_channels=128,
        patch_size=16,
        dim_out=256,
        embedding_dim=128,
        num_heads=4,
        norm="linear",
    ):
        super(mViT, self).__init__()
        self.norm = norm
        self.n_query_channels = n_query_channels
        self.patch_transformer = PatchTransformerEncoder(
            in_channels, patch_size, embedding_dim, num_heads
        )
        self.dot_product_layer = PixelWiseDotProduct()

        self.conv3x3 = nn.Conv2d(
            in_channels, embedding_dim, kernel_size=3, stride=1, padding=1
        )
        self.regressor = nn.Sequential(
            nn.Linear(embedding_dim, 256),
            nn.LeakyReLU(),
            nn.Linear(256, 256),
            nn.LeakyReLU(),
            nn.Linear(256, dim_out),
        )

    def forward(self, x):
        # n, c, h, w = x.size()
        tgt = self.patch_transformer(x.clone())  # .shape = S, N, E

        x = self.conv3x3(x)

        regression_head, queries = tgt[0, ...], tgt[1 : self.n_query_channels + 1, ...]

        # Change from S, N, E to N, S, E
        queries = queries.permute(1, 0, 2)
        range_attention_maps = self.dot_product_layer(
            x, queries
        )  # .shape = n, n_query_channels, h, w

        y = self.regressor(regression_head)  # .shape = N, dim_out
        if self.norm == "linear":
            y = torch.relu(y)
            eps = 0.1
            y = y + eps
        elif self.norm == "softmax":
            return torch.softmax(y, dim=1), range_attention_maps
        else:
            y = torch.sigmoid(y)
        y = y / y.sum(dim=1, keepdim=True)
        return y, range_attention_maps
