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

import comfy.model_management
import comfy.ops
from comfy.image_encoders.dino2 import Dino2PatchEmbeddings, Dinov2Model
from comfy.ldm.modules.attention import optimized_attention_for_device
from comfy.ldm.modules.diffusionmodules.model import Decoder

ops = comfy.ops.disable_weight_init

DINO_CONFIG = {"num_hidden_layers": 12, "hidden_size": 768, "num_attention_heads": 12, "layer_norm_eps": 1e-6, "use_swiglu_ffn": False, "use_mask_token": False}
DINO_MEAN = (0.485, 0.456, 0.406)
DINO_STD = (0.229, 0.224, 0.225)
DINO_PATCH_SIZE = 14
DINO_GRID = 37
LATENT_DOWNSAMPLE_FACTOR = 16
ENCODER_LAYERS = 6


def dino_hidden_states(dino, pixels, patch_embeddings):
    embeddings = dino.embeddings
    x = patch_embeddings(pixels)
    pos = comfy.model_management.cast_to(embeddings.position_embeddings, dtype=torch.float32, device=x.device)
    patch_pos = pos[:, 1:].reshape(1, DINO_GRID, DINO_GRID, -1).permute(0, 3, 1, 2)
    grid = (pixels.shape[-2] // DINO_PATCH_SIZE, pixels.shape[-1] // DINO_PATCH_SIZE)
    patch_pos = F.interpolate(patch_pos, size=grid, mode="bicubic", antialias=True).flatten(2).transpose(1, 2)
    pos = torch.cat([pos[:, :1], patch_pos], dim=1).to(x.dtype)
    cls = comfy.model_management.cast_to(embeddings.cls_token, dtype=x.dtype, device=x.device)
    x = torch.cat([cls.expand(x.shape[0], -1, -1), x], dim=1) + pos

    optimized_attention = optimized_attention_for_device(x.device, False, small_input=True)
    hidden_states = []
    for layer in dino.encoder.layer:
        x = layer(x, optimized_attention)
        hidden_states.append(x)
    return hidden_states


class Nanosaur2VAE(nn.Module):

    def __init__(self, latent_channels=64):
        super().__init__()
        embed_dim = DINO_CONFIG["hidden_size"]
        self.encoder = Dinov2Model(DINO_CONFIG, None, None, ops)
        self.semantic_patch_embeddings = Dino2PatchEmbeddings(embed_dim, patch_size=DINO_PATCH_SIZE, operations=ops)
        self.feature_norms = nn.ModuleList([ops.LayerNorm(embed_dim) for _ in range(ENCODER_LAYERS)])
        self.encoder_projection = ops.Conv2d(embed_dim * ENCODER_LAYERS, latent_channels // 2, 1)
        self.semantic_projection = ops.Conv2d(embed_dim, latent_channels // 2, 1)
        self.decoder = Decoder(ch=128, out_ch=3, ch_mult=(1, 1, 2, 2, 4), num_res_blocks=2, attn_resolutions=(16,), in_channels=3, resolution=256, z_channels=latent_channels, tanh_out=True)
        self.register_buffer("latent_mean", torch.empty(1, latent_channels, 1, 1))
        self.register_buffer("latent_std", torch.empty(1, latent_channels, 1, 1))

    def encode(self, x):
        b, _, h, w = x.shape
        grid_h = max(1, (h + LATENT_DOWNSAMPLE_FACTOR // 2) // LATENT_DOWNSAMPLE_FACTOR)
        grid_w = max(1, (w + LATENT_DOWNSAMPLE_FACTOR // 2) // LATENT_DOWNSAMPLE_FACTOR)
        x = F.interpolate(x, size=(grid_h * DINO_PATCH_SIZE, grid_w * DINO_PATCH_SIZE), mode="bicubic", align_corners=False, antialias=True)
        mean = torch.tensor(DINO_MEAN, device=x.device, dtype=x.dtype).view(1, 3, 1, 1)
        std = torch.tensor(DINO_STD, device=x.device, dtype=x.dtype).view(1, 3, 1, 1)
        x = (x.add(1.0).mul(0.5) - mean) / std

        features = dino_hidden_states(self.encoder, x, self.encoder.embeddings.patch_embeddings)[-ENCODER_LAYERS:]
        features = [norm(f[:, 1:]) for f, norm in zip(features, self.feature_norms)]
        features = torch.cat(features, dim=-1).transpose(1, 2).reshape(b, -1, grid_h, grid_w)
        semantic = self.encoder.layernorm(dino_hidden_states(self.encoder, x, self.semantic_patch_embeddings)[-1][:, 1:])
        semantic = semantic.transpose(1, 2).reshape(b, -1, grid_h, grid_w)
        z = torch.cat([self.encoder_projection(features), self.semantic_projection(semantic)], dim=1)
        return (z - comfy.ops.cast_to_input(self.latent_mean, z)) / comfy.ops.cast_to_input(self.latent_std, z)

    def decode(self, z):
        z = z * comfy.ops.cast_to_input(self.latent_std, z) + comfy.ops.cast_to_input(self.latent_mean, z)
        return self.decoder(z)