from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn as nn
from einops import rearrange
from jaxtyping import Float
from torch import Tensor

from sf3d.models.tokenizers.dinov2 import Dinov2Model
from sf3d.models.transformers.attention import Modulation
from sf3d.models.utils import BaseModule


class DINOV2SingleImageTokenizer(BaseModule):
    @dataclass
    class Config(BaseModule.Config):
        pretrained_model_name_or_path: str = "facebook/dinov2-large"
        width: int = 512
        height: int = 512
        modulation_cond_dim: int = 768

    cfg: Config

    def configure(self) -> None:
        self.model = Dinov2Model.from_pretrained(self.cfg.pretrained_model_name_or_path)

        for p in self.model.parameters():
            p.requires_grad_(False)
        self.model.eval()

        self.model.set_gradient_checkpointing(False)

        # add modulation
        modulations = []
        for layer in self.model.encoder.layer:
            norm1_modulation = Modulation(
                self.model.config.hidden_size,
                self.cfg.modulation_cond_dim,
                zero_init=True,
                single_layer=True,
            )
            norm2_modulation = Modulation(
                self.model.config.hidden_size,
                self.cfg.modulation_cond_dim,
                zero_init=True,
                single_layer=True,
            )
            layer.register_ada_norm_modulation(norm1_modulation, norm2_modulation)
            modulations += [norm1_modulation, norm2_modulation]
        self.modulations = nn.ModuleList(modulations)

        self.register_buffer(
            "image_mean",
            torch.as_tensor([0.485, 0.456, 0.406]).reshape(1, 1, 3, 1, 1),
            persistent=False,
        )
        self.register_buffer(
            "image_std",
            torch.as_tensor([0.229, 0.224, 0.225]).reshape(1, 1, 3, 1, 1),
            persistent=False,
        )

    def forward(
        self,
        images: Float[Tensor, "B *N C H W"],
        modulation_cond: Optional[Float[Tensor, "B *N Cc"]],
        **kwargs,
    ) -> Float[Tensor, "B *N Ct Nt"]:
        model = self.model

        packed = False
        if images.ndim == 4:
            packed = True
            images = images.unsqueeze(1)
            if modulation_cond is not None:
                assert modulation_cond.ndim == 2
                modulation_cond = modulation_cond.unsqueeze(1)

        batch_size, n_input_views = images.shape[:2]
        images = (images - self.image_mean) / self.image_std
        out = model(
            rearrange(images, "B N C H W -> (B N) C H W"),
            modulation_cond=rearrange(modulation_cond, "B N Cc -> (B N) Cc")
            if modulation_cond is not None
            else None,
        )
        local_features = out.last_hidden_state
        local_features = local_features.permute(0, 2, 1)
        local_features = rearrange(
            local_features, "(B N) Ct Nt -> B N Ct Nt", B=batch_size
        )
        if packed:
            local_features = local_features.squeeze(1)

        return local_features

    def detokenize(self, *args, **kwargs):
        raise NotImplementedError