mtrcnn.py

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "huggingface-hub",
#     "jinja2",
#     "numpy",
#     "scipy",
#     "soundfile",
#     "torch",
# ]
# ///
"""Train a compact domain-generalized mosquito species classifier."""

# pylint: disable=too-few-public-methods,too-many-arguments,too-many-instance-attributes
# pylint: disable=too-many-lines,too-many-locals,too-many-positional-arguments

import argparse
import json
import math
import os
import random
import re
import shutil
import statistics
import urllib.request
import zipfile
from collections import Counter
from pathlib import Path

import numpy as np
import soundfile
import torch
from huggingface_hub import (
    EvalResult,
    HfApi,
    ModelCard,
    ModelCardData,
    hf_hub_download,
    snapshot_download,
)
from huggingface_hub.errors import RepositoryNotFoundError
from scipy.signal import resample_poly
from torch import nn
from torch.nn import functional
from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler

DATASET = "aptemvs/mosquitoes-biodcase2026-task5"
REPO = "jgalego/mosquito-mtrcnn-dg"
ZENODO = "https://zenodo.org/api/records/20478577/files/Development_data.zip/content"
SAMPLE_RATE = 8_000
N_FFT = 512
HOP_LENGTH = 80
N_MELS = 64
# Mains hum and handling noise sit below 200 Hz; wingbeat fundamentals are higher.
MIN_HZ = 200
LEVEL_RMS = 0.05
# Every D5 clip lasts 0.625 s; longer crops would pad them with telltale silence.
SECONDS = 0.625
SPECIES = [
    "Aedes aegypti",
    "Aedes albopictus",
    "Culex quinquefasciatus",
    "Anopheles gambiae",
    "Anopheles arabiensis",
    "Anopheles dirus",
    "Culex pipiens",
    "Anopheles minimus",
    "Anopheles stephensi",
]
DOMAINS = ["D1", "D2", "D3", "D4", "D5"]
FILE_ID = re.compile(r"^S_(\d+)_D_(\d+)_(\d+)$")
SYNTHETIC_DOMAINS = [
    [0, 3, 4],
    [0, 2, 4],
    [3, 4],
    [4],
    [4],
    [0],
    [0, 1, 4],
    [0, 2, 3],
    [0, 1, 2],
]


def parse_file_id(file_id):
    """Return zero-based species and domain indices from an official file ID."""
    match = FILE_ID.fullmatch(file_id)
    if not match:
        raise ValueError(f"invalid BioDCASE file ID: {file_id}")
    species, domain, _ = map(int, match.groups())
    if not (1 <= species <= len(SPECIES) and 1 <= domain <= len(DOMAINS)):
        raise ValueError(f"out-of-range BioDCASE file ID: {file_id}")
    return species - 1, domain - 1


def held_domain_split(ids, fold, seed=42, singleton_fraction=0.1):
    """Hold out a rare domain, or rows when only one domain exists, per species."""
    groups = {}
    for file_id in ids:
        groups.setdefault(parse_file_id(file_id), []).append(file_id)
    held_domains = {}
    validation_ids = []
    for species, species_name in enumerate(SPECIES):
        available = [domain for domain in range(len(DOMAINS)) if groups.get((species, domain))]
        candidates = sorted(available)
        if len(candidates) == 1:
            held_domains[species] = None
            rows = groups[species, candidates[0]].copy()
            random.Random(seed + species).shuffle(rows)
            validation_size = min(
                len(rows) - 1,
                max(1, round(singleton_fraction * len(rows))),
            )
            validation_ids.extend(rows[:validation_size])
            continue
        if fold >= len(candidates):
            raise ValueError(
                f"fold {fold} unavailable for {species_name}: "
                f"only {len(candidates)} domains"
            )
        held_domains[species] = candidates[fold]
        validation_ids.extend(groups[species, candidates[fold]])
    validation_set = set(validation_ids)
    training_ids = [file_id for file_id in ids if file_id not in validation_set]
    training_species = {parse_file_id(file_id)[0] for file_id in training_ids}
    validation_species = {parse_file_id(file_id)[0] for file_id in validation_ids}
    expected_species = set(range(len(SPECIES)))
    if training_species != expected_species or validation_species != expected_species:
        raise ValueError("held-domain split does not cover every species")
    return training_ids, validation_ids, held_domains


def load_ids(split="Training"):
    """Download and read one official ID list."""
    path = hf_hub_download(DATASET, f"data/metadata/{split}_ids.txt", repo_type="dataset")
    with open(path, encoding="utf-8") as handle:
        return [line.strip() for line in handle if line.strip()]


def load_unseen_domains():
    """Return the official unseen domain index for each species index."""
    path = hf_hub_download(DATASET, "data/metadata/split_summary.json", repo_type="dataset")
    mapping = json.loads(Path(path).read_text(encoding="utf-8"))["unseen_domain_by_species"]
    return {SPECIES.index(species): DOMAINS.index(domain) for species, domain in mapping.items()}


def download(url, path):
    """Download a large file atomically with coarse progress output."""
    partial = path.with_suffix(path.suffix + ".part")
    with urllib.request.urlopen(url) as response, open(partial, "wb") as output:
        total = int(response.headers.get("Content-Length", 0))
        downloaded = 0
        report_at = 0
        while chunk := response.read(1 << 20):
            output.write(chunk)
            downloaded += len(chunk)
            if downloaded >= report_at:
                print(
                    f"downloaded {downloaded / 1e9:.1f}/{total / 1e9:.1f} GB",
                    flush=True,
                )
                report_at += 250_000_000
    partial.rename(path)


def prepare_data(data):
    """Download and flatten the official Zenodo development archive."""
    destination = Path(data)
    marker = destination / ".complete"
    if marker.exists():
        print(f"audio already prepared at {destination}")
        return
    destination.mkdir(parents=True, exist_ok=True)
    archive = destination.parent / "Development_data.zip"
    if not archive.exists():
        download(ZENODO, archive)
    with zipfile.ZipFile(archive) as source:
        members = [
            name
            for name in source.namelist()
            if name.lower().endswith(".wav")
            and not name.startswith("__MACOSX")
            and not Path(name).name.startswith("._")
        ]
        print(f"extracting {len(members)} waveforms", flush=True)
        for index, name in enumerate(members, 1):
            with source.open(name) as src, open(
                destination / Path(name).name, "wb"
            ) as dst:
                shutil.copyfileobj(src, dst)
            if index % 25_000 == 0:
                print(f"extracted {index}/{len(members)}", flush=True)
    marker.write_text(str(len(members)), encoding="utf-8")
    archive.unlink()
    print(f"prepared {len(members)} waveforms at {destination}")


def select_ids(ids, size, seed):
    """Select a seeded species-domain-stratified subset."""
    if not size or size >= len(ids):
        return ids
    groups = {}
    for file_id in ids:
        groups.setdefault(parse_file_id(file_id), []).append(file_id)
    rng = random.Random(seed)
    for rows in groups.values():
        rng.shuffle(rows)
    selected = []
    while len(selected) < size:
        added = False
        for key in sorted(groups):
            if groups[key] and len(selected) < size:
                selected.append(groups[key].pop())
                added = True
        if not added:
            break
    rng.shuffle(selected)
    return selected


