import copy
import math
from typing import Optional

import einops
import torch
import torch.nn.functional as F
from einops import rearrange
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.attention import CrossAttention, PlainAttention, RoPE
from src.model.dino import (
    DINOv3ViTAttention,
    DINOv3ViTDropPath,
    DINOv3ViTEmbeddings,
    DINOv3ViTGatedMLP,
    DINOv3ViTLayer,
    DINOv3ViTLayerScale,
    DINOv3ViTMLP,
    DINOv3ViTRopePositionEmbedding,
)
from src.model.dit import modulate
from src.model.resnet import ResBlock


def create_coordinate(h, w, start=0, end=1, device="cuda:1", dtype=torch.float32):
    # Create a grid of coordinates
    x = torch.linspace(start, end, h, device=device, dtype=dtype)
    y = torch.linspace(start, end, w, device=device, dtype=dtype)
    # Create a 2D map using meshgrid
    xx, yy = torch.meshgrid(x, y, indexing="ij")
    # Stack the x and y coordinates to create the final map
    coord_map = torch.stack([xx, yy], dim=-1)[None, ...]
    coords = rearrange(coord_map, "b h w c -> b (h w) c", h=h, w=w)
    return coords


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 DinoEncoderLayer(DINOv3ViTLayer):
    def __init__(self, config: DINOv3ViTConfig):
        super().__init__(config)

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

        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,
        **kwargs,
    ) -> torch.Tensor:
        assert position_embeddings is not None

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

        hidden_states = (
            self.drop_path(
                self.layer_scale_cond(
                    self.cond(
                        hidden_states,
                        self.norm_cond(conditioning_input),
                    )
                )
            )
            + hidden_states
        )

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

        return hidden_states


class DinoDecoderLayer(DINOv3ViTLayer):
    def __init__(self, config: DINOv3ViTConfig, depth: int):
        super().__init__(config)

        hidden_size = config.hidden_size // (16**depth)
        hacky_config = copy.copy(config)
        hacky_config.hidden_size = hidden_size
        hacky_config.intermediate_size = hacky_config.intermediate_size // (16**depth)

        self.norm1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)
        self.attention = PlainAttention(
            hidden_size, config.num_attention_heads // (4**depth)
        )  # head scaling law?
        self.layer_scale1 = DINOv3ViTLayerScale(hacky_config)
        self.drop_path = (
            DINOv3ViTDropPath(hacky_config.drop_path_rate)
            if hacky_config.drop_path_rate > 0.0
            else nn.Identity()
        )

        self.norm2 = nn.LayerNorm(hidden_size, eps=hacky_config.layer_norm_eps)

        if config.use_gated_mlp:
            self.mlp = DINOv3ViTGatedMLP(hacky_config)
        else:
            self.mlp = DINOv3ViTMLP(hacky_config)
        self.layer_scale2 = DINOv3ViTLayerScale(hacky_config)

        # adaln
        self.adaln = nn.Sequential(
            nn.SiLU(),
            nn.Linear(config.hidden_size, 6 * hidden_size, bias=True),
        )

    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,
        **kwargs,
    ) -> torch.Tensor:
        conditioning_input = conditioning_input.squeeze(1)  # type: ignore
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln(
            conditioning_input
        ).chunk(6, dim=-1)

        hidden_states = (
            self.drop_path(
                gate_msa.unsqueeze(1)
                * self.layer_scale1(
                    self.attention(
                        modulate(self.norm1(hidden_states), shift_msa, scale_msa),
                        position_embeddings=position_embeddings,
                    )
                )
            )
            + hidden_states
        )

        hidden_states = (
            self.drop_path(
                gate_mlp.unsqueeze(1)
                * self.layer_scale2(
                    self.mlp(modulate(self.norm2(hidden_states), shift_mlp, scale_mlp))
                )
            )
            + hidden_states
        )

        return hidden_states


