"""
VAE Wrappers — compatible VAE interfaces for LiquidFlow.

Supports two VAE backends:
1. TAESD (Tiny AutoEncoder for SD): < 1M params, extremely fast, perfect for mobile
2. SD-VAE (Stability AI VAE): Higher quality, 84M params, standard for SD pipelines

TAESD is the DEFAULT for LiquidFlow — it's designed to be lightweight and
fast enough for Colab/Kaggle free tier.

Paper reference: "Tiny AutoEncoder for Stable Diffusion" (madebyollin/taesd)
Model: madebyollin/taesd (335K downloads on HF)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional


class TAESDWrapper:
    """
    Wrapper for Tiny AutoEncoder for Stable Diffusion (TAESD).
    
    TAESD properties:
    - ~1M parameters (vs 84M for SD VAE)
    - Latent dim: 4 channels @ 8x compression
    - Extremely fast encode/decode
    - Works on CPU — no GPU needed
    - Perfect for Colab/Kaggle free tier
    
    Model on HF: madebyollin/taesd
    """
    
    def __init__(self, device='cpu'):
        self.device = device
        self.model = None
        
    @staticmethod
    def is_available():
        """Check if TAESD can be loaded."""
        try:
            from diffusers import AutoencoderTiny
            return True
        except ImportError:
            return False
    
    @staticmethod
    def load(device='cpu'):
        """Load TAESD model."""
        from diffusers import AutoencoderTiny
        model = AutoencoderTiny.from_pretrained(
            "madebyollin/taesd",
            torch_dtype=torch.float32,
        )
        model = model.to(device)
        model.eval()
        return model
    
    @staticmethod
    def get_latent_shape(image_size):
        """Get latent spatial size given image size (8x compression)."""
        return image_size // 8
    
    @staticmethod
    def encode(vae, x):
        """
        Encode image to latent.
        Args:
            vae: TAESD model
            x: [B, 3, H, W] images in [-1, 1]
        Returns:
            z: [B, 4, H/8, W/8]
        """
        with torch.no_grad():
            posterior = vae.encode(x).latent_dist
            z = posterior.sample()
            z = z * vae.config.scaling_factor
        return z
    
    @staticmethod
    def decode(vae, z):
        """
        Decode latent to image.
        Args:
            vae: TAESD model  
            z: [B, 4, H/8, W/8]
        Returns:
            x: [B, 3, H, W] images in [-1, 1]
        """
        with torch.no_grad():
            z = z / vae.config.scaling_factor
            x = vae.decode(z).sample
        return x


class SDVAEWrapper:
    """
    Wrapper for Stability AI VAE (sd-vae-ft-mse).
    
    Properties:
    - ~84M parameters
    - Latent dim: 4 channels @ 8x compression
    - Higher quality reconstruction than TAESD
    - Requires GPU for reasonable speed
    
    Model on HF: stabilityai/sd-vae-ft-mse
    """
    
    def __init__(self, device='cpu'):
        self.device = device
        self.model = None
    
    @staticmethod
    def load(device='cpu'):
        """Load SD VAE model."""
        from diffusers import AutoencoderKL
        model = AutoencoderKL.from_pretrained(
            "stabilityai/sd-vae-ft-mse",
            torch_dtype=torch.float32,
        )
        model = model.to(device)
        model.eval()
        return model
    
    @staticmethod
    def encode(vae, x):
        """Encode image to latent."""
        with torch.no_grad():
            posterior = vae.encode(x).latent_dist
            z = posterior.sample()
            z = z * vae.config.scaling_factor
        return z
    
    @staticmethod
    def decode(vae, z):
        """Decode latent to image."""
        with torch.no_grad():
            z = z / vae.config.scaling_factor
            x = vae.decode(z).sample
        return x