# examples/latentdiffusion.py
"""Offline latent diffusion mechanics with a random, frozen KL-style interface.

Authored reduction of Rombach et al. (2112.10752v2), Sections 3.1-3.3,
Appendices B/G. Not a trained perceptual compressor, U-Net, or text encoder.
Official interface provenance: CompVis/latent-diffusion at
 a506df5756472e2ebaf9078affdde2c4f1502cd4 (autoencoder.py, ddpm.py, ddim.py).
Sampling uses deterministic DDIM (eta=0) over every reduced schedule entry.
Import neither creates a model nor trains, downloads, or samples anything.
"""

import argparse
import hashlib
import json
from pathlib import Path

import torch
from torch import nn


def require_tensor(value, ndim, name):
    """Reject empty, nonfinite, non-CPU, or non-float32 numerical inputs."""
    if not isinstance(value, torch.Tensor):
        raise ValueError(f"{name} must be a tensor")
    if (
        value.ndim != ndim
        or min(value.shape) < 1
        or value.dtype != torch.float32
        or value.device.type != "cpu"
        or not torch.isfinite(value).all()
    ):
        raise ValueError(
            f"{name} must be nonempty finite CPU float32 rank {ndim}"
        )


# region: fixture_and_first_stage
class FrozenKLFirstStage(nn.Module):
    """Random Gaussian encoder/decoder interface, not perceptual compression.

    Attributes:
        downsampling: Spatial factor two, independent of latent channel count.
        latent_channels: Number of channels in a sampled diagonal Gaussian.
    """

    def __init__(self, image_channels=3, latent_channels=2):
        super().__init__()
        if min(image_channels, latent_channels) < 1:
            raise ValueError("Channel counts must be positive")
        self.image_channels = image_channels
        self.latent_channels = latent_channels
        self.downsampling = 2
        self.encoder = nn.Conv2d(
            image_channels, 2 * latent_channels, 2, stride=2
        )
        self.decoder = nn.ConvTranspose2d(
            latent_channels, image_channels, 2, stride=2
        )
        self.requires_grad_(False)
        self.eval()


def synthetic_fixture(generator, batch=2, height=8, width=8, tokens=3):
    """Return random images and supplied numerical conditions, without semantics."""
    if min(batch, height, width, tokens) < 1 or height % 2 or width % 2:
        raise ValueError(
            "Positive dimensions and even image sizes are required"
        )
    images = torch.randn(batch, 3, height, width, generator=generator)
    condition = torch.randn(batch, tokens, 3, generator=generator)
    return images, condition


# endregion: fixture_and_first_stage


# region: encode_and_freeze
@torch.no_grad()
def encode(first_stage, images, generator):
    """Sample z = mu + exp(logvar/2)*xi from the frozen KL-style interface."""
    require_tensor(images, 4, "images")
    if (
        images.shape[1] != first_stage.image_channels
        or images.shape[2] % 2
        or images.shape[3] % 2
    ):
        raise ValueError(
            "Image channels or spatial downsampling are incompatible"
        )
    moments = first_stage.encoder(images)
    mean, logvar = moments.chunk(2, dim=1)
    logvar = logvar.clamp(-30.0, 20.0)
    xi = torch.randn(mean.shape, generator=generator)
    latent = mean + (0.5 * logvar).exp() * xi
    return latent, mean, logvar


# endregion: encode_and_freeze


def validate_schedule(alpha_bar):
    """Validate positive, decreasing cumulative alphas for the reduced chain."""
    require_tensor(alpha_bar, 1, "alpha_bar")
    if (
        not ((alpha_bar > 0) & (alpha_bar < 1)).all()
        or not (alpha_bar[1:] < alpha_bar[:-1]).all()
    ):
        raise ValueError(
            "Cumulative alphas must decrease strictly within (0,1)"
        )


# region: latent_noising
def noise_latent(latent, timesteps, noise, alpha_bar):
    """Apply the closed-form forward process; indices 0..T-1 mean paper t=1..T."""
    require_tensor(latent, 4, "latent")
    require_tensor(noise, 4, "noise")
    validate_schedule(alpha_bar)
    if (
        noise.shape != latent.shape
        or not isinstance(timesteps, torch.Tensor)
        or timesteps.dtype != torch.long
        or timesteps.device.type != "cpu"
        or timesteps.shape != (latent.shape[0],)
        or not ((timesteps >= 0) & (timesteps < len(alpha_bar))).all()
    ):
        raise ValueError(
            "Noise shape or per-example timestep indices are invalid"
        )
    coefficients = alpha_bar[timesteps].reshape(-1, 1, 1, 1)
    noisy = coefficients.sqrt() * latent + (1 - coefficients).sqrt() * noise
    return noisy


# endregion: latent_noising


# region: spatial_queries
def spatial_queries(latent, projection):
    """Flatten N=h*w locations and retain shape for the inverse transform."""
    require_tensor(latent, 4, "latent")
    batch, channels, height, width = latent.shape
    if channels != projection.in_features:
        raise ValueError("Query input width must match latent channels")
    spatial = latent.flatten(2).transpose(1, 2)
    queries = projection(spatial)
    return queries, (batch, height, width)