def cpus():
    """Return the CPU quota visible to this process."""
    try:
        quota, period = Path("/sys/fs/cgroup/cpu.max").read_text(
            encoding="utf-8"
        ).split()
        if quota != "max":
            return max(1, int(quota) // int(period))
    except OSError:
        pass
    return len(os.sched_getaffinity(0))


def load_waveform(path, seconds=None, training=False):
    """Read a mono waveform at the model sample rate, or only a random or central crop."""
    start, frames = 0, -1
    if seconds is not None:
        info = soundfile.info(path)
        frames = math.ceil(seconds * info.samplerate)
        gap = max(0, info.frames - frames)
        start = random.randrange(gap + 1) if training else gap // 2
    waveform, sample_rate = soundfile.read(
        path, start=start, frames=frames, dtype="float32", always_2d=False
    )
    if waveform.ndim > 1:
        waveform = waveform.mean(axis=1)
    if sample_rate != SAMPLE_RATE:
        waveform = resample_poly(waveform, SAMPLE_RATE, sample_rate).astype(np.float32)
    return waveform


def normalize_level(waveforms, level_rms):
    """Scale each waveform to one RMS level so recording gain cannot reveal its domain."""
    if level_rms <= 0:
        return waveforms.astype(np.float32)
    rms = np.sqrt(np.mean(np.square(waveforms), axis=-1, keepdims=True))
    return (waveforms * (level_rms / np.maximum(rms, 1e-6))).astype(np.float32)


def split_windows(waveform, samples):
    """Cover a complete clip with fixed-length windows, the last aligned to its end."""
    starts = list(range(0, max(len(waveform) - samples + 1, 1), samples))
    if starts[-1] + samples < len(waveform):
        starts.append(len(waveform) - samples)
    return np.stack(
        [fixed_length(waveform[start : start + samples], samples, False) for start in starts]
    )


def fixed_length(waveform, samples, training):
    """Randomly or centrally crop and zero-pad a waveform."""
    if len(waveform) < samples:
        gap = samples - len(waveform)
        left = random.randrange(gap + 1) if training else gap // 2
        waveform = np.pad(waveform, (left, gap - left))
    if len(waveform) > samples:
        gap = len(waveform) - samples
        start = random.randrange(gap + 1) if training else gap // 2
        waveform = waveform[start : start + samples]
    return waveform.astype(np.float32)


class MosquitoDataset(Dataset):
    """Fold-filtered 8 kHz audio with species-domain sampling weights."""

    def __init__(
        self, root, ids, seconds, training, balance_power, synthetic=False, level_rms=LEVEL_RMS
    ):
        self.root = Path(root) if root else None
        self.ids = ids
        self.samples = round(seconds * SAMPLE_RATE)
        self.training = training
        self.synthetic = synthetic
        self.level_rms = level_rms
        self.labels = [parse_file_id(file_id) for file_id in ids]
        counts = Counter(self.labels)
        self.sample_weights = [counts[label] ** -balance_power for label in self.labels]

    def __len__(self):
        return len(self.ids)

    def _synthetic_waveform(self, index):
        species, domain = self.labels[index]
        row = int(self.ids[index].rsplit("_", maxsplit=1)[1])
        time = np.arange(self.samples, dtype=np.float32) / SAMPLE_RATE
        phase = (row % 7) * math.pi / 7
        return 0.2 * np.sin(2 * np.pi * (180 + 35 * species + domain) * time + phase)

    def __getitem__(self, index):
        if self.synthetic:
            waveform = self._synthetic_waveform(index)
        else:
            waveform = load_waveform(
                self.root / f"{self.ids[index]}.wav",
                self.samples / SAMPLE_RATE,
                self.training,
            )
            waveform = fixed_length(waveform, self.samples, self.training)
        waveform = normalize_level(waveform, self.level_rms)
        if self.training:
            waveform = np.roll(waveform, random.randrange(len(waveform)))
            waveform = waveform * 10 ** random.uniform(-0.3, 0.3)
            waveform = waveform + np.random.normal(
                0, random.uniform(0, 0.005), len(waveform)
            )
        species, domain = self.labels[index]
        return {
            "file_id": self.ids[index],
            "waveform": waveform.astype(np.float32),
            "species": species,
            "domain": domain,
        }


class AudioCollator:
    """Stack fixed-length waveforms and labels."""

    def __call__(self, rows):
        waveforms = torch.from_numpy(np.stack([row["waveform"] for row in rows]))
        return {
            "file_ids": [row["file_id"] for row in rows],
            "waveforms": waveforms,
            "lengths": torch.full((len(rows),), waveforms.size(1), dtype=torch.long),
            "species": torch.tensor([row["species"] for row in rows]),
            "domains": torch.tensor([row["domain"] for row in rows]),
        }


class ClipDataset(MosquitoDataset):
    """Complete clips split into training-length evaluation windows."""

    def __getitem__(self, index):
        if self.synthetic:
            waveform = self._synthetic_waveform(index)
        else:
            waveform = load_waveform(self.root / f"{self.ids[index]}.wav")
        species, domain = self.labels[index]
        return {
            "windows": normalize_level(split_windows(waveform, self.samples), self.level_rms),
            "species": species,
            "domain": domain,
        }


class WindowCollator:
    """Concatenate every clip's windows and record which clip owns each."""

    def __call__(self, rows):
        counts = torch.tensor([len(row["windows"]) for row in rows])
        waveforms = torch.from_numpy(np.concatenate([row["windows"] for row in rows]))
        return {
            "waveforms": waveforms,
            "lengths": torch.full((len(waveforms),), waveforms.size(1), dtype=torch.long),
            "owners": torch.repeat_interleave(torch.arange(len(rows)), counts),
            "species": torch.tensor([row["species"] for row in rows]),
            "domains": torch.tensor([row["domain"] for row in rows]),
        }


def synthetic_ids(rows_per_cell):
    """Build official-shaped IDs with the real species-domain availability pattern."""
    return [
        f"S_{species + 1}_D_{domain + 1}_{row + 1}"
        for species, domains in enumerate(SYNTHETIC_DOMAINS)
        for domain in domains
        for row in range(rows_per_cell)
    ]


def make_datasets(args):
    """Create one pseudo-unseen fold, or for fold -1 all training rows and the official
    validation split."""
    ids = synthetic_ids(args.synthetic_rows) if args.synthetic else load_ids()
    if args.fold < 0:
        training_ids = ids
        validation_ids = ids if args.synthetic else load_ids("Validation")
        held_domains = dict.fromkeys(range(len(SPECIES)))
    else:
        training_ids, validation_ids, held_domains = held_domain_split(
            ids, args.fold, seed=args.seed
        )
    training_ids = select_ids(training_ids, args.train_size, args.seed)
    validation_ids = select_ids(validation_ids, args.validation_size, args.seed)
    training = MosquitoDataset(
        args.data,
        training_ids,
        args.seconds,
        True,
        args.balance_power,
        args.synthetic,
        args.level_rms,
    )
    validation = MosquitoDataset(
        args.data,
        validation_ids,
        args.seconds,
        False,
        args.balance_power,
        args.synthetic,
        args.level_rms,
    )
    expected = set(range(len(SPECIES)))
    if {label[0] for label in training.labels} != expected:
        raise ValueError("training size limit removed a species")
    if {label[0] for label in validation.labels} != expected:
        raise ValueError("validation size limit removed a species")
    return training, validation, held_domains


def seed_worker(_worker_id):
    """Seed NumPy and Python from a DataLoader worker's PyTorch seed."""
    worker_seed = torch.initial_seed() % (2**32)
    random.seed(worker_seed)
    np.random.seed(worker_seed)


def make_loader(dataset, batch_size, workers, training, seed, collator=None):
    """Create a deterministic loader, balanced by species-domain cell for training."""
    generator = torch.Generator().manual_seed(seed)
    sampler = None
    if training:
        sampler = WeightedRandomSampler(
            dataset.sample_weights,
            num_samples=len(dataset),
            replacement=True,
            generator=generator,
        )
    return DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,
        shuffle=False,
        num_workers=workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=workers > 0,
        worker_init_fn=seed_worker,
        generator=generator,
        collate_fn=collator or AudioCollator(),
    )


