"""The small networks every lecture reuses.

Four networks cover the whole course. A lecture changes the loss and the sampler and keeps the
network fixed, so any difference in the samples comes from the mathematics alone.

| Network            | Input                      | Output                | Used by |
| ------------------ | -------------------------- | --------------------- | ------- |
| `TimeEmbedding`    | a time or step index       | a vector of width D   | all     |
| `SequenceDenoiser` | tokens [B, L] and a time   | logits [B, L, K]      | 4, 5, 6 |
| `VelocityMLP`      | a point [B, D] and a time  | a vector [B, D]       | 2, 3, 6, 7 |
| `TwoTimeMLP`       | a point, a start, an end   | a vector [B, D]       | 6 |
| `SmallUNet`        | an image [B, 1, 28, 28]    | an image of the same  | 2, 3 |
"""

from __future__ import annotations

import math

import torch
import torch.nn as nn


# %% 01. Embedding a continuous time
#
# Every network in the course has a time input as well as a state. A raw scalar in [0, 1] is a poor
# input to a linear layer, so we first spread it over many frequencies, exactly as in Section 3.28
# of the notes, and then pass it through a small MLP.


def sinusoidal_embedding(t: torch.Tensor, width: int) -> torch.Tensor:
    """Spread a time [B] over ``width`` sine and cosine frequencies. Returns [B, width]."""
    half = width // 2  # Half the coordinates are sines and half cosines.
    frequencies = torch.exp(
        -math.log(10000.0) * torch.arange(half, device=t.device).to(torch.get_default_dtype()) / half
    )  # [half] geometrically spaced frequencies.
    angles = t.to(torch.get_default_dtype())[:, None] * frequencies[None, :]  # [B, half].
    return torch.cat([torch.sin(angles), torch.cos(angles)], dim=1)  # [B, width].


class TimeEmbedding(nn.Module):
    """Sinusoidal features of the time, followed by a two-layer MLP."""

    def __init__(self, width: int = 128, scale: float = 1000.0):
        super().__init__()
        self.width = width  # Width of the embedding the rest of the network receives.
        self.scale = scale  # We multiply a time in [0, 1] by this before the sines.
        self.mlp = nn.Sequential(
            nn.Linear(width, width), nn.SiLU(), nn.Linear(width, width),
        )

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        """t: [B] -> [B, width]."""
        return self.mlp(sinusoidal_embedding(t * self.scale, self.width))  # [B, width].


# %% 02. A denoiser over token sequences
#
# Lectures 4, 5 and 6 all train the same shape of network: it reads a corrupted sequence and a time,
# and produces a distribution over the vocabulary at every position. The notes call it the denoiser
# d_theta(. | z, t). Attention lets the prediction at one position depend on every other position,
# so the factorized reverse step of Chapter 4 is worth taking.


