from typing import Union
import json
import os
import random
from pathlib import Path

import numpy as np
import torch


def ensure_dir(path: Union[str, Path]) -> None:
    """
    Input:
        path: Directory path.
    Output:
        Creates the directory if it does not exist.
    """
    Path(path).mkdir(parents=True, exist_ok=True)


def save_json(obj: dict, path: Union[str, Path]) -> None:
    """
    Input:
        obj: Python dictionary.
        path: Output JSON path.
    Output:
        Saves dictionary as formatted JSON.
    """
    with open(path, "w", encoding="utf-8") as f:
        json.dump(obj, f, indent=2, ensure_ascii=False)


def load_json(path: Union[str, Path]) -> dict:
    """
    Input:
        path: JSON file path.
    Output:
        Loaded Python dictionary.
    """
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def set_seed(seed: int) -> None:
    """
    Input:
        seed: Integer seed.
    Output:
        Makes Python, NumPy, and PyTorch randomness more reproducible.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def is_rank_zero() -> bool:
    """
    Output:
        True only for the main distributed process.
    """
    return int(os.environ.get("RANK", "0")) == 0


def get_world_size() -> int:
    """
    Output:
        Number of distributed processes.
    """
    return int(os.environ.get("WORLD_SIZE", "1"))