import compress_pickle as cp
import numpy as np
import os

def convert_bytes(num):
    """
    this function will convert bytes to MB.... GB... etc
    """
    for x in ['bytes', 'KB', 'MB', 'GB', 'TB']:
        if num < 1024.0:
            return "%3.1f %s" % (num, x)
        num /= 1024.0


def file_size(file_path):
    """
    this function will return the file size
    """
    if os.path.isfile(file_path):
        file_info = os.stat(file_path)
        return convert_bytes(file_info.st_size)
    

if __name__ == '__main__':
    # Create a random array
    uv_mask = np.load('custom_utils/mask_256x512.npy')
    arr = np.stack([np.random.rand(512, 256, 3) * uv_mask[:,:,None] for _ in range(5)])
    cp.dump(arr, 'temp/test_pickle', 'bz2')
    np.savez_compressed('temp/test_numpy.npz', arr)

    print('Pickle file size: ', file_size('temp/test_pickle.' + 'bz'))
    print('Numpy file size: ', file_size('temp/test_numpy.npz'))