from __future__ import division, print_function, absolute_import, unicode_literals
'''
Conventional MRC2014 and on-the-fly compressed MRCZ file interface

CCPEM MRC2014 specification:

http://www.ccpem.ac.uk/mrc_format/mrc2014.php

IMOD specification:

http://bio3d.colorado.edu/imod/doc/mrc_format.txt

Testing:

http://emportal.nysbc.org/mrc2014/

Tested output on: Gatan GMS, IMOD, Chimera, Relion, MotionCorr, UnBlur
'''

import os, os.path, sys
import numpy as np
import threading
import struct
try:
    from concurrent.futures import ThreadPoolExecutor
except ImportError as e:
    if sys.version_info > (3,0):
        raise ImportError('Get the backport for `concurrent.futures` for Py2.7 as `pip install futures`')
    raise e
from mrcz.__version__ import __version__

import logging
logger = logging.getLogger('MRCZ')
try: 
    import blosc
    bloscPresent = True
    # For async operations we want to release the GIL in blosc operations and 
    # file IO operations.
    blosc.set_releasegil(True)
except:
    bloscPresent = False
    logger.info( '`blosc` meta-compression library not found, file compression disabled.' )
try: 
    import rapidjson as json
except ImportError:
    import json
    logger.info( '`python-rapidjson` not found, using builtin `json` instead.' )

# Buffer for file I/O
# Quite arbitrary, in bytes (hand-optimized)
BUFFERSIZE = 2**20
BLOSC_BLOCK = 2**16

DEFAULT_HEADER_LEN = 1024

# ENUM dicts for our various Python to MRC constant conversions
COMPRESSOR_ENUM = { 0:None, 1:'blosclz', 2:'lz4', 3:'lz4hc', 4:'snappy', 5:'zlib', 6:'zstd' }
REVERSE_COMPRESSOR_ENUM = { None:0, 'blosclz':1, 'lz4':2, 'lz4hc':2, 'snappy':4, 'zlib':5, 'zstd':6 }

MRC_COMP_RATIO = 1000  
CCPEM_ENUM = { 0: 'i1', 1:'i2', 2:'f4', 4:'c8', 6:'u2', 101:'u1' }
EMAN2_ENUM = { 1: 'i1', 2:'u1', 3:'i2', 4:'u2', 5:'i4', 6:'u4', 7:'f4' }
REVERSE_CCPEM_ENUM = { 'int8':0, 'i1':0, 
                      'uint4':101, 
                      'int16':1, 'i2':1,
                      'uint16':6, 'u2':6,
                      'float64':2, 'f8':2, 'float32':2, 'f4':2,
                      'complex128':4, 'c16':4, 'complex64':4, 'c8':4 }


_asyncExecutor = ThreadPoolExecutor( max_workers = 1 )
def _setAsyncWorkers( N_workers ):
    '''
    _setAsyncWorkers( N_workers )

    *This function is protected as there appears to be little value in using more 
    than one worker. It may be subject to removal in the future.*

    Sets the maximum number of background workers that can be used for reading 
    or writing with the functions.  Defaults to 1. Generally when writing 
    to hard drives use 1 worker. For random-access drives there may be 
    advantages to using multiple workers.
    
    Some test results, 30 files, each 10 x 2048 x 2048 x float-32, on a CPU
    with 4 physical cores:

        HD, 1 worker:   42.0 s
        HD, 2 workers:  50.0 s
        HD, 4 workers:  50.4 s
        SSD, 1 worker:  12.6 s
        SSD, 2 workers: 11.6 s
        SSD, 4 workers: 8.9 s
        SSD, 8 workers: 16.9 s

    In g
    '''
    if N_workers <= 0:
        raise ValueError('N_workers must be greater than 0')
    if _asyncExecutor._max_workers == N_workers:
        return
    _asyncExecutor._max_workers = N_workers
    _asyncExecutor._adjust_thread_count()

def defaultHeader():
    '''
    Returns a default MRC header dictionary with all fields with default values.
    '''
    header = {}
    header['fileConvention'] = 'ccpem'
    header['endian'] = 'le'
    header['MRCtype'] = 0
    header['dimensions'] = np.array( [0,0,0], dtype=int )
    header['dtype'] = 'u1'
    
    header['compressor'] = None
    header['packedBytes'] = 0
    header['clevel'] = 1
    
    header['maxImage'] = 1.0
    header['minImage'] = 0.0
    header['meanImage'] = 0.0
    
    header['pixelsize'] = 0.1 
    header['pixelunits'] = u'nm' # Can be '\\AA' for Angstroms
    header['voltage'] = 300.0 # kV
    header['C3'] = 2.7 # mm
    header['gain'] = 1.0 # counts/electron
    
    if bloscPresent:
        header['n_threads'] = blosc.detect_number_of_cores()
    
    return header
    
    