def deterministic_view(dataset, ids=None):
    """Return an unaugmented view over a dataset's IDs."""
    return MosquitoDataset(
        dataset.root,
        dataset.ids if ids is None else ids,
        dataset.samples / SAMPLE_RATE,
        training=False,
        balance_power=0.0,
        synthetic=dataset.synthetic,
        level_rms=dataset.level_rms,
    )


@torch.no_grad()
def compute_normalization(dataset, size, batch_size, workers, device, seed, min_hz, lifter=0):
    """Estimate per-mel mean and variance from fold-training rows only."""
    ids = select_ids(dataset.ids, size, seed)
    loader = make_loader(
        deterministic_view(dataset, ids), batch_size, workers, False, seed
    )
    frontend = LogMelFrontend(min_hz=min_hz, lifter=lifter).to(device)
    total = torch.zeros(N_MELS, dtype=torch.float64, device=device)
    squared = torch.zeros_like(total)
    frames = 0
    for batch in loader:
        waveforms = batch["waveforms"].to(device, non_blocking=True)
        lengths = batch["lengths"].to(device, non_blocking=True)
        features, frame_lengths = frontend(waveforms, lengths)
        positions = torch.arange(features.size(1), device=device)[None, :]
        mask = positions < frame_lengths[:, None]
        selected = features[mask].double()
        total += selected.sum(0)
        squared += selected.square().sum(0)
        frames += selected.size(0)
    mean = total / frames
    variance = squared / frames - mean.square()
    return mean.float().cpu(), variance.clamp_min(1e-6).sqrt().float().cpu()


def balanced_accuracy(gold, predicted):
    """Return mean per-species recall over labels present in gold."""
    gold = np.asarray(gold)
    predicted = np.asarray(predicted)
    recalls = [
        np.mean(predicted[gold == label] == label)
        for label in sorted(set(gold.tolist()))
    ]
    return float(np.mean(recalls)) if recalls else 0.0


def prediction_metrics(labels, predictions):
    """Return balanced accuracy, per-species recall and a gold-by-predicted confusion."""
    labels = labels.cpu().numpy()
    predictions = predictions.cpu().numpy()
    confusion = np.bincount(
        labels * len(SPECIES) + predictions, minlength=len(SPECIES) ** 2
    ).reshape(len(SPECIES), len(SPECIES))
    return {
        "balanced_accuracy": balanced_accuracy(labels, predictions),
        "rows": len(labels),
        "per_species": {
            SPECIES[label]: float(np.mean(predictions[labels == label] == label))
            for label in sorted(set(labels.tolist()))
        },
        "confusion": confusion.tolist(),
    }


def official_metrics(labels, predictions, domains, unseen_domains):
    """Score the BioDCASE seen and unseen species-domain partitions."""
    unseen = torch.tensor(
        [
            unseen_domains[label.item()] == domain.item()
            for label, domain in zip(labels, domains)
        ],
        dtype=torch.bool,
    )
    return {
        "seen": prediction_metrics(labels[~unseen], predictions[~unseen]),
        "unseen": prediction_metrics(labels[unseen], predictions[unseen]),
    }


def fold_prediction_metrics(labels, predictions, domains, held_domains):
    """Split fold metrics into true domain holdouts and singleton reserves."""
    metrics = prediction_metrics(labels, predictions)
    domain_holdout = torch.tensor(
        [held_domains[label.item()] is not None for label in labels],
        dtype=torch.bool,
    )
    metrics["domain_holdout"] = prediction_metrics(
        labels[domain_holdout], predictions[domain_holdout]
    )
    metrics["same_domain"] = prediction_metrics(
        labels[~domain_holdout], predictions[~domain_holdout]
    )
    expected_domains = torch.tensor(
        [
            held_domains[label.item()]
            if held_domains[label.item()] is not None
            else domain.item()
            for label, domain in zip(labels, domains)
        ]
    )
    if not torch.equal(domains, expected_domains):
        raise ValueError("validation rows do not match the declared fold policy")
    return metrics


@torch.no_grad()
def collect_outputs(model, frontend, dataset, batch_size, workers, device, seed):
    """Collect deterministic logits and embeddings for one dataset."""
    model.eval()
    loader = make_loader(dataset, batch_size, workers, False, seed)
    collected = {"labels": [], "domains": [], "logits": [], "embeddings": []}
    for batch in loader:
        waveforms = batch["waveforms"].to(device, non_blocking=True)
        lengths = batch["lengths"].to(device, non_blocking=True)
        domains = batch["domains"].to(device, non_blocking=True)
        features, frame_lengths = frontend(waveforms, lengths)
        outputs = model(features, frame_lengths, domains)
        collected["labels"].append(batch["species"])
        collected["domains"].append(batch["domains"])
        collected["logits"].append(outputs["species_logits"].cpu())
        collected["embeddings"].append(outputs["embedding"].cpu())
    return {name: torch.cat(values) for name, values in collected.items()}


@torch.no_grad()
def collect_clip_outputs(model, frontend, dataset, batch_size, workers, device, seed):
    """Average window probabilities and embeddings over each complete clip."""
    model.eval()
    loader = make_loader(dataset, batch_size, workers, False, seed, WindowCollator())
    collected = {"labels": [], "domains": [], "probabilities": [], "embeddings": []}
    for batch in loader:
        owners = batch["owners"].to(device, non_blocking=True)
        domains = batch["domains"].to(device, non_blocking=True)
        window_outputs = {"probabilities": [], "embeddings": []}
        for start in range(0, len(owners), batch_size):
            window = slice(start, start + batch_size)
            features, frame_lengths = frontend(
                batch["waveforms"][window].to(device, non_blocking=True),
                batch["lengths"][window].to(device, non_blocking=True),
            )
            outputs = model(features, frame_lengths, domains[owners[window]])
            window_outputs["probabilities"].append(
                functional.softmax(outputs["species_logits"], dim=1)
            )
            window_outputs["embeddings"].append(outputs["embedding"])
        counts = torch.bincount(owners, minlength=len(domains)).unsqueeze(1)
        for name, values in window_outputs.items():
            values = torch.cat(values)
            totals = values.new_zeros(len(domains), values.size(1))
            collected[name].append((totals.index_add_(0, owners, values) / counts).cpu())
        collected["labels"].append(batch["species"])
        collected["domains"].append(batch["domains"])
    return {name: torch.cat(values) for name, values in collected.items()}


