"""Deterministic seeding across python / numpy / torch."""
from __future__ import annotations

import os
import random


def set_seed(seed: int, deterministic: bool = True) -> None:
    """Seed all RNGs used in training/eval.

    Mirrors the upstream Transolver/Geo-FNO seeding (``torch.manual_seed``,
    ``np.random.seed``, cudnn deterministic) and extends it to python's ``random``
    and the ``PYTHONHASHSEED`` env var so runs are reproducible across seeds {0,1,2}.
    """
    import numpy as np
    import torch

    os.environ["PYTHONHASHSEED"] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        # NOTE: torch.use_deterministic_algorithms(True) was tried but it coincided with a worse
        # Transolver-baseline result on GPU (baseline eager ~0.0090 vs ~0.0068 without it); we did
        # not isolate the cause and do not force it. Reproducibility is handled by averaging seeds
        # and disclosing run-to-run variance (README + model card). See PART 6 caveat 6.