source/src/bgc_retrieval/training.py

"""One-GPU Phase 1 and Phase 2 training with immutable artifacts."""

from __future__ import annotations

import copy
import random
from collections.abc import Iterable
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader

from .artifacts import write_json_immutable
from .checkpoints import load_checkpoint, save_checkpoint
from .data import BGCEmbeddingDataset, collate_bgcs
from .losses import augment_gene_sets, masked_gene_loss, supervised_contrastive_loss
from .metrics import expected_tie_aware_metrics
from .model import LeakageFreeBGCSetNet, MaskedGenePredictionHead, ModelConfig, WeightedPfamJaccard
from .sampling import UniqueGroupBatchSampler


def set_reproducible(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


def choose_device() -> torch.device:
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def _move_batch(batch: dict[str, object], device: torch.device) -> tuple[torch.Tensor, ...]:
    return (
        batch["gene_embeddings"].to(device),
        batch["relative_positions"].to(device),
        batch["padding_mask"].to(device),
        batch["pfam_tokens"].to(device),
    )


def _encode_genes(
    model: nn.Module,
    embeddings: torch.Tensor,
    positions: torch.Tensor,
    padding_mask: torch.Tensor,
    pfam_tokens: torch.Tensor,
) -> torch.Tensor:
    if getattr(model, "uses_pfam", False):
        return model.encode_genes(embeddings, positions, padding_mask, pfam_tokens)
    return model.encode_genes(embeddings, positions, padding_mask)


def _encode_bgc(
    model: nn.Module,
    embeddings: torch.Tensor,
    positions: torch.Tensor,
    padding_mask: torch.Tensor,
    pfam_tokens: torch.Tensor,
) -> torch.Tensor:
    if getattr(model, "uses_pfam", False):
        return model(embeddings, positions, padding_mask, pfam_tokens)
    return model(embeddings, positions, padding_mask)


def _mask_gene_batch(
    embeddings: torch.Tensor,
    padding_mask: torch.Tensor,
    probability: float,
    seed: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    generator = torch.Generator(device="cpu").manual_seed(seed)
    random_values = torch.rand(padding_mask.shape, generator=generator)
    masked = (random_values < probability) & ~padding_mask.cpu()
    for row in range(masked.shape[0]):
        if not masked[row].any():
            valid = torch.nonzero(~padding_mask[row].cpu(), as_tuple=False).flatten()
            masked[row, valid[torch.randint(len(valid), (1,), generator=generator)]] = True
    masked = masked.to(embeddings.device)
    model_input = embeddings.clone()
    model_input[masked] = 0.0
    return model_input, masked


def _phase1_epoch(
    model: LeakageFreeBGCSetNet,
    head: MaskedGenePredictionHead,
    loader: DataLoader,
    device: torch.device,
    mask_probability: float,
    seed: int,
    optimizer: torch.optim.Optimizer | None,
    mixed_precision: bool,
) -> float:
    training = optimizer is not None
    model.train(training)
    head.train(training)
    losses: list[float] = []
    scaler = torch.amp.GradScaler("cuda", enabled=mixed_precision and device.type == "cuda")
    context = torch.enable_grad() if training else torch.no_grad()
    with context:
        for batch_index, batch in enumerate(loader):
            embeddings, positions, padding_mask, pfam_tokens = _move_batch(batch, device)
            model_input, masked = _mask_gene_batch(
                embeddings, padding_mask, mask_probability, seed + batch_index
            )
            if optimizer:
                optimizer.zero_grad(set_to_none=True)
            with torch.amp.autocast(device_type=device.type, enabled=mixed_precision and device.type == "cuda"):
                contextual = _encode_genes(
                    model, model_input, positions, padding_mask, pfam_tokens
                )
                prediction = head(contextual)
                loss = masked_gene_loss(prediction, embeddings, masked)
            if not torch.isfinite(loss):
                raise FloatingPointError("Non-finite Phase 1 loss")
            if optimizer:
                scaler.scale(loss).backward()
                scaler.unscale_(optimizer)
                nn.utils.clip_grad_norm_(list(model.parameters()) + list(head.parameters()), 1.0)
                scaler.step(optimizer)
                scaler.update()
            losses.append(float(loss.detach().cpu()))
    return float(np.mean(losses))


@torch.no_grad()
def encode_dataset(
    model: LeakageFreeBGCSetNet,
    dataset: BGCEmbeddingDataset,
    device: torch.device,
    num_workers: int = 0,
) -> dict[str, torch.Tensor]:
    loader = DataLoader(
        dataset, batch_size=32, shuffle=False, num_workers=num_workers, collate_fn=collate_bgcs
    )
    model.eval()
    result: dict[str, torch.Tensor] = {}
    for batch in loader:
        embeddings, positions, padding_mask, pfam_tokens = _move_batch(batch, device)
        encoded = _encode_bgc(model, embeddings, positions, padding_mask, pfam_tokens).cpu()
        result.update(zip(batch["bgc_ids"], encoded))
    return result


def embedding_validation_metric(
    embeddings: dict[str, torch.Tensor], assignments: pd.DataFrame, cutoff: int = 50
) -> float:
    group_by_bgc = assignments.set_index("bgc_id")["group_id"].astype(str).to_dict()
    identifiers = sorted(embeddings)
    values: list[float] = []
    for reference in identifiers:
        relevant = {
            item for item in identifiers if item != reference and group_by_bgc[item] == group_by_bgc[reference]
        }
        if not relevant:
            continue
        scores = {
            candidate: float(F.cosine_similarity(embeddings[reference], embeddings[candidate], dim=0))
            for candidate in identifiers
            if candidate != reference
        }
        values.append(expected_tie_aware_metrics(scores, relevant, recall_at=(cutoff,), ndcg_at=())[f"recall@{cutoff}"])
    if not values:
        raise ValueError("Validation split has no positive retrieval queries")
    return float(np.mean(values))


def train_phase1(
    model: LeakageFreeBGCSetNet,
    model_config: ModelConfig,
    train_dataset: BGCEmbeddingDataset,
    validation_dataset: BGCEmbeddingDataset,
    split_path: str | Path,
    input_paths: list[str | Path],
    output_dir: str | Path,
    training_config: dict[str, Any],
    seed: int,
) -> Path:
    set_reproducible(seed)
    device = choose_device()
    model.to(device)
    head = MaskedGenePredictionHead(model_config).to(device)
    optimizer = torch.optim.AdamW(
        list(model.parameters()) + list(head.parameters()),
        lr=float(training_config["learning_rate"]),
        weight_decay=float(training_config["weight_decay"]),
    )
    generator = torch.Generator().manual_seed(seed)
    train_loader = DataLoader(
        train_dataset,
        batch_size=int(training_config["batch_groups"]),
        shuffle=True,
        generator=generator,
        num_workers=int(training_config["num_workers"]),
        collate_fn=collate_bgcs,
    )
    validation_loader = DataLoader(
        validation_dataset,
        batch_size=int(training_config["batch_groups"]),
        shuffle=False,
        num_workers=int(training_config["num_workers"]),
        collate_fn=collate_bgcs,
    )
    history: list[dict[str, float | int]] = []
    best_loss = float("inf")
    best_state: dict[str, Any] | None = None
    patience = 0
    for epoch in range(int(training_config["phase1_epochs"])):
        train_loss = _phase1_epoch(
            model, head, train_loader, device, float(training_config["mask_probability"]),
            seed + epoch * 10000, optimizer, bool(training_config["mixed_precision"]),
        )
        validation_loss = _phase1_epoch(
            model, head, validation_loader, device, float(training_config["mask_probability"]),
            seed + 900000, None, bool(training_config["mixed_precision"]),
        )
        history.append({"epoch": epoch, "train_loss": train_loss, "validation_loss": validation_loss})
        if validation_loss < best_loss:
            best_loss = validation_loss
            best_state = {
                "model": copy.deepcopy(model.state_dict()),
                "head": copy.deepcopy(head.state_dict()),
                "epoch": epoch,
            }
            patience = 0
        else:
            patience += 1
            if patience >= int(training_config["patience"]):
                break
    if best_state is None:
        raise RuntimeError("Phase 1 did not produce a valid checkpoint")
    model.load_state_dict(best_state["model"])
    output = Path(output_dir)
    checkpoint_path = output / "phase1_best.pt"
    save_checkpoint(
        checkpoint_path, model, model_config, split_path, input_paths,
        {"stage": "phase1", "best_epoch": best_state["epoch"], "best_loss": best_loss,
         "head_state": best_state["head"]}, optimizer,
    )
    write_json_immutable(output / "phase1_history.json", history)
    return checkpoint_path


def weighted_jaccard_contrastive_loss(
    similarities: torch.Tensor,
    group_ids: list[str],
    temperature: float,
) -> torch.Tensor:
    identity = torch.eye(len(group_ids), dtype=torch.bool, device=similarities.device)
    positives = torch.tensor(
        [[left == right for right in group_ids] for left in group_ids],
        dtype=torch.bool,
        device=similarities.device,
    ) & ~identity
    positive_counts = positives.sum(dim=1)
    if torch.any(positive_counts == 0):
        raise ValueError("Every weighted-Pfam batch item must have a positive")
    logits = similarities / temperature
    logits = logits.masked_fill(identity, float("-inf"))
    log_probabilities = logits - torch.logsumexp(logits, dim=1, keepdim=True)
    positive_log_probability = log_probabilities.masked_fill(~positives, 0.0).sum(dim=1)
    positive_log_probability = positive_log_probability / positive_counts
    return -positive_log_probability.mean()


@torch.no_grad()
def weighted_jaccard_validation_metric(
    model: WeightedPfamJaccard,
    dataset: BGCEmbeddingDataset,
    assignments: pd.DataFrame,
    cutoff: int = 50,
) -> float:
    device = next(model.parameters()).device
    vocabulary = model.raw_weights.shape[0]
    presence = torch.zeros(len(dataset), vocabulary, device=device)
    for index in range(len(dataset)):
        tokens = dataset[index]["pfam_tokens"].to(device)
        if len(tokens):
            presence[index, tokens] = 1.0
    presence[:, 0] = 0.0
    weighted = presence * model.domain_weights().to(device)
    totals = weighted.sum(dim=1)
    intersection = weighted @ presence.T
    union = totals[:, None] + totals[None, :] - intersection
    similarities = (intersection / union.clamp_min(1e-8)).cpu()
    group_by_bgc = assignments.set_index("bgc_id")["group_id"].astype(str).to_dict()
    values: list[float] = []
    for index, identifier in enumerate(dataset.bgc_ids):
        relevant = {
            other_index
            for other_index, other_id in enumerate(dataset.bgc_ids)
            if other_index != index and group_by_bgc[other_id] == group_by_bgc[identifier]
        }
        if not relevant:
            continue
        scores = {
            other_id: float(similarities[index, other_index])
            for other_index, other_id in enumerate(dataset.bgc_ids)
            if other_index != index
        }
        values.append(
            expected_tie_aware_metrics(
                scores, {dataset.bgc_ids[item] for item in relevant}, recall_at=(cutoff,), ndcg_at=()
            )[f"recall@{cutoff}"]
        )
    if not values:
        raise ValueError("Validation split has no positive weighted-Pfam queries")
    return float(np.mean(values))


def train_weighted_pfam(
    model: WeightedPfamJaccard,
    model_config: ModelConfig,
    train_dataset: BGCEmbeddingDataset,
    validation_dataset: BGCEmbeddingDataset,
    validation_assignments: pd.DataFrame,
    split_path: str | Path,
    input_paths: list[str | Path],
    output_dir: str | Path,
    training_config: dict[str, Any],
    seed: int,
) -> Path:
    set_reproducible(seed)
    device = choose_device()
    model.to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=float(training_config["learning_rate"]),
        weight_decay=float(training_config["weight_decay"]),
    )
    sampler = UniqueGroupBatchSampler(
        [train_dataset.group_by_bgc[item] for item in train_dataset.bgc_ids],
        groups_per_batch=int(training_config["batch_groups"]),
        examples_per_group=int(training_config["examples_per_group"]),
        seed=seed,
    )
    loader = DataLoader(
        train_dataset,
        batch_sampler=sampler,
        num_workers=int(training_config["num_workers"]),
        collate_fn=collate_bgcs,
    )
    best_metric = -float("inf")
    best_state: dict[str, Any] | None = None
    history: list[dict[str, float | int]] = []
    patience = 0
    regularization = float(training_config.get("pfam_weight_regularization", 0.001))
    for epoch in range(int(training_config["phase2_epochs"])):
        sampler.set_epoch(epoch)
        model.train()
        epoch_losses: list[float] = []
        for batch in loader:
            tokens = batch["pfam_tokens"].to(device)
            optimizer.zero_grad(set_to_none=True)
            similarities = model.pairwise_jaccard(tokens)
            loss = weighted_jaccard_contrastive_loss(
                similarities, batch["group_ids"], float(training_config["temperature"])
            )
            log_weights = torch.log(model.domain_weights()[1:].clamp_min(1e-8))
            loss = loss + regularization * torch.mean(log_weights.square())
            if not torch.isfinite(loss):
                raise FloatingPointError("Non-finite weighted-Pfam loss")
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            epoch_losses.append(float(loss.detach().cpu()))
        validation_recall = weighted_jaccard_validation_metric(
            model, validation_dataset, validation_assignments, cutoff=50
        )
        history.append(
            {"epoch": epoch, "train_loss": float(np.mean(epoch_losses)),
             "validation_recall@50": validation_recall}
        )
        if validation_recall > best_metric:
            best_metric = validation_recall
            best_state = {"model": copy.deepcopy(model.state_dict()), "epoch": epoch}
            patience = 0
        else:
            patience += 1
            if patience >= int(training_config["patience"]):
                break
    if best_state is None:
        raise RuntimeError("Weighted Pfam training did not produce a valid checkpoint")
    model.load_state_dict(best_state["model"])
    output = Path(output_dir)
    checkpoint_path = output / "phase2_best.pt"
    save_checkpoint(
        checkpoint_path, model, model_config, split_path, input_paths,
        {"stage": "weighted_pfam", "best_epoch": best_state["epoch"],
         "best_validation_recall@50": best_metric}, optimizer,
    )
    write_json_immutable(output / "phase2_history.json", history)
    return checkpoint_path


def train_phase2(
    model: LeakageFreeBGCSetNet,
    model_config: ModelConfig,
    train_dataset: BGCEmbeddingDataset,
    validation_dataset: BGCEmbeddingDataset,
    validation_assignments: pd.DataFrame,
    split_path: str | Path,
    input_paths: list[str | Path],
    output_dir: str | Path,
    training_config: dict[str, Any],
    seed: int,
    phase1_checkpoint: str | Path | None = None,
) -> Path:
    set_reproducible(seed)
    device = choose_device()
    if phase1_checkpoint:
        load_checkpoint(phase1_checkpoint, model, split_path)
    model.to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=float(training_config["learning_rate"]),
        weight_decay=float(training_config["weight_decay"]),
    )
    sampler = UniqueGroupBatchSampler(
        [train_dataset.group_by_bgc[item] for item in train_dataset.bgc_ids],
        groups_per_batch=int(training_config["batch_groups"]),
        examples_per_group=int(training_config["examples_per_group"]), seed=seed,
    )
    loader = DataLoader(
        train_dataset, batch_sampler=sampler, num_workers=int(training_config["num_workers"]),
        collate_fn=collate_bgcs,
    )
    scaler = torch.amp.GradScaler(
        "cuda", enabled=bool(training_config["mixed_precision"]) and device.type == "cuda"
    )
    best_metric = -float("inf")
    best_state: dict[str, Any] | None = None
    history: list[dict[str, float | int]] = []
    patience = 0
    for epoch in range(int(training_config["phase2_epochs"])):
        sampler.set_epoch(epoch)
        model.train()
        epoch_losses: list[float] = []
        for batch in loader:
            embeddings, positions, padding_mask, pfam_tokens = _move_batch(batch, device)
            embeddings, positions, padding_mask = augment_gene_sets(
                embeddings, positions, padding_mask,
                float(training_config["gene_dropout"]), float(training_config["position_jitter"]),
            )
            optimizer.zero_grad(set_to_none=True)
            with torch.amp.autocast(
                device_type=device.type,
                enabled=bool(training_config["mixed_precision"]) and device.type == "cuda",
            ):
                encoded = _encode_bgc(
                    model, embeddings, positions, padding_mask, pfam_tokens
                )
                loss = supervised_contrastive_loss(
                    encoded, batch["group_ids"], float(training_config["temperature"])
                )
            if not torch.isfinite(loss):
                raise FloatingPointError("Non-finite Phase 2 loss")
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(optimizer)
            scaler.update()
            epoch_losses.append(float(loss.detach().cpu()))
        validation_embeddings = encode_dataset(
            model, validation_dataset, device, int(training_config["num_workers"])
        )
        validation_recall = embedding_validation_metric(
            validation_embeddings, validation_assignments, cutoff=50
        )
        history.append(
            {"epoch": epoch, "train_loss": float(np.mean(epoch_losses)),
             "validation_recall@50": validation_recall}
        )
        if validation_recall > best_metric:
            best_metric = validation_recall
            best_state = {"model": copy.deepcopy(model.state_dict()), "epoch": epoch}
            patience = 0
        else:
            patience += 1
            if patience >= int(training_config["patience"]):
                break
    if best_state is None:
        raise RuntimeError("Phase 2 did not produce a valid checkpoint")
    model.load_state_dict(best_state["model"])
    output = Path(output_dir)
    checkpoint_path = output / "phase2_best.pt"
    save_checkpoint(
        checkpoint_path, model, model_config, split_path, input_paths,
        {"stage": "phase2", "best_epoch": best_state["epoch"],
         "best_validation_recall@50": best_metric, "phase1_checkpoint": str(phase1_checkpoint)},
        optimizer,
    )
    write_json_immutable(output / "phase2_history.json", history)
    return checkpoint_path