def readMRC( MRCfilename, idx=None, endian='le', 
              pixelunits=u'\\AA', fileConvention='ccpem', 
              useMemmap=False, n_threads=None, asList=False ):
    '''
    readMRC( MRCfilename, idx=None, endian='le', 
              pixelunits=u'\\AA', fileConvention='ccpem', 
              useMemmap=False, n_threads=None, asList=False )
    
    Imports an MRC/Z file as a NumPy array and a meta-data dict.  
    
    Usage::
        
        [image, meta] = readMRC( MRCfilename, idx=None, endian='le', 
              pixelunits=u'\\AA', fileConvention = 'ccpem', 
              useMemmap=False, n_threads = None, asList=False  ):
    
    * ``image`` is a NumPy array.
    * ``meta`` is a dict with various fields relating to the MRC header information 
      and extended metadata.
    * ``idx = (first, last)`` is a tuple with first (inclusive) and last (not 
      inclusive) indices of images to be read from the stack. Index of first image 
      is 0. Negative indices can be used to count backwards. A singleton integer 
      can be provided to read only one image. If omitted, will read whole file. 
      Compression is currently not supported with this option.
    * ``pixelunits`` can be ``'AA' (Angstoms), 'nm', '\mum', or 'pm'``.  Internally pixel 
      sizes are always encoded in Angstroms in the MRC file.
    * ``fileConvention`` can be ``'ccpem'`` (equivalent to IMOD) or ``'eman2'``, which is
      only partially supported at present.
    * ``endian`` can be big-endian as ``'be'`` or little-endian as ``'le'``. Defaults 
      to `'le'` as the vast majority of modern computers are little-endian.
    * ``n_threads`` is the number of threads to use for decompression, defaults to 
       use all virtual cores.
    * ``useMemmap=True`` returns a ``numpy.memmap`` instead of a ``numpy.ndarray``
    * ``asList=True`` returns a `list` of 2D NumPy arrays, with the list indices 
      representing the Z-axis. Since the data is non-contigious, very large 
      arrays can fit more easily into memory. Ignored if ``useMemmap=True``.


    '''

    with open(MRCfilename, 'rb', buffering=BUFFERSIZE) as f:
        # Read in header as a dict
        
        header = readMRCHeader(MRCfilename, endian=endian, fileConvention = fileConvention, pixelunits=pixelunits)
        # Support for compressed data in MRCZ

        
        if ( (header['compressor'] in REVERSE_COMPRESSOR_ENUM) 
            and (REVERSE_COMPRESSOR_ENUM[header['compressor']] > 0) 
            and idx == None ):
            return __MRCZImport(f, header, endian=endian, fileConvention=fileConvention, 
                                n_threads=n_threads, asList=asList)
        # Else load as uncompressed MRC file

        if idx != None:
        # If specific images were requested:
        # TO DO: add support to read all images within a range at once

            if header['compressor'] != None:
                raise RuntimeError( 'Reading from arbitrary positions not supported for compressed files. Compressor = %s'%header['compressor'] )
            if np.isscalar( idx ):
                indices = np.array([idx, idx], dtype='int')
            else:
                indices = np.array(idx, dtype='int')

            # Convert to old way:
            idx = indices[0]
            n = indices[1] - indices[0] + 1

            if idx < 0:
                # Convert negative index to equivalent positive index:
                idx = header['dimensions'][0] + idx

            # Just check if the desired image is within the stack range:
            if idx < 0 or idx >= header['dimensions'][0]:
                raise ValueError( 'Error: image or slice index out of range. idx = %d, z_dimension = %d'%(idx, header['dimensions'][0]) )
            elif idx + n > header['dimensions'][0]:
            	raise ValueError( 'Error: image or slice index out of range. idx + n = %d, z_dimension = %d'%(idx + n, header['dimensions'][0]) )
            elif n < 1:
            	raise ValueError( 'Error: n must be >= 1. n = %d'%n )
            else:
                # We adjust the dimensions of the returned image in the header:
                header['dimensions'][0] = n

                # This offset will be applied to f.seek():
                offset = idx * np.product(header['dimensions'][1:])*np.dtype(header['dtype']).itemsize

        else:
            offset = 0

        f.seek(DEFAULT_HEADER_LEN + header['extendedBytes'] + offset)
        if bool(useMemmap):
            image = np.memmap( f, dtype=header['dtype'], 
                              mode='c', 
                              shape=tuple(dim for dim in header['dimensions']) )
        else: # Load entire file into memory
            dims = header['dimensions']
            if asList: # List of NumPy 2D-arrays
                XYsize = np.product(dims[1:])
                dtype=header['dtype']

                # np.fromfile advances the file pointer `f` for us.
                image = [ np.squeeze(np.fromfile(f, dtype=dtype, count=XYsize).reshape(dims[1:])) for imSlice in range(dims[0]) ]
               
            else: # NumPy 3D-array
                image = np.fromfile(f, dtype=header['dtype'], count=np.product(dims))
                if header['MRCtype'] == 101:
                    # Seems the 4-bit is interlaced ...
                    interlaced_image = image
                    
                    image = np.empty( np.product(dims), dtype=header['dtype'] )

                    image[0::2] = np.left_shift(interlaced_image,4) / 15
                    image[1::2] = np.right_shift(interlaced_image,4)

                image = np.squeeze( image.reshape(dims) )



        return image, header

       
