"""HSIGene AutoencoderKL - nn.Module, no Lightning. Loss = Identity."""

import torch
import torch.nn as nn

from .vae_blocks import Encoder, Decoder, DiagonalGaussianDistribution


class AutoencoderKL(nn.Module):
    """
    AutoencoderKL - nn.Module (not Lightning).
    Uses Encoder, Decoder, quant_conv, post_quant_conv.
    encode() returns posterior, decode() takes z.
    Loss = Identity (no-op).
    """

    def __init__(
        self,
        ddconfig,
        embed_dim=4,
        lossconfig=None,
        **kwargs,
    ):
        super().__init__()
        self.encoder = Encoder(**ddconfig)
        self.decoder = Decoder(**ddconfig)
        assert ddconfig.get("double_z", True)
        z_channels = ddconfig["z_channels"]
        self.quant_conv = nn.Conv2d(2 * z_channels, 2 * embed_dim, 1)
        self.post_quant_conv = nn.Conv2d(embed_dim, z_channels, 1)
        self.embed_dim = embed_dim
        self.loss = nn.Identity()

    def encode(self, x):
        h = self.encoder(x)
        moments = self.quant_conv(h)
        posterior = DiagonalGaussianDistribution(moments, deterministic=True)
        return posterior

    def decode(self, z):
        z = self.post_quant_conv(z)
        return self.decoder(z)

    def forward(self, input, sample_posterior=True):
        posterior = self.encode(input)
        if sample_posterior:
            z = posterior.sample()
        else:
            z = posterior.mode()
        dec = self.decode(z)
        return dec, posterior


class HSIGeneAutoencoderKL(AutoencoderKL):
    """
    HSIGene VAE with diffusers-style config.
    Accepts in_channels, out_channels, latent_channels, block_out_channels.
    """

    def __init__(
        self,
        in_channels: int = 48,
        out_channels: int = 48,
        latent_channels: int = 96,
        embed_dim: int = 4,
        block_out_channels: tuple = (64, 128, 256),
        num_res_blocks: int = 4,
        attn_resolutions: tuple = (16, 32, 64),
        dropout: float = 0.0,
        double_z: bool = True,
        resolution: int = 256,
        **kwargs,
    ):
        ch = block_out_channels[0]
        ch_mult = tuple(
            block_out_channels[i] // ch for i in range(len(block_out_channels))
        )
        ddconfig = dict(
            double_z=double_z,
            z_channels=latent_channels,
            resolution=resolution,
            in_channels=in_channels,
            out_ch=out_channels,
            ch=ch,
            ch_mult=list(ch_mult),
            num_res_blocks=num_res_blocks,
            attn_resolutions=list(attn_resolutions),
            dropout=dropout,
        )
        super().__init__(ddconfig=ddconfig, embed_dim=embed_dim, **kwargs)