# endregion: spatial_queries


# region: condition_keys_values
def condition_keys_values(condition, key_projection, value_projection):
    """Project supplied numerical embeddings, without claiming text semantics."""
    require_tensor(condition, 3, "condition")
    if condition.shape[-1] != key_projection.in_features:
        raise ValueError("Condition width must match the key projection")
    if condition.shape[-1] != value_projection.in_features:
        raise ValueError("Condition width must match the value projection")
    keys = key_projection(condition)
    values = value_projection(condition)
    return keys, values


# endregion: condition_keys_values


# region: cross_attention
def cross_attention(queries, keys, values, spatial_shape, output_projection):
    """Normalize over M condition tokens, aggregate values, then restore space."""
    for name, tensor in (
        ("queries", queries),
        ("keys", keys),
        ("values", values),
    ):
        require_tensor(tensor, 3, name)
    batch, height, width = spatial_shape
    if (
        queries.shape[:2] != (batch, height * width)
        or keys.shape[0] != batch
        or values.shape[:2] != keys.shape[:2]
        or queries.shape[-1] != keys.shape[-1]
        or values.shape[-1] != output_projection.in_features
    ):
        raise ValueError(
            "Attention batch, token, feature, or spatial shapes mismatch"
        )
    scores = queries @ keys.transpose(-1, -2) / queries.shape[-1] ** 0.5
    weights = scores.softmax(dim=-1)
    attended = output_projection(weights @ values)
    residual = attended.transpose(1, 2).reshape(batch, -1, height, width)
    return residual, weights


# endregion: cross_attention


# region: epsilon_objective
class ConditionalDenoiser(nn.Module):
    """Reduced conditional epsilon network, not the paper's multiscale U-Net."""

    def __init__(self, latent_channels=2, condition_width=3, attention_width=4):
        super().__init__()
        if min(latent_channels, condition_width, attention_width) < 1:
            raise ValueError("Model widths must be positive")
        self.query = nn.Linear(latent_channels, attention_width, bias=False)
        self.key = nn.Linear(condition_width, attention_width, bias=False)
        self.value = nn.Linear(condition_width, attention_width, bias=False)
        self.output = nn.Linear(attention_width, latent_channels, bias=False)
        self.time = nn.Linear(1, latent_channels)
        self.conv = nn.Conv2d(latent_channels, latent_channels, 3, padding=1)

    def forward(self, noisy, timesteps, condition, total_steps):
        """Predict noise from current state, normalized time, and condition only."""
        require_tensor(noisy, 4, "noisy")
        if (
            not isinstance(total_steps, int)
            or total_steps < 1
            or not isinstance(timesteps, torch.Tensor)
            or timesteps.dtype != torch.long
            or timesteps.device.type != "cpu"
            or timesteps.shape != (noisy.shape[0],)
            or not ((timesteps >= 0) & (timesteps < total_steps)).all()
        ):
            raise ValueError("Denoiser time indices are invalid")
        queries, shape = spatial_queries(noisy, self.query)
        keys, values = condition_keys_values(condition, self.key, self.value)
        residual, weights = cross_attention(
            queries, keys, values, shape, self.output
        )
        time = ((timesteps.float() + 1) / total_steps).unsqueeze(1)
        time = self.time(time).unsqueeze(-1).unsqueeze(-1)
        predicted = self.conv(torch.tanh(noisy + residual + time))
        return predicted, weights


def epsilon_loss(predicted, target):
    """Use squared-vector norm per example, then batch mean (paper Eq. 3)."""
    require_tensor(predicted, 4, "predicted")
    require_tensor(target, 4, "target")
    if predicted.shape != target.shape:
        raise ValueError(
            "Noise target and prediction must have identical shapes"
        )
    return (predicted - target).square().flatten(1).sum(1).mean()


def train_step(
    model, optimizer, noisy, timesteps, condition, target, total_steps
):
    """Backpropagate noise regression and update only optimizer-owned parameters."""
    optimizer.zero_grad(set_to_none=True)
    predicted, weights = model(noisy, timesteps, condition, total_steps)
    loss = epsilon_loss(predicted, target)
    loss.backward()
    optimizer.step()
    return loss.detach(), predicted.detach(), weights.detach()


# endregion: epsilon_objective


