src/pino/train.py

from __future__ import annotations

import argparse
import json
import logging
import random
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from datasets import DatasetDict
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from pino.heads import PIMTHeads
from pino.embeddings import COMBINED_DIM_REAL_POM, MORGAN2048_TOKEN_DIM
from pino.pimt_model import (
    DEFAULT_EMBEDDING_DIM,
    FragranceTrajectoryDataset,
    PhysicsInformedMixtureTransformer,
    objective_targets_to_pyramid,
)
from pino.thermo.losses import AdaptiveLossBalancer
from pino.upload_data import DEFAULT_TRAIN_RATIO, create_molecule_disjoint_split_from_records

logger = logging.getLogger("pino.train")

SMOKE_TEST_CONFIG = {
    "hidden_dim": 64,
    "num_layers": 2,
    "num_heads": 2,
    "batch_size": 4,
    "max_batches_per_epoch": 20,
    "lr": 1e-3,
    "weight_decay": 1e-2,
    "checkpoint_name": "pimt_smoke_test.pt",
}

# Idea3 physics-off ablation: when True, pad_trajectory_collate zeroes the
# physics tensor for every batch (train + val). Set from --physics-off-train.
_PHYSICS_OFF_TRAIN = False
# Factorial physics-channel ablation (idea-3 follow-up): which channel carries
# the drydown gain? Channels are [x_liquid, log10 OAV, log10 gamma_prop].
#   full           : all channels live (default; physics_gamma arm)
#   off            : all zeroed (physics_off arm; legacy --physics-off-train)
#   gamma_only     : keep channel 2 (log gamma), zero 0-1 -- isolates the
#                    non-redundant blend-interaction channel
#   headspace_only : keep channels 0-1, zero channel 2 -- the E2-redundant part
_PHYSICS_MODE = "full"

def pad_trajectory_collate(batch: list[dict[str, Any]]) -> dict[str, Any]:
    """
    Collate variable-length formulations into a single padded batch.

    Because recipes contain a variable number of ingredients (K), this function
    finds K_max in the batch and pads shorter token/physics tensors with zeros.
    It also emits a `src_key_padding_mask` of shape (B, K_max) where True marks
    padding positions and False marks valid chemical tokens.
    """
    max_molecules = max(item["tokens"].size(0) for item in batch)
    max_timesteps = max(item["physics"].size(0) for item in batch)
    bsz = len(batch)

    emb_dim = batch[0]["tokens"].size(1)
    # physics channel count from the batch (2 = x_liq/logOAV, 3 = + gamma interaction)
    phys_dim = batch[0]["physics"].size(-1)
    # objective target width from the batch itself (138 for pyramid, 575 for all-tags)
    obj_dim = batch[0]["target_obj"].size(-1)
    tokens = torch.zeros(bsz, max_molecules, emb_dim, dtype=torch.float32)
    physics = torch.zeros(bsz, max_timesteps, max_molecules, phys_dim, dtype=torch.float32)
    src_key_padding_mask = torch.ones(bsz, max_molecules, dtype=torch.bool)
    target_obj = torch.zeros(bsz, 3, obj_dim, dtype=torch.float32)
    target_sub_seasonality = torch.zeros(bsz, 4, dtype=torch.float32)
    target_sub_gender = torch.zeros(bsz, dtype=torch.float32)
    target_sub_wearability = torch.zeros(bsz, 2, dtype=torch.float32)
    target_substantivity = torch.zeros(bsz, dtype=torch.float32)
    target_substantivity_mask = torch.zeros(bsz, dtype=torch.float32)
    genre_labels = torch.zeros(bsz, dtype=torch.int64)

    for i, item in enumerate(batch):
        n_mol = item["tokens"].size(0)
        t_steps = item["physics"].size(0)
        tokens[i, :n_mol] = item["tokens"]
        physics[i, :t_steps, :n_mol] = item["physics"]
        src_key_padding_mask[i, :n_mol] = False
        target_obj[i] = objective_targets_to_pyramid(item["target_obj"]).float()
        sub = item["target_sub"]
        target_sub_seasonality[i] = sub[[0,1,2,3]]
        target_sub_gender[i] = sub[4]
        target_sub_wearability[i] = sub[[5,6]]
        target_substantivity[i] = item.get("target_substantivity", torch.tensor(0.0))
        target_substantivity_mask[i] = item.get("target_substantivity_mask", torch.tensor(0.0))
        genre_labels[i] = item.get("genre_label", 0)

    # Force padding positions to a neutral log10_OAV sentinel so the Synergistic
    # Objective Head treats them as zero-contribution even without the mask.
    physics[..., 1] = physics[..., 1].masked_fill(src_key_padding_mask.unsqueeze(1), -6.0)

    # Idea3 physics-off ablation arm: zero the entire physics state so the FiLM
    # block and physics-informed heads receive no physics signal during BOTH
    # training and validation. This is a fair trained baseline (not an
    # inference-time shock), isolating the physics channel's contribution.
    if _PHYSICS_OFF_TRAIN or _PHYSICS_MODE == "off":
        physics = torch.zeros_like(physics)
    elif _PHYSICS_MODE == "gamma_only":
        # Keep only channel 2 (log gamma); zero x_liquid and log10 OAV.
        physics = physics.clone()
        physics[..., 0] = 0.0
        physics[..., 1] = 0.0
    elif _PHYSICS_MODE == "headspace_only":
        # Keep channels 0-1 (x_liquid, log10 OAV); zero the gamma channel.
        physics = physics.clone()
        physics[..., 2] = 0.0

    return {
        "tokens": tokens,
        "physics": physics,
        "src_key_padding_mask": src_key_padding_mask,
        "target_obj": target_obj,
        "target_sub_seasonality": target_sub_seasonality,
        "target_sub_gender": target_sub_gender,
        "target_sub_wearability": target_sub_wearability,
        "target_substantivity": target_substantivity,
        "target_substantivity_mask": target_substantivity_mask,
        "genre_label": genre_labels,
    }