def __MRCZImport( f, header, endian='le', fileConvention='ccpem', 
        returnHeader=False, n_threads=None, asList=False ):
    '''
    Equivalent to MRCImport, but for compressed data using the blosc library.
    
    The following compressors are supported: 
        ``'zlib'``
        ``'zstd'``
        ``'lz4'``
    
    Memory mapping is not possible in this case at present.  
    '''
    if not bloscPresent:
        logger.error( 'blosc not present, cannot compress files.' )
        return
        
    if n_threads == None:
        blosc.nthreads = blosc.detect_number_of_cores()
    else:
        blosc.nthreads = n_threads

    dims = header['dimensions']
    dtype = header['dtype']
    if asList:
        image = []
    else:
        image = np.empty(dims, dtype=dtype )
    
    blosc_chunk_pos = DEFAULT_HEADER_LEN + header['extendedBytes']
    for J in range(dims[0]):
        f.seek( blosc_chunk_pos )
        ( (nbytes, blockSize, ctbytes ), (ver_info) ) = readBloscHeader(f)
        f.seek(blosc_chunk_pos)
        # blosc includes the 16 header bytes in ctbytes
        if asList:
            image.append( np.squeeze(
                np.frombuffer(blosc.decompress(f.read(ctbytes)), 
                                dtype=dtype).reshape(dims[1:])
            ))
        else:
            image[J,:,:] = np.frombuffer(blosc.decompress(f.read(ctbytes)), 
                                dtype=dtype).reshape(dims[1:])
            
        blosc_chunk_pos += (ctbytes)
        pass
    
    
    if header['MRCtype'] == 101:
        # Seems the 4-bit is interlaced 
        if asList:
            raise NotImplementedError('MRC type 101 (uint4) not supported with return as `list`')
        interlaced_image = image
            
        image = np.empty( np.product(header['dimensions']), dtype=dtype )
        # Bit-shift and Bit-and to seperate decimated pixels
        image[0::2] = np.left_shift(interlaced_image,4) / 15
        image[1::2] = np.right_shift(interlaced_image,4)

    if not asList:
        image = np.squeeze( image )
    
    return image, header
    

    
def readBloscHeader( filehandle ):
    '''
    Reads in the 16 byte header file from a blosc chunk. Blosc header format 
    for each chunk is as follows::
    
        |-0-|-1-|-2-|-3-|-4-|-5-|-6-|-7-|-8-|-9-|-A-|-B-|-C-|-D-|-E-|-F-|
        ^   ^   ^   ^ |     nbytes    |   blocksize   |    ctbytes    |
        |   |   |   |
        |   |   |   +--typesize
        |   |   +------flags
        |   +----------versionlz
        +--------------version
    '''
    # Probably 
    [version, versionlz, flags, typesize] = np.fromfile( filehandle, dtype='uint8', count=4 )
    [nbytes, blocksize, ctbytes] = np.fromfile( filehandle, dtype='uint32', count=3 )
    return ( [nbytes, blocksize, ctbytes], [version, versionlz, flags, typesize] )
    
    
