"""Адаптер на сторонний VAE (AsymmetricAutoencoderKL из sdxs-micro) под API нашего AutoEncoder.

Что адаптер прячет от остального кода:
  * нормировка латента своя — `latents_mean`/`latents_std` из vae_sdxs/config.json (≈ -0.04 / 1.58),
    а не наши ≈ -4.1 / 1.03; снаружи всё так же: encode(x) → нормированный латент,
    `to_raw()` в модели возвращает сырой;
  * декодер асимметричный: энкодер f8, декодер f16, то есть латент 32×32 отдаёт картинку 512²
    (вдвое больше входа) — это и есть смысл этого VAE;
  * нет «хвоста» — поточного пиксельного декодера, поэтому пиксельный лосс при обучении не считается
    (`--pix-weight 0`), а сэмплы декодируются один раз встроенным декодером VAE.

API повторяет `vae/ae/model.py::AutoEncoder` в той части, которую использует train.py/sample.py:
`encode`, `decode_raw`, `latent_mean`, `latent_std`, `config.latent_channels`, `config.patch_size`,
`compute_dtype`, `.eval()`, `.to(device)`.
"""
from __future__ import annotations

import torch
import torch.nn as nn

PATCH = 8                      # энкодер сжимает в 8 раз; декодер потом растягивает в 16 → 2× на выходе


class SdxsAE(nn.Module):
    def __init__(self, vae, device: str = "cuda"):
        super().__init__()
        self.vae = vae
        self.device_name = device
        self.compute_dtype = torch.float32
        cfg = vae.config
        mean = torch.tensor(getattr(cfg, "latents_mean"), dtype=torch.float32)
        std = torch.tensor(getattr(cfg, "latents_std"), dtype=torch.float32)
        self.register_buffer("latent_mean", mean.view(1, -1, 1, 1))
        self.register_buffer("latent_std", std.view(1, -1, 1, 1))
        self.config = type("Cfg", (), {"latent_channels": int(cfg.latent_channels),
                                       "patch_size": PATCH})()

    @classmethod
    def from_pretrained(cls, path, device: str = "cuda"):
        from diffusers import AsymmetricAutoencoderKL
        vae = AsymmetricAutoencoderKL.from_pretrained(str(path), torch_dtype=torch.float32)
        vae = vae.to(device).eval()
        vae.requires_grad_(False)
        return cls(vae, device).to(device).eval()

    @property
    def device(self):
        return self.latent_mean.device

    def train(self, mode: bool = True):
        super().train(mode)
        self.vae.eval()                     # VAE заморожен всегда
        return self

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """x в -1..1 → нормированный латент (B, 32, H/8, W/8)."""
        with torch.no_grad():
            raw = self.vae.encode(x.to(self.compute_dtype)).latent_dist.mode()
        return ((raw - self.latent_mean) / self.latent_std).to(torch.float32)

    def decode_raw(self, z_raw: torch.Tensor, steps: int = 0, generator=None) -> torch.Tensor:
        """Сырой латент → картинка в -1..1. `steps`/`generator` — для совместимости с нашим API."""
        with torch.no_grad():
            return self.vae.decode(z_raw.to(self.compute_dtype), return_dict=False)[0].clamp(-1, 1)