class SequenceDenoiser(nn.Module):
    """A small transformer that maps a token sequence and a time to per-position logits."""

    def __init__(
        self,
        vocab_size: int,
        length: int,
        width: int = 128,
        depth: int = 2,
        heads: int = 4,
        num_classes: int = 0,
    ):
        super().__init__()
        self.vocab_size = vocab_size  # K, including a mask token when the lecture adds one.
        self.length = length  # L, the number of positions.
        self.width = width  # D, the model width.
        self.token_embedding = nn.Embedding(vocab_size, width)  # One vector per token id.
        self.position_embedding = nn.Embedding(length, width)  # One vector per position.
        self.time_embedding = TimeEmbedding(width)  # Shared across positions.
        # A label embedding, used only by the conditional examples. Index num_classes is the null
        # label that classifier-free guidance drops to.
        self.label_embedding = (
            nn.Embedding(num_classes + 1, width) if num_classes else None
        )
        layer = nn.TransformerEncoderLayer(
            d_model=width, nhead=heads, dim_feedforward=4 * width,
            batch_first=True, dropout=0.0, norm_first=True,
        )
        self.transformer = nn.TransformerEncoder(
            layer, num_layers=depth, enable_nested_tensor=False,
        )
        self.head = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, vocab_size))

    def forward(
        self, tokens: torch.Tensor, t: torch.Tensor, labels: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """tokens: [B, L], t: [B], labels: [B] or None -> logits [B, L, K]."""
        positions = torch.arange(tokens.shape[1], device=tokens.device)  # [L].
        h = (
            self.token_embedding(tokens)  # [B, L, D] what each position currently holds.
            + self.position_embedding(positions)[None, :, :]  # [1, L, D] where it sits.
            + self.time_embedding(t)[:, None, :]  # [B, 1, D] how far along the path we are.
        )
        if labels is not None and self.label_embedding is not None:
            h = h + self.label_embedding(labels)[:, None, :]  # [B, 1, D] the condition.
        h = self.transformer(h)  # [B, L, D] every position now depends on every other.
        return self.head(h)  # [B, L, K] logits over the vocabulary at each position.


# %% 03. A velocity network on a vector state
#
# Lectures 2, 3, 6 and 7 learn a field on R^d: a velocity, a score, a noise prediction or a drift.
# The same MLP serves for all of them, because only the target the loss regresses onto differs.


class VelocityMLP(nn.Module):
    """An MLP from a point and a time to a vector of the same width as the point."""

    def __init__(self, dim: int = 2, width: int = 128, depth: int = 3, num_classes: int = 0):
        super().__init__()
        self.dim = dim  # d, the width of the state.
        self.time_embedding = TimeEmbedding(width, scale=1.0)  # Times here already live in [0, 1].
        self.label_embedding = (
            nn.Embedding(num_classes + 1, width) if num_classes else None
        )
        self.input = nn.Linear(dim, width)  # Lift the state to the hidden width.
        layers: list[nn.Module] = []
        for _ in range(depth):
            layers += [nn.Linear(width, width), nn.SiLU()]
        self.body = nn.Sequential(*layers)
        self.output = nn.Linear(width, dim)  # Back down to one vector per coordinate.

    def forward(
        self, x: torch.Tensor, t: torch.Tensor, labels: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """x: [B, d], t: [B], labels: [B] or None -> [B, d]."""
        h = self.input(x) + self.time_embedding(t)  # [B, width] state and time together.
        if labels is not None and self.label_embedding is not None:
            h = h + self.label_embedding(labels)  # [B, width] the condition.
        return self.output(self.body(h))  # [B, d].


# %% 04. A two-time network, for the flow maps of Lecture 6
#
# A flow map sends a state at a start time s to the state at an end time t, so its network has both
# times as inputs. The identity axiom X_{s,s} = id holds by construction when we write the map as
# the state plus the elapsed time times a network output, as `flow_map` below does.


class TwoTimeMLP(nn.Module):
    """An MLP from a point, a start time and an end time to a vector."""

    def __init__(self, dim: int = 2, width: int = 128, depth: int = 3):
        super().__init__()
        self.dim = dim
        self.start_embedding = TimeEmbedding(width, scale=1.0)  # Embedding of s.
        self.end_embedding = TimeEmbedding(width, scale=1.0)  # Embedding of t.
        self.input = nn.Linear(dim, width)
        layers: list[nn.Module] = []
        for _ in range(depth):
            layers += [nn.Linear(width, width), nn.SiLU()]
        self.body = nn.Sequential(*layers)
        self.output = nn.Linear(width, dim)

    def forward(self, x: torch.Tensor, s: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """The average velocity v_{s,t}(x). x: [B, d], s: [B], t: [B] -> [B, d]."""
        h = self.input(x) + self.start_embedding(s) + self.end_embedding(t)  # [B, width].
        return self.output(self.body(h))  # [B, d].

    def flow_map(self, x: torch.Tensor, s: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """X_{s,t}(x) = x + (t - s) v_{s,t}(x), so the identity axiom holds by construction."""
        elapsed = (t - s)[:, None]  # [B, 1] elapsed time, broadcast over coordinates.
        return x + elapsed * self.forward(x, s, t)  # [B, d].


# %% 05. A small U-Net for 28-by-28 images
#
# Lectures 2 and 3 generate images. The network is deliberately small, about half a million
# parameters, so that a few epochs on a laptop CPU already produce recognizable digits.


def _conv_block(in_channels: int, out_channels: int) -> nn.Sequential:
    """Two 3-by-3 convolutions with a SiLU after each, keeping the height and width."""
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, 3, padding=1), nn.SiLU(),
        nn.Conv2d(out_channels, out_channels, 3, padding=1), nn.SiLU(),
    )


class SmallUNet(nn.Module):
    """A three-scale U-Net that maps an image and a time to an image of the same shape.

    The time enters as an extra constant channel, the simplest conditioning that works and the one
    the Lecture 2 script runs on.
    """

    def __init__(self, base_channels: int = 32, in_channels: int = 1):
        super().__init__()
        c = base_channels  # Widths are c, 2c, 4c.
        self.encoder1 = _conv_block(in_channels + 1, c)  # Image plus the time channel.
        self.encoder2 = _conv_block(c, 2 * c)
        self.middle = _conv_block(2 * c, 4 * c)
        self.decoder2 = _conv_block(6 * c, 2 * c)  # Upsampled 4c joined with the 2c skip.
        self.decoder1 = _conv_block(3 * c, c)  # Upsampled 2c joined with the c skip.
        self.output = nn.Conv2d(c, in_channels, 1)  # One prediction per pixel.
        self.pool = nn.MaxPool2d(2)  # Halve the height and width.

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """x: [B, 1, 28, 28], t: [B] -> [B, 1, 28, 28]."""
        import torch.nn.functional as F  # Local import keeps the module header short.

        time_channel = t[:, None, None, None].expand(x.shape[0], 1, *x.shape[2:])  # [B, 1, H, W].
        h = torch.cat([x, time_channel], dim=1)  # [B, 2, H, W].
        skip1 = self.encoder1(h)  # [B, c, 28, 28].
        skip2 = self.encoder2(self.pool(skip1))  # [B, 2c, 14, 14].
        h = self.middle(self.pool(skip2))  # [B, 4c, 7, 7] the compressed representation.
        h = F.interpolate(h, size=skip2.shape[-2:], mode="nearest")  # 7 to 14.
        h = self.decoder2(torch.cat([h, skip2], dim=1))  # Rejoin the 14-by-14 detail.
        h = F.interpolate(h, size=skip1.shape[-2:], mode="nearest")  # 14 to 28.
        h = self.decoder1(torch.cat([h, skip1], dim=1))  # Rejoin the 28-by-28 detail.
        return self.output(h)  # [B, 1, 28, 28].


def count_parameters(model: nn.Module) -> int:
    """Number of trainable weights and biases, printed by every script at startup."""
    return sum(p.numel() for p in model.parameters() if p.requires_grad)