def mahalanobis_fit(reference_embeddings, reference_labels):
    """Return class means and the precision of one regularized shared covariance matrix."""
    if set(reference_labels.tolist()) != set(range(len(SPECIES))):
        raise ValueError("Mahalanobis reference data must cover every species")
    means = torch.stack(
        [
            reference_embeddings[reference_labels == label].mean(0)
            for label in range(len(SPECIES))
        ]
    )
    residuals = reference_embeddings - means[reference_labels]
    covariance = residuals.T @ residuals / max(1, len(residuals) - len(SPECIES))
    regularization = covariance.diagonal().mean().clamp_min(1e-6) * 1e-3
    covariance = covariance + regularization * torch.eye(covariance.size(0))
    return means, torch.linalg.pinv(covariance)


def mahalanobis_predict(means, precision, embeddings):
    """Assign each embedding to the nearest class mean in Mahalanobis distance."""
    differences = embeddings[:, None, :] - means[None, :, :]
    distances = torch.einsum("ncd,de,nce->nc", differences, precision, differences)
    return distances.argmin(1)


def sampled_priors(dataset, device):
    """Return species priors under the balanced training sampler for logit adjustment."""
    totals = torch.zeros(len(SPECIES), dtype=torch.float64).index_add_(
        0,
        torch.tensor([species for species, _domain in dataset.labels]),
        torch.tensor(dataset.sample_weights, dtype=torch.float64),
    )
    return (totals / totals.sum()).float().to(device)


def grl_strength(step, total_steps):
    """Return the standard sigmoid gradient-reversal warmup."""
    progress = step / max(1, total_steps - 1)
    return 2 / (1 + math.exp(-10 * progress)) - 1


def learning_rate_schedule(total_steps, warmup_ratio):
    """Return a linear-warmup cosine-decay multiplier."""
    warmup_steps = max(1, round(total_steps * warmup_ratio))

    def multiplier(step):
        if step < warmup_steps:
            return (step + 1) / warmup_steps
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return 0.5 * (1 + math.cos(math.pi * min(1.0, progress)))

    return multiplier


def train_epoch(model, frontend, loader, optimizer, scheduler, priors, epoch, state):
    """Run one balanced training epoch or stop at the configured step limit."""
    model.train()
    totals = {}
    batches = 0
    for batch in loader:
        if 0 < state["max_steps"] <= state["step"]:
            break
        waveforms = batch["waveforms"].to(state["device"], non_blocking=True)
        lengths = batch["lengths"].to(state["device"], non_blocking=True)
        species = batch["species"].to(state["device"], non_blocking=True)
        domains = batch["domains"].to(state["device"], non_blocking=True)
        with torch.no_grad():
            features, frame_lengths = frontend(waveforms, lengths)
        optimizer.zero_grad(set_to_none=True)
        outputs = model(
            features,
            frame_lengths,
            domains,
            domain_strength=grl_strength(state["step"], state["total_steps"]),
            augment=True,
        )
        losses = training_loss(outputs, species, domains, priors, epoch + 1)
        losses["loss"].backward()
        nn.utils.clip_grad_norm_(model.parameters(), state["gradient_clip"])
        optimizer.step()
        scheduler.step()
        state["step"] += 1
        batches += 1
        for name, value in losses.items():
            totals[name] = totals.get(name, 0.0) + value.detach().item()
    if not batches:
        return None
    return {name: value / batches for name, value in totals.items()}


def checkpoint_state(model, frontend, optimizer, scheduler, epoch, state, metrics, args):
    """Build a resumable best-checkpoint payload."""
    return {
        "model": model.state_dict(),
        "frontend": frontend.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "epoch": epoch,
        "step": state["step"],
        "metrics": metrics,
        "arguments": {key: value for key, value in vars(args).items() if key != "run"},
    }


def final_metrics(model, frontend, training, validation, held_domains, args, device):
    """Score softmax and shared-covariance Mahalanobis inference."""
    reference_ids = select_ids(training.ids, args.reference_size, args.seed + 1)
    reference = deterministic_view(training, reference_ids)
    reference_outputs = collect_outputs(
        model,
        frontend,
        reference,
        args.eval_batch_size,
        args.workers,
        device,
        args.seed,
    )
    validation_outputs = collect_outputs(
        model,
        frontend,
        validation,
        args.eval_batch_size,
        args.workers,
        device,
        args.seed,
    )
    softmax = fold_prediction_metrics(
        validation_outputs["labels"],
        validation_outputs["logits"].argmax(1),
        validation_outputs["domains"],
        held_domains,
    )
    mahalanobis = fold_prediction_metrics(
        validation_outputs["labels"],
        mahalanobis_predict(
            *mahalanobis_fit(reference_outputs["embeddings"], reference_outputs["labels"]),
            validation_outputs["embeddings"],
        ),
        validation_outputs["domains"],
        held_domains,
    )
    return softmax, mahalanobis, reference_outputs


