import os, pickle, yaml, json
from datetime import datetime
import pytz

def str_datetime(fmt=None, tz=pytz.UTC):
    fmt = fmt or '%m-%d-%Y_%I-%M-%S-%p'
    return datetime.now(tz).strftime(fmt)

def get_world_size():
    world_size = int(os.getenv('WORLD_SIZE', 1))
    return world_size

def get_global_rank():
    node_rank = int(os.getenv('NODE_RANK', 0))
    local_rank = int(os.getenv('LOCAL_RANK', 0))
    world_size = get_world_size()
    return world_size * node_rank + local_rank

def is_rank_zero():
    node_rank = int(os.getenv('NODE_RANK', 0))
    local_rank = int(os.getenv('LOCAL_RANK', 0))
    return (local_rank + node_rank) == 0

def set_cuda_visible_devices(ids):
    if ids is None:
        return
    os.environ['CUDA_VISIBLE_DEVICES'] = ids if isinstance(ids, str) else ','.join(ids)

def get_cuda_visible_devices():
    if  'CUDA_VISIBLE_DEVICES' not in os.environ:
        return None
    devices = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
    ids = [int(device_id) for device_id in devices if len(device_id)]
    return ids

def pickle_save(obj, path, **kwargs):
    with open(path, 'wb') as write_buffer:
        pickle.dump(obj, write_buffer, **kwargs)


def pickle_load(path):
    with open(path, 'rb') as read_buffer:
        obj = pickle.load(read_buffer)
    return obj

def yaml_save(data, path, **kwargs):
    with open(path, 'w') as file:
        yaml.dump(data, file, **kwargs)

def yaml_load(path, **kwargs):
    with open(path, 'r') as file:
        data = yaml.safe_load(file, **kwargs)
    return data

def json_save(data, path, **kwargs):
    with open(path, 'w') as file:
        json.dump(data, file, **kwargs)

def json_load(path, **kwargs):
    with open(path, 'r') as file:
        data = json.load(file, **kwargs)
    return data