source/src/bgc_retrieval/residual.py

"""Baseline-preserving gene weighting for species-scoped BGC retrieval."""

from __future__ import annotations

import copy
from collections.abc import Mapping, Sequence
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 .baselines import cosine_scores, ensemble_scores, pfam_jaccard_scores
from .checkpoints import save_checkpoint
from .data import BGCEmbeddingDataset, collate_bgcs
from .evaluation import evaluate_retrieval
from .losses import augment_gene_sets
from .model import ModelConfig
from .sampling import UniqueGroupBatchSampler
from .training import choose_device, embedding_validation_metric, set_reproducible


class ResidualGeneWeightingEncoder(nn.Module):
    """Reweight genes in frozen ESM space while preserving mean pooling."""

    input_names = ("gene_embeddings", "relative_positions", "padding_mask")

    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.config = config
        dimension = config.hidden_dimension
        self.embedding_norm = nn.LayerNorm(config.esm_dimension)
        self.gene_projection = nn.Linear(config.esm_dimension, dimension)
        self.position_projection = nn.Sequential(
            nn.Linear(1, dimension),
            nn.GELU(),
            nn.Linear(dimension, dimension),
        )
        self.gate = nn.Sequential(
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(dimension, 1),
        )
        nn.init.zeros_(self.gate[-1].weight)
        nn.init.zeros_(self.gate[-1].bias)

    def components(
        self,
        gene_embeddings: torch.Tensor,
        relative_positions: torch.Tensor,
        padding_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if padding_mask is None:
            padding_mask = torch.zeros(
                gene_embeddings.shape[:2],
                dtype=torch.bool,
                device=gene_embeddings.device,
            )
        valid = ~padding_mask
        if torch.any(valid.sum(dim=1) == 0):
            raise ValueError("Every BGC must contain at least one unmasked gene")

        counts = valid.sum(dim=1, keepdim=True).to(gene_embeddings.dtype)
        raw = (gene_embeddings * valid.unsqueeze(-1)).sum(dim=1) / counts

        features = self.gene_projection(self.embedding_norm(gene_embeddings))
        features = features + self.position_projection(relative_positions.unsqueeze(-1))
        logits = self.gate(features).squeeze(-1)
        logits = logits.masked_fill(padding_mask, float("-inf"))
        weights = torch.softmax(logits, dim=1)
        weighted = (gene_embeddings * weights.unsqueeze(-1)).sum(dim=1)
        return (
            F.normalize(raw, p=2, dim=-1),
            F.normalize(weighted, p=2, dim=-1),
            weights,
        )

    def forward(
        self,
        gene_embeddings: torch.Tensor,
        relative_positions: torch.Tensor,
        padding_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        return self.components(gene_embeddings, relative_positions, padding_mask)[1]

    @staticmethod
    def combine(
        raw_embeddings: torch.Tensor,
        learned_embeddings: torch.Tensor,
        alpha: float,
    ) -> torch.Tensor:
        if not 0.0 <= alpha <= 1.0:
            raise ValueError("Residual alpha must be in [0, 1]")
        raw = F.normalize(raw_embeddings, p=2, dim=-1)
        learned = F.normalize(learned_embeddings, p=2, dim=-1)
        return torch.cat(
            [
                np.sqrt(1.0 - alpha) * raw,
                np.sqrt(alpha) * learned,
            ],
            dim=-1,
        )


def pfam_teacher_matrix(
    bgc_ids: Sequence[str],
    pfam_sets: Mapping[str, set[str]],
    device: torch.device,
) -> torch.Tensor:
    values = torch.empty((len(bgc_ids), len(bgc_ids)), dtype=torch.float32)
    for row, left_id in enumerate(bgc_ids):
        left = pfam_sets.get(str(left_id), set())
        for column, right_id in enumerate(bgc_ids):
            right = pfam_sets.get(str(right_id), set())
            union = left | right
            values[row, column] = len(left & right) / len(union) if union else 0.0
    return values.to(device)


def soft_relational_loss(
    learned_embeddings: torch.Tensor,
    raw_embeddings: torch.Tensor,
    group_ids: Sequence[str],
    teacher_similarities: torch.Tensor,
    student_temperature: float,
    teacher_temperature: float,
    preservation_weight: float,
) -> torch.Tensor:
    """Match soft Pfam/family geometry without hard-negative false assumptions."""
    count = len(group_ids)
    if learned_embeddings.shape[0] != count or teacher_similarities.shape != (count, count):
        raise ValueError("Residual loss inputs have inconsistent batch dimensions")
    if student_temperature <= 0.0 or teacher_temperature <= 0.0:
        raise ValueError("Temperatures must be positive")
    if preservation_weight < 0.0:
        raise ValueError("Preservation weight must be non-negative")

    learned = F.normalize(learned_embeddings, dim=1)
    raw = F.normalize(raw_embeddings, dim=1)
    identity = torch.eye(count, dtype=torch.bool, device=learned.device)
    same_group = torch.tensor(
        [[left == right for right in group_ids] for left in group_ids],
        dtype=torch.bool,
        device=learned.device,
    )
    teacher = teacher_similarities.to(learned.device).clone()
    teacher = torch.where(same_group & ~identity, torch.ones_like(teacher), teacher)
    teacher_logits = (teacher / teacher_temperature).masked_fill(identity, float("-inf"))
    teacher_probabilities = torch.softmax(teacher_logits, dim=1)

    student_logits = (learned @ learned.T / student_temperature).masked_fill(
        identity, float("-inf")
    )
    student_log_probabilities = torch.log_softmax(student_logits, dim=1)
    student_log_probabilities = student_log_probabilities.masked_fill(identity, 0.0)
    relational = F.kl_div(
        student_log_probabilities,
        teacher_probabilities,
        reduction="batchmean",
    )
    preservation = F.mse_loss(
        (learned @ learned.T)[~identity],
        (raw @ raw.T)[~identity],
    )
    return relational + preservation_weight * preservation


@torch.no_grad()
def encode_residual_components(
    model: ResidualGeneWeightingEncoder,
    dataset: BGCEmbeddingDataset,
    device: torch.device,
    num_workers: int = 0,
) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
    loader = DataLoader(
        dataset,
        batch_size=32,
        shuffle=False,
        num_workers=num_workers,
        collate_fn=collate_bgcs,
    )
    model.eval()
    raw_result: dict[str, torch.Tensor] = {}
    learned_result: dict[str, torch.Tensor] = {}
    for batch in loader:
        embeddings = batch["gene_embeddings"].to(device)
        positions = batch["relative_positions"].to(device)
        padding_mask = batch["padding_mask"].to(device)
        raw, learned, _ = model.components(embeddings, positions, padding_mask)
        raw_result.update(zip(batch["bgc_ids"], raw.cpu()))
        learned_result.update(zip(batch["bgc_ids"], learned.cpu()))
    return raw_result, learned_result


def train_residual_gene_weighting(
    model: ResidualGeneWeightingEncoder,
    model_config: ModelConfig,
    train_dataset: BGCEmbeddingDataset,
    validation_dataset: BGCEmbeddingDataset,
    validation_assignments: pd.DataFrame,
    pfam_sets: Mapping[str, set[str]],
    split_path: str | Path,
    input_paths: list[str | Path],
    output_dir: str | Path,
    training_config: dict[str, Any],
    residual_config: dict[str, Any],
    seed: int,
) -> Path:
    set_reproducible(seed)
    device = choose_device()
    model.to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=float(residual_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,
    )
    mixed_precision = bool(training_config["mixed_precision"]) and device.type == "cuda"
    scaler = torch.amp.GradScaler("cuda", enabled=mixed_precision)
    history: list[dict[str, float | int]] = []
    best_metric = -float("inf")
    best_state: dict[str, Any] | None = None
    patience = 0

    for epoch in range(int(residual_config["epochs"])):
        sampler.set_epoch(epoch)
        model.train()
        epoch_losses: list[float] = []
        for batch in loader:
            embeddings = batch["gene_embeddings"].to(device)
            positions = batch["relative_positions"].to(device)
            padding_mask = batch["padding_mask"].to(device)
            embeddings, positions, padding_mask = augment_gene_sets(
                embeddings,
                positions,
                padding_mask,
                float(training_config["gene_dropout"]),
                float(training_config["position_jitter"]),
            )
            teacher = pfam_teacher_matrix(batch["bgc_ids"], pfam_sets, device)
            optimizer.zero_grad(set_to_none=True)
            with torch.amp.autocast(
                device_type=device.type,
                enabled=mixed_precision,
            ):
                raw, learned, _ = model.components(embeddings, positions, padding_mask)
                loss = soft_relational_loss(
                    learned,
                    raw,
                    batch["group_ids"],
                    teacher,
                    float(residual_config["student_temperature"]),
                    float(residual_config["teacher_temperature"]),
                    float(residual_config["preservation_weight"]),
                )
            if not torch.isfinite(loss):
                raise FloatingPointError("Non-finite residual training 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_residual_components(
            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(residual_config["patience"]):
                break

    if best_state is None:
        raise RuntimeError("Residual training did not produce a valid checkpoint")
    model.load_state_dict(best_state["model"])
    output = Path(output_dir)
    checkpoint_path = output / "residual_best.pt"
    save_checkpoint(
        checkpoint_path,
        model,
        model_config,
        split_path,
        input_paths,
        {
            "stage": "residual_gene_weighting",
            "best_epoch": best_state["epoch"],
            "best_validation_recall@50": best_metric,
            "objective": "soft_pfam_family_relational_distillation",
        },
        optimizer,
    )
    write_json_immutable(output / "residual_history.json", history)
    return checkpoint_path


class RetrievalScoreCache:
    """Cache three base scorers once per deterministic retrieval draw."""

    def __init__(
        self,
        raw_embeddings: Mapping[str, torch.Tensor],
        learned_embeddings: Sequence[Mapping[str, torch.Tensor]],
        pfam_sets: Mapping[str, set[str]],
    ) -> None:
        self.raw_embeddings = raw_embeddings
        self.learned_embeddings = list(learned_embeddings)
        self.pfam_sets = pfam_sets
        self.key: tuple[tuple[str, ...], tuple[str, ...]] | None = None
        self.values: tuple[dict[str, float], dict[str, float], dict[str, float]] | None = None

    def components(
        self,
        candidates: Sequence[str],
        references: Sequence[str],
    ) -> tuple[dict[str, float], dict[str, float], dict[str, float]]:
        key = (tuple(candidates), tuple(references))
        if key == self.key and self.values is not None:
            return self.values
        raw = cosine_scores(candidates, references, self.raw_embeddings, "mean")
        learned_by_model = [
            cosine_scores(candidates, references, embeddings, "mean")
            for embeddings in self.learned_embeddings
        ]
        learned = {
            identifier: float(
                np.mean([scores[identifier] for scores in learned_by_model])
            )
            for identifier in candidates
        }
        pfam = pfam_jaccard_scores(candidates, references, self.pfam_sets, "max")
        self.key = key
        self.values = (raw, learned, pfam)
        return self.values

    def residual(
        self,
        candidates: Sequence[str],
        references: Sequence[str],
        alpha: float,
    ) -> dict[str, float]:
        raw, learned, _ = self.components(candidates, references)
        return {
            identifier: (1.0 - alpha) * raw[identifier] + alpha * learned[identifier]
            for identifier in candidates
        }

    def hybrid(
        self,
        candidates: Sequence[str],
        references: Sequence[str],
        alpha: float,
        beta: float,
    ) -> dict[str, float]:
        residual = self.residual(candidates, references, alpha)
        _, _, pfam = self.components(candidates, references)
        return ensemble_scores(residual, pfam, beta)


def validation_grid(
    assignments: pd.DataFrame,
    cache: RetrievalScoreCache,
    alphas: Sequence[float],
    betas: Sequence[float],
    reference_size: int,
    draws: int,
    seed: int,
    recall_at: Sequence[int],
    ndcg_at: Sequence[int],
) -> pd.DataFrame:
    methods: dict[str, Any] = {}
    for alpha in alphas:
        methods[f"residual_a{alpha:g}"] = (
            lambda candidates, references, weight=float(alpha): cache.residual(
                candidates, references, weight
            )
        )
        for beta in betas:
            methods[f"hybrid_a{alpha:g}_b{beta:g}"] = (
                lambda candidates, references, a=float(alpha), b=float(beta): cache.hybrid(
                    candidates, references, a, b
                )
            )
    return evaluate_retrieval(
        assignments,
        "validation",
        methods,
        reference_size,
        draws,
        seed,
        recall_at,
        ndcg_at,
    )


def select_validation_weights(
    results: pd.DataFrame,
    metric: str,
    prefix: str,
) -> tuple[float, float | None]:
    selected = results[results["method"].str.startswith(prefix)].copy()
    if selected.empty:
        raise ValueError(f"No validation methods start with {prefix}")
    means = selected.groupby("method")[metric].mean()
    best = float(means.max())
    candidates = sorted(str(name) for name, value in means.items() if np.isclose(value, best))
    parsed: list[tuple[float, float | None, str]] = []
    for name in candidates:
        alpha = float(name.split("_a", 1)[1].split("_", 1)[0])
        beta = float(name.rsplit("_b", 1)[1]) if "_b" in name else None
        parsed.append((alpha, beta, name))
    if prefix == "hybrid_":
        alpha, beta, _ = min(
            parsed,
            key=lambda item: (
                float(item[1]),
                item[0],
            ),
        )
    else:
        alpha, beta, _ = min(parsed, key=lambda item: item[0])
    return alpha, beta