def train(args):
    """Train and select MTRCNN-DG on one training-only pseudo-unseen fold."""
    set_seed(args.seed)
    if not args.synthetic and not (Path(args.data) / ".complete").exists():
        prepare_data(args.data)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    training, validation, held_domains = make_datasets(args)
    mean, std = compute_normalization(
        training,
        args.stats_size,
        args.eval_batch_size,
        args.workers,
        device,
        args.seed,
        args.min_hz,
        args.lifter,
    )
    frontend = LogMelFrontend(mean, std, args.min_hz, args.lifter).to(device)
    model = MTRCNNDG(dropout=args.dropout).to(device)
    loader = make_loader(
        training, args.batch_size, args.workers, True, args.seed
    )
    natural_steps = args.epochs * len(loader)
    total_steps = min(natural_steps, args.max_steps) if args.max_steps > 0 else natural_steps
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay
    )
    scheduler = torch.optim.lr_scheduler.LambdaLR(
        optimizer, learning_rate_schedule(total_steps, args.warmup_ratio)
    )
    priors = sampled_priors(training, device)
    output = Path(args.output)
    output.mkdir(parents=True, exist_ok=True)
    checkpoint = output / "checkpoint.pt"
    state = {
        "device": device,
        "step": 0,
        "total_steps": total_steps,
        "max_steps": args.max_steps,
        "gradient_clip": args.gradient_clip,
    }
    history = []
    best_score = -1.0
    best_epoch = 0
    stale_epochs = 0
    for epoch in range(args.epochs):
        losses = train_epoch(
            model, frontend, loader, optimizer, scheduler, priors, epoch, state
        )
        if losses is None:
            break
        validation_outputs = collect_outputs(
            model,
            frontend,
            validation,
            args.eval_batch_size,
            args.workers,
            device,
            args.seed,
        )
        metrics = fold_prediction_metrics(
            validation_outputs["labels"],
            validation_outputs["logits"].argmax(1),
            validation_outputs["domains"],
            held_domains,
        )
        record = {"epoch": epoch + 1, "step": state["step"], "losses": losses, **metrics}
        history.append(record)
        print(json.dumps(record), flush=True)
        # Without a held-out domain (fold -1) the last epoch is kept.
        score = metrics["domain_holdout"]["balanced_accuracy"] if args.fold >= 0 else epoch
        if score > best_score:
            best_score = score
            best_epoch = epoch + 1
            stale_epochs = 0
            torch.save(
                checkpoint_state(
                    model,
                    frontend,
                    optimizer,
                    scheduler,
                    epoch + 1,
                    state,
                    metrics,
                    args,
                ),
                checkpoint,
            )
        else:
            stale_epochs += 1
        reached_limit = 0 < args.max_steps <= state["step"]
        if reached_limit or stale_epochs >= args.patience:
            break
    saved = torch.load(checkpoint, map_location=device, weights_only=True)
    model.load_state_dict(saved["model"])
    frontend.load_state_dict(saved["frontend"])
    softmax, mahalanobis, reference = final_metrics(
        model, frontend, training, validation, held_domains, args, device
    )
    result = {
        "device": str(device),
        "fold": args.fold,
        "seed": args.seed,
        "training_rows": len(training),
        "validation_rows": len(validation),
        "normalization_rows": len(select_ids(training.ids, args.stats_size, args.seed)),
        "reference_rows": len(reference["labels"]),
        "held_domains": {
            SPECIES[species]: DOMAINS[domain] if domain is not None else "same-domain"
            for species, domain in held_domains.items()
        },
        "best_epoch": best_epoch,
        "steps": state["step"],
        "softmax": softmax,
        "mahalanobis": mahalanobis,
        "history": history,
        "checkpoint": str(checkpoint),
        "test": test_metrics(model, frontend, reference, args, device) if args.test else None,
    }
    (output / "results.json").write_text(
        json.dumps(result, indent=1), encoding="utf-8"
    )
    print(json.dumps(result, indent=1))


def test_metrics(model, frontend, reference, args, device):
    """Score the official development test split with softmax and Mahalanobis inference."""
    test_ids = (
        [
            f"S_{species + 1}_D_{domain + 1}_{row + 1}"
            for species in range(len(SPECIES))
            for domain in range(len(DOMAINS))
            for row in range(args.synthetic_rows)
        ]
        if args.synthetic
        else load_ids("Test")
    )
    test = ClipDataset(
        args.data,
        select_ids(test_ids, args.test_size, args.seed),
        args.seconds,
        False,
        0.0,
        args.synthetic,
        args.level_rms,
    )
    outputs = collect_clip_outputs(
        model, frontend, test, args.eval_batch_size, args.workers, device, args.seed
    )
    unseen_domains = load_unseen_domains()
    mahalanobis = mahalanobis_predict(
        *mahalanobis_fit(reference["embeddings"], reference["labels"]), outputs["embeddings"]
    )
    return {
        "rows": len(test),
        "softmax": official_metrics(
            outputs["labels"],
            outputs["probabilities"].argmax(1),
            outputs["domains"],
            unseen_domains,
        ),
        "mahalanobis": official_metrics(
            outputs["labels"], mahalanobis, outputs["domains"], unseen_domains
        ),
    }


