# examples/vae.py
"""Seeded offline CPU Bernoulli VAE mechanics, not trained image semantics.

Authored from Kingma and Welling, Auto-Encoding Variational Bayes, v11:
https://arxiv.org/html/1312.6114v11 (Sections 2.2-2.4, 3; Appendices B, C).
Uses one posterior sample, feature-summed NLL and latent-summed KL followed
by a batch mean. Beta is exactly one. No checkpoint, network or import run.
"""

import argparse
import hashlib
import json
import math
from pathlib import Path

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


def check_matrix(value, name, width=None):
    """Reject empty, nonfinite, nonfloating or non-CPU rank-two tensors."""
    if not isinstance(value, torch.Tensor):
        raise ValueError(f"{name} must be a tensor")
    if (
        value.ndim != 2
        or 0 in value.shape
        or not value.is_floating_point()
        or value.device.type != "cpu"
        or not bool(torch.isfinite(value).all())
    ):
        raise ValueError(f"{name} must be a finite nonempty CPU float matrix")
    if width is not None and value.shape[1] != width:
        raise ValueError(f"{name} width must be {width}")


def check_pair(left, right, left_name, right_name):
    """Validate matching tensor shape and dtype without broadcasting."""
    check_matrix(left, left_name)
    check_matrix(right, right_name)
    if left.shape != right.shape or left.dtype != right.dtype:
        raise ValueError(
            f"{left_name} and {right_name} must match shape and dtype"
        )


# region: encoder
class Encoder(nn.Module):
    """Map [B,Dx] inputs to [B,Dz] means and log variances (Appendix C.2)."""

    def __init__(self, input_dim, hidden_dim, latent_dim):
        super().__init__()
        self.hidden = nn.Linear(input_dim, hidden_dim)
        self.mu_head = nn.Linear(hidden_dim, latent_dim)
        self.logvar_head = nn.Linear(hidden_dim, latent_dim)

    def forward(self, x):
        """Return diagonal Gaussian parameters; logvar is log(sigma squared)."""
        check_matrix(x, "input", self.hidden.in_features)
        if x.dtype != self.hidden.weight.dtype:
            raise ValueError("input dtype must match encoder parameters")
        h = torch.tanh(self.hidden(x))
        return self.mu_head(h), self.logvar_head(h)


# endregion: encoder


# region: reparameterize
def reparameterize(mu, logvar, epsilon):
    """Use supplied standard-normal noise, preserving encoder gradients.

    Args:
        mu: Posterior means, shaped [B,Dz].
        logvar: Log variances, with the same shape and dtype as mu.
        epsilon: Independent standard-normal noise, shaped like mu.

    Returns:
        A differentiable [B,Dz] posterior sample.
    """
    check_pair(mu, logvar, "mu", "logvar")
    check_pair(mu, epsilon, "mu", "epsilon")
    std = torch.exp(0.5 * logvar)
    z = mu + std * epsilon
    check_matrix(z, "latent sample")
    return z


# endregion: reparameterize


# region: decoder
class Decoder(nn.Module):
    """Return [B,Dx] Bernoulli logits from [B,Dz] codes (Appendix C.1)."""

    def __init__(self, latent_dim, hidden_dim, input_dim):
        super().__init__()
        self.hidden = nn.Linear(latent_dim, hidden_dim)
        self.output = nn.Linear(hidden_dim, input_dim)

    def forward(self, z):
        """Return logits, rather than thresholded pixels or sampled images."""
        check_matrix(z, "latent sample", self.hidden.in_features)
        if z.dtype != self.hidden.weight.dtype:
            raise ValueError("latent dtype must match decoder parameters")
        return self.output(torch.tanh(self.hidden(z)))


# endregion: decoder


# region: reconstruction_nll
def reconstruction_nll(logits, x):
    """Return feature-summed Bernoulli negative log likelihood per sample."""
    check_pair(logits, x, "logits", "input")
    if not bool(((x == 0) | (x == 1)).all()):
        raise ValueError("Bernoulli observations must be binary")
    element_nll = F.binary_cross_entropy_with_logits(
        logits, x, reduction="none"
    )
    return element_nll.sum(dim=1)


# endregion: reconstruction_nll


# region: analytic_kl
def analytic_kl(mu, logvar):
    """Sum diagonal Gaussian-to-standard-normal KL over latent dimensions."""
    check_pair(mu, logvar, "mu", "logvar")
    element_kl = 0.5 * (mu.square() + logvar.exp() - 1 - logvar)
    per_sample = element_kl.sum(dim=1)
    if not bool(torch.isfinite(per_sample).all()):
        raise ValueError("KL overflowed")
    return per_sample


# endregion: analytic_kl


# region: negative_elbo
def negative_elbo(logits, x, mu, logvar):
    """Minimize a one-sample negative ELBO; beta=1 and batch reduction=mean.

    Args:
        logits: Decoder Bernoulli logits, shaped [B,Dx].
        x: Binary observations, shaped like logits.
        mu: Posterior means, shaped [B,Dz].
        logvar: Posterior log variances, shaped like mu.

    Returns:
        Scalar loss, per-sample feature-summed NLL, and latent-summed KL.
    """
    nll = reconstruction_nll(logits, x)
    kl = analytic_kl(mu, logvar)
    if nll.shape != kl.shape or logits.dtype != mu.dtype:
        raise ValueError("likelihood and posterior must share batch and dtype")
    loss = (nll + kl).mean()
    if not bool(torch.isfinite(loss)):
        raise ValueError("objective overflowed")
    return loss, nll, kl


# endregion: negative_elbo


