import random
from pathlib import Path

import numpy as np
import torch


def set_seeds(seed: int) -> np.random.Generator:
    """Set seeds for reproducibility across various libraries."""

    # Python random module
    random.seed(seed)

    # Torch & Cuda reproducibility
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    # NumPy generator
    np.random.seed(seed)
    return np.random.default_rng(seed)


def get_seed_sequence(rng: np.random.Generator, num_seeds: int) -> list[int]:
    """Get a sequence of seeds from a given RNG."""
    return rng.integers(low=0, high=2**32 - 1, size=num_seeds, dtype=np.uint32).tolist()


def resolve_source_dir(arg_path: str, project_root: Path, current_path: Path) -> Path:
    input_path = Path(arg_path)

    candidates = [
        input_path,  # Absolute path
        project_root / input_path,  # Project-relative path
        current_path / input_path,  # Script-relative path
    ]

    for candidate in candidates:
        if candidate.is_dir():
            return candidate

    raise FileNotFoundError(f"Error: Directory '{arg_path}' not found.")