def evaluate(args):
    """Score a checkpoint on its validation rows and on the official development test split."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    saved = torch.load(args.checkpoint, map_location=device, weights_only=True)
    options = argparse.Namespace(**{"lifter": 0, **saved["arguments"]})
    options.data = args.data
    options.eval_batch_size = args.eval_batch_size
    options.workers = args.workers
    options.test_size = args.test_size
    set_seed(options.seed)
    if not options.synthetic and not (Path(args.data) / ".complete").exists():
        prepare_data(args.data)
    training, validation, held_domains = make_datasets(options)
    frontend = LogMelFrontend(min_hz=options.min_hz, lifter=options.lifter).to(device)
    frontend.load_state_dict(saved["frontend"])
    model = MTRCNNDG(dropout=options.dropout).to(device)
    model.load_state_dict(saved["model"])
    fold_softmax, fold_mahalanobis, reference = final_metrics(
        model, frontend, training, validation, held_domains, options, device
    )
    result = {
        "checkpoint": args.checkpoint,
        "epoch": saved["epoch"],
        "fold": options.fold,
        "seed": options.seed,
        "arguments": saved["arguments"],
        "test": test_metrics(model, frontend, reference, options, device),
        "fold_softmax": fold_softmax,
        "fold_mahalanobis": fold_mahalanobis,
    }
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(result, indent=1), encoding="utf-8")
    print(json.dumps(result, indent=1))
    if args.export:
        export(model, frontend, reference, saved["arguments"], result, Path(args.export))
        if args.repo:
            HfApi().upload_folder(
                repo_id=args.repo,
                folder_path=args.export,
                commit_message="Add MTRCNN-DG weights and evaluation",
            )


def export(model, frontend, reference, arguments, result, folder):
    """Write weights, Mahalanobis statistics, this script and the evaluation to folder."""
    means, precision = mahalanobis_fit(reference["embeddings"], reference["labels"])
    (folder / "results").mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "model": model.state_dict(),
            "frontend": frontend.state_dict(),
            "arguments": arguments,
            "means": means,
            "precision": precision,
        },
        folder / "model.pt",
    )
    shutil.copy(__file__, folder / "mtrcnn.py")
    (folder / "results" / "eval.json").write_text(json.dumps(result, indent=1), encoding="utf-8")


def predict(args):
    """Print the Mahalanobis species prediction for each WAV file."""
    path = Path(args.model)
    if not path.is_file():
        path = hf_hub_download(args.model, "model.pt")
    saved = torch.load(path, map_location="cpu", weights_only=True)
    options = saved["arguments"]
    frontend = LogMelFrontend(min_hz=options["min_hz"], lifter=options.get("lifter", 0))
    frontend.load_state_dict(saved["frontend"])
    model = MTRCNNDG().eval()
    model.load_state_dict(saved["model"])
    samples = round(options["seconds"] * SAMPLE_RATE)
    for file in args.files:
        windows = split_windows(load_waveform(file), samples)
        waveforms = torch.from_numpy(normalize_level(windows, options["level_rms"]))
        with torch.no_grad():
            features, lengths = frontend(waveforms, torch.full((len(waveforms),), samples))
            embedding = model(features, lengths, torch.zeros(len(waveforms), dtype=torch.long))[
                "embedding"
            ].mean(0, keepdim=True)
        species = mahalanobis_predict(saved["means"], saved["precision"], embedding).item()
        print(json.dumps({"file": file, "species": SPECIES[species]}))


def card(args):
    """Render card.jinja into card/README.md with the results stored in the model repo."""
    try:
        folder = Path(snapshot_download(args.repo, allow_patterns="results/*.json"))
        paths = folder.glob("results/*.json")
    except RepositoryNotFoundError:
        paths = []
    results = {path.stem: json.loads(path.read_text(encoding="utf-8")) for path in paths}
    runs = {}
    for name, result in sorted(results.items()):
        if name.startswith("train-"):
            runs.setdefault(name.split("-")[1], []).append(result["test"])
    seeds = {
        config: {
            inference: {
                partition: [run[inference][partition]["balanced_accuracy"] * 100 for run in tests]
                for partition in ("seen", "unseen")
            }
            for inference in ("softmax", "mahalanobis")
        }
        for config, tests in runs.items()
    }
    evaluation = results.get("eval")
    data = ModelCardData(
        model_name=args.repo.split("/")[1],
        datasets=[DATASET],
        license="cc-by-4.0",
        tags=["audio-classification", "bioacoustics", "mosquito", "domain-generalization"],
        eval_results=[
            EvalResult(
                task_type="audio-classification",
                dataset_type=DATASET,
                dataset_name="BioDCASE 2026 Task 5 development test, unseen domains",
                metric_type="balanced_accuracy",
                metric_name="Unseen-domain balanced accuracy",
                metric_value=round(
                    evaluation["test"]["mahalanobis"]["unseen"]["balanced_accuracy"] * 100, 2
                ),
            )
        ]
        if evaluation
        else None,
    )
    rendered = ModelCard.from_template(
        data,
        template_path=Path(__file__).parent / "card.jinja",
        repo=args.repo,
        dataset=DATASET,
        species=SPECIES,
        evaluation=evaluation,
        seeds=seeds,
        mean=statistics.mean,
        stdev=statistics.stdev,
    )
    output = Path(__file__).parent / "card"
    output.mkdir(exist_ok=True)
    rendered.save(output / "README.md")


def set_seed(seed):
    """Seed Python, NumPy and PyTorch."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def mel_filter_bank(min_hz):
    """Return a triangular Slaney-style mel filter bank from min_hz to Nyquist."""
    frequencies = torch.linspace(0, SAMPLE_RATE / 2, N_FFT // 2 + 1)
    mel_min = 2595 * math.log10(1 + min_hz / 700)
    mel_max = 2595 * math.log10(1 + (SAMPLE_RATE / 2) / 700)
    mel_points = torch.linspace(mel_min, mel_max, N_MELS + 2)
    hz_points = 700 * (torch.pow(10, mel_points / 2595) - 1)
    lower = (frequencies[:, None] - hz_points[:-2]) / (
        hz_points[1:-1] - hz_points[:-2]
    )
    upper = (hz_points[2:] - frequencies[:, None]) / (
        hz_points[2:] - hz_points[1:-1]
    )
    return torch.minimum(lower, upper).clamp_min(0)


def lifter_projection(lifter):
    """Return the matrix that removes each log-mel frame's lifter smoothest DCT components.

    A recording device colours the spectrum with a smooth gain curve, which lands in these
    components; the wingbeat harmonics vary faster across mel bins and survive.
    """
    bins = torch.arange(N_MELS)
    basis = torch.cos(math.pi * torch.arange(lifter)[:, None] * (bins + 0.5) / N_MELS)
    basis = basis / basis.norm(dim=1, keepdim=True)
    return torch.eye(N_MELS) - basis.T @ basis


class LogMelFrontend(nn.Module):
    """Fixed 8 kHz, 64-bin log-mel frontend with optional liftering and global normalization."""

    def __init__(self, mean=None, std=None, min_hz=MIN_HZ, lifter=0):
        super().__init__()
        self.register_buffer("window", torch.hann_window(N_FFT), persistent=False)
        self.register_buffer("mel_filters", mel_filter_bank(min_hz), persistent=False)
        self.register_buffer("projection", lifter_projection(lifter), persistent=False)
        self.register_buffer(
            "mean",
            torch.zeros(N_MELS) if mean is None else torch.as_tensor(mean).float(),
        )
        self.register_buffer(
            "std",
            torch.ones(N_MELS) if std is None else torch.as_tensor(std).float(),
        )

    def forward(self, waveforms, sample_lengths):
        """Convert padded waveforms to normalized log-mel frames."""
        spectrum = torch.stft(
            waveforms,
            n_fft=N_FFT,
            hop_length=HOP_LENGTH,
            win_length=N_FFT,
            window=self.window,
            center=True,
            pad_mode="reflect",
            return_complex=True,
        )
        power = spectrum.abs().square().transpose(1, 2)
        features = 10 * torch.log10((power @ self.mel_filters).clamp_min(1e-10))
        features = features @ self.projection
        features = (features - self.mean) / self.std.clamp_min(1e-6)
        frame_lengths = torch.div(sample_lengths, HOP_LENGTH, rounding_mode="floor") + 1
        return features, frame_lengths


def different_domain_partners(domains):
    """Choose a random different-domain partner for every available sample."""
    partners = []
    for index, domain in enumerate(domains):
        candidates = torch.nonzero(domains != domain, as_tuple=False).flatten()
        if candidates.numel():
            choice = torch.randint(candidates.numel(), (), device=domains.device)
            partners.append(candidates[choice])
        else:
            partners.append(torch.tensor(index, device=domains.device))
    return torch.stack(partners)


class FourierMix(nn.Module):
    """Mix log-mel Fourier amplitudes between recording domains."""

    def __init__(self, probability=0.5):
        super().__init__()
        self.probability = probability

    def forward(self, features, domains):
        """Optionally replace each sample's amplitude style across domains."""
        if not self.training or torch.rand((), device=features.device) >= self.probability:
            return features
        partners = different_domain_partners(domains)
        spectrum = torch.fft.rfft2(features, dim=(-2, -1), norm="ortho")
        amplitude = spectrum.abs()
        phase = spectrum / amplitude.clamp_min(1e-8)
        weight = torch.rand(features.size(0), 1, 1, device=features.device)
        mixed_amplitude = (1 - weight) * amplitude + weight * amplitude[partners]
        return torch.fft.irfft2(
            mixed_amplitude * phase,
            s=features.shape[-2:],
            dim=(-2, -1),
            norm="ortho",
        )


class MixStyle(nn.Module):
    """Mix channel statistics between domains after a convolutional stage."""

    def __init__(self, probability=0.5, alpha=0.1):
        super().__init__()
        self.probability = probability
        self.beta = torch.distributions.Beta(alpha, alpha)

    def forward(self, features, domains):
        """Optionally mix feature statistics across recording domains."""
        if not self.training or torch.rand((), device=features.device) >= self.probability:
            return features
        mean = features.mean(dim=(2, 3), keepdim=True).detach()
        std = (features.var(dim=(2, 3), keepdim=True, unbiased=False) + 1e-6).sqrt().detach()
        normalized = (features - mean) / std
        partners = different_domain_partners(domains)
        weight = self.beta.sample((features.size(0), 1, 1, 1)).to(features.device)
        mixed_mean = weight * mean + (1 - weight) * mean[partners]
        mixed_std = weight * std + (1 - weight) * std[partners]
        return normalized * mixed_std + mixed_mean


class ConvStage(nn.Module):
    """One convolution, normalization, pooling and dropout stage."""

    def __init__(self, in_channels, out_channels, kernel, dilation, padding, dropout):
        super().__init__()
        self.kernel = kernel
        self.dilation = dilation
        self.padding = padding
        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel,
            dilation=dilation,
            padding=padding,
            bias=False,
        )
        self.batch_norm = nn.BatchNorm2d(out_channels)
        self.pool = nn.AvgPool2d(2)
        self.dropout = nn.Dropout2d(dropout)

    def forward(self, features):
        """Transform and downsample one feature map."""
        features = functional.relu(self.batch_norm(self.conv(features)), inplace=True)
        return self.dropout(self.pool(features))

    def output_lengths(self, lengths):
        """Propagate valid time lengths through convolution and pooling."""
        convolved = lengths + 2 * self.padding[0] - self.dilation[0] * (
            self.kernel[0] - 1
        )
        return torch.div(convolved.clamp_min(0), 2, rounding_mode="floor")