def readMRCHeader( MRCfilename, endian='le', fileConvention = 'ccpem', pixelunits=u'\\AA' ):
    '''
    Reads in the first 1024 bytes from an MRC file and parses it into a Python dictionary, yielding 
    header information.
    '''
    if endian == 'le':
        endchar = '<'
    else:
        endchar = '>'
        
    header = {}
    with open( MRCfilename, 'rb' ) as f:
        # diagStr = ''
        # Get dimensions, in format [nz, ny, nx] (stored as [nx,ny,nz] in the file)
        header['dimensions'] = np.flipud( np.fromfile( f, dtype=endchar+'i4', count=3 ) )
    
        header['MRCtype'] = int( np.fromfile( f, dtype=endchar+'i4', count=1 )[0] )
        # Hack to fix lack of standard endian indication in the file header
        if header['MRCtype'] > 16000000: 
            # Endianess found to be backward
            header['MRCtype'] = int( np.asarray( header['MRCtype'] ).byteswap()[0] )
            header['dimensions'] = header['dimensions'].byteswap()
            if endchar == '<':
                endchar = '>' 
            else:
                endchar = '<'
        
        # Extract compressor from dtype > MRC_COMP_RATIO
        header['compressor'] = COMPRESSOR_ENUM[ np.floor_divide(header['MRCtype'], MRC_COMP_RATIO) ]
        header['MRCtype'] = np.mod( header['MRCtype'], MRC_COMP_RATIO )
        logger.info( 'compressor: %s, MRCtype: %s' % (str(header['compressor']),str(header['MRCtype'])) )
        
        fileConvention = fileConvention.lower()
        # if fileConvention == 'ccpem':
        #     diagStr += ('ioMRC.readMRCHeader: MRCtype: %s, compressor: %s, dimensions %s' % 
        #         (CCPEM_ENUM[header['MRCtype']],header['compressor'], header['dimensions'] ) )
        # elif fileConvention == 'eman2':
        #     diagStr += ( 'ioMRC.readMRCHeader: MRCtype: %s, compressor: %s, dimensions %s' % 
        #         (EMAN2_ENUM[header['MRCtype']],header['compressor'], header['dimensions'] ) )

        
        if fileConvention == 'eman2':
            try:
                header['dtype'] = EMAN2_ENUM[ header['MRCtype'] ]
            except:
                raise ValueError( 'Error: unrecognized EMAN2-MRC data type = ' + str(header['MRCtype']) )
                
        elif fileConvention == 'ccpem': # Default is CCPEM
            try:
                header['dtype'] = CCPEM_ENUM[ header['MRCtype'] ]
            except:
                raise ValueError( 'Error: unrecognized CCPEM-MRC data type = ' + str(header['MRCtype']) )
        else:
            raise ValueError( 'Error: unrecognized MRC file convention: {}'.format(fileConvention) )
                
        # Apply endian-ness to NumPy dtype
        header['dtype'] = endchar + header['dtype']
        # Read in pixelsize
        f.seek(40)
        cellsize = np.fromfile( f, dtype=endchar + 'f4', count=3 )
        header['pixelsize'] = np.flipud( cellsize ) / header['dimensions']
        # MRC is Angstroms by convention

        header['pixelunits'] = pixelunits
            
        # '\AA' will eventually be deprecated, please cease using it.
        if header['pixelunits'] == u'\\AA' or header['pixelunits']==u'\AA':
            pass
        elif header['pixelunits'] == u'\mum':
            header['pixelsize'] *= 1E-5
        elif header['pixelunits'] == u'pm':
            header['pixelsize'] *= 100.0
        else: # Default to nm
            header['pixelsize'] *= 0.1
            
        # Read in [X,Y,Z] array ordering
        # Currently I don't use this
        # f.seek(64)
        # axesTranpose = np.fromfile( f, dtype=endchar + 'i4', count=3 ) - 1
        
        # Read in statistics
        f.seek(76)
        (header['minImage'], header['maxImage'], header['meanImage']) = np.fromfile( f, dtype=endchar + 'f4', count=3 )

        f.seek(92)
        header['extendedBytes'] = int( np.fromfile( f, dtype=endchar + 'i4', count=1) )
        if header['extendedBytes'] > 0:
            # diagStr += ', extended header %d' % header['extendedBytes']

            f.seek(104)
            header['metaId'] = f.read(4)
            if header['metaId'] == b'json':
                f.seek(DEFAULT_HEADER_LEN)
                header.update( json.loads( f.read(header['extendedBytes'] ).decode('utf-8') ) )

        # Read in kV, C3, and gain
        f.seek(132)
        header['voltage']  = np.fromfile( f, dtype=endchar + 'f4', count=1 )
        header['C3']  = np.fromfile( f, dtype=endchar + 'f4', count=1 )
        header['gain']  = np.fromfile( f, dtype=endchar + 'f4', count=1 )
        
        #diagStr += ', voltage: %.1f, C3: %.2f, gain: %.2f' % (header['voltage'], header['C3'], header['gain']) 

        # Read in size of packed data
        f.seek(144)
        # Have to convert to Python int to avoid index warning.
        header['packedBytes'] = struct.unpack( 'q', f.read(8) )
        # header['packedBytes'] = int( np.fromfile( f, dtype=endchar + 'i8', count=1) )
        # if header['packedBytes'] > 0:
        #     diagStr += ', packedBytes: %d' % header['packedBytes']
        
        

        # How many bytes in an MRC
        return header
        