def multitask_loss(
    pred_obj: torch.Tensor,
    target_obj: torch.Tensor,
    pred_sub: dict[str, torch.Tensor],
    target_sub: dict[str, torch.Tensor],
    pred_alignment: torch.Tensor,
    loss_balancer: AdaptiveLossBalancer,
) -> tuple[torch.Tensor, dict[str, float]]:
    """
    Multi-task loss for pyramid targets.

    The objective task is now binary multi-label classification over three static
    tiers (top, middle, base notes).  Sigmoid BCE replaces MSE on the 49-step
    trajectory, while subjective and alignment tasks remain unchanged.
    """
    # Pyramid objective: binary multi-label classification (B, 3, 138).
    # Clamp targets into [0,1]: some all-tags labels carry a float32 rounding
    # artifact (1.00000012) that trips the CUDA BCE device-side assert on GPU.
    obj_loss = nn.functional.binary_cross_entropy(pred_obj, target_obj.clamp(0.0, 1.0), reduction="mean")

    season_loss = nn.functional.cross_entropy(pred_sub["seasonality"], target_sub["seasonality"].argmax(dim=-1))
    gender_loss = nn.functional.mse_loss(pred_sub["gender_profile"].squeeze(-1), target_sub["gender_profile"])
    wear_loss = nn.functional.binary_cross_entropy_with_logits(pred_sub["wearability"], target_sub["wearability"])
    wear_loss = wear_loss + gender_loss
    sub_mask = target_sub.get("substantivity_mask")
    if sub_mask is not None and float(sub_mask.sum().detach().cpu()) > 0:
        sub_pred = pred_sub["substantivity"].squeeze(-1)
        sub_loss = ((sub_pred - target_sub["substantivity"]) ** 2 * sub_mask).sum() / sub_mask.sum().clamp_min(1.0)
    else:
        sub_loss = pred_obj.new_tensor(0.0)

    # Continuous cosine alignment between predicted 138-D descriptor profile and
    # the mean of the true pyramid tiers over the three tiers.  This keeps the
    # overall descriptor profile aligned even as the pyramid specialises.
    true_profile = target_obj.mean(dim=1)  # (B, 138)
    alignment_loss = 1.0 - nn.functional.cosine_similarity(pred_alignment, true_profile, dim=-1).mean()

    loss_payload = {
        "obj": obj_loss,
        "season": season_loss,
        "wear": wear_loss,
        "substantivity": sub_loss,
        "alignment": alignment_loss,
    }

    total = loss_balancer(loss_payload)
    components = {
        "obj": obj_loss.item(),
        "season": season_loss.item(),
        "wear": wear_loss.item(),
        "substantivity": sub_loss.item(),
        "alignment": alignment_loss.item(),
        "total": total.item(),
    }
    return total, components


