import math
from functools import lru_cache
from typing import Optional

import torch
from huggingface_hub import hf_hub_download
from safetensors import safe_open
from torch import nn
from transformers.models.dinov3_vit.configuration_dinov3_vit import DINOv3ViTConfig

from src.model.dino import (
    DINOv3ViTAttention,
    DINOv3ViTEmbeddings,
    DINOv3ViTLayer,
    DINOv3ViTLayerScale,
    DINOv3ViTRopePositionEmbedding,
)


class TimestepEmbedder(nn.Module):
    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size),
        )
        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def timestep_embedding(t, dim, max_period=10000):
        half = dim // 2
        freqs = torch.exp(
            -math.log(max_period) * torch.arange(start=0, end=half) / half
        ).to(t.device)
        args = t[:, None] * freqs[None]
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2:
            embedding = torch.cat(
                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
            )
        return embedding

    def forward(self, t):
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(
            dtype=next(self.parameters()).dtype
        )
        t_emb = self.mlp(t_freq)
        return t_emb


class LabelEmbedder(nn.Module):
    def __init__(self, num_classes, hidden_size, dropout_prob):
        super().__init__()
        use_cfg_embedding = int(dropout_prob > 0)
        self.embedding_table = nn.Embedding(
            num_classes + use_cfg_embedding, hidden_size
        )
        self.num_classes = num_classes
        self.dropout_prob = dropout_prob

    def token_drop(self, labels, force_drop_ids=None):
        if force_drop_ids is None:
            drop_ids = torch.rand(labels.shape[0]) < self.dropout_prob
            drop_ids = drop_ids.cuda()
            drop_ids = drop_ids.to(labels.device)
        else:
            drop_ids = force_drop_ids == 1
        labels = torch.where(drop_ids, self.num_classes, labels)
        return labels

    def forward(self, labels, train, force_drop_ids=None):
        use_dropout = self.dropout_prob > 0
        if (train and use_dropout) or (force_drop_ids is not None):
            labels = self.token_drop(labels, force_drop_ids)
        embeddings = self.embedding_table(labels)
        return embeddings


