"""Datasets for the course.

Everything here runs offline in under a second, so a fresh clone and the test suite run without a
network. MNIST and ESM-2 embeddings download on first use, and only the scripts with a
``--dataset`` flag reach them.

The discrete lectures all run on the same synthetic DNA, so a loss from Lecture 4 and a loss from
Lecture 5 are comparable.
"""

from __future__ import annotations

from pathlib import Path

import torch

# %% 01. The DNA alphabet
#
# The vocabulary is the four DNA letters, because Chapters 4 and 5 run every worked example on
# DNA. A lecture that needs a mask token appends it as index 4 in its own file, since masked
# diffusion has a mask and uniform diffusion does not.

DNA = "ACGT"  # Token 0 is A, 1 is C, 2 is G, 3 is T.
DNA_SIZE = len(DNA)  # K = 4, the alphabet size the chapters call |V|.
MOTIF = "ACGT"  # The planted motif, used as the label in the guidance examples.


def encode_dna(sequences: list[str]) -> torch.Tensor:
    """Turn a list of strings over ACGT into a long tensor of shape [N, L]."""
    index = {letter: i for i, letter in enumerate(DNA)}  # Letter to token id.
    rows = [[index[letter] for letter in sequence] for sequence in sequences]
    return torch.tensor(rows, dtype=torch.long)  # [N, L] token ids.


def decode_dna(tokens: torch.Tensor) -> list[str]:
    """Turn a long tensor of shape [N, L] back into strings over ACGT."""
    return ["".join(DNA[int(i)] for i in row) for row in tokens]  # One string per row.


# %% 02. Synthetic DNA
#
# We plant two pieces of structure in these sequences: a GC-rich background, and the motif ACGT
# at a random position in a fraction of them. If we trained on uniform random sequences the
# sampler would look the same before and after. The GC fraction gives the label for the
# conditional examples, and the motif gives the property to steer toward.