def writeMRC( input_image, MRCfilename, meta=None, endian='le', dtype=None, 
               pixelsize=[0.1,0.1,0.1], pixelunits=u'\\AA', shape=None, 
               voltage=0.0, C3=0.0, gain=1.0,
               compressor=None, clevel = 1, n_threads=None, quickStats=True, idx = None ):
    '''
    writeMRC( input_image, MRCfilename, meta=None, idx=None, 
               endian='le', dtype=None, 
               pixelsize=[0.1,0.1,0.1], pixelunits=u'\\AA', shape=None, 
               voltage=0.0, C3=0.0, gain=1.0,
               compressor=None, clevel=1, n_threads=None, quickStats=True, 
               )
    
    Given a ``numpy`` 2-D or 3-D array ``input_image`` write it has an MRC file ``MRCfilename``.

    * ``meta`` is a Python dict{} which will be serialized by JSON and written 
        into the extended header.
    * ``dtype`` will cast the data before writing it.
    * ``pixelsize`` is [z,y,x] pixel size (singleton values are ok for square/cubic pixels)
    * ``pixelunits`` is ``'\\AA'`` for Angstroms, ``'pm'`` for picometers, ``'\mum'`` for micrometers, 
      or ``'nm'`` for nanometers.  MRC standard is always Angstroms, so pixelsize 
      is converted internally from nm to Angstroms if necessary
    * ``shape`` is only used if you want to later append to the file, such as 
      merging together Relion particles for Frealign.  Not recommended and 
      only present for legacy reasons.
    * ``voltage`` is accelerating potential in keV, defaults to 300.0
    * ``C3`` is spherical aberration in mm, defaults to 2.7 mm
    * ``gain`` is detector gain (counts/primary electron), defaults to 1.0 (for counting camera)
    * ``compressor`` is a choice of ``'lz4'``, ``'zlib'``, or ``'zstd'``, plus ``'blosclz'``, ``'lz4hc'`` 
        - ``'lz4'`` is  generally the fastest.
        - ``'zstd'`` generally gives the best compression performance, and is still almost 
        as fast as 'lz4' with clevel = 1
    * ``clevel`` is the compression level, 1 is fastest, 9 is slowest.  The compression
      ratio will rise slowly with clevel.
    * ``n_threads`` is number of threads to use for blosc compression
    * ``quickStats=True`` estimates the image mean, min, max from the first frame only,
      which saves a lot of computational time for stacks.
    * ``idx`` can be used to write an image or set of images starting at a 
      specific position in the MRC file (which may already exist). Index of 
      first image is 0. A negative index can be used to count backwards. If 
      omitted, will write whole stack to file. If writing to an existing 
      file, compression or extended MRC2014 headers are currently not 
      supported with this option.
    
    *Note: MRC definitions are not consistent. Generally we support the CCPEM schema.*
    '''



    # For dask, we don't want to import dask, but we can still work-around how to 
    # check its type without isinstance()
    image_type = type(input_image)
    if image_type.__module__ == 'dask.array.core' and image_type.__name__ == 'Array':
        # Ideally it would be faster to iterate over the chunks and pass each one 
        # to blosc but that likely requires c-blosc2
        input_image = input_image.__array__()
        dims = input_image.shape

    if isinstance(input_image, (tuple,list)):
        asList = True
        if input_image[0].ndim != 2:
            raise ValueError( "List of arrays must have 2D arrays as elements" )
        dims = np.array( [len(input_image), input_image[0].shape[0], input_image[0].shape[1]])

        # Verify that each image in the list is the same 2D shape and dtype
        for z_slice in input_image:
            assert( np.all(z_slice.shape == dims[1:]) )
            assert( z_slice.dtype == input_image[0].dtype )

    else: # Array-'like' object
        asList = False
        dims = input_image.shape
        if len( input_image.shape ) == 2:
            # If it's a 2D image we force it to 3D - this makes life easier later:
            input_image = input_image.reshape( ( 1, input_image.shape[0], input_image.shape[1] ) )

    # We will need this regardless if writing to an existing file or not:
    if endian == 'le':
        endchar = '<'
    else:
        endchar = '>'

    # We now check if we have to create a new header (i.e. new file) or not. If 
    # the file exists, but idx is 'None', it will be replaced by a new file 
    # with new header anyway:
    if os.path.isfile( MRCfilename ):
        if idx == None:
            idxnewfile = True
        else:
            idxnewfile = False
    else:
        idxnewfile = True

    
    if idxnewfile:
        if dtype == 'uint4' and compressor != None:
            raise TypeError( 'uint4 packing is not compatible with compression, use int8 datatype.' )
            
        header = {'meta': meta }
        if dtype == None:
            if asList:
                header['dtype'] = endchar + input_image[0].dtype.descr[0][1].strip( '<>|' )
            else:
                header['dtype'] = endchar + input_image.dtype.descr[0][1].strip( '<>|' )
        else:
            header['dtype'] = dtype
            
        # Now we need to filter dtype to make sure it's actually acceptable to MRC
        if not header['dtype'].strip( '<>|' ) in REVERSE_CCPEM_ENUM:
            raise TypeError( 'ioMRC.MRCExport: Unsupported dtype cast for MRC %s' % header['dtype'] )
            
        header['dimensions'] = dims
        
        header['pixelsize'] = pixelsize
        header['pixelunits'] = pixelunits
        header['compressor'] = compressor
        header['clevel'] = clevel
        header['shape'] = shape
        
        # This overhead calculation is annoying but many 3rd party tools that use 
        # MRC require these statistical parameters.
        if bool(quickStats):
            if asList:
                first_image = input_image[0]
            else:
                first_image = input_image[0,:,:]

            imMin = first_image.real.min(); imMax = first_image.real.max()
            header['maxImage'] = imMax
            header['minImage'] =  imMin
            header['meanImage'] = 0.5*(imMax + imMin)
        else:
            if asList:
                header['maxImage'] = np.max( [z_slice.real.max() for z_slice in input_image] )
                header['minImage'] = np.min( [z_slice.real.min() for z_slice in input_image] )
                header['meanImage'] = np.mean( [z_slice.real.mean() for z_slice in input_image] )
            else:
                header['maxImage'] = input_image.real.max()
                header['minImage'] = input_image.real.min()
                header['meanImage'] = input_image.real.mean()
        
        header['voltage'] = voltage
        if not bool( header['voltage'] ):
            header['voltage'] = 0.0
        header['C3'] = C3
        if not bool( header['C3'] ):
            header['C3'] = 0.0
        header['gain'] = gain
        if not bool( header['gain'] ):
            header['gain'] = 1.0
        
        header['compressor'] = compressor
        header['clevel'] = clevel
        if n_threads == None and bloscPresent:
            n_threads = blosc.detect_number_of_cores()
        header['n_threads'] = n_threads
        
        # TODO: can we detect the number of cores without adding a heavy dependancy?
        
        if dtype == 'uint4':
            if asList:
                raise NotImplementedError('Saving of lists of arrays not supported for `dtype=uint4`')
            # Decimate to packed 4-bit
            input_image = input_image.astype('uint8')
            input_image = input_image[:,:,::2] + np.left_shift(input_image[:,:,1::2],4)

    else: # We are going to append to an already existing file:
        # So we try to figure out its header with 'CCPEM' or 'eman2' file conventions:
        try:
            header = readMRCHeader( MRCfilename, endian, fileConvention = 'CCPEM', pixelunits=pixelunits )

        except ValueError:
            try:
                header = readMRCHeader( MRCfilename, endian, fileConvention = 'eman2', pixelunits=pixelunits )
            except ValueError:
            # If neither 'CCPEM' nor 'eman2' formats satisfy:
                raise ValueError( 'Error: unrecognized MRC type for file: %s ' % MRCfilename )

        # No support for extended headers in arbitrary appending mode:
        # RAM: should work now
        # if header['extendedBytes'] > 0:
        #     raise ValueError( 'Error: MRC2014 files with extended headers not supported for writing: %s = %d' % ('extendedBytes', header['extendedBytes'] ) )

        # If the file already exists, its X,Y dimensions must be consistent with the current image to be written:
        if np.any( header['dimensions'][1:] != input_image.shape[1:] ):
            raise ValueError( 'Error: x,y dimensions of image do not match that of MRC file: %s ' % MRCfilename )
            # TO DO: check also consistency of dtype?

        if 'meta' not in header.keys():
            header['meta'] = meta

    # Now that we have a proper header, we go into the details of writing to a specific position:
    if idx != None:
        if header['compressor'] != None:
            raise RuntimeError( 'Writing at arbitrary positions not supported for compressed files. Compressor = %s' % header['compressor'] )

        idx = int(idx)
        # Force 2D to 3D dimensions:
        if len( header['dimensions'] ) == 2:
            header['dimensions'] = np.array( [1, header['dimensions'][0], header['dimensions'][1]] )

        # Convert negative index to equivalent positive index:
        if idx < 0:
            idx = header['dimensions'][0] + idx

        # Just check if the desired image is within the stack range:
        # In principle we could write to a position beyond the limits of the file (missing slots would be filled with zeros), but let's avoid that the user writes a big file with zeros by mistake. So only positions within or immediately consecutive to the stack are allowed:
        if idx < 0 or idx > header['dimensions'][0]:
            raise ValueError( 'Error: image or slice index out of range. idx = %d, z_dimension = %d' % (idx, header['dimensions'][0]) )

        # The new Z dimension may be larger than that of the existing file, or even of the new file, if an index larger than the current stack is specified:
        newZ = idx + input_image.shape[0]
        if newZ > header['dimensions'][0]:
            header['dimensions'] = np.array( [idx + input_image.shape[0], header['dimensions'][1], header['dimensions'][2]] )

        # This offset will be applied to f.seek():
        offset = idx * np.product( header['dimensions'][1:] ) * np.dtype( header['dtype'] ).itemsize

    else:
        offset = 0
        
    __MRCExport(input_image, header, MRCfilename, endchar, offset, idxnewfile, asList=asList)
 
        