# region: latent_reverse_and_decode
@torch.no_grad()
def reverse_and_decode(model, first_stage, initial, condition, alpha_bar):
    """Perform all T deterministic DDIM eta=0 updates, then decode once.

    The trace includes the initial z_T followed by each z_{t-1}. The final
    previous cumulative alpha is one, yielding the predicted clean latent.
    """
    require_tensor(initial, 4, "initial")
    validate_schedule(alpha_bar)
    if initial.shape[1] != first_stage.latent_channels:
        raise ValueError(
            "Latent channel count is incompatible with the decoder"
        )
    latent = initial.clone()
    trace = [latent.clone()]
    for index in reversed(range(len(alpha_bar))):
        timesteps = torch.full((latent.shape[0],), index, dtype=torch.long)
        epsilon, _ = model(latent, timesteps, condition, len(alpha_bar))
        current = alpha_bar[index]
        previous = (
            alpha_bar[index - 1] if index > 0 else current.new_tensor(1.0)
        )
        clean = (latent - (1 - current).sqrt() * epsilon) / current.sqrt()
        latent = previous.sqrt() * clean + (1 - previous).sqrt() * epsilon
        trace.append(latent.clone())
    decoded = first_stage.decoder(latent)
    return decoded, torch.stack(trace)


# endregion: latent_reverse_and_decode


def smoke(seed=7):
    """Execute seeded CPU mechanics and return actual numerical observations."""
    with torch.random.fork_rng():
        torch.manual_seed(seed)
        generator = torch.Generator().manual_seed(seed)
        first_stage = FrozenKLFirstStage()
        model = ConditionalDenoiser()
        images, condition = synthetic_fixture(generator)
        frozen_before = [p.detach().clone() for p in first_stage.parameters()]
        model_before = [p.detach().clone() for p in model.parameters()]
        latent, mean, logvar = encode(first_stage, images, generator)
        alpha_bar = torch.cumprod(1 - torch.linspace(0.05, 0.2, 4), 0)
        timesteps = torch.randint(
            len(alpha_bar), (images.shape[0],), generator=generator
        )
        noise = torch.randn(latent.shape, generator=generator)
        noisy = noise_latent(latent, timesteps, noise, alpha_bar)
        optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
        loss, predicted, weights = train_step(
            model, optimizer, noisy, timesteps, condition, noise, len(alpha_bar)
        )
        gradient_norm = sum(
            p.grad.square().sum() for p in model.parameters()
        ).sqrt()
        with torch.no_grad():
            updated, _ = model(noisy, timesteps, condition, len(alpha_bar))
            changed_condition, _ = model(
                noisy, timesteps, condition + 0.5, len(alpha_bar)
            )
            initial = torch.randn(latent.shape, generator=generator)
            decoded, trace = reverse_and_decode(
                model, first_stage, initial, condition, alpha_bar
            )
        return {
            "source_sha256": hashlib.sha256(
                Path(__file__).read_bytes()
            ).hexdigest(),
            "seed": seed,
            "reductions": {
                "epsilon_loss": "sum Cz*h*w per example, then batch mean",
                "attention": "softmax over condition tokens M",
            },
            "shapes": {
                name: list(tensor.shape)
                for name, tensor in (
                    ("images", images),
                    ("latent", latent),
                    ("condition", condition),
                    ("noisy", noisy),
                    ("predicted_epsilon", predicted),
                    ("attention", weights),
                    ("reverse_trace", trace),
                    ("decoded", decoded),
                )
            },
            "configuration": {
                "first_stage": "random frozen KL-style diagonal Gaussian",
                "posterior": "sample, no latent scaling",
                "downsampling": 2,
                "latent_channels": 2,
                "reverse": "DDIM eta=0, all 4 entries",
                "updates": 1,
                "learning_rate": 0.001,
            },
            "observations": {
                "loss": loss.item(),
                "post_update_loss": epsilon_loss(updated, noise).item(),
                "denoiser_gradient_norm": gradient_norm.item(),
                "denoiser_parameter_delta": sum(
                    (a - b).abs().sum().item()
                    for a, b in zip(model.parameters(), model_before)
                ),
                "first_stage_parameter_delta": sum(
                    (a - b).abs().sum().item()
                    for a, b in zip(first_stage.parameters(), frozen_before)
                ),
                "first_stage_gradient_count": sum(
                    p.grad is not None for p in first_stage.parameters()
                ),
                "latent_requires_grad": latent.requires_grad,
                "attention_row_sum_error": (weights.sum(-1) - 1)
                .abs()
                .max()
                .item(),
                "condition_output_delta": (updated - changed_condition)
                .abs()
                .max()
                .item(),
                "mean_abs": mean.abs().mean().item(),
                "logvar_range": [logvar.min().item(), logvar.max().item()],
                "decoded_range": [decoded.min().item(), decoded.max().item()],
            },
            "timesteps_zero_based": timesteps.tolist(),
            "alpha_bar": alpha_bar.tolist(),
            "reverse_trace": {
                "paper_times": [4, 3, 2, 1, 0],
                "states": trace.tolist(),
            },
            "boundary": "Random numerical shapes only; no trained compression, text semantics, or image quality.",
        }


def main():
    """Print strict finite smoke JSON only when explicitly requested."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if not args.smoke:
        parser.error("Pass --smoke to execute the offline numerical reference")
    torch.set_num_threads(1)
    print(json.dumps(smoke(args.seed), allow_nan=False))


if __name__ == "__main__":
    main()