def masked_mean_max(features, lengths):
    """Pool valid time frames by adding their mean and maximum."""
    lengths = lengths.clamp(min=1, max=features.size(2))
    positions = torch.arange(features.size(2), device=features.device).view(1, 1, -1, 1)
    mask = positions < lengths.view(-1, 1, 1, 1)
    mean = (features * mask).sum(2) / lengths.view(-1, 1, 1)
    maximum = features.masked_fill(~mask, float("-inf")).max(2).values
    return mean + torch.where(torch.isfinite(maximum), maximum, torch.zeros_like(maximum))


class MTRCNNBranch(nn.Module):
    """One temporal-resolution branch of MTRCNN."""

    def __init__(self, stage_specs, dropout):
        super().__init__()
        channels = [1, 16, 32, 64]
        self.stages = nn.ModuleList(
            ConvStage(channels[index], channels[index + 1], *specification, dropout)
            for index, specification in enumerate(stage_specs)
        )
        self.mix_style = MixStyle()
        self.frequency_projection = nn.Linear(self._frequency_bins(), 1)

    def _frequency_bins(self):
        with torch.no_grad():
            features = torch.zeros(1, 1, 128, N_MELS)
            for stage in self.stages:
                features = stage(features)
        return features.shape[-1]

    def forward(self, features, lengths, domains, augment):
        """Encode one temporal-resolution branch."""
        for index, stage in enumerate(self.stages):
            features = stage(features)
            lengths = stage.output_lengths(lengths)
            if index == 0 and augment:
                features = self.mix_style(features, domains)
        pooled = masked_mean_max(features, lengths)
        return functional.relu(self.frequency_projection(pooled).squeeze(-1), inplace=True)


class GradientReversal(torch.autograd.Function):
    """Identity forward pass with a sign-reversed backward pass."""

    @staticmethod
    def forward(ctx, features, strength):
        """Return features unchanged and retain the reversal strength."""
        ctx.strength = strength
        return features.view_as(features)

    @staticmethod
    def backward(ctx, gradient):
        """Reverse and scale the upstream embedding gradient."""
        return -ctx.strength * gradient, None


class MTRCNNDG(nn.Module):
    """MTRCNN with spectral style, adversarial, and conditional domain heads."""

    def __init__(self, dropout=0.2):
        super().__init__()
        self.fourier_mix = FourierMix()
        self.input_batch_norm = nn.BatchNorm2d(N_MELS)
        specifications = [
            [((3, 3), (1, 1), (1, 1)), ((3, 3), (2, 1), (2, 0)), ((3, 3), (3, 1), (3, 0))],
            [((5, 5), (1, 1), (2, 2)), ((5, 5), (2, 1), (4, 1)), ((5, 5), (3, 1), (6, 1))],
            [((7, 7), (1, 1), (3, 3)), ((7, 7), (2, 1), (6, 2)), ((7, 7), (3, 1), (9, 2))],
        ]
        self.branches = nn.ModuleList(MTRCNNBranch(specs, dropout) for specs in specifications)
        self.embedding = nn.Linear(64 * 3, 32)
        self.species_classifier = nn.Linear(32, len(SPECIES))
        self.domain_classifier = nn.Linear(32, len(DOMAINS))
        self.aux_domain_classifier = nn.Linear(32, len(DOMAINS))

    def forward(self, features, lengths, domains, domain_strength=0.0, augment=False):
        """Return species, domain, auxiliary, and embedding outputs."""
        if augment:
            features = self.fourier_mix(features, domains)
        features = features.unsqueeze(1).transpose(1, 3)
        features = self.input_batch_norm(features).transpose(1, 3)
        branch_outputs = [
            branch(features, lengths, domains, augment) for branch in self.branches
        ]
        embedding = functional.gelu(self.embedding(torch.cat(branch_outputs, dim=1)))
        reversed_embedding = GradientReversal.apply(embedding, domain_strength)
        auxiliary = self.aux_domain_classifier
        return {
            "embedding": embedding,
            "species_logits": self.species_classifier(embedding),
            "domain_logits": self.domain_classifier(reversed_embedding),
            "aux_domain_logits": auxiliary(embedding),
            "frozen_aux_domain_logits": functional.linear(
                embedding,
                auxiliary.weight.detach(),
                auxiliary.bias.detach(),
            ),
        }


def supervised_contrastive(embedding, labels, temperature=0.01):
    """Pull embeddings of the same species together."""
    normalized = functional.normalize(embedding, dim=1)
    similarity = normalized @ normalized.T / temperature
    identity = torch.eye(len(labels), dtype=torch.bool, device=labels.device)
    positives = labels[:, None].eq(labels[None, :]) & ~identity
    similarity = similarity.masked_fill(identity, float("-inf"))
    log_probability = similarity - torch.logsumexp(similarity, dim=1, keepdim=True)
    log_probability = torch.where(positives, log_probability, torch.zeros_like(log_probability))
    counts = positives.sum(1)
    valid = counts > 0
    if not valid.any():
        return embedding.sum() * 0
    return -(log_probability.sum(1)[valid] / counts[valid]).mean()


def conditional_domain_losses(logits, species_labels):
    """Return class-conditional domain entropy and balance losses."""
    probabilities = functional.softmax(logits, dim=1).clamp_min(1e-8)
    entropy_losses = []
    balance_losses = []
    for species in species_labels.unique():
        class_probabilities = probabilities[species_labels == species]
        entropy_losses.append((class_probabilities * class_probabilities.log()).sum(1).mean())
        mean_probability = class_probabilities.mean(0)
        balance_losses.append(
            (mean_probability * (mean_probability.log() + math.log(len(DOMAINS)))).sum()
        )
    return torch.stack(entropy_losses).mean(), torch.stack(balance_losses).mean()


def training_loss(outputs, species_labels, domain_labels, priors, epoch):
    """Compute the full MTRCNN-DG objective and its components."""
    adjusted_logits = outputs["species_logits"] + priors.clamp_min(1e-8).log()
    species_loss = functional.cross_entropy(adjusted_logits, species_labels)
    domain_loss = functional.cross_entropy(outputs["domain_logits"], domain_labels)
    auxiliary_loss = functional.cross_entropy(outputs["aux_domain_logits"], domain_labels)
    contrastive_loss = supervised_contrastive(outputs["embedding"], species_labels)
    ccde_loss, cdb_loss = conditional_domain_losses(
        outputs["frozen_aux_domain_logits"], species_labels
    )
    conditional_weight = 0.05 if epoch >= 5 else 0.0
    total = (
        species_loss
        + domain_loss
        + auxiliary_loss
        + 0.5 * contrastive_loss
        + conditional_weight * (ccde_loss + cdb_loss)
    )
    return {
        "loss": total,
        "species": species_loss,
        "domain": domain_loss,
        "aux_domain": auxiliary_loss,
        "contrastive": contrastive_loss,
        "ccde": ccde_loss,
        "cdb": cdb_loss,
    }


