from json_tricks.np import dump, load
import pickle
import logging
log = logging.getLogger(__name__)


def save_json(obj, outfile, allow_nan=True, compression=True):
    """Save an ssbio object as a JSON file using json_tricks"""
    if compression:
        with open(outfile, 'wb') as f:
            dump(obj, f, allow_nan=allow_nan, compression=compression)
    else:
        with open(outfile, 'w') as f:
            dump(obj, f, allow_nan=allow_nan, compression=compression)
    log.info('Saved {} (id: {}) to {}'.format(type(obj), obj.id, outfile))


def load_json(file, new_root_dir=None, decompression=True):
    """Load a JSON file using json_tricks"""
    if decompression:
        with open(file, 'rb') as f:
            my_object = load(f, decompression=decompression)
    else:
        with open(file, 'r') as f:
            my_object = load(f, decompression=decompression)
    if new_root_dir:
        my_object.root_dir = new_root_dir

    return my_object


def save_pickle(obj, outfile, protocol=2):
    """Save the object as a pickle file

    Args:
        outfile (str): Filename
        protocol (int): Pickle protocol to use. Default is 2 to remain compatible with Python 2

    Returns:
        str: Path to pickle file

    """
    with open(outfile, 'wb') as f:
        pickle.dump(obj, f, protocol=protocol)

    return outfile


def load_pickle(file):
    """Load a pickle file.
    
    Args:
        file (str): Path to pickle file

    Returns:
        object: Loaded object from pickle file

    """
    with open(file, 'rb') as f:
       return pickle.load(f)