def train_epoch(
    model: nn.Module,
    heads: nn.Module,
    loader: DataLoader,
    optimizer: optim.Optimizer,
    device: torch.device,
    loss_balancer: AdaptiveLossBalancer,
    max_batches: int | None = None,
) -> dict[str, float]:
    model.train()
    heads.train()
    total_loss = 0.0
    totals = {"obj": 0.0, "season": 0.0, "wear": 0.0, "substantivity": 0.0, "alignment": 0.0}
    count = 0

    for batch_idx, batch in enumerate(loader, start=1):
        if max_batches is not None and batch_idx > max_batches:
            break
        tokens = batch["tokens"].to(device)
        physics = batch["physics"].to(device)
        src_key_padding_mask = batch["src_key_padding_mask"].to(device)
        target_obj = batch["target_obj"].to(device)
        target_sub = {
            "seasonality": batch["target_sub_seasonality"].to(device),
            "gender_profile": batch["target_sub_gender"].to(device),
            "wearability": batch["target_sub_wearability"].to(device),
            "substantivity": batch["target_substantivity"].to(device),
            "substantivity_mask": batch["target_substantivity_mask"].to(device),
        }
        optimizer.zero_grad()
        latent = model(tokens, physics, src_key_padding_mask)  # (B, T, S, H)
        # Compute pyramid objective predictions with physics-informed routing.
        pred = heads(latent, physics, src_key_padding_mask)
        loss, comps = multitask_loss(
            pred["objective"], target_obj, pred["subjective"], target_sub,
            pred["alignment"], loss_balancer,
        )
        loss.backward()
        torch.nn.utils.clip_grad_norm_(list(model.parameters()) + list(heads.parameters()), 1.0)
        optimizer.step()

        total_loss += comps["total"]
        for k in totals:
            totals[k] += comps[k]
        count += 1
        weights = getattr(loss_balancer, "last_weights", [0.35, 0.12, 0.12, 0.16, 0.25])
        logger.info(
            "Batch %d | total=%.4f | obj=%.4f | season=%.4f | wear=%.4f | substantivity=%.4f | alignment=%.4f | weights=%s",
            batch_idx, comps["total"], comps["obj"], comps["season"], comps["wear"], comps["substantivity"], comps["alignment"],
            [round(float(w), 3) for w in weights],
        )

    return {k: v / max(count, 1) for k, v in {**totals, "total": total_loss}.items()}


@torch.no_grad()
def validate(
    model: nn.Module,
    heads: nn.Module,
    loader: DataLoader,
    device: torch.device,
    loss_balancer: AdaptiveLossBalancer,
) -> dict[str, float]:
    model.eval()
    heads.eval()
    total_loss = 0.0
    totals = {"obj": 0.0, "season": 0.0, "wear": 0.0, "substantivity": 0.0, "alignment": 0.0}
    count = 0

    for batch in loader:
        tokens = batch["tokens"].to(device)
        physics = batch["physics"].to(device)
        src_key_padding_mask = batch["src_key_padding_mask"].to(device)
        target_obj = batch["target_obj"].to(device)
        target_sub = {
            "seasonality": batch["target_sub_seasonality"].to(device),
            "gender_profile": batch["target_sub_gender"].to(device),
            "wearability": batch["target_sub_wearability"].to(device),
            "substantivity": batch["target_substantivity"].to(device),
            "substantivity_mask": batch["target_substantivity_mask"].to(device),
        }
        latent = model(tokens, physics, src_key_padding_mask)
        pred = heads(latent, physics, src_key_padding_mask)
        loss, comps = multitask_loss(
            pred["objective"], target_obj, pred["subjective"], target_sub,
            pred["alignment"], loss_balancer,
        )
        total_loss += comps["total"]
        for k in totals:
            totals[k] += comps[k]
        count += 1

    return {k: v / max(count, 1) for k, v in {**totals, "total": total_loss}.items()}


