"""
This module implements functions for dealing with the state of the program.
Standard torch.distributed is used to handle distributed training and inference.
"""

from __future__ import annotations

import contextlib
import os
import sys

import torch
import torch.distributed as dist
import torch.multiprocessing as mp

__all__ = [
    "barrier",
    "check_main_process",
    "cpus_available",
    "get_device",
    "get_interactive",
    "get_process_count",
    "get_process_index",
]


###################
# Multiprocessing #
###################


def cpus_available():
    r"""
    Returns the number of CPUs available to the current process.

    This is preferred over `os.cpu_count()` as it is more reliable in HPC environments.
    """
    with contextlib.suppress(KeyError, ValueError, TypeError):
        return int(os.environ["SLURM_CPUS_PER_GPU"])

    with contextlib.suppress(Exception):
        return len(os.sched_getaffinity(0))

    return max(mp.cpu_count(), 1)


####################
# Unipercept state #
####################


def get_interactive():
    if os.getenv("UNIPERCEPT_INTERACTIVE", "").lower() in ("1", "true"):
        return True
    if sys.stdin.isatty():  # Terminal
        return True
    if hasattr(sys, "ps1"):  # IPython, Python shell, etc.
        return True
    if sys.flags.interactive:  # Python launched with -i
        return True
    return os.isatty(sys.stdout.fileno())  # Redirected output


##################################
# Data and multiprocessing utils #
##################################


def get_process_index(local: bool = False) -> int:
    """Returns the rank of the current process."""
    if not dist.is_available() or not dist.is_initialized():
        return 0
    if local:
        return int(os.environ.get("LOCAL_RANK", 0))
    return dist.get_rank()


def get_process_count() -> int:
    """Returns the total number of processes."""
    if not dist.is_available() or not dist.is_initialized():
        return 1
    return dist.get_world_size()


def check_main_process(local: bool = False) -> bool:
    """Returns True if the current process is the main process."""
    return get_process_index(local=local) == 0


def get_device() -> torch.device:
    """Returns the device assigned to the current process."""
    if torch.cuda.is_available():
        if dist.is_available() and dist.is_initialized():
            local_rank = int(os.environ.get("LOCAL_RANK", 0))
            return torch.device(f"cuda:{local_rank}")
        return torch.device("cuda:0")
    return torch.device("cpu")


def barrier():
    """Wait for all processes to reach this point."""
    if dist.is_available() and dist.is_initialized():
        dist.barrier()