def synthetic_dna(
    num_sequences: int = 512,
    length: int = 16,
    motif_fraction: float = 0.75,
    gc_bias: float = 0.6,
    seed: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Draw synthetic DNA with a planted motif.

    Returns the tokens as a long tensor of shape [N, L] and a binary label of shape [N] that is
    one when the sequence has GC fraction at least 0.6. The label is a deterministic function of
    the sequence, with no biological meaning.
    """
    generator = torch.Generator().manual_seed(seed)  # Local generator; global seed untouched.

    # Background letters: C and G each get gc_bias/2, A and T share the rest.
    background = torch.tensor([
        (1.0 - gc_bias) / 2.0,  # A
        gc_bias / 2.0,  # C
        gc_bias / 2.0,  # G
        (1.0 - gc_bias) / 2.0,  # T
    ])
    tokens = torch.multinomial(
        background.expand(num_sequences * length, DNA_SIZE),  # One draw per position.
        num_samples=1,
        replacement=True,
        generator=generator,
    ).view(num_sequences, length)  # [N, L] background sequence.

    # Plant the motif in a fraction of the sequences, at a position drawn uniformly.
    motif = encode_dna([MOTIF])[0]  # [4] token ids of ACGT.
    carries_motif = torch.rand(num_sequences, generator=generator) < motif_fraction  # [N] bool.
    # The motif starts at one of three places, so a small model can learn the pattern from a few
    # hundred sequences while the position still varies. We keep a start only when the whole motif
    # fits, so a short sequence has fewer of them.
    room = length - len(MOTIF)  # Last index at which the motif still fits.
    choices = torch.tensor([s for s in (2, 6, 10) if s <= room] or [0])
    starts = choices[torch.randint(0, len(choices), (num_sequences,), generator=generator)]
    for row in torch.nonzero(carries_motif).flatten().tolist():
        start = int(starts[row])
        tokens[row, start:start + len(MOTIF)] = motif  # Overwrite those positions.

    gc = ((tokens == 1) | (tokens == 2)).to(torch.get_default_dtype()).mean(dim=1)  # [N] fraction of C and G.
    labels = (gc >= 0.6).long()  # [N] the toy label used by the conditional examples.
    return tokens, labels


def dna_splits(
    num_sequences: int = 512, length: int = 16, seed: int = 0, validation: float = 0.2,
) -> dict[str, torch.Tensor]:
    """The same synthetic DNA, split into a training and a validation set."""
    tokens, labels = synthetic_dna(num_sequences, length, seed=seed)
    cut = int(num_sequences * (1.0 - validation))  # Index where validation begins.
    return {
        "train_tokens": tokens[:cut], "train_labels": labels[:cut],
        "valid_tokens": tokens[cut:], "valid_labels": labels[cut:],
    }


def write_dna_tsv(path: str | Path, tokens: torch.Tensor, labels: torch.Tensor) -> None:
    """Write sequences and labels as a two-column TSV, the format ``--data`` reads."""
    lines = ["sequence\tlabel"]
    lines += [f"{s}\t{int(y)}" for s, y in zip(decode_dna(tokens), labels)]
    Path(path).write_text("\n".join(lines) + "\n")


def read_dna_tsv(path: str | Path) -> tuple[torch.Tensor, torch.Tensor]:
    """Read a two-column TSV of sequences and labels."""
    rows = Path(path).read_text().strip().split("\n")[1:]  # Skip the header line.
    sequences = [row.split("\t")[0] for row in rows]
    labels = [int(row.split("\t")[1]) for row in rows]
    return encode_dna(sequences), torch.tensor(labels, dtype=torch.long)


# %% 03. Two-dimensional densities
#
# The continuous lectures run on four Gaussians at the corners of a square, a target that plots in
# two dimensions and has an exact density. The modes are far enough apart that a collapsed sampler
# is obvious. The log density has a closed form the reports can quote.

GAUSSIAN_CENTERS = torch.tensor([[2.0, 2.0], [2.0, -2.0], [-2.0, 2.0], [-2.0, -2.0]])
GAUSSIAN_STD = 0.35  # Standard deviation of each component.


def gaussian_mixture_2d(num_points: int = 2048, seed: int = 0) -> torch.Tensor:
    """Draw from four equally weighted Gaussians at the corners of a square. Returns [N, 2]."""
    generator = torch.Generator().manual_seed(seed)
    which = torch.randint(0, len(GAUSSIAN_CENTERS), (num_points,), generator=generator)  # [N].
    noise = torch.randn(num_points, 2, generator=generator) * GAUSSIAN_STD  # [N, 2].
    return GAUSSIAN_CENTERS[which] + noise  # [N, 2] points.


def gaussian_mixture_log_density(points: torch.Tensor) -> torch.Tensor:
    """Exact log density of the four-Gaussian mixture at each row of ``points`` ([N, 2] -> [N])."""
    diff = points[:, None, :] - GAUSSIAN_CENTERS.to(points)  # [N, 4, 2] offset to each center.
    squared = (diff ** 2).sum(dim=-1)  # [N, 4] squared distance to each center.
    log_component = (
        -0.5 * squared / GAUSSIAN_STD ** 2  # Gaussian exponent.
        - torch.log(torch.tensor(2.0 * torch.pi * GAUSSIAN_STD ** 2))  # Normalizer in 2-D.
    )  # [N, 4].
    weight = torch.log(torch.tensor(1.0 / len(GAUSSIAN_CENTERS)))  # Equal mixture weights.
    return torch.logsumexp(log_component + weight, dim=-1)  # [N] mixture log density.


def two_moons(num_points: int = 2048, noise: float = 0.08, seed: int = 0) -> torch.Tensor:
    """Two interleaving half circles, the standard test for a sampler that must bend. [N, 2]."""
    generator = torch.Generator().manual_seed(seed)
    half = num_points // 2
    angle = torch.rand(half, generator=generator) * torch.pi  # [half] angle along each arc.
    upper = torch.stack([torch.cos(angle), torch.sin(angle)], dim=1)  # Upper half circle.
    lower = torch.stack([1.0 - torch.cos(angle), 0.5 - torch.sin(angle)], dim=1)  # Lower arc.
    points = torch.cat([upper, lower], dim=0)  # [2*half, 2].
    points = points + noise * torch.randn(points.shape, generator=generator)  # Blur the arcs.
    return (points - points.mean(0)) / points.std(0)  # Standardize to roughly unit scale.


def checkerboard(num_points: int = 2048, seed: int = 0) -> torch.Tensor:
    """A checkerboard of eight squares, a density with disconnected support. [N, 2]."""
    generator = torch.Generator().manual_seed(seed)
    x = torch.rand(num_points, generator=generator) * 4.0 - 2.0  # [N] in [-2, 2].
    y = torch.rand(num_points, generator=generator) - torch.randint(
        0, 2, (num_points,), generator=generator,
    ).to(torch.get_default_dtype()) * 2.0  # [N] in [-2, 1] before the shift below.
    y = y + torch.floor(x) % 2.0  # Shift alternate columns so the squares interleave.
    return torch.stack([x, y * 2.0], dim=1)  # [N, 2] points.


# %% 04. Short phrases
#
# Chapter 6 runs its categorical flow maps on four-token phrases. The grammar is small enough that
# a reader can tell at a glance whether a generated phrase is in the language.

_SUBJECTS = ["red", "blue", "green", "small"]
_OBJECTS = ["circle", "square", "arrow", "dot"]
_VERBS = ["moves", "turns", "drifts", "jumps"]
_DIRECTIONS = ["left", "right", "up", "down"]


def phrase_vocabulary() -> list[str]:
    """The token vocabulary of the phrase language, in a fixed order."""
    return _SUBJECTS + _OBJECTS + _VERBS + _DIRECTIONS  # 16 tokens, four per slot.


def phrases(num_phrases: int = 512, seed: int = 0) -> tuple[torch.Tensor, list[str]]:
    """Draw four-token phrases such as ``red circle moves left``. Returns [N, 4] and the vocabulary.

    Each slot draws from its own quarter of the vocabulary, so a correct phrase is one whose
    position i holds a token from group i. The training support is then easy to check.
    """
    generator = torch.Generator().manual_seed(seed)
    vocabulary = phrase_vocabulary()
    groups = [_SUBJECTS, _OBJECTS, _VERBS, _DIRECTIONS]
    columns = []
    for slot, group in enumerate(groups):
        offset = slot * len(group)  # Where this slot's tokens start in the vocabulary.
        picks = torch.randint(0, len(group), (num_phrases,), generator=generator)
        columns.append(picks + offset)  # [N] token ids inside the full vocabulary.
    return torch.stack(columns, dim=1), vocabulary  # [N, 4] and the 16 token strings.


# %% 05. MNIST and ESM-2
#
# These two download only when a script runs with --dataset mnist or --dataset esm2, and every
# such script also runs on the offline data above.


def mnist_loader(batch_size: int = 128, root: str | Path = "data", train: bool = True):
    """MNIST as a DataLoader of images normalized to [-1, 1]. Downloads on first use."""
    from torch.utils.data import DataLoader  # Imported here so the module loads without torchvision.
    from torchvision import datasets
    from torchvision.transforms import v2

    transform = v2.Compose([
        v2.ToImage(),  # Channel-first image: [1, 28, 28].
        v2.ToDtype(torch.float32, scale=True),  # Integer pixels 0-255 to floats in [0, 1].
        v2.Normalize(mean=(0.5,), std=(0.5,)),  # Then [0, 1] to [-1, 1].
    ])
    dataset = datasets.MNIST(root=root, train=train, download=True, transform=transform)
    return DataLoader(dataset, batch_size=batch_size, shuffle=train, num_workers=0)


def toy_images(num_images: int = 256, size: int = 28, seed: int = 0) -> torch.Tensor:
    """Images shaped like MNIST, drawn from a two-mode distribution of bars and blobs.

    These stand in for MNIST wherever a test has to run without a download. They have the same
    shape and the same [-1, 1] range, so the identical training and sampling code paths run.
    """
    generator = torch.Generator().manual_seed(seed)
    images = torch.full((num_images, 1, size, size), -1.0)  # Background at the low end of [-1, 1].
    grid = torch.arange(size)
    for i in range(num_images):
        if i % 2 == 0:  # A vertical bar.
            column = int(torch.randint(6, size - 6, (1,), generator=generator))
            images[i, 0, 4:size - 4, column - 2:column + 2] = 1.0
        else:  # A round blob.
            cy = int(torch.randint(8, size - 8, (1,), generator=generator))
            cx = int(torch.randint(8, size - 8, (1,), generator=generator))
            dy = (grid[:, None] - cy) ** 2
            dx = (grid[None, :] - cx) ** 2
            images[i, 0][(dy + dx) <= 25] = 1.0
    images = images + 0.05 * torch.randn(images.shape, generator=generator)  # Light pixel noise.
    return images.clamp(-1.0, 1.0)  # [N, 1, size, size] in the training range.


PEPTIDES = [
    "KWKLFKKIGAVLKVL", "GIGKFLHSAKKFGKA", "FLPIIAKLLSGLL", "ILPWKWPWWPWRR",
    "KRWWKWWRR", "GLLSKLWKKVFKAFKKFLKK", "ACDEFGHIKLMNPQR", "RRWWRRWRR",
    "LLGDFFRKSKEKIGKE", "KIAKVALKAL", "VRRFPWWWPFLRR", "GWLKKIGKKIERVGQ",
]  # Short antimicrobial-peptide-like sequences, used when ESM-2 is unavailable.


def esm2_embeddings(
    sequences: list[str] | None = None,
    model_name: str = "facebook/esm2_t33_650M_UR50D",
    max_length: int = 32,
    device: torch.device | str = "cpu",
):
    """Residue embeddings from a frozen ESM-2 encoder. Downloads the checkpoint on first use.

    Returns the embeddings [N, L, D], the attention mask [N, L], the tokenizer and the masked
    language-model head that decodes generated latents back to amino acids in Lecture 3.
    """
    from transformers import AutoModel, AutoModelForMaskedLM, AutoTokenizer

    sequences = sequences or PEPTIDES
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    encoder = AutoModel.from_pretrained(model_name).to(device).eval()
    head = AutoModelForMaskedLM.from_pretrained(model_name).lm_head.to(device).eval()
    for parameter in encoder.parameters():
        parameter.requires_grad = False  # The encoder is frozen; only the denoiser trains.

    batch = tokenizer(
        sequences, padding="max_length", truncation=True,
        max_length=max_length, return_tensors="pt",
    ).to(device)
    with torch.no_grad():
        embeddings = encoder(**batch).last_hidden_state  # [N, L, D] residue embeddings.
    return embeddings, batch["attention_mask"], tokenizer, head


def toy_embeddings(
    num_sequences: int = 128, length: int = 16, width: int = 32, seed: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Embeddings shaped like ESM-2 output, for the offline path of the latent examples.

    Each sequence is a smooth low-frequency curve along the residue axis plus noise, so neighbouring
    positions are correlated the way real residue embeddings are. Returns [N, L, D] and a mask.
    """
    generator = torch.Generator().manual_seed(seed)
    phase = torch.rand(num_sequences, 1, width, generator=generator) * 2.0 * torch.pi
    frequency = torch.rand(num_sequences, 1, width, generator=generator) * 0.6 + 0.1
    position = torch.arange(length).view(1, length, 1).to(torch.get_default_dtype())  # [1, L, 1].
    signal = torch.sin(frequency * position + phase)  # [N, L, D] smooth along the residue axis.
    noise = 0.2 * torch.randn(num_sequences, length, width, generator=generator)
    mask = torch.ones(num_sequences, length)  # No padding in the toy set.
    return signal + noise, mask


def standardize(embeddings: torch.Tensor, mask: torch.Tensor):
    """Center and scale each embedding coordinate over the valid residues.

    Returns the standardized embeddings and the mean and standard deviation, used to undo the
    transformation before the decoder reads amino acids.
    """
    valid = mask.bool().unsqueeze(-1)  # [N, L, 1] true at a real residue.
    count = valid.sum()  # Number of valid residue-coordinate entries.
    mean = (embeddings * valid).sum(dim=(0, 1)) / count  # [D] per-coordinate mean.
    variance = ((embeddings - mean) ** 2 * valid).sum(dim=(0, 1)) / count  # [D].
    std = variance.sqrt().clamp(min=1e-6)  # [D], clamped so a dead coordinate cannot divide by zero.
    return (embeddings - mean) / std, mean, std