def main():
    parser = argparse.ArgumentParser(description="Train the PINO Physics-Informed Mixture Transformer")
    parser.add_argument(
        "--data",
        default="data/empirical_dataset_v9_plus_wisemoor.jsonl",
        help="Path to the empirical dataset JSONL (default: v9 plus resolved WiseMoor formulas).",
    )
    parser.add_argument("--epochs", type=int, default=5)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--hidden-dim", type=int, default=256)
    parser.add_argument("--num-layers", type=int, default=4)
    parser.add_argument("--num-heads", type=int, default=4)
    parser.add_argument("--workers", type=int, default=0, help="DataLoader workers; 0 for local CPU check cycles")
    parser.add_argument("--seed", type=int, default=2026)
    parser.add_argument("--output-dir", default="models")
    parser.add_argument("--log-dir", default="runs")
    parser.add_argument("--checkpoint-name", default="pimt_v9.pt", help="Filename for the best checkpoint")
    parser.add_argument("--train-ratio", type=float, default=DEFAULT_TRAIN_RATIO, help="Molecule-disjoint train molecule ratio; default keeps a 15% molecule holdout")
    parser.add_argument("--structural-source", choices=["morgan", "morgan_2048_rp", "morgan_2048_direct", "openpom_256", "pom_alltags", "disjoint_256"], default="morgan", help="Structural embedding block + objective target: 'morgan' (138-dim Morgan, input 151, 138-dim pyramid) | 'morgan_2048_rp' (2048-bit Morgan token randomly projected to 269 for input-parameter parity with the OpenPOM arm, 138-dim pyramid) | 'morgan_2048_direct' (2048-bit Morgan token fed directly to a learned projection, input 2071 — round-10 control for the RP subspace constraint) | 'openpom_256' (genuine 256-dim OpenPOM, input 269, 138-dim pyramid) | 'disjoint_256' (256-dim MPNN encoder retrained on the benchmark-disjoint corpus, input 269, 138-dim pyramid — leakage-free OpenPOM-style arm) | 'pom_alltags' (genuine 256-dim OpenPOM, input 269, 575-dim molequles all-tags objective)")
    parser.add_argument("--smoke-test", action="store_true", help="Run lightweight CPU smoke test")
    parser.add_argument("--use-gamma", action="store_true", help="Add the UNIFAC blend-interaction (gamma) 3rd physics channel; requires a gamma-augmented dataset (physics_gamma field)")
    parser.add_argument("--split-strategy", choices=["molecule", "formula"], default="molecule", help="molecule = strict molecule-disjoint holdout (default, used for frozen-task headline results); formula = random formula-level holdout (E1 physics-gamma substantivity experiment; discloses molecule leakage for measured-Poucher eval power)")
    parser.add_argument("--physics-off-train", action="store_true", help="Idea3 ablation: zero the physics tensor for every batch (train+val) so the model is a fair no-physics baseline")
    parser.add_argument("--physics-mode", choices=["full", "off", "gamma_only", "headspace_only"], default="full", help="Factorial channel ablation: full = all channels; off = all zeroed; gamma_only = keep only log-gamma channel (zero x_liquid/OAV); headspace_only = keep x_liquid/OAV (zero gamma). Applied during BOTH train and val so each arm is a fair trained baseline.")
    args = parser.parse_args()

    global _PHYSICS_OFF_TRAIN, _PHYSICS_MODE
    _PHYSICS_OFF_TRAIN = bool(getattr(args, "physics_off_train", False))
    _PHYSICS_MODE = getattr(args, "physics_mode", "full")
    if _PHYSICS_OFF_TRAIN:
        _PHYSICS_MODE = "off"

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    logger.info("Training PIMT with fallback embeddings (use_fallback=True)")
    logger.info("Embedding backend: deterministic 128-bit Morgan + 10 physicochemical descriptors")

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    random.seed(args.seed)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info("Device: %s", device)

    logger.info("Structural source: %s", args.structural_source)
    use_pom = args.structural_source in ("openpom_256", "pom_alltags", "disjoint_256")
    if use_pom or args.structural_source == "morgan_2048_rp":
        embedding_dim = COMBINED_DIM_REAL_POM
    elif args.structural_source == "morgan_2048_direct":
        embedding_dim = MORGAN2048_TOKEN_DIM
    else:
        embedding_dim = DEFAULT_EMBEDDING_DIM
    objective_dim = 575 if args.structural_source == "pom_alltags" else 138
    # engine only knows the two structural blocks; the tag target is selected in the dataset via objective_dim
    engine_source = "openpom_256" if args.structural_source == "pom_alltags" else args.structural_source
    logger.info("Embedding input dim: %d | objective target dim: %d", embedding_dim, objective_dim)

    cfg = SMOKE_TEST_CONFIG if args.smoke_test else {
        "hidden_dim": args.hidden_dim,
        "num_layers": args.num_layers,
        "num_heads": args.num_heads,
        "batch_size": args.batch_size,
        "max_batches_per_epoch": None,
        "lr": args.lr,
        "weight_decay": 1e-2,
        "checkpoint_name": args.checkpoint_name,
    }
    logger.info("Configuration: %s", cfg)

    # Ingest and split the dataset with strict active-molecule isolation.
    # NOTE (leakage review): a connected-components split over the formula-molecule
    # co-occurrence graph was evaluated and REJECTED for this corpus. The graph is
    # dominated by a single giant component (5,592/5,707 records; shared carriers +
    # genre ingredient pools link nearly all molecules), so a strict component split
    # is mathematically impossible, and the hub-excluded variant yields a validation
    # set of only singleton/duplicate formulas (degenerate, unrepresentative).
    # Molecule-disjoint splitting is therefore the correct instrument here: it gives
    # a representative held-out set with a strict zero-active-molecule guarantee.
    # See scripts/split_leakage_analysis.py for the full evidence + assertion.
    print("⏳ Splicing empirical bootstrap data streams...")
    with open(args.data, "r") as f:
        all_records = [json.loads(line) for line in f]

    control_records = [r for r in all_records if r.get("is_control", False)]
    print(f"📊 Extracted {len(control_records)} pure single-molecule anchors.")

    if getattr(args, "split_strategy", "molecule") == "formula":
        # Formula-level random holdout. Used by the E1 physics-gamma substantivity
        # experiment, where strict molecule-disjointness starves the measured-
        # Poucher eval (n~10) because Poucher molecules are ubiquitous. Trades a
        # weaker generalization claim (molecules may leak) for eval power; the
        # leakage is disclosed in the run summary. The frozen-task headline
        # results in the paper still use the molecule-disjoint path below.
        rng = np.random.default_rng(args.seed)
        n_total = len(all_records)
        n_val = max(1, int(round(n_total * (1.0 - args.train_ratio))))
        perm = rng.permutation(n_total)
        val_idx = set(perm[:n_val].tolist())
        train_records = [all_records[i] for i in range(n_total) if i not in val_idx]
        val_records = [all_records[i] for i in sorted(val_idx)]
        print(f"📊 Formula-level split (E1; leakage disclosed): train={len(train_records)} val={len(val_records)}")
    else:
        split = create_molecule_disjoint_split_from_records(
            all_records,
            train_ratio=args.train_ratio,
            seed=args.seed,
        )
        train_records = split["train"]
        val_records = split["validation"]

    if getattr(args, "split_strategy", "molecule") == "molecule":
        train_cas = set(split["train_compounds"])
        val_cas = set(split["validation_compounds"])
        overlap = train_cas & val_cas
        print(f"🔒 Molecule isolation check: {len(overlap)} active CAS shared between splits.")
        print(f"🚫 Excluded mixed-boundary records: {len(split['excluded_indices'])}")
        assert len(overlap) == 0, f"Molecule leakage detected: {sorted(overlap)}"
        # Persist the zero-intersection assertion for the paper's reproducibility ledger.
        try:
            leakage_report = {
                "split_strategy": "molecule_disjoint",
                "seed": args.seed,
                "train_ratio": args.train_ratio,
                "n_train_records": len(train_records),
                "n_validation_records": len(val_records),
                "n_excluded_records": len(split["excluded_indices"]),
                "n_train_molecules": len(train_cas),
                "n_validation_molecules": len(val_cas),
                "shared_molecule_count": len(overlap),
                "zero_intersection_assertion": "PASSED" if not overlap else "FAILED",
            }
            Path(args.output_dir).mkdir(parents=True, exist_ok=True)
            report_name = f"split_leakage_report_seed{args.seed}.json"
            with open(Path(args.output_dir) / report_name, "w") as rf:
                json.dump(leakage_report, rf, indent=2)
            print(f"📝 Split leakage report: {report_name} (zero-intersection: {leakage_report['zero_intersection_assertion']})")
        except Exception as e:  # noqa: BLE001
            logger.warning("Could not write split leakage report: %s", e)
    if not train_records or not val_records:
        raise RuntimeError(
            "Split produced an empty train or validation set. "
            "Adjust --train-ratio or --seed."
        )

    total_rows = len(train_records) + len(val_records)
    print(f"📊 Dataset fully loaded. Total rows: {total_rows} | Training samples: {len(train_records)} | Validation samples: {len(val_records)}")

    # Use the genre column for both stratification and the contrastive loss.
    all_genres = ["citrus_cologne", "fougere", "floral_woody", "amber_oriental", "wildcard"]
    genre_map = {genre: idx for idx, genre in enumerate(all_genres)}

    train_dataset = FragranceTrajectoryDataset(
        records=train_records,
        use_embedding_fallback=True,
        structural_source=engine_source,
        genre_map=genre_map,
        objective_dim=objective_dim,
        use_gamma=args.use_gamma,
    )
    val_dataset = FragranceTrajectoryDataset(
        records=val_records,
        use_embedding_fallback=True,
        structural_source=engine_source,
        genre_map=genre_map,
        objective_dim=objective_dim,
        use_gamma=args.use_gamma,
    )

    train_loader = DataLoader(
        train_dataset, batch_size=cfg["batch_size"], shuffle=True, num_workers=args.workers, collate_fn=pad_trajectory_collate
    )
    val_loader = DataLoader(
        val_dataset, batch_size=cfg["batch_size"], shuffle=False, num_workers=args.workers, collate_fn=pad_trajectory_collate
    )

    state_dim = 3 if args.use_gamma else 2
    model = PhysicsInformedMixtureTransformer(
        embedding_dim=embedding_dim,
        state_dim=state_dim,
        hidden_dim=cfg["hidden_dim"],
        num_heads=cfg["num_heads"],
        num_layers=cfg["num_layers"],
    ).to(device)
    heads = PIMTHeads(hidden_dim=cfg["hidden_dim"], objective_dim=objective_dim).to(device)

    params = list(model.parameters()) + list(heads.parameters())
    optimizer = optim.AdamW(params, lr=cfg["lr"], weight_decay=cfg["weight_decay"])
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)

    # Multi-task loss balancer.  Fix objective high so the pyramid head is not
    # starved by the easier subjective tasks.
    loss_weights = torch.tensor([0.35, 0.12, 0.12, 0.16, 0.25])
    loss_balancer = AdaptiveLossBalancer(
        num_tasks=5, fixed_weights=loss_weights
    ).to(device)

    writer = SummaryWriter(log_dir=args.log_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    checkpoint_path = output_dir / cfg["checkpoint_name"]

    best_val = float("inf")
    for epoch in range(1, args.epochs + 1):
        train_metrics = train_epoch(model, heads, train_loader, optimizer, device, loss_balancer, max_batches=cfg.get("max_batches_per_epoch"))
        val_metrics = validate(model, heads, val_loader, device, loss_balancer)
        scheduler.step()

        for split, metrics in [("train", train_metrics), ("val", val_metrics)]:
            for k, v in metrics.items():
                writer.add_scalar(f"{split}/{k}", v, epoch)

        logger.info(
            "Epoch %d/%d | train_total=%.4f | val_total=%.4f | obj=%.4f | season=%.4f | wear=%.4f | substantivity=%.4f | alignment=%.4f",
            epoch, args.epochs, train_metrics["total"], val_metrics["total"],
            val_metrics["obj"], val_metrics["season"], val_metrics["wear"], val_metrics["substantivity"], val_metrics["alignment"],
        )

        if val_metrics["total"] < best_val:
            best_val = val_metrics["total"]
            torch.save({
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "heads_state_dict": heads.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "val_loss": best_val,
            }, checkpoint_path)
            logger.info("Saved best checkpoint to %s", checkpoint_path)

    val_final_metrics = validate(model, heads, val_loader, device, loss_balancer)
    logger.info(
        "Final Val | total=%.4f | obj=%.4f | season=%.4f | wear=%.4f | substantivity=%.4f | alignment=%.4f",
        val_final_metrics["total"], val_final_metrics["obj"], val_final_metrics["season"], val_final_metrics["wear"], val_final_metrics["substantivity"], val_final_metrics["alignment"],
    )
    writer.close()
    return checkpoint_path


if __name__ == "__main__":
    main()