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

from roma.utils.utils import get_grid
from .layers.block import Block
from .layers.attention import MemEffAttention
from .dinov2 import vit_large

device = "cuda" if torch.cuda.is_available() else "cpu"


class TransformerDecoder(nn.Module):
    def __init__(
        self,
        blocks,
        hidden_dim,
        out_dim,
        is_classifier=False,
        *args,
        amp=False,
        pos_enc=True,
        learned_embeddings=False,
        embedding_dim=None,
        **kwargs
    ) -> None:
        super().__init__(*args, **kwargs)
        self.blocks = blocks
        self.to_out = nn.Linear(hidden_dim, out_dim)
        self.hidden_dim = hidden_dim
        self.out_dim = out_dim
        self._scales = [16]
        self.is_classifier = is_classifier
        self.amp = amp
        if torch.cuda.is_available():
            if torch.cuda.is_bf16_supported():
                self.amp_dtype = torch.bfloat16
            else:
                self.amp_dtype = torch.float16
        else:
            self.amp_dtype = torch.float32

        self.pos_enc = pos_enc
        self.learned_embeddings = learned_embeddings
        if self.learned_embeddings:
            self.learned_pos_embeddings = nn.Parameter(
                nn.init.kaiming_normal_(
                    torch.empty((1, hidden_dim, embedding_dim, embedding_dim))
                )
            )

    def scales(self):
        return self._scales.copy()

    def forward(self, gp_posterior, features, old_stuff, new_scale):
        with torch.autocast(device, dtype=self.amp_dtype, enabled=self.amp):
            B, C, H, W = gp_posterior.shape
            x = torch.cat((gp_posterior, features), dim=1)
            B, C, H, W = x.shape
            grid = get_grid(B, H, W, x.device).reshape(B, H * W, 2)
            if self.learned_embeddings:
                pos_enc = (
                    F.interpolate(
                        self.learned_pos_embeddings,
                        size=(H, W),
                        mode="bilinear",
                        align_corners=False,
                    )
                    .permute(0, 2, 3, 1)
                    .reshape(1, H * W, C)
                )
            else:
                pos_enc = 0
            tokens = x.reshape(B, C, H * W).permute(0, 2, 1) + pos_enc
            z = self.blocks(tokens)
            out = self.to_out(z)
            out = out.permute(0, 2, 1).reshape(B, self.out_dim, H, W)
            warp, certainty = out[:, :-1], out[:, -1:]
            return warp, certainty, None