def smoke(args):
    """Exercise frontend, augmentations, all heads, losses, and backward pass."""
    set_seed(args.seed)
    multi_domain_species = {0, 1, 2, 6, 7, 8}
    fold_ids = [
        f"S_{species + 1}_D_{domain + 1}_{row + 1}"
        for species in range(len(SPECIES))
        for domain in range(2 if species in multi_domain_species else 1)
        for row in range(2)
    ]
    training_ids, validation_ids, held_domains = held_domain_split(
        fold_ids, 0, seed=args.seed
    )
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dataset = MosquitoDataset(
        None,
        synthetic_ids(4),
        args.seconds,
        training=True,
        balance_power=0.5,
        synthetic=True,
    )
    sampler = WeightedRandomSampler(
        dataset.sample_weights,
        num_samples=args.batch_size,
        replacement=True,
        generator=torch.Generator().manual_seed(args.seed),
    )
    batch = next(
        iter(
            DataLoader(
                dataset,
                batch_size=args.batch_size,
                sampler=sampler,
                collate_fn=AudioCollator(),
            )
        )
    )
    waveforms = batch["waveforms"].to(device)
    lengths = batch["lengths"].to(device)
    labels = batch["species"].to(device)
    domains = batch["domains"].to(device)
    mean, std = compute_normalization(
        dataset,
        size=min(20, len(dataset)),
        batch_size=args.batch_size,
        workers=0,
        device=device,
        seed=args.seed,
        min_hz=MIN_HZ,
    )
    frontend = LogMelFrontend(mean, std).to(device)
    model = MTRCNNDG().to(device).train()
    features, frame_lengths = frontend(waveforms, lengths)
    outputs = model(features, frame_lengths, domains, domain_strength=0.3, augment=True)
    priors = torch.full((len(SPECIES),), 1 / len(SPECIES), device=device)
    losses = training_loss(outputs, labels, domains, priors, epoch=5)
    losses["loss"].backward()
    assert outputs["species_logits"].shape == (args.batch_size, len(SPECIES))
    assert outputs["domain_logits"].shape == (args.batch_size, len(DOMAINS))
    assert all(torch.isfinite(value) for value in losses.values())
    assert all(
        parameter.grad is None or torch.isfinite(parameter.grad).all()
        for parameter in model.parameters()
    )
    result = {
        "device": str(device),
        "features": list(features.shape),
        "embedding": list(outputs["embedding"].shape),
        "parameters": sum(parameter.numel() for parameter in model.parameters()),
        "normalization": {
            "mean": [round(mean.min().item(), 4), round(mean.max().item(), 4)],
            "std": [round(std.min().item(), 4), round(std.max().item(), 4)],
        },
        "fold": {
            "train": len(training_ids),
            "validation": len(validation_ids),
            "held_domains": {
                SPECIES[species]: DOMAINS[domain] if domain is not None else "same-domain"
                for species, domain in held_domains.items()
            },
        },
        "losses": {name: round(value.item(), 6) for name, value in losses.items()},
    }
    print(json.dumps(result, indent=1))


def add_train_parser(commands):
    """Register the train command and its options."""
    train_parser = commands.add_parser("train")
    train_parser.add_argument("--data", default="/tmp/mosquito-audio")
    train_parser.add_argument("--output", default="/tmp/mosquito-mtrcnn")
    train_parser.add_argument("--fold", type=int, default=0)
    train_parser.add_argument("--epochs", type=int, default=20)
    train_parser.add_argument("--max-steps", type=int, default=-1)
    train_parser.add_argument("--train-size", type=int, default=None)
    train_parser.add_argument("--validation-size", type=int, default=None)
    train_parser.add_argument("--stats-size", type=int, default=10_000)
    train_parser.add_argument("--reference-size", type=int, default=20_000)
    train_parser.add_argument("--batch-size", type=int, default=128)
    train_parser.add_argument("--eval-batch-size", type=int, default=256)
    train_parser.add_argument("--learning-rate", type=float, default=1e-3)
    train_parser.add_argument("--weight-decay", type=float, default=1e-4)
    train_parser.add_argument("--warmup-ratio", type=float, default=0.05)
    train_parser.add_argument("--gradient-clip", type=float, default=5.0)
    train_parser.add_argument("--dropout", type=float, default=0.2)
    train_parser.add_argument("--balance-power", type=float, default=1.0)
    train_parser.add_argument("--seconds", type=float, default=SECONDS)
    train_parser.add_argument("--min-hz", type=float, default=MIN_HZ)
    train_parser.add_argument("--level-rms", type=float, default=LEVEL_RMS)
    train_parser.add_argument("--lifter", type=int, default=0)
    train_parser.add_argument("--patience", type=int, default=4)
    train_parser.add_argument("--workers", type=int, default=max(0, cpus() - 1))
    train_parser.add_argument("--seed", type=int, default=42)
    train_parser.add_argument("--synthetic", action="store_true")
    train_parser.add_argument("--synthetic-rows", type=int, default=4)
    train_parser.add_argument("--test", action="store_true")
    train_parser.add_argument("--test-size", type=int, default=None)
    train_parser.set_defaults(run=train)


def main():
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    smoke_parser = commands.add_parser("smoke")
    smoke_parser.add_argument("--batch-size", type=int, default=10)
    smoke_parser.add_argument("--seconds", type=float, default=SECONDS)
    smoke_parser.add_argument("--seed", type=int, default=42)
    smoke_parser.set_defaults(run=smoke)
    add_train_parser(commands)
    evaluate_parser = commands.add_parser("evaluate")
    evaluate_parser.add_argument("--checkpoint", required=True)
    evaluate_parser.add_argument("--data", default="/tmp/mosquito-audio")
    evaluate_parser.add_argument("--output", default="/tmp/mosquito-mtrcnn-eval.json")
    evaluate_parser.add_argument("--test-size", type=int, default=None)
    evaluate_parser.add_argument("--eval-batch-size", type=int, default=128)
    evaluate_parser.add_argument("--workers", type=int, default=max(0, cpus() - 1))
    evaluate_parser.add_argument("--export", default=None)
    evaluate_parser.add_argument("--repo", default=None)
    evaluate_parser.set_defaults(run=evaluate)
    predict_parser = commands.add_parser("predict")
    predict_parser.add_argument("--model", default=REPO)
    predict_parser.add_argument("files", nargs="+")
    predict_parser.set_defaults(run=predict)
    card_parser = commands.add_parser("card")
    card_parser.add_argument("--repo", default=REPO)
    card_parser.set_defaults(run=card)
    args = parser.parse_args()
    args.run(args)


if __name__ == "__main__":
    main()