"""Inference-only VQ-VAE using the released three-level SPAC checkpoint."""

import json
from pathlib import Path
from typing import Mapping

import torch
import torch.nn as nn
from .sparse import SparseTensor
from .nn.encoder import SparseSDFEncoder
from .nn.decoder import SparseSDFDecoder


class SparseVectorQuantizer(nn.Module):
    def __init__(self, num_embeddings: int, embedding_dim: int, use_ema_update: bool = True):
        super().__init__()
        self.embeddings = nn.Embedding(num_embeddings, embedding_dim)
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        if use_ema_update:
            # Retain checkpoint buffers while omitting all training updates.
            self.register_buffer("ema_cluster_size", torch.zeros(num_embeddings))
            self.register_buffer("ema_w", torch.zeros(num_embeddings, embedding_dim))
            self.register_buffer("_ema_initialized", torch.tensor(False))

    def forward(self, value: SparseTensor) -> SparseTensor:
        distances = torch.cdist(value.feats.to(self.embeddings.weight.dtype), self.embeddings.weight)
        indices = torch.argmin(distances, dim=1)
        return value.replace(indices.unsqueeze(-1).float())


class SparseSDFVQVAE(nn.Module):
    def __init__(
        self, resolution: int = 32, model_channels: int = 512, latent_channels: int = 16,
        num_blocks: int = 8, num_embeddings: int = 8192, vq_group_size: int = 1,
        num_heads: int = 8, num_head_channels: int = 64, mlp_ratio: float = 4.0,
        attn_mode: str = "swin", window_size: int = 8, pe_mode: str = "ape",
        use_ema_update: bool = True, legacy_three_level: bool = True,
    ):
        super().__init__()
        if not legacy_three_level or vq_group_size != 1:
            raise ValueError("This package contains only the released three-level, single-voxel SPAC VQ-VAE")
        self.resolution = resolution
        self.embed_dim = latent_channels
        self.vq_group_size = 1
        self.vq_block_side = 1
        arguments = dict(
            resolution=resolution, model_channels=model_channels,
            latent_channels=latent_channels, num_blocks=num_blocks,
            num_heads=num_heads, num_head_channels=num_head_channels,
            mlp_ratio=mlp_ratio, attn_mode=attn_mode, window_size=window_size, pe_mode=pe_mode,
        )
        self.encoder = SparseSDFEncoder(**arguments)
        self.decoder = SparseSDFDecoder(**arguments)
        self.vq = SparseVectorQuantizer(num_embeddings, latent_channels, use_ema_update)

    @torch.inference_mode()
    def Encode(self, batch) -> SparseTensor:
        if isinstance(batch, Mapping):
            features = batch["sparse_sdf"]
            if features.ndim == 1:
                features = features.unsqueeze(-1)
            coordinates = torch.cat([batch["batch_idx"].unsqueeze(-1), batch["sparse_index"]], dim=-1).int()
            value = SparseTensor(features.float(), coordinates)
            factor = batch.get("factor")
        else:
            value = batch
            factor = None
        return self.vq(self.encoder(value, factor))

    @torch.inference_mode()
    def Decode(self, indices: SparseTensor, band_prune=None, gt_prune=None) -> SparseTensor:
        ids = indices.feats.long().squeeze(-1)
        if not ids.numel() or ids.min() < 0 or ids.max() >= self.vq.num_embeddings:
            raise ValueError("A nonempty sequence of valid VQ codebook indices is required")
        latent = indices.replace(self.vq.embeddings(ids))
        return self.decoder(latent, band_prune=band_prune, gt_prune=gt_prune)


def load_vae(checkpoint, config, device) -> SparseSDFVQVAE:
    """Load every released encoder, decoder and codebook tensor strictly in FP32."""
    configuration = json.loads(Path(config).read_text()) if not isinstance(config, Mapping) else dict(config)
    arguments = configuration.get("args", configuration)
    model = SparseSDFVQVAE(**arguments)
    state = torch.load(checkpoint, map_location="cpu", weights_only=True)
    if isinstance(state, Mapping) and isinstance(state.get("state_dict"), Mapping):
        state = state["state_dict"]
    state = {key.removeprefix("module."): value for key, value in state.items()}
    model.load_state_dict(state, strict=True)
    return model.float().eval().requires_grad_(False).to(device)