class DinoConditionedLayer(DINOv3ViTLayer):
    def __init__(self, config: DINOv3ViTConfig, is_encoder: bool = False):
        super().__init__(config)
        self.is_encoder = is_encoder

        self.norm_cond = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.cond = DINOv3ViTAttention(config)
        self.layer_scale_cond = DINOv3ViTLayerScale(config)

        # no init zeros!
        if is_encoder:
            nn.init.constant_(self.layer_scale_cond.lambda1, 0)
            self.norm1.requires_grad_(False)
            self.norm2.requires_grad_(False)
            self.attention.requires_grad_(False)
            self.mlp.requires_grad_(False)
            self.layer_scale1.requires_grad_(False)
            self.layer_scale2.requires_grad_(False)

    def forward(
        self,
        hidden_states: torch.Tensor,
        *,
        conditioning_input: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        do_condition: bool = True,
        **kwargs,
    ) -> torch.Tensor:
        assert position_embeddings is not None
        assert conditioning_input is not None or not do_condition

        residual = hidden_states
        hidden_states = self.norm1(hidden_states)
        hidden_states = self.attention(
            hidden_states,
            attention_mask=attention_mask,
            position_embeddings=position_embeddings,
        )
        hidden_states = self.layer_scale1(hidden_states)
        hidden_states = self.drop_path(hidden_states) + residual

        if do_condition:
            residual = hidden_states
            hidden_states = self.norm_cond(hidden_states)
            hidden_states = self.cond(
                hidden_states,
                conditioning_input,
            )
            hidden_states = self.layer_scale_cond(hidden_states)
            hidden_states = self.drop_path(hidden_states) + residual

        residual = hidden_states
        hidden_states = self.norm2(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = self.layer_scale2(hidden_states)
        hidden_states = self.drop_path(hidden_states) + residual

        return hidden_states


# class DinoV3ViTDecoder(nn.Module):
#     def __init__(self, config: DINOv3ViTConfig):
#         super().__init__()
#         self.config = config
#         self.num_channels_out = config.num_channels

#         self.projection = nn.Linear(
#             config.hidden_size,
#             self.num_channels_out * config.patch_size * config.patch_size,
#             bias=True,
#         )

#     def forward(self, x: torch.Tensor, image_size: tuple[int, int]) -> torch.Tensor:
#         batch_size = x.shape[0]

#         num_special_tokens = 1 + self.config.num_register_tokens
#         patch_tokens = x[:, num_special_tokens:, :]

#         projected_tokens = self.projection(patch_tokens)

#         p = self.config.patch_size
#         c = self.num_channels_out
#         h_grid = image_size[0] // p
#         w_grid = image_size[1] // p

#         assert patch_tokens.shape[1] == h_grid * w_grid, (
#             "Number of patches does not match image size."
#         )

#         x_reshaped = projected_tokens.reshape(batch_size, h_grid, w_grid, p, p, c)

#         x_permuted = torch.einsum("nhwpqc->nchpwq", x_reshaped)

#         reconstructed_image = x_permuted.reshape(batch_size, c, h_grid * p, w_grid * p)

#         return reconstructed_image

# lets try conv decoder


class DinoV3ViTDecoder(nn.Module):
    def __init__(self, config: DINOv3ViTConfig):
        super().__init__()
        self.config = config
        self.num_channels_out = config.num_channels
        self.patch_size = config.patch_size

        self.projection = nn.Linear(
            config.hidden_size, config.num_channels * (self.patch_size**2), bias=True
        )
        self.pixel_shuffle = nn.PixelShuffle(self.patch_size)

        nn.init.zeros_(self.projection.weight)
        nn.init.zeros_(
            self.projection.bias
        ) if self.projection.bias is not None else None

    def forward(self, x: torch.Tensor, image_size: tuple[int, int]) -> torch.Tensor:
        batch_size = x.shape[0]

        x = x[:, 1 + self.config.num_register_tokens :, :]

        p = self.config.patch_size
        h_grid = image_size[0] // p
        w_grid = image_size[1] // p

        assert x.shape[1] == h_grid * w_grid
        x = self.projection(x)
        x = x.reshape(batch_size, h_grid, w_grid, -1).permute(0, 3, 1, 2)
        x = self.pixel_shuffle(x)

        return x


# how about transposed conv decoderclass DinoV3ViTDecoder(nn.Module):
# class DinoV3ViTDecoder(nn.Module):
#     def __init__(self, config: DINOv3ViTConfig):
#         super().__init__()
#         self.config = config
#         self.num_channels_out = config.num_channels
#         self.patch_size = config.patch_size

#         intermediate_channels = config.hidden_size // 4

#         self.decoder_block = nn.Sequential(
#             nn.ConvTranspose2d(
#                 in_channels=config.hidden_size,
#                 out_channels=intermediate_channels,
#                 kernel_size=self.patch_size,
#                 stride=self.patch_size,
#                 bias=True,
#             ),
#             nn.LayerNorm(intermediate_channels),
#             nn.GELU(),
#             nn.Conv2d(
#                 in_channels=intermediate_channels,
#                 out_channels=config.num_channels,
#                 kernel_size=1,
#                 bias=True,
#             ),
#         )

#         nn.init.zeros_(self.decoder_block[-1].weight)
#         nn.init.zeros_(self.decoder_block[-1].bias)

#     def forward(self, x: torch.Tensor, image_size: tuple[int, int]) -> torch.Tensor:
#         batch_size = x.shape[0]

#         x = x[:, 1 + self.config.num_register_tokens :, :]

#         p = self.config.patch_size
#         h_grid = image_size[0] // p
#         w_grid = image_size[1] // p
#         assert x.shape[1] == h_grid * w_grid

#         x = x.transpose(1, 2).reshape(
#             batch_size, self.config.hidden_size, h_grid, w_grid
#         )
#         x = self.decoder_block(x)

#         return x


class NerfEmbedder(nn.Module):
    def __init__(self, in_channels, hidden_size_input, max_freqs):
        super().__init__()
        self.max_freqs = max_freqs
        self.hidden_size_input = hidden_size_input
        self.embedder = nn.Sequential(
            nn.Linear(in_channels + max_freqs**2, hidden_size_input, bias=True),
        )

    @lru_cache
    def fetch_pos(self, patch_size, device, dtype):
        pos_x = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
        pos_y = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
        pos_y, pos_x = torch.meshgrid(pos_y, pos_x, indexing="ij")
        pos_x = pos_x.reshape(-1, 1, 1)
        pos_y = pos_y.reshape(-1, 1, 1)

        freqs = torch.linspace(
            0, self.max_freqs, self.max_freqs, dtype=dtype, device=device
        )
        freqs_x = freqs[None, :, None]
        freqs_y = freqs[None, None, :]
        coeffs = (1 + freqs_x * freqs_y) ** -1
        dct_x = torch.cos(pos_x * freqs_x * torch.pi)
        dct_y = torch.cos(pos_y * freqs_y * torch.pi)
        dct = (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs**2)
        return dct

    def forward(self, inputs):
        target_dtype = self.embedder[0].weight.dtype
        inputs = inputs.to(dtype=target_dtype)
        B, P2, C = inputs.shape
        patch_size = int(P2**0.5)
        device = inputs.device
        dtype = inputs.dtype
        dct = self.fetch_pos(patch_size, device, dtype)
        dct = dct.repeat(B, 1, 1)
        inputs = torch.cat([inputs, dct], dim=-1)
        inputs = self.embedder(inputs)
        return inputs


class NerfBlock(nn.Module):
    def __init__(self, hidden_size_s: int, hidden_size_x: int, mlp_ratio: int = 4):
        super().__init__()
        self.param_generator1 = nn.Sequential(
            nn.Linear(hidden_size_s, 2 * hidden_size_x**2 * mlp_ratio, bias=True),
        )
        self.norm = nn.RMSNorm(hidden_size_x, eps=1e-6)
        self.mlp_ratio = mlp_ratio

    def forward(self, x, s):
        batch_size, num_x, hidden_size_x = x.shape
        mlp_params1 = self.param_generator1(s)
        fc1_param1, fc2_param1 = mlp_params1.chunk(2, dim=-1)
        fc1_param1 = fc1_param1.view(
            batch_size, hidden_size_x, hidden_size_x * self.mlp_ratio
        )
        fc2_param1 = fc2_param1.view(
            batch_size, hidden_size_x * self.mlp_ratio, hidden_size_x
        )

        # normalize fc1
        normalized_fc1_param1 = torch.nn.functional.normalize(fc1_param1, dim=-2)
        # normalize fc2
        normalized_fc2_param1 = torch.nn.functional.normalize(fc2_param1, dim=-2)
        # mlp 1
        res_x = x
        x = self.norm(x)
        x = torch.bmm(x, normalized_fc1_param1)
        x = torch.nn.functional.silu(x)
        x = torch.bmm(x, normalized_fc2_param1)
        x = x + res_x
        return x


class NerfFinalLayer(nn.Module):
    def __init__(self, hidden_size, out_channels):
        super().__init__()
        self.norm = nn.RMSNorm(hidden_size, eps=1e-6)
        self.linear = nn.Linear(hidden_size, out_channels, bias=True)

    def forward(self, x):
        x = self.norm(x)
        x = self.linear(x)
        return x


class UTransformer(nn.Module):
    def __init__(
        self,
        config: DINOv3ViTConfig,
        num_classes: int,
        nerf_patch=16,
        nerf_hidden=64,
        scale_factor: int = 4,
    ):
        super().__init__()
        self.config = config
        self.scale_factor = scale_factor
        self.nerf_patch_size = nerf_patch

        assert config.num_hidden_layers % scale_factor == 0

        self.embeddings = DINOv3ViTEmbeddings(config)
        self.rope_embeddings = DINOv3ViTRopePositionEmbedding(config)
        self.t_embedder = TimestepEmbedder(config.hidden_size)

        self.encoder_layers = nn.ModuleList(
            [
                DinoConditionedLayer(config, True)
                for _ in range(config.num_hidden_layers)
            ]
        )
        self.encoder_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

        self.decoder_layers = nn.ModuleList(
            [
                DinoConditionedLayer(config, False)
                for _ in range(config.num_hidden_layers // scale_factor)
            ]
        )
        self.residual_merger = nn.ModuleList(
            [
                nn.Sequential(
                    nn.SiLU(), nn.Linear(config.hidden_size, 2 * config.hidden_size)
                )
                for _ in range(config.num_hidden_layers // scale_factor)
            ]
        )
        self.rest_decoder = nn.ModuleList(
            [DinoConditionedLayer(config, False) for _ in range(4)]
        )

        # nerf!
        self.nerf_encoder = NerfEmbedder(3, nerf_hidden, 8)  # (rgb, hidden, freq)
        self.nerf_decoder = nn.ModuleList(
            [NerfBlock(self.config.hidden_size, nerf_hidden) for _ in range(12)]
        )
        self.final_layer = NerfFinalLayer(nerf_hidden, 3)

        # freeze pretrained
        self.embeddings.requires_grad_(False)
        self.rope_embeddings.requires_grad_(False)
        self.encoder_norm.requires_grad_(False)

    def forward(
        self,
        pixel_values: torch.Tensor,
        time: torch.Tensor,
        # cond: torch.Tensor,
        bool_masked_pos: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
    ):
        if time.dim() == 0:
            time = time.repeat(pixel_values.shape[0])

        # resolution config
        B = pixel_values.shape[0]
        dino_h = pixel_values.shape[-2] // self.config.patch_size
        dino_w = pixel_values.shape[-1] // self.config.patch_size
        nerf_h = pixel_values.shape[-2] // self.nerf_patch_size
        nerf_w = pixel_values.shape[-1] // self.nerf_patch_size

        pixel_values = pixel_values.to(self.embeddings.patch_embeddings.weight.dtype)
        position_embeddings = self.rope_embeddings(pixel_values)

        x = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
        t = self.t_embedder(time).unsqueeze(1)
        # y = self.y_embedder(cond, self.training).unsqueeze(1)
        # conditioning_input = t.to(x.dtype) + y.to(x.dtype)
        conditioning_input = t.to(x.dtype)

        residual = []
        for i, layer_module in enumerate(self.encoder_layers):
            if i % self.scale_factor == 0:
                residual.append(x)
            layer_head_mask = head_mask[i] if head_mask is not None else None
            x = layer_module(
                x,
                conditioning_input=conditioning_input,
                attention_mask=layer_head_mask,
                position_embeddings=position_embeddings,
            )

        x = self.encoder_norm(x)

        reversed_residual = residual[::-1]
        for i, layer_module in enumerate(self.decoder_layers):
            layer_head_mask = head_mask[i] if head_mask is not None else None
            x = layer_module(
                x,
                conditioning_input=conditioning_input,
                attention_mask=layer_head_mask,
                position_embeddings=position_embeddings,
                do_condition=False,
            )
            shift, scale = self.residual_merger[i](reversed_residual[i]).chunk(
                2, dim=-1
            )
            x = x * (1 + scale) + shift

        for i, layer_module in enumerate(self.rest_decoder):
            layer_head_mask = head_mask[i] if head_mask is not None else None
            x = layer_module(
                x,
                conditioning_input=conditioning_input,
                attention_mask=layer_head_mask,
                position_embeddings=position_embeddings,
                do_condition=False,
            )  # (batch, image // patch^2, 1024)

        x = x[:, 1 + self.config.num_register_tokens :, :]

        nerf_cond = nn.functional.silu(t + x)  # (batch, image // patch^2, 1024)
        nerf_cond = nerf_cond.reshape(
            B, dino_h, dino_w, self.config.hidden_size
        ).permute(0, 3, 1, 2)  # (batch, 1024, image // patch, image // patch)
        # nerf_cond = nn.functional.interpolate(
        #     nerf_cond, size=(nerf_h, nerf_w), mode="bilinear", align_corners=False
        # )
        nerf_cond = (
            nerf_cond.permute(0, 2, 3, 1)
            .reshape(-1, nerf_h * nerf_w, self.config.hidden_size)
            .view(-1, self.config.hidden_size)
        )

        # nerf
        x_nerf = nn.functional.unfold(
            pixel_values, self.nerf_patch_size, stride=self.nerf_patch_size
        ).transpose(1, 2)
        x_nerf = x_nerf.reshape(
            B * x_nerf.shape[1], -1, self.nerf_patch_size**2
        ).transpose(1, 2)
        x_nerf = self.nerf_encoder(x_nerf)

        for module in self.nerf_decoder:
            x_nerf = module(x_nerf, nerf_cond)

        x_nerf = self.final_layer(x_nerf)

        num_patches = nerf_h * nerf_w
        x_nerf = x_nerf.reshape(
            B * num_patches, -1
        )  # (B*num_patches, 48): flatten pixels+RGB per patch
        x_nerf = (
            x_nerf.view(B, num_patches, -1).transpose(1, 2).contiguous()
        )  # (B, 48, num_patches)

        res = nn.functional.fold(
            x_nerf,
            (pixel_values.shape[-2], pixel_values.shape[-1]),
            kernel_size=self.nerf_patch_size,
            stride=self.nerf_patch_size,
        )
        return res

    def get_residual(
        self,
        pixel_values: torch.Tensor,
        time: Optional[torch.Tensor],
        do_condition: bool,
    ):
        pixel_values = pixel_values.to(self.embeddings.patch_embeddings.weight.dtype)
        position_embeddings = self.rope_embeddings(pixel_values)

        x = self.embeddings(pixel_values, bool_masked_pos=None)

        if do_condition:
            t = self.t_embedder(time).unsqueeze(1)
            # y = self.y_embedder(cond, self.training).unsqueeze(1)
            # conditioning_input = t.to(x.dtype) + y.to(x.dtype)
            conditioning_input = t.to(x.dtype)
        else:
            conditioning_input = None

        residual = []
        for i, layer_module in enumerate(self.encoder_layers):
            if i % self.scale_factor == 0:
                residual.append(x)

            x = layer_module(
                x,
                conditioning_input=conditioning_input,
                attention_mask=None,
                position_embeddings=position_embeddings,
                do_condition=do_condition,
            )

        return residual

    @staticmethod
    def from_pretrained_backbone(name: str):
        config = DINOv3ViTConfig.from_pretrained(name)
        instance = UTransformer(config, 0)

        weight_dict = {}
        with safe_open(
            hf_hub_download(name, "model.safetensors"), framework="pt", device="cuda:1"
        ) as f:
            for key in f.keys():
                new_key = key.replace("layer.", "encoder_layers.").replace(
                    "norm.", "encoder_norm."
                )

                weight_dict[new_key] = f.get_tensor(key)

        instance.load_state_dict(weight_dict, strict=False)

        return instance
