import torch.nn as nn

from ...models.dino.head import DINOHead


class iBOTHead(DINOHead):
    def __init__(
        self,
        *args,
        patch_out_dim=8192,
        n_layers=3,
        hidden_dim=2048,
        bottleneck_dim=256,
        norm_last_layer=True,
        shared_head=False,
        **kwargs
    ):
        super(iBOTHead, self).__init__(
            *args,
            n_layers=n_layers,
            hidden_dim=hidden_dim,
            bottleneck_dim=bottleneck_dim,
            norm_last_layer=norm_last_layer,
            **kwargs
        )

        if not shared_head:
            if bottleneck_dim > 0:
                self.last_layer2 = nn.utils.weight_norm(
                    nn.Linear(bottleneck_dim, patch_out_dim, bias=False)
                )
                self.last_layer2.weight_g.data.fill_(1)
                if norm_last_layer:
                    self.last_layer2.weight_g.requires_grad = False
            else:
                self.mlp2 = nn.Linear(hidden_dim, patch_out_dim)
                self.last_layer2 = None

        else:
            if bottleneck_dim > 0:
                self.last_layer2 = self.last_layer
            else:
                self.mlp2 = self.mlp[-1]
                self.last_layer2 = None

    def forward(self, x):
        if len(x.shape) == 2:
            return super(iBOTHead, self).forward(x)

        if self.last_layer is not None:
            x = self.mlp(x)
            x = nn.functional.normalize(x, dim=-1, p=2)
            x1 = self.last_layer(x[:, 0])
            x2 = self.last_layer2(x[:, 1:])
        else:
            x = self.mlp[:-1](x)
            x1 = self.mlp[-1](x[:, 0])
            x2 = self.mlp2(x[:, 1:])

        return x1, x2

    def _build_norm(self, norm, hidden_dim, **kwargs):
        if norm == "bn":
            norm = nn.BatchNorm1d(hidden_dim, **kwargs)
        elif norm == "syncbn":
            norm = nn.SyncBatchNorm(hidden_dim, **kwargs)
        elif norm == "ln":
            norm = nn.LayerNorm(hidden_dim, **kwargs)
        else:
            assert norm is None, "unknown norm type {}".format(norm)
        return norm

    def _build_act(self, act):
        if act == "relu":
            act = nn.ReLU()
        elif act == "gelu":
            act = nn.GELU()
        else:
            assert False, "unknown act type {}".format(act)
        return act