# region: training_step
def training_step(encoder, decoder, x, epsilon, learning_rate=0.01):
    """Backpropagate both networks and take one SGD step with explicit noise.

    Args:
        encoder: A float CPU Encoder module.
        decoder: A compatible Decoder module.
        x: Binary observations, shaped [B,Dx].
        epsilon: Independent standard-normal noise, shaped [B,Dz].
        learning_rate: Positive finite SGD learning rate.

    Returns:
        Detached pre-update samples and objective terms, recorded gradients,
        and actual parameter-change evidence after the SGD update.
    """
    if not math.isfinite(learning_rate) or learning_rate <= 0:
        raise ValueError("learning_rate must be finite and positive")
    mu, logvar = encoder(x)
    z = reparameterize(mu, logvar, epsilon)
    logits = decoder(z)
    loss, nll, kl = negative_elbo(logits, x, mu, logvar)
    named = [("encoder." + k, p) for k, p in encoder.named_parameters()]
    named += [("decoder." + k, p) for k, p in decoder.named_parameters()]
    before = {name: p.detach().clone() for name, p in named}
    optimizer = torch.optim.SGD([p for _, p in named], lr=learning_rate)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    gradients = {name: p.grad.detach().clone() for name, p in named}
    if not all(bool(torch.isfinite(g).all()) for g in gradients.values()):
        raise ValueError("nonfinite gradients; parameters were not updated")
    optimizer.step()
    delta = sum(
        float((p.detach() - before[name]).square().sum()) for name, p in named
    )
    return {
        "mu": mu.detach(),
        "logvar": logvar.detach(),
        "z": z.detach(),
        "logits": logits.detach(),
        "objective": loss.detach(),
        "reconstruction_nll": nll.detach(),
        "kl": kl.detach(),
        "gradients": gradients,
        "parameter_delta_l2": math.sqrt(delta),
        "finite_parameters": all(
            bool(torch.isfinite(p).all()) for _, p in named
        ),
        "both_networks_updated": all(
            any(
                not torch.equal(p.detach(), before[name])
                for name, p in named
                if name.startswith(prefix)
            )
            for prefix in ("encoder.", "decoder.")
        ),
    }


# endregion: training_step


def make_models(seed, input_dim=4, latent_dim=2, hidden_dim=5):
    """Construct seeded float64 CPU networks without changing global RNG."""
    for dimension in (input_dim, latent_dim, hidden_dim):
        if (
            isinstance(dimension, bool)
            or not isinstance(dimension, int)
            or dimension < 1
        ):
            raise ValueError("model dimensions must be positive integers")
    with torch.random.fork_rng(devices=[]):
        torch.manual_seed(seed)
        encoder = Encoder(input_dim, hidden_dim, latent_dim).double()
        decoder = Decoder(latent_dim, hidden_dim, input_dim).double()
    return encoder, decoder


def synthetic_fixture(seed):
    """Return three binary inputs and independent Gaussian noise."""
    x = torch.tensor(
        [[0, 1, 0, 1], [1, 0, 1, 0], [1, 1, 0, 0]], dtype=torch.float64
    )
    generator = torch.Generator(device="cpu").manual_seed(seed + 1)
    epsilon = torch.randn((3, 2), generator=generator, dtype=torch.float64)
    return x, epsilon


def smoke(seed):
    """Execute an update and serialize pre-update samples and evidence."""
    encoder, decoder = make_models(seed)
    x, epsilon = synthetic_fixture(seed)
    result = training_step(encoder, decoder, x, epsilon)
    gradients = result["gradients"]

    def gradient_norm(prefix):
        return math.sqrt(
            sum(
                float(g.square().sum())
                for name, g in gradients.items()
                if name.startswith(prefix)
            )
        )

    tensors = {"input": x, "epsilon": epsilon}
    tensors.update(
        {
            key: result[key]
            for key in (
                "mu",
                "logvar",
                "z",
                "logits",
                "reconstruction_nll",
                "kl",
                "objective",
            )
        }
    )
    return {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "seed": seed,
        "boundary": (
            "offline CPU float64, synthetic binary data, one SGD step; "
            "no learned semantics"
        ),
        "reductions": {
            "reconstruction": "sum_features_per_sample",
            "kl": "sum_latents_per_sample",
            "objective": "mean_batch",
            "monte_carlo_samples": 1,
            "beta": 1,
        },
        "shapes": {key: list(value.shape) for key, value in tensors.items()},
        "objective": float(result["objective"]),
        "reconstruction_nll": result["reconstruction_nll"].tolist(),
        "kl": result["kl"].tolist(),
        "latent_samples": result["z"].tolist(),
        "bernoulli_probabilities": result["logits"].sigmoid().tolist(),
        "finite_outputs": all(
            bool(torch.isfinite(t).all()) for t in tensors.values()
        ),
        "finite_gradients": all(
            bool(torch.isfinite(g).all()) for g in gradients.values()
        ),
        "gradient_l2": {
            "encoder": gradient_norm("encoder."),
            "mu_head": gradient_norm("encoder.mu_head."),
            "logvar_head": gradient_norm("encoder.logvar_head."),
            "decoder": gradient_norm("decoder."),
        },
        "parameter_delta_l2": result["parameter_delta_l2"],
        "finite_parameters": result["finite_parameters"],
        "both_networks_updated": result["both_networks_updated"],
        "learning_rate": 0.01,
    }


def main():
    """Print finite JSON only when an explicit smoke run is requested."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--seed", type=int, default=7)
    arguments = parser.parse_args()
    if arguments.smoke:
        print(json.dumps(smoke(arguments.seed), indent=2, allow_nan=False))
    else:
        parser.print_help()


if __name__ == "__main__":
    main()