"""Original latent-diffusion subsystem lesson with text cross-attention.

An untrained deterministic convolutional autoencoder stands in for the paper's
pretrained KL/VQ first stage. Width-eight denoising U-Net; 12 beta steps.
No pretrained weights, tokenizer, perceptual training or image-quality claim.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


class ImageEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.down = nn.Conv2d(3, 8, 3, stride=2, padding=1)
        self.latent = nn.Conv2d(8, 2, 3, padding=1)

    def forward(self, image):
        return self.latent(F.silu(self.down(image)))


class ImageDecoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.up = nn.ConvTranspose2d(2, 8, 2, stride=2)
        self.rgb = nn.Conv2d(8, 3, 3, padding=1)

    def forward(self, latent):
        return self.rgb(F.silu(self.up(latent)))


class LatentNoise(nn.Module):
    def __init__(self):
        super().__init__()
        beta = torch.cat(
            [
                torch.zeros(1, dtype=torch.float64),
                torch.linspace(0.02, 0.18, 12, dtype=torch.float64),
            ]
        )
        self.register_buffer("beta", beta)
        self.register_buffer("alpha_bar", (1 - beta).cumprod(0))

    def forward(self, latent, noise, timestep):
        a = self.alpha_bar[timestep].to(latent.dtype).reshape(-1, 1, 1, 1)
        return a.sqrt() * latent + (1 - a).sqrt() * noise


class TextConditioner(nn.Module):
    def __init__(self):
        super().__init__()
        self.tokens = nn.Embedding(16, 8)
        self.positions = nn.Parameter(torch.randn(1, 3, 8) * 0.02)

    def forward(self, tokens):
        return self.tokens(tokens) + self.positions


class CrossAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.query = nn.Linear(8, 8)
        self.key = nn.Linear(8, 8)
        self.value = nn.Linear(8, 8)
        self.softmax = nn.Softmax(dim=-1)
        self.project = nn.Linear(8, 8)

    def weights(self, image, context):
        return self.softmax(
            self.query(image) @ self.key(context).transpose(-1, -2) / (8 ** 0.5)
        )

    def forward(self, image, context):
        q = self.query(image)
        k = self.key(context)
        v = self.value(context)
        weights = self.softmax(q @ k.transpose(-1, -2) / (8 ** 0.5))
        return self.project(weights @ v)


class ConditionalDenoiser(nn.Module):
    def __init__(self):
        super().__init__()
        self.time_embedding = nn.Embedding(13, 8)
        self.stem = nn.Conv2d(2, 8, 3, padding=1)
        self.down = nn.Conv2d(8, 8, 3, stride=2, padding=1)
        self.cross_attention = CrossAttention()
        self.up = nn.ConvTranspose2d(8, 8, 2, stride=2)
        self.merge = nn.Conv2d(16, 8, 3, padding=1)
        self.noise = nn.Conv2d(8, 2, 3, padding=1)

    def forward(self, noisy, timestep, context):
        fine = F.silu(
            self.stem(noisy) + self.time_embedding(timestep)[:, :, None, None]
        )
        coarse = F.silu(self.down(fine))
        tokens = coarse.flatten(2).transpose(1, 2)
        tokens = tokens + self.cross_attention(tokens, context)
        conditioned = tokens.transpose(1, 2).reshape_as(coarse)
        return self.noise(
            F.silu(self.merge(torch.cat([fine, self.up(conditioned)], 1)))
        )


class EstimateCleanLatent(nn.Module):
    def forward(self, noisy, epsilon, alpha_bar):
        return (noisy - (1 - alpha_bar).sqrt() * epsilon) / alpha_bar.sqrt()


class TinyLatentDiffusion(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = ImageEncoder()
        self.corrupt = LatentNoise()
        self.conditioner = TextConditioner()
        self.denoiser = ConditionalDenoiser()
        self.estimate = EstimateCleanLatent()
        self.decoder = ImageDecoder()
        for parameter in self.encoder.parameters():
            parameter.requires_grad_(False)
        for parameter in self.decoder.parameters():
            parameter.requires_grad_(False)

    def forward(self, image, tokens, noise, timestep):
        latent = self.encoder(image)
        noisy = self.corrupt(latent, noise, timestep)
        context = self.conditioner(tokens)
        epsilon = self.denoiser(noisy, timestep, context)
        a = self.corrupt.alpha_bar[timestep].to(image.dtype).reshape(-1, 1, 1, 1)
        estimated = self.estimate(noisy, epsilon, a)
        return self.decoder(estimated)


def setup():
    torch.manual_seed(43)
    model = TinyLatentDiffusion().double()
    image = torch.randn(1, 3, 8, 8, dtype=torch.float64)
    tokens = torch.tensor([[1, 3, 5]])
    noise = torch.randn(1, 2, 4, 4, dtype=torch.float64)
    return model, image, tokens, noise


def verify():
    model, image, tokens, noise = setup()
    latent = model.encoder(image)
    assert image.numel() == 192 and latent.numel() == 32
    t = torch.tensor([6])
    noisy = model.corrupt(latent, noise, t)
    context = model.conditioner(tokens)
    pred = model.denoiser(noisy, t, context)
    assert not torch.allclose(
        pred, model.denoiser(noisy, t, model.conditioner(torch.tensor([[7, 9, 11]])))
    )
    attention = model.denoiser.cross_attention
    q = torch.randn(1, 4, 8, dtype=torch.float64)
    torch.testing.assert_close(
        attention(q, context), attention(q, context[:, [2, 0, 1]])
    )
    torch.testing.assert_close(
        attention.weights(q, context).sum(-1), torch.ones(1, 4, dtype=torch.float64)
    )
    a = model.corrupt.alpha_bar[6]
    torch.testing.assert_close(model.estimate(noisy, noise, a), latent)
    F.mse_loss(pred, noise).backward()
    assert all(p.grad is None for p in model.encoder.parameters())
    assert all(p.grad is None for p in model.decoder.parameters())
    assert model.conditioner.tokens.weight.grad.norm() > 0
    assert model.denoiser.cross_attention.key.weight.grad.norm() > 0
    return [
        "The 192-value image is compressed to 32 latent values and decoded back to its original shape",
        "Text changes epsilon predictions while the cached image latent remains independent of text",
        "Cross-attention rows normalize over text tokens and are invariant to joint K/V ordering",
        "Known epsilon recovers the clean latent; diffusion-loss gradients reach the conditioner and denoiser while the first stage stays frozen",
    ]


def experiment():
    model, image, tokens, noise = setup()
    cases = []
    with torch.no_grad():
        latent = model.encoder(image)
        for identity, label, t, ids in [
            ("early", "Same text · t = 1", 1, tokens),
            ("middle", "Same text · t = 6", 6, tokens),
            ("late", "Same text · t = 12", 12, tokens),
            ("text", "Changed text · t = 6", 6, torch.tensor([[7, 9, 11]])),
        ]:
            time = torch.tensor([t])
            context = model.conditioner(ids)
            noisy = model.corrupt(latent, noise, time)
            pred = model.denoiser(noisy, time, context)
            estimated = model.estimate(noisy, pred, model.corrupt.alpha_bar[t])
            decoded = model.decoder(estimated)
            fine = F.silu(
                model.denoiser.stem(noisy)
                + model.denoiser.time_embedding(time)[:, :, None, None]
            )
            coarse = F.silu(model.denoiser.down(fine))
            weights = model.denoiser.cross_attention.weights(
                coarse.flatten(2).transpose(1, 2), context
            )
            cases.append(
                dict(
                    id=identity,
                    label=label,
                    target="denoiser.cross_attention.softmax",
                    note="The same encoded image and Gaussian noise are reused. Queries come from four coarse spatial locations; keys and values come from three text tokens. The changed-text case keeps t = 6 fixed. The one-step clean estimate is decoded only for inspection and is not a complete generation.",
                    matrices=[
                        dict(
                            label="Clean latent · channel 0",
                            values=latent[0, 0].tolist(),
                        ),
                        dict(
                            label="Noisy latent · channel 0",
                            values=noisy[0, 0].tolist(),
                        ),
                        dict(
                            label="Spatial query × text attention",
                            values=weights[0].tolist(),
                        ),
                        dict(
                            label="Decoded estimate · red channel",
                            values=decoded[0, 0].tolist(),
                        ),
                    ],
                    vectors=[dict(label="Text IDs", values=ids[0].tolist())],
                    metrics=[
                        dict(label="Latent values", value=32),
                        dict(label="Pixel-space values", value=192),
                        dict(
                            label="Noise-prediction MSE",
                            value=F.mse_loss(pred, noise).item(),
                        ),
                    ],
                )
            )
        current = torch.randn_like(latent)
        context = model.conditioner(tokens)
        for t in range(12, 0, -1):
            pred = model.denoiser(current, torch.tensor([t]), context)
            b = model.corrupt.beta[t]
            a = 1 - b
            abar = model.corrupt.alpha_bar[t]
            mean = (current - b / (1 - abar).sqrt() * pred) / a.sqrt()
            var = b * (1 - model.corrupt.alpha_bar[t - 1]) / (1 - abar)
            current = mean + (var.sqrt() * torch.randn_like(current) if t > 1 else 0)
        decoded = model.decoder(current)
        cases.append(
            dict(
                id="sample",
                label="After all 12 latent reverse steps",
                target="decoder",
                note="Generation starts from a new Gaussian latent and reuses the same text context for 12 denoiser calls. The image encoder is not called in this loop. After the last step, the decoder is called once. All weights are untrained; the output has no semantic image-quality claim.",
                matrices=[
                    dict(
                        label="Final latent · channel 0", values=current[0, 0].tolist()
                    ),
                    dict(
                        label="Decoded result · red channel",
                        values=decoded[0, 0].tolist(),
                    ),
                ],
                vectors=[
                    dict(label="Reverse timesteps", values=list(range(12, 0, -1)))
                ],
                metrics=[
                    dict(label="Denoiser calls", value=12),
                    dict(label="Image-encoder calls during generation", value=0),
                    dict(label="Decoder calls after sampling", value=1),
                ],
            )
        )
    return dict(
        kind="matrices",
        title="Denoise in latent space; decode once.",
        description="Recorded 4 × 4 latent diffusion with text conditioning. A frozen random autoencoder illustrates compression; it does not reproduce a trained KL/VQ first stage.",
        controlLabel="Latent diffusion case",
        cases=cases,
    )