class ResidualUpscaler(nn.Module):
    def __init__(
        self, config: DINOv3ViTConfig, depth: list[int], bottleneck_dim: int = 128
    ):  # max depth 2 (4**2 = 16 = patch size)
        super().__init__()

        def build_encoder(in_dim, num_layers=2):
            return nn.Sequential(
                nn.Conv2d(
                    in_dim,
                    bottleneck_dim,
                    kernel_size=1,
                    padding=0,
                    padding_mode="reflect",
                    bias=False,
                ),
                *[
                    ResBlock(
                        bottleneck_dim,
                        bottleneck_dim,
                        kernel_size=1,
                        num_groups=8,
                        pad_mode="reflect",
                        norm_fn=nn.GroupNorm,
                        activation_fn=nn.SiLU,
                        use_conv_shortcut=False,
                    )
                    for _ in range(num_layers)
                ],
            )

        self.config = config
        self.depth = depth
        self.global_encode = nn.Linear(
            config.hidden_size * len(depth),
            bottleneck_dim * 2,
        )
        self.local_encode = nn.ModuleList(
            [
                nn.Linear(config.hidden_size, bottleneck_dim * 2)
                if d != 0
                else nn.Identity()
                for d in depth
            ]
        )
        self.q_norm = nn.ModuleList(
            [nn.LayerNorm(bottleneck_dim) if d != 0 else nn.Identity() for d in depth]
        )
        self.k_norm = nn.ModuleList(
            [nn.LayerNorm(bottleneck_dim) if d != 0 else nn.Identity() for d in depth]
        )
        self.v_downsample = nn.ModuleList(
            [
                nn.Linear(config.hidden_size, config.hidden_size // (16**d))
                if d != 0
                else nn.Identity()
                for d in depth
            ]
        )
        self.cross_attn = nn.ModuleList(
            [
                CrossAttention(
                    bottleneck_dim,
                    bottleneck_dim,
                    config.hidden_size // (16**d),
                )
                if d != 0
                else nn.Identity()
                for d in depth
            ]
        )

        self.image_encoder = build_encoder(3)
        self.q_encoder = build_encoder(bottleneck_dim)
        self.k_encoder = build_encoder(bottleneck_dim)

        self.rope = RoPE(bottleneck_dim)

    def forward(self, pixel_values: torch.Tensor, residuals: list[torch.Tensor]):
        # residual[0] => deepest, -1 => shallowest; pixel values (b, 3, h, w) / residuals [(b, seq, d), (b, seq, d)]
        # objective: say we have (1024, 1024, 512) residual. we want to make multi head attention query well
        assert self.config.patch_size is not None
        image_h, image_w = pixel_values.shape[-2], pixel_values.shape[-1]

        global_shift, global_scale = self.global_encode(
            einops.rearrange(
                torch.stack(residuals, dim=1), "b depth s h -> b s (depth h)"
            )
        ).chunk(2, dim=-1)  # patch-level global btw

        image_residual = self.image_encoder(pixel_values)  # messy; todo: cleanup
        coords = create_coordinate(pixel_values.shape[-2], pixel_values.shape[-1])
        image_residual = rearrange(image_residual, "b c h w -> b (h w) c")
        image_residual = self.rope(image_residual, coords)
        image_residual = rearrange(
            image_residual, "b (h w) c -> b c h w", h=pixel_values.shape[-2]
        )
        q = self.q_encoder(image_residual)
        k = self.k_encoder(image_residual)

        reformed_residual = []
        for i, (depth, residual) in enumerate(zip(self.depth, residuals)):
            if depth == 0:
                reformed_residual.append(residual)
                continue

            local_shift, local_scale = self.local_encode[i](residual).chunk(2, dim=-1)
            local_q = self.q_norm[i](
                einops.rearrange(
                    F.adaptive_avg_pool2d(
                        q,
                        output_size=(
                            image_h // self.config.patch_size * (4**depth),
                            image_w // self.config.patch_size * (4**depth),
                        ),
                    ),
                    "b c h w -> b (h w) c",
                )
            )
            local_k = (1 + local_scale) * (
                (1 + global_scale)
                * self.k_norm[i](
                    einops.rearrange(
                        F.adaptive_avg_pool2d(
                            k,
                            output_size=(
                                image_h // self.config.patch_size,
                                image_w // self.config.patch_size,
                            ),
                        ),
                        "b c h w -> b (h w) c",
                    )
                )
                + global_shift
            ) + local_shift
            local_v = self.v_downsample[i](residual)

            reformed_residual.append(self.cross_attn[i](local_q, local_k, local_v))

        return reformed_residual


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 // 16 // 16, config.num_channels, bias=True
        )

        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:
        return self.projection(
            einops.rearrange(
                x, "b (h w) d -> b h w d", h=image_size[0], w=image_size[1]
            )
        ).permute(0, 3, 1, 2)


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

        assert config.num_hidden_layers % scale_factor % 3 == 0

        self.embeddings = DINOv3ViTEmbeddings(config)
        self.rope_embeddings = DINOv3ViTRopePositionEmbedding(config)

        def gen_rope(depth: int):
            hidden_size = config.hidden_size // (16**depth)
            hacky_config = copy.copy(config)
            hacky_config.hidden_size = hidden_size
            hacky_config.intermediate_size = hacky_config.intermediate_size // (
                16**depth
            )
            hacky_config.num_attention_heads = hacky_config.num_attention_heads // (
                4**depth
            )
            return DINOv3ViTRopePositionEmbedding(hacky_config)

        self.decode_ropes = nn.ModuleList([gen_rope(i + 1) for i in range(2)])
        self.t_embedder = TimestepEmbedder(config.hidden_size)

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

        self.residual_upscaler = ResidualUpscaler(
            config,
            [0, 0, 1, 1, 2, 2],  # hardcoded, sorry
            bottleneck_dim=128,
        )
        self.decoder_layers = nn.ModuleList(
            [
                nn.ModuleList(
                    [
                        DinoDecoderLayer(config, depth)
                        for _ in range((config.num_hidden_layers // scale_factor) // 3)
                    ]
                )
                for depth in range(3)
            ]
        )
        self.residual_merger = nn.ModuleList(
            [
                nn.ModuleList(
                    [
                        nn.Sequential(
                            nn.SiLU(),
                            nn.Linear(
                                config.hidden_size // (16**depth),
                                2 * config.hidden_size // (16**depth),
                            ),
                        )
                        for _ in range((config.num_hidden_layers // scale_factor) // 3)
                    ]
                )
                for depth in range(3)
            ]
        )
        self.upsample = nn.ModuleList([nn.PixelShuffle(4) for _ in range(2)])
        self.rest_decoder = nn.ModuleList(
            [DinoDecoderLayer(config, 2) for _ in range(4)]
        )
        self.decoder_norm = nn.LayerNorm(
            (config.hidden_size // (16**2)), eps=config.layer_norm_eps
        )
        self.decoder = DinoV3ViTDecoder(config)

        # 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])

        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[:, 1 + self.config.num_register_tokens :])
            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)
        x = x[:, 1 + self.config.num_register_tokens :]
        reversed_residual = self.residual_upscaler(pixel_values, residual[::-1])

        residual_idx = 0
        for depth, layers in enumerate(self.decoder_layers):
            for i, layer_module in enumerate(layers):  # type: ignore
                x = layer_module(
                    x,
                    conditioning_input=conditioning_input,
                    attention_mask=None,
                    position_embeddings=position_embeddings,
                )
                shift, scale = self.residual_merger[depth][i](  # type: ignore
                    reversed_residual[residual_idx]
                ).chunk(2, dim=-1)
                x = x * (1 + scale) + shift
                residual_idx += 1
            x = (
                rearrange(
                    self.upsample[depth](
                        rearrange(
                            x,
                            "b (h w) d -> b d h w",
                            h=pixel_values.shape[-2]
                            // (self.config.patch_size)
                            * (4**depth),
                        )
                    ),
                    "b d h w -> b (h w) d",
                )
                if depth != 2
                else x
            )
            position_embeddings = (
                self.decode_ropes[depth](
                    torch.zeros(
                        (
                            1,
                            1,
                            pixel_values.shape[-2] * (4 ** (depth + 1)),
                            pixel_values.shape[-1] * (4 ** (depth + 1)),
                        ),
                        device=x.device,
                    ).to(self.embeddings.patch_embeddings.weight.dtype)
                )
                if depth != 2
                else position_embeddings
            )

        for i, layer_module in enumerate(self.rest_decoder):
            x = layer_module(
                x,
                conditioning_input=conditioning_input,
                attention_mask=None,
                position_embeddings=position_embeddings,
            )

        x = self.decoder_norm(x)

        return self.decoder(x, image_size=pixel_values.shape[-2:]), residual

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

        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