def __MRCExport(input_image, header, MRCfilename, endchar='<', offset=0, idxnewfile=True, asList=False):
    '''
    MRCExport private interface with a dictionary rather than a mess of function 
    arguments.
    '''

    if idxnewfile: # If forcing a new file we truncate it even if it already exists:
        fmode = 'wb'
    else: # Otherwise we'll just update its header and append images as required:
        fmode = 'rb+'

    with open(MRCfilename, fmode, buffering=BUFFERSIZE) as f:
        extendedBytes = writeMRCHeader(f, header, endchar)
        f.seek(DEFAULT_HEADER_LEN + extendedBytes + offset)

        dtype = header['dtype']
        if ('compressor' in header) \
                and (header['compressor'] in REVERSE_COMPRESSOR_ENUM) \
                and (REVERSE_COMPRESSOR_ENUM[header['compressor']]) > 0:
            # compressed MRCZ
            logger.info( 'Compressing %s with compressor %s%d' %
                    (MRCfilename, header['compressor'], header['clevel'] ) )
            

            applyCast = False
            if asList:
                chunkSize = input_image[0].size
                typeSize = input_image[0].dtype.itemsize
                if dtype != 'uint4' and input_image[0].dtype != dtype: applyCast = True
            else:
                chunkSize = input_image[0,:,:].size
                typeSize = input_image.dtype.itemsize
                if dtype != 'uint4' and input_image.dtype != dtype: applyCast = True
                
            blosc.set_nthreads( header['n_threads'] )
            blosc.set_blocksize( BLOSC_BLOCK )
            
            header['packedBytes'] = 0
            
            zCount = len(input_image) if asList else input_image.shape[0]
            for J in range( zCount ):
                # print( 'Slice %d: Compressing address at: %d of %d:' % (J, int(J*typeSize*blockSize), input_image.nbytes) )
                
                # # Looks like I have problem for typesize > 1?
                # if int(J*typeSize*chunkSize) >= input_image.nbytes:
                #     raise MemoryError( 'MRCExport: Tried to reference past end of ndarray %d > %d' % (int(J*typeSize*chunkSize), input_image.nbytes ) )
                    
                if asList:
                    if applyCast:
                        compressedData = blosc.compress( input_image[J].astype(dtype).tobytes(),
                                    typeSize, 
                                    clevel=header['clevel'], 
                                    shuffle=blosc.BITSHUFFLE,
                                    cname=header['compressor'] )
                    else:
                        compressedData = blosc.compress( input_image[J].tobytes(),
                                    typeSize, 
                                    clevel=header['clevel'], 
                                    shuffle=blosc.BITSHUFFLE,
                                    cname=header['compressor'] )
                else: # Array-like
                    if applyCast:
                        compressedData = blosc.compress( input_image[J,:,:].astype(dtype).tobytes(),
                                    typeSize, 
                                    clevel=header['clevel'], 
                                    shuffle=blosc.BITSHUFFLE,
                                    cname=header['compressor'] )
                    else:
                        compressedData = blosc.compress( input_image[J,:,:].tobytes(),
                                    typeSize, 
                                    clevel=header['clevel'], 
                                    shuffle=blosc.BITSHUFFLE,
                                    cname=header['compressor'] )
                f.write( compressedData )
                    
                header['packedBytes'] += len(compressedData)
                
            # Rewind and write out the total compressed size
            f.seek(144)
            np.int64( header['packedBytes'] ).astype( endchar + 'i8' ).tofile(f)

            
        else: # vanilla MRC
            
            if asList: 
                if dtype != 'uint4' and dtype != input_image.dtype:
                    [z_slice.astype(dtype).tofile(f) for z_slice in input_image]
                else:
                    [z_slice.tofile(f) for z_slice in input_image]
            else:
                if dtype != 'uint4' and dtype != input_image.dtype:
                    input_image = input_image.astype(dtype)
                input_image.tofile(f)
            
            
    return 
    
