"""Conditional VQ-VAE-2 (2 дискретных codebook-уровня, conditioning на класс+elevation).

Архитектура вынесена из обучающего скрипта Анастасии (conditional-vq-vae-2.py),
БЕЗ кода обучения/данных — только nn.Module-классы, чтобы загрузить
`best_conditional_vqvae2.pt` (model_state) и посчитать метрики тем же пайплайном,
что VAE/β-VAE/CVAE. Конфиг чекпоинта: base=56, z=64, cond=24, top=768, bottom=1024.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

COMMITMENT_COST = 0.25
EMA_DECAY = 0.99


class ResBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        groups = min(8, channels)
        while channels % groups != 0:
            groups -= 1
        self.net = nn.Sequential(
            nn.GroupNorm(groups, channels),
            nn.SiLU(inplace=True),
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.GroupNorm(groups, channels),
            nn.SiLU(inplace=True),
            nn.Conv2d(channels, channels, 3, padding=1),
        )

    def forward(self, x):
        return x + self.net(x)


class DownBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 4, stride=2, padding=1),
            nn.GroupNorm(8 if out_ch % 8 == 0 else 1, out_ch),
            nn.SiLU(inplace=True),
            ResBlock(out_ch),
        )

    def forward(self, x):
        return self.net(x)


class UpBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.net = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.GroupNorm(8 if out_ch % 8 == 0 else 1, out_ch),
            nn.SiLU(inplace=True),
            ResBlock(out_ch),
        )

    def forward(self, x):
        return self.net(x)


class VectorQuantizerEMA(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, commitment_cost=0.25, decay=0.99, eps=1e-5):
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.commitment_cost = commitment_cost
        self.decay = decay
        self.eps = eps
        embed = torch.randn(num_embeddings, embedding_dim) / math.sqrt(embedding_dim)
        self.register_buffer('embedding', embed)
        self.register_buffer('cluster_size', torch.zeros(num_embeddings))
        self.register_buffer('embed_avg', embed.clone())

    def forward(self, inputs):
        # inputs: [B, D, H, W]
        b, d, h, w = inputs.shape
        flat = inputs.permute(0, 2, 3, 1).contiguous().view(-1, d)
        flat_float = flat.float()
        embed_float = self.embedding.float()
        distances = (
            flat_float.pow(2).sum(1, keepdim=True)
            - 2 * flat_float @ embed_float.t()
            + embed_float.pow(2).sum(1).unsqueeze(0)
        )
        indices = torch.argmin(distances, dim=1)
        encodings = F.one_hot(indices, self.num_embeddings).type(flat_float.dtype)
        quantized = self.embedding[indices].view(b, h, w, d).permute(0, 3, 1, 2).contiguous()

        if self.training:
            with torch.no_grad():
                cluster_size = encodings.sum(0)
                embed_sum = encodings.t() @ flat_float
                self.cluster_size.mul_(self.decay).add_(cluster_size, alpha=1 - self.decay)
                self.embed_avg.mul_(self.decay).add_(embed_sum, alpha=1 - self.decay)
                n = self.cluster_size.sum()
                cluster_size = (self.cluster_size + self.eps) / (n + self.num_embeddings * self.eps) * n
                embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
                self.embedding.copy_(embed_normalized)

        loss = self.commitment_cost * F.mse_loss(inputs.float(), quantized.detach().float())
        quantized = quantized.to(inputs.dtype)
        quantized = inputs + (quantized - inputs).detach()
        avg_probs = encodings.float().mean(0)
        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
        indices = indices.view(b, h, w)
        return quantized, loss, perplexity, indices

    def decode_indices(self, indices):
        # indices: [B, H, W]
        b, h, w = indices.shape
        quantized = self.embedding[indices.reshape(-1)].view(b, h, w, self.embedding_dim)
        return quantized.permute(0, 3, 1, 2).contiguous()


class ConditionalVQVAE2(nn.Module):
    def __init__(self, image_size=256, base=56, z_dim=64, cond_dim=24, class_emb_dim=12, elev_emb_dim=12,
                 top_codes=768, bottom_codes=1024):
        super().__init__()
        self.image_size = image_size
        self.z_dim = z_dim
        self.cond_dim = cond_dim
        self.class_emb = nn.Embedding(3, class_emb_dim)
        self.elev_mlp = nn.Sequential(nn.Linear(1, 32), nn.SiLU(), nn.Linear(32, elev_emb_dim), nn.SiLU())
        self.cond_mlp = nn.Sequential(nn.Linear(class_emb_dim + elev_emb_dim, 64), nn.SiLU(), nn.Linear(64, cond_dim), nn.SiLU())

        self.bottom_encoder = nn.Sequential(
            nn.Conv2d(1 + cond_dim, base, 3, padding=1),
            nn.GroupNorm(8 if base % 8 == 0 else 1, base),
            nn.SiLU(inplace=True),
            DownBlock(base, base),      # 256 -> 128
            DownBlock(base, base * 2),  # 128 -> 64
            ResBlock(base * 2),
            ResBlock(base * 2),
        )
        self.top_encoder = nn.Sequential(
            DownBlock(base * 2 + cond_dim, base * 4),  # 64 -> 32
            ResBlock(base * 4),
            ResBlock(base * 4),
            nn.Conv2d(base * 4, z_dim, 1),
        )
        self.top_vq = VectorQuantizerEMA(top_codes, z_dim, COMMITMENT_COST, EMA_DECAY)
        self.top_decoder = nn.Sequential(
            nn.Conv2d(z_dim + cond_dim, base * 2, 3, padding=1),
            ResBlock(base * 2),
            UpBlock(base * 2, z_dim),  # 32 -> 64
        )
        self.bottom_pre_vq = nn.Sequential(
            nn.Conv2d(base * 2 + z_dim + cond_dim, base * 2, 3, padding=1),
            ResBlock(base * 2),
            nn.Conv2d(base * 2, z_dim, 1),
        )
        self.bottom_vq = VectorQuantizerEMA(bottom_codes, z_dim, COMMITMENT_COST, EMA_DECAY)
        self.decoder = nn.Sequential(
            nn.Conv2d(z_dim + z_dim + cond_dim, base * 2, 3, padding=1),
            ResBlock(base * 2),
            UpBlock(base * 2, base),   # 64 -> 128
            UpBlock(base, base),       # 128 -> 256
            nn.GroupNorm(8 if base % 8 == 0 else 1, base),
            nn.SiLU(inplace=True),
            nn.Conv2d(base, 1, 3, padding=1),
            nn.Tanh(),
        )

    def condition_vec(self, labels, elev):
        c = torch.cat([self.class_emb(labels), self.elev_mlp(elev)], dim=1)
        return self.cond_mlp(c)

    @staticmethod
    def cond_map(cond, h, w):
        return cond[:, :, None, None].expand(-1, -1, h, w)

    def encode_quantize(self, x, labels, elev):
        cond = self.condition_vec(labels, elev)
        x_in = torch.cat([x, self.cond_map(cond, x.shape[-2], x.shape[-1])], dim=1)
        bottom_feat = self.bottom_encoder(x_in)
        c64 = self.cond_map(cond, bottom_feat.shape[-2], bottom_feat.shape[-1])
        top_z = self.top_encoder(torch.cat([bottom_feat, c64], dim=1))
        top_q, top_loss, top_ppx, top_idx = self.top_vq(top_z)
        c32 = self.cond_map(cond, top_q.shape[-2], top_q.shape[-1])
        top_dec = self.top_decoder(torch.cat([top_q, c32], dim=1))
        bottom_z = self.bottom_pre_vq(torch.cat([bottom_feat, top_dec, c64], dim=1))
        bottom_q, bottom_loss, bottom_ppx, bottom_idx = self.bottom_vq(bottom_z)
        vq_loss = top_loss + bottom_loss
        info = {
            'vq_loss': vq_loss,
            'top_perplexity': top_ppx.detach(),
            'bottom_perplexity': bottom_ppx.detach(),
            'top_indices': top_idx,
            'bottom_indices': bottom_idx,
        }
        return top_q, bottom_q, top_dec, cond, info

    def decode_quantized(self, top_q, bottom_q, labels, elev):
        cond = self.condition_vec(labels, elev)
        c32 = self.cond_map(cond, top_q.shape[-2], top_q.shape[-1])
        top_dec = self.top_decoder(torch.cat([top_q, c32], dim=1))
        c64 = self.cond_map(cond, bottom_q.shape[-2], bottom_q.shape[-1])
        recon = self.decoder(torch.cat([bottom_q, top_dec, c64], dim=1))
        return recon

    def decode_from_indices(self, top_idx, bottom_idx, labels, elev):
        top_q = self.top_vq.decode_indices(top_idx.to(next(self.parameters()).device))
        bottom_q = self.bottom_vq.decode_indices(bottom_idx.to(next(self.parameters()).device))
        return self.decode_quantized(top_q, bottom_q, labels, elev)

    def forward(self, x, labels, elev):
        top_q, bottom_q, top_dec, cond, info = self.encode_quantize(x, labels, elev)
        c64 = self.cond_map(cond, bottom_q.shape[-2], bottom_q.shape[-1])
        recon = self.decoder(torch.cat([bottom_q, top_dec, c64], dim=1))
        info['recon'] = recon
        return info