def writeMRCHeader( f, header, endchar = '<' ):
    '''
    writeMRCHeader( f, header, endchar = '<' )

    Usage::

        writeMRCHeader( f, header )
        
    Writes a  header to the file-like object ``f``, requires a dict 
    called ``header`` to parse the appropriate fields. Use ``defaultHeader()`` to 
    retrieve an example with all potential fields.
    '''
        
    f.seek(0)
    # Write dimensions
    if len(header['dimensions']) == 2: # force to 3-D
        dimensions = np.array( [1, header['dimensions'][0], header['dimensions'][1]] )
    else: 
        dimensions = np.array( header['dimensions'] )
        
    # Flip to Fortran order
    dimensions = np.flipud( dimensions )
    dimensions.astype(endchar+'i4').tofile( f )
    
    # 64-bit floats are automatically down-cast
    dtype = header['dtype'].lower().strip( '<>|' )
    try:
        MRCmode = np.int32( REVERSE_CCPEM_ENUM[dtype] ).astype( endchar + 'i4' )
    except:
        raise ValueError( 'Warning: Unknown dtype for MRC encountered = ' + str(dtype) )
        
    # Add 1000 * COMPRESSOR_ENUM to the dtype for compressed data
    if ('compressor' in header 
                and header['compressor'] in REVERSE_COMPRESSOR_ENUM 
                and REVERSE_COMPRESSOR_ENUM[header['compressor']] > 0):
        header['compressor'] = header['compressor'].lower()
        MRCmode += MRC_COMP_RATIO * REVERSE_COMPRESSOR_ENUM[ header['compressor'] ]
        
        # How many bytes in an MRCZ file, so that the file can be appended-to.
        try:
            f.seek(144)
            np.int32( header['packedBytes'] ).astype( endchar + 'i8' ).tofile(f)
        except:
            # This is written afterward so we don't try to keep the entire compressed file in RAM
            pass
            
        
    f.seek(12)
    MRCmode.tofile(f)
    
    # Print NXSTART, NYSTART, NZSTART
    np.array( [0,0,0], dtype=endchar+'i4' ).tofile(f)
    # Print MX, MY, MZ, the number of pixels
    np.array( dimensions, dtype=endchar+'i4' ).tofile(f)
    # Print cellsize = pixelsize * dimensions
    # '\AA' will eventually be deprecated, please cease using it.
    if header['pixelunits'] == '\\AA' or header['pixelunits'] == '\AA':
        AApixelsize = np.array(header['pixelsize'])
    elif header['pixelunits'] == '\mum':
        AApixelsize = np.array(header['pixelsize'])*10000.0
    elif header['pixelunits'] == 'pm':
        AApixelsize = np.array(header['pixelsize'])/100.0
    else: # Default is nm
        AApixelsize = np.array(header['pixelsize'])*10.0   
    
    # The above AApixelsize insures we cannot have an array of len=1 here
    if isinstance( AApixelsize, np.ndarray ) and AApixelsize.size == 1:
        cellsize = np.array( [AApixelsize,AApixelsize,AApixelsize]  ) * dimensions
    elif not isinstance( AApixelsize, (list,tuple,np.ndarray) ):
        cellsize = np.array( [AApixelsize,AApixelsize,AApixelsize]  ) * dimensions
    elif len(AApixelsize) == 2:
        # Default to z-axis pixelsize of 10.0 Angstroms
        cellsize = np.flipud(np.array( [10.0, AApixelsize[0], AApixelsize[1]] )) * dimensions
    else:
        cellsize = np.flipud(np.array( AApixelsize )) *  dimensions
        
    f.seek(40)
    np.array( cellsize, dtype=endchar+'f4' ).tofile(f)
    # Print default cell angles
    np.array( [90.0,90.0,90.0], dtype=endchar+'f4' ).tofile(f)
    
    # Print axis associations (we use C ordering internally in all Python code)
    np.array( [1,2,3], dtype=endchar+'i4').tofile(f)
    
    # Print statistics (if available)
    f.seek(76)
    if 'minImage' in header:
        np.float32( header['minImage'] ).astype(endchar+'f4').tofile(f)
    else:
        np.float32( 0.0 ).astype(endchar+'f4').tofile(f)
    if 'maxImage' in header:
        np.float32( header['maxImage'] ).astype(endchar+'f4').tofile(f)
    else:
        np.float32( 1.0 ).astype(endchar+'f4').tofile(f)
    if 'meanImage' in header:
        np.float32( header['meanImage'] ).astype(endchar+'f4').tofile(f)
    else:
        np.float32( 0.0 ).astype(endchar+'f4').tofile(f)
        
    # We'll put the compressor info and number of compressed bytes in 132-204
    # and new metadata
    # RESERVED: 132: 136: 140 : 144 for voltage, C3, and gain
    f.seek(132)
    if 'voltage' in header:
        np.float32( header['voltage'] ).astype(endchar+'f4').tofile(f)
    if 'C3' in header:
         np.float32( header['C3'] ).astype(endchar+'f4').tofile(f)
    if 'gain' in header:
        np.float32( header['gain'] ).astype(endchar+'f4').tofile(f)
        

    # Magic MAP_ indicator that tells us this is in-fact an MRC file
    f.seek( 208 )
    f.write( b'MAP ' )
    # Write a machine stamp, '17,17' for big-endian or '68,65'
    # Note that the MRC format doesn't indicate the endianness of the endian 
    # identifier............................................................
    f.seek( 212 )
    if endchar == '<':
        f.write( struct.pack( b'BB', 68, 65 ) )
        # np.array( [68,65], dtype='uint8' ).tofile(f)
    else:
        f.write( struct.pack( b'BB', 17, 17 ) )
        #np.array( [17,17], dtype='uint8' ).tofile(f)

    # Write b'MRCZ<version>' into labels
    f.seek(220)
    f.write( struct.pack( b'i', 1) ) # We have one label
    f.write( b'MRCZ' + __version__.encode('ascii') )
    
    # Extended header, if meta is not None
    if isinstance(header['meta'], dict):
        # Encode metadata as bytes with UTF-8
        for key, value in header['meta'].items():
            # Sanitize types that `python-rapidjson` and `json` don't understand.
            # Should we make a copy of `meta` in this case?  We are modifying 
            # the passed data.
            if type(value) == np.ndarray:
                header['meta'][key] = value.tolist()
                
        jsonMeta = json.dumps( header['meta'] ).encode('utf-8')
        jsonLen = len(jsonMeta)
        # Length of extended header
        f.seek(92)
        f.write( struct.pack( endchar+'i', jsonLen ) )

        # 4-byte char ID string of extended metadata type
        f.seek(104)
        f.write( b'json' )

        # Go to the extended header
        f.seek(DEFAULT_HEADER_LEN)
        f.write( jsonMeta )
        return jsonLen
    
    # No extended header
    return 0

def asyncReadMRC( *args, **kwargs ):
    '''
    asyncReadMRC( *args, **kwargs )

    Returns a ``concurrent.futures.Future`` object, call the method ``result()`` to get the image 
    and metadata.

    Usage:: 

        worker = asyncReadMRC( 'someones_file.mrc' )
        # Do some work
        mrcImage, mrcMeta = worker.result()

    Valid arguments are as for ``readMRC()``.
    '''
    return _asyncExecutor.submit( readMRC, *args, **kwargs )

def asyncWriteMRC( *args, **kwargs ):
    '''
    asyncWriteMRC( *args, **kwargs )

    Calls writeMRC in a seperate thread and executes it in the background. Returns the thread 
    object, so if necessary you can call ``concurrent.futures.Future.result()`` to wait for the write to finish, 
    or check with ``concurrent.futures.Future.done()``.  Usage::

        worker = asyncWriteMRC( npImageData, 'my_mrcfile.mrc' )
        # Do some work
        if not worker.done():
            time.sleep(0.001)
        # File is written to disk

    Most of the time you can ignore the return and let the system finish writing 
    when it finishes.  
 
    Valid arguments are as for writeMRC().
    '''
    return _asyncExecutor.submit( writeMRC, *args, **kwargs )