import warnings
def __warning_format(message, category, filename, lineno, file=None, line=None):
        return '%s:%s: %s:%s\n' % (filename, lineno, category.__name__, message)
warnings.formatwarning = __warning_format

class UCF:
    """UCF low-level access class"""
    def __init__(self,file=None,verbosity=False,debug=False):
        self.__initializeConstants()
        self.__resetPublicProperties()
        self.__resetPrivateProperties()
        self.__resetCurrentStep()
        self.__resetCurrentSet()
        if file is not None:
            self.open(file)
        self.Debug = debug
        self.Verbosity = verbosity
        
    def open(self,file):
        """Opens an input stream for reading access. The variable 'file' can be of the following types:
           str                    opens a file on disk whose path is specified by 'file'
           bytes                  data which is already located in memory as a bytes or bytearray object (can be used for streams)
           """
        from io import BytesIO
        import numpy
        # Check what class 'file' belongs to and treat it accordingly
        if isinstance(file,str):
            self.File = file
            self.__fileobj = open(self.File,'rb')
            self.__inputAvailable = True
        elif isinstance(file,bytes):
            self.File = 'byte-stream'
            self.__fileobj = BytesIO(file)
            self.__inputAvailable = True
        # Determine file size
        self.__fileobj.seek(0,2)
        self.FileSize = self.__fileobj.tell()
        self.__fileobj.seek(self.__fileBeg,0)
        # Read the header of the file
        self.__readHeaderFile()
        # Scan through file to get the basic structure (steps/sets)
        self.__timeStep      = numpy.zeros(self.__scanBuffSize,dtype=numpy.float64)
        self.__posStep       = numpy.zeros(self.__scanBuffSize,dtype=numpy.int64)
        self.__numSetPerStep = numpy.zeros(self.__scanBuffSize,dtype=numpy.int32)
        istep = 0;
        while self.__fileobj.tell()<self.FileSize:
            self.__readHeaderStep();
            self.__timeStep[istep]      = self.__currentStepTime
            self.__posStep[istep]       = self.__fileobj.tell()
            self.__numSetPerStep[istep] = self.__currentStepNumSet
            istep = istep+1
            if self.__currentStepSize==-1:
              break
            else:
              self.__fileobj.seek(self.__currentStepSize,1)
        nstep = istep
        # Truncate buffered arrays
        if nstep>self.__scanBuffSize:
            warnings.warn('Buffer overflow detected: increase scanBuffSize.')
        self.__timeStep      = self.__timeStep[0:nstep]
        self.__posStep       = self.__posStep[0:nstep]
        self.__numSetPerStep = self.__numSetPerStep[0:nstep]
        # Set some internal variables
        self.NumDataset  = numpy.max(self.__numSetPerStep)
        self.NumTimestep = nstep
        self.__isFileHeaderWritten = True
        self.__isStepHeaderWritten = True

    def close(self):
        """Closes input file object"""
        self.__fileobj.close()
        self.__init__()
  
    def addFileHeaderToBuffer(self,rank=0,rankijk=(0,0,0),ftype=1999):
        """Initialize a buffer to generate a new UCF file."""
        self.__bufNumSteps  = 0
        self.__bufStep      = []
        self.__bufParams    = []
        self.__bufData      = []
        self.__bufRank      = rank
        self.__bufRankijk   = rankijk
        self.__bufFileType  = ftype
        self.__bufAvailable = True

    def addStepToBuffer(self,step=1,time=0.0):
        """Add a new step to buffer."""
        if not self.__bufAvailable:
            raise BufferError('Buffer has not been initialized.')
        if step>self.__bufNumSteps:
            self.__bufStep.extend([None] for ii in range(self.__bufNumSteps,step))
            self.__bufParams.extend([] for ii in range(self.__bufNumSteps,step))
            self.__bufData.extend([] for ii in range(self.__bufNumSteps,step))
            self.__bufNumSteps = step
        self.__bufStep[step-1] = time

    def addDatasetToBuffer(self,data,params=None,step=1,dset=1):
        import numpy
        """Add a new dataset to specified step of buffer."""
        if not self.__bufAvailable:
            raise BufferError('Buffer has not been initialized.')
        if step>self.__bufNumSteps:
            raise ValueError('Requested step does not exist.')
        if not hasattr(data,'dtype'):
            raise TypeError('Cannot determine datatype of provided data')
        if not hasattr(data,'nbytes'):
            raise TypeError('Cannot determine number of bytes of provided data')
        if not hasattr(data,'tobytes'):
            raise TypeError('Cannot convert provided data to bytes')
        if not hasattr(data,'shape'):
            raise TypeError('Cannot determine shape of provided data')
        if params is not None and not all(numpy.issubdtype(type(ii),numpy.integer) for ii in params):
            raise TypeError('Parameters must be provided as integer')
        nset = len(self.__bufData[step-1])
        if dset>nset:
            self.__bufParams[step-1].extend(None for ii in range(nset,dset))
            self.__bufData[step-1].extend(None for ii in range(nset,dset))
        self.__bufParams[step-1][dset-1] = params
        self.__bufData[step-1][dset-1] = data

    def copyFileHeaderToBuffer(self):
        if not self.__inputAvailable:
            raise IOError('No input file available')
        self.addFileHeaderToBuffer(rank=self.IORank[0],rankijk=self.IORank[1:],ftype=self.__typeID)

    def copyStepToBuffer(self,step_in,step_out=1,recursive=False,singlePrecision=False):
        """Copy a step from an input file to output buffer. If recursive copying is activated, all datasets
           within the step will be copied, otherwise only the step header is copied without datasets.
           If datasets are copied, the precision can be reduced using the 'singlePrecision' flag."""
        if not self.__inputAvailable:
            raise IOError('No input file available')
        if not self.__bufAvailable:
            raise BufferError('Buffer has not been initialized.')
        self.addStepToBuffer(step=step_out,time=self.__timeStep[step_in-1])
        if recursive:
            for dset in range(0,self.__numSetPerStep[step_in-1]):
                self.copyDatasetToBuffer(step_in,dset+1,step_out=step_out,dset_out=dset+1,singlePrecision=singlePrecision)

    def copyDatasetToBuffer(self,step_in,dset_in,step_out=1,dset_out=1,singlePrecision=False):
        """Copy a dataset from an input file to output buffer at specified step. The precision of the
           dataset can be reduced using the 'singlePrecision' flag."""
        import numpy
        if not self.__inputAvailable:
            raise IOError('No input file available')
        if not self.__bufAvailable:
            raise BufferError('Buffer has not been initialized.')
        (data,params) = self.readSet(step_in,dset_in)
        if singlePrecision:
            if data.dtype==numpy.dtype('float64'):
                data = numpy.float32(data)
            elif data.dtype==numpy.dtype('int64'):
                data = numpy.int32(data)
        self.addDatasetToBuffer(data,params=params,step=step_out,dset=dset_out)

    def flushBuffer(self):
        """Returns the buffer as a bytes object, which can be written to a file using a file object."""
        from time import time
        from sys import stderr
        from struct import pack
        import numpy as np
        # Sanity check and size gathering
        sizeStep = []
        sizeSet  = [[]]
        for step in range(0,self.__bufNumSteps):
            nset = len(self.__bufData[step])
            tmpSizeStep = 0
            if nset==0:
                warnings.warn('Step #{} in buffer does not contain any dataset.'.format(step+1),RuntimeWarning)
            for dset in range(0,nset):
                if self.__bufData[step][dset] is None:
                    raise ValueError('Step #{}, dataset #{} does not contain any data.'.format(step+1,dset+1))
                if self.__bufParams[step][dset] is None:
                    warnings.warn('No parameters were provided for step #{}, dataset #{}.'.format(step+1,dset+1),RuntimeWarning)
                    nparam==0
                else:
                    nparam = len(self.__bufParams[step][dset])
                sizeSet[step].append(self.__bufData[step][dset].nbytes)
                tmpSizeStep += (self.__nHeaderSet+nparam)*self.__nByteHeaderSet
                tmpSizeStep += self.__bufData[step][dset].nbytes 
            sizeStep.append(tmpSizeStep)
        # Create output buffer
        obuff = b''
        # Build file header
        magicFile   = self.__magicFile
        fileVersion = 2
        unixTime    = int(time())
        fileType    = self.__bufFileType
        rank        = self.__bufRank
        (iproc,jproc,kproc) = self.__bufRankijk
        if self.Debug:
            print('Write the following file header at {} bytes'.format(len(obuff)),file=stderr)
            print((magicFile,fileVersion,unixTime,fileType,rank,iproc,jproc,kproc),file=stderr)
        obuff += pack('qqqqqqqq',magicFile,fileVersion,unixTime,fileType,rank,iproc,jproc,kproc)
        # Build step header
        for step in range(0,self.__bufNumSteps):
            if self.Verbosity:
                print('Adding step #{} to output buffer'.format(step+1),file=stderr)
            magicStep = self.__magicStep
            stepBytes = sizeStep[step]
            stepTime  = self.__bufStep[step]
            nset      = len(self.__bufData[step])
            if self.Debug:
                print('Write the following step header at {} bytes'.format(len(obuff)),file=stderr)
                print((magicStep,stepBytes,stepTime,nset),file=stderr)
            obuff += pack('qqdq',magicStep,stepBytes,stepTime,nset)
            # Build dataset headers + attach data
            for dset in range(0,nset):
                if self.Verbosity:
                   print('  dataset #{}'.format(dset+1),file=stderr)
                magicSet = self.__magicSet
                setSize  = sizeSet[step][dset]
                nptype = self.__bufData[step][dset].dtype
                if nptype==numpy.dtype('int32'):
                    dtEncoded = 11
                elif nptype==numpy.dtype('int64'):
                    dtEncoded = 12
                elif nptype==numpy.dtype('float32'):
                    dtEncoded = 21
                elif nptype==numpy.dtype('float64'):
                    dtEncoded = 22
                else:
                    raise TypeError('Data at step #{}, dataset #{} has an invalid datatype.'.format(step+1,dset+1))
                if self.__bufParams[step][dset] is None:
                    nparam = 0
                else:
                    nparam = len(self.__bufParams[step][dset])
                if self.Debug:
                    print('Write the following set header at {} bytes'.format(len(obuff)),file=stderr)
                    print((magicSet,setSize,dtEncoded,nparam),file=stderr)
                    print('with parameters:',file=stderr)
                    print(self.__bufParams[step][dset],file=stderr)

                obuff += pack('qqqq',magicSet,setSize,dtEncoded,nparam)
                if nparam!=0:
                    obuff += pack(nparam*'q',*self.__bufParams[step][dset])
                obuff += self.__bufData[step][dset].tobytes('F')
        # Return bytes
        return obuff

    def getTime(self,step=1):
        assert(step-1<self.NumTimestep)
        return self.__timeStep[step-1]

    def readSet(self,step=1,dset=1,memmap=False):
        """Read a dataset from input file. If 'memmap' is activated, the file will only be read partially on demand."""
        import numpy
        from io import BytesIO
        if not self.__inputAvailable:
            raise IOError('No input file available')
        self.__fileobj.seek(self.__findSet(step,dset),0)
        self.__readHeaderSet();
        params = self.__currentSetParams
        if isinstance(self.__fileobj,BytesIO):
            data = numpy.frombuffer(self.__fileobj.read(self.__currentSetSize),dtype=self.__currentSetDatatype) 
        else:
            data = numpy.fromfile(self.__fileobj,dtype=self.__currentSetDatatype,count=self.__currentSetNumElements)
        return (data,params)
    
    def __readHeaderFile(self):
        from datetime import datetime
        from sys import stderr
        from struct import unpack
        import numpy
        self.__fileobj.seek(self.__fileBeg,0);
        # Determine endianess
        mfmt = "<>"
        buff = self.__fileobj.read(8)
        for fmt in mfmt:
            currentMagic = unpack("%sq"%fmt,buff)[0]
            if currentMagic==self.__magicFile:
                break
        if currentMagic!=self.__magicFile:
            raise ValueError('Magic mismatch: invalid file header. {}'.format(currentMagic))
        self.Endian = fmt
        # Read header
        self.__fileobj.seek(self.__fileBeg,0);
        buff = self.__fileobj.read(self.__nHeaderFile*8)
        header = unpack("%s%dq"%(self.Endian,8),buff)
        if self.Debug:
            print('Read the following file header at 0 bytes',file=stderr)
            print(header,file=stderr)
        # Parse version
        self.__versionMajor = numpy.floor(header[1]/self.__factorMajor)
        self.__versionMinor = numpy.floor(numpy.mod(header[1],self.__factorMajor)/self.__factorMinor)
        self.__versionPatch = numpy.floor(numpy.mod(header[1],self.__factorMinor)/self.__factorPatch)
        self.CodeVersion  = "%d.%d.%d" %(self.__versionMajor,self.__versionMinor,self.__versionPatch)
        self.UCFVersion   = numpy.mod(header[1],self.__factorPatch);
        # Parse time stamp (UTC)
        self.__creationTimeUnix = header[2];
        self.CreationTime = datetime.utcfromtimestamp(self.__creationTimeUnix).strftime('%Y-%m-%d %H:%M:%S')
        #Parse file type
        self.__typeID = header[3];
        typeDict = {
               0: "grid",
              10: "processor grid",
            1000: "fluid snapshot",
            1010: "scalar snapshot",
            1999: "matlab field data",
            2000: "particle snapshot",
            2001: "particle append",
            2011: "particle lagrange",
            2021: "particle balancing",
            2999: "matlab particle data",
            3000: "statistics fluid",
            3010: "statistics fluid pure",
            3020: "statistics scalar"
            }
        self.Type = typeDict.get(self.__typeID,"unkmown")
        # Parse file class
        classDict = {
            1: "field",
            2: "particle",
            3: "statistics"
            }
        self.Class = classDict.get(numpy.floor(self.__typeID/self.__factorTypeIDClass),"unknown")
        # Parse IO rank
        self.IORank = header[4:8]
    
    def __readHeaderStep(self):
        from sys import stderr
        from struct import unpack
        # Read and parse
        self.__currentStepPosHeader = self.__fileobj.tell()
        buff = self.__fileobj.read(self.__nHeaderStep*8)
        header = unpack("%s%dq"%(self.Endian,self.__nHeaderStep),buff)
        self.__currentStepPosData = self.__fileobj.tell()
        currentMagic = header[0]
        self.__currentStepSize   = header[1]
        self.__currentStepTime   = unpack("%sd"%self.Endian,buff[16:24])[0]
        self.__currentStepNumSet = header[3]
        if self.Debug:
            print("Read the following step header at %d bytes" % self.__currentStepPosHeader,file=stderr)
            print("%d,%d,%f,%d" % (currentMagic,self.__currentStepSize,self.__currentStepTime,self.__currentStepNumSet),file=stderr) 
        # Check if magic is correct
        if currentMagic!=self.__magicStep:
            raise ValueError("Magic mismatch: invalid step header. %d" % currentMagic);
    
    def __readHeaderSet(self):
        from sys import stderr
        from struct import unpack
        import numpy
        # Read and parse
        self.__currentSetPosHeader = self.__fileobj.tell()
        buff = self.__fileobj.read(self.__nHeaderSet*8)
        header = unpack("%s%dq"%(self.Endian,self.__nHeaderSet),buff)
        self.__currentSetPosData = self.__fileobj.tell()
        currentMagic = header[0]
        self.__currentSetSize = header[1]
        self.__currentSetDatatypeNumeric = header[2]
        dtSizeDict = {
            11: 4,
            12: 8,
            21: 4,
            22: 8
            }
        dtNameDict = {
            11: "%si4" % self.Endian,
            12: "%si8" % self.Endian,
            21: "%sf4" % self.Endian,
            22: "%sf8" % self.Endian
        }
        self.__currentSetSizeof      = dtSizeDict[self.__currentSetDatatypeNumeric]
        self.__currentSetDatatype    = dtNameDict[self.__currentSetDatatypeNumeric]
        self.__currentSetNumParams   = header[3]
        self.__currentSetNumElements = numpy.around(self.__currentSetSize/self.__currentSetSizeof).astype(numpy.int32)
        if self.Debug:
            print("Read the following set header at %d bytes" % self.__currentSetPosHeader,file=stderr)
            print(header,file=stderr)
        # Check if magic is correct
        if currentMagic!=self.__magicSet:
            raise ValueError("Magic mismatch: invalid dataset header. %d" % currentMagic)
        # Read variable number of parameters
        buff = self.__fileobj.read(self.__currentSetNumParams*8)
        self.__currentSetParams = unpack("%s%dq"%(self.Endian,self.__currentSetNumParams),buff)
        if self.Debug:
            print('with parameters:',file=stderr)
            print(self.__currentSetParams,file=stderr)
    
    def __findSet(self,tstep,dset):
        from sys import stderr
        # Check input
        if tstep>self.NumTimestep:
            raise ValueError("Out of bounds: timestep. %d, %d" %(tstep,self.NumTimestep))
        if dset>self.__numSetPerStep[tstep-1]:
            raise ValueError("Out of bounds: dataset. %d, %d" % (dset,self.NumDataset))
        # Navigate to correct set
        self.__fileobj.seek(self.__posStep[tstep-1],0)
        for iset in range(0,dset-1):
            self.__readHeaderSet()
            self.__fileobj.seek(self.__currentSetSize,1)
        posHeader = self.__fileobj.tell()
        if self.Debug:
            print("Found step #%d, set #%d at position %d" % (tstep,dset,posHeader),file=stderr)
        return posHeader
    
    def __initializeConstants(self):
        self.__magicFile                 = 81985529216486895;
        self.__magicStep                 = 11944304052957;
        self.__magicSet                  = 240217520921210;
        self.__nHeaderFile               = 8;
        self.__nHeaderStep               = 4;
        self.__nHeaderSet                = 4;
        self.__nByteHeaderFile           = 8;
        self.__nByteHeaderStep           = 8;
        self.__nByteHeaderSet            = 8;
        self.__nSetParamsField           = 10;
        self.__nSetParamsParticle        = 16;
        self.__factorMajor               = 1000000000;
        self.__factorMinor               = 1000000;
        self.__factorPatch               = 1000;
        self.__factorTypeIDClass         = 1000;
        self.__factorTypeIDKind          = 10;
        self.__typeIDmatlabField         = 1999;
        self.__typeIDmatlabParticle      = 2999;
        self.__scanBuffSize              = 131072;
    def __resetPublicProperties(self):
        self.File         = '' # file name
        self.Type         = '' # file type
        self.Class        = '' # file class
        self.Endian       = '' # endianess
        self.CodeVersion  = '' # version of the simulation code
        self.UCFVersion   = '' # version of the data format ("unified container format")
        self.NumDataset   = 0  # maximum number of datasets in this file (over all time steps)
        self.NumTimestep  = 0  # number of time steps in this file
        self.FileSize     = 0  # file size
        self.CreationTime = 0  # time of creation
        self.IOMode       = '' # file opened in read-only or read-write mode?
        self.IORank       = 0  # rank of processor + col,row,pln
        self.Verbosity    = 0  # verbose output?
        self.Debug        = 0  # debug information?
    def __resetPrivateProperties(self):
        self.__fileobj                   = None
        self.__fileBeg                   = 0
        self.__typeID                    = 0
        self.__creationTimeUnix          = ''
        self.__versionMajor              = 0
        self.__versionMinor              = 0
        self.__versionPatch              = 0
        self.__versionFile               = 0
        self.__posStep                   = 0
        self.__numSetPerStep             = 0
        self.__timeStep                  = 0
        self.__inputAvailable            = False
        self.__bufAvailable              = False
    def __resetCurrentStep(self):
        self.__currentStep               = 0
        self.__currentStepPosHeader      = 0
        self.__currentStepPosData        = 0
        self.__currentStepSize           = 0
        self.__currentStepTime           = 0
        self.__currentStepNumSet         = 0
    def __resetCurrentSet(self):
        self.__currentSet                = 0
        self.__currentSetPosHeader       = 0
        self.__currentSetPosData         = 0
        self.__currentSetSize            = 0
        self.__currentSetDatatype        = 0
        self.__currentSetDatatypeNumeric = 0
        self.__currentSetSizeof          = 0
        self.__currentSetNumParams       = 0
        self.__currentSetParams          = 0
        self.__currentSetNumElements     = 0

# TBD: refactor -> UCFCollection, UCFTarFile constructors etc
class UCFTarFile:
    '''Superclass for instances of ucf.tar files.'''
    def __init__(self,file_tar,file_index=None):
        self.handler = Ustar(file_tar,file_index)
        self.verbose = False
        self.debug   = False
        self._buffer_grid = None
        self._buffer_proc = None
        self._buffer_params = None
    def read_grid(self,key=None):
        self._initialize_buffer_grid()
        if key is not None:
            gidx = self.gidx_from_key(key)
            return self._buffer_grid[gidx]
        return self._buffer_grid
    def read_procgrid(self,key=None):
        self._initialize_buffer_proc()
        if key is not None:
            gidx = self.gidx_from_key(key)
            return self._buffer_proc[gidx]
        return self._buffer_proc
    def read_parameters(self):
        self._initialize_buffer_params()
        return self._buffer_params
    def bounds(self,axis=None):
        params = self.read_parameters()
        keys = ('a','b','c','d','e','f')
        if axis is None:
            return tuple(params.getfloat('geometry',key) for key in keys)
        else: 
            assert axis<3, "'axis' must be smaller than 3"
            return tuple(params.getfloat('geometry',key) for key in keys[2*axis:2*(axis+1)])
    def periodicity(self,axis=None):
        params = self.read_parameters()
        keys = ('xperiodic','yperiodic','zperiodic')
        if axis is None:
            return tuple(params.getboolean('geometry',key) for key in keys)
        else:
            assert axis<3, "'axis' must be smaller than 3"
            return params.getboolean('geometry',keys[axis])
    def period(self,axis=None):
        if axis is None:
            return tuple(self.period(axis=ii) for ii in range(0,3))
        else:
            assert axis<3, "'axis' must be smaller than 3"
            params = self.read_parameters()
            key_bound  = ('b','d','f')
            key_period = ('xperiodic','yperiodic','zperiodic')
            period     = params.getfloat('geometry',key_bound[axis])
            isPeriodic = params.getboolean('geometry',key_period[axis]) 
            return period if isPeriodic else None
    def procgrid(self):
        self._initialize_buffer_proc()
        proc = {}
        offset = 0
        for key in ('u','v','w','p','s'):
            proc[key] = tuple(self._buffer_proc[offset])
            offset += 1
        return proc
    def origin(self):
        self._initialize_buffer_grid()
        origin = {}
        offset = 0
        for key in ('u','v','w','p','s'):
            origin[key] = (
                self._buffer_grid[offset][0][0],
                self._buffer_grid[offset][1][0],
                self._buffer_grid[offset][2][0])
            offset +=1
        return origin
    def spacing(self):
        self._initialize_buffer_grid()
        spacing = (
                self._buffer_grid[0][0][1]-self._buffer_grid[0][0][0],
                self._buffer_grid[0][1][1]-self._buffer_grid[0][1][0],
                self._buffer_grid[0][2][1]-self._buffer_grid[0][2][0],
            )
        return spacing
    def gidx_from_key(self,key):
        if key[0]=='u': return 0
        elif key[0]=='v': return 1
        elif key[0]=='w': return 2
        elif key[0]=='p': return 3
        elif key[0]=='s': return 4
        else: raise ValueError("Invalid value of 'key'.")
    def _initialize_buffer_grid(self):
        if self._buffer_grid is None:
            data = self.handler.read('grid.bin')
            self._buffer_grid = read_grid(data,verbosity=self.verbose,debug=self.debug)
        return 
    def _initialize_buffer_proc(self):
        if self._buffer_proc is None:
            data = self.handler.read('proc.bin')
            self._buffer_proc = read_procgrid(data,verbosity=self.verbose,debug=self.debug)
        return
    def _initialize_buffer_params(self):
        if self._buffer_params is None:
            data = self.handler.read('parameters.asc')
            self._buffer_params = read_parameters(data)
        return

class UCFSnapshot(UCFTarFile):
    '''Handles a snapshot.ucf.tar file.'''
    def __init__(self,file_tar,file_index=None):
        super(UCFSnapshot,self).__init__(file_tar,file_index=file_index)
        file_name_string = '\t'.join(self.handler.file_name)
        if 'uvwp' in file_name_string:
            self.type = 'uvwp'
        elif 'scal' in file_name_string:
            self.type = 'scal'
        else:
            raise ValueError("Archive does not contain 'uvwp' nor 'scal' files.")
        self._nproc = sum(self.type in s for s in self.handler.file_name)
        self._nprocijk = None
    def read_particles(self):
        data = self.handler.read('particles.bin')
        return read_particles(data,step=1,verbosity=self.verbose,debug=self.debug)
    def read_chunk(self,rank,dset=None,keep_ghost=True):
        file_target = self.type+'.{:05d}'.format(rank)
        data = self.handler.read(file_target)
        return read_chunk(data,step=1,dset=dset,keep_ghost=keep_ghost,
                    verbosity=self.verbose,debug=self.debug)
    def field_chunk(self,rank,key,keep_ghost=True):
        '''Returns chunk data as Field3d object.'''
        from .field import Field3d
        # Support a list/tuple of keys: this is not the most efficient 
        # implementation but convenient
        if isinstance(key,(list,tuple)):
            return [self.field_chunk(rank,k,keep_ghost=keep_ghost) 
                    for k in key]
        dset   = self.dset_from_key(key)
        chunk  = self.read_chunk(rank,dset=dset,keep_ghost=keep_ghost)
        grid   = self.read_grid(key=key)
        dx,dy,dz = (grid[0][1]-grid[0][0],
                    grid[1][1]-grid[1][0],
                    grid[2][1]-grid[2][0])
        xo,yo,zo = (grid[0][chunk['ibeg']-1]-(chunk['ighost'])*dx,
                    grid[1][chunk['jbeg']-1]-(chunk['ighost'])*dy,
                    grid[2][chunk['kbeg']-1]-(chunk['ighost'])*dz)
        return Field3d(chunk['data'],(xo,yo,zo),(dx,dy,dz))
    def particles(self,select_col=None):
        '''Returns particle data as Particles object.'''
        from .particle import Particles
        pp,col,time = self.read_particles()
        return Particles.from_array(pp,col,time,self.period(),select_col=select_col)
    def nproc(self,axis=None):
        if axis is None:
            return self._nproc
        else:
            assert axis<3, "'axis' must be smaller than 3"
            self._initialize_nprocijk()
            return self._nprocijk[axis]
    def ijk_from_rank(self,rank):
        assert rank<self.nproc(), "'rank' exceeds highest rank."
        ijk = (
            rank//(self.nproc(axis=1)*self.nproc(axis=2)), 
            (rank//self.nproc(axis=2))%self.nproc(axis=1), 
            rank%self.nproc(axis=2)
            )
        return ijk
    def rank_from_ijk(self,ijk):
        assert len(ijk)==3, "len(ijk) must be 3."
        assert all([ijk[ii]<self.nproc(axis=ii) for ii in range(0,3)])
        return ijk[0]*self.nproc(axis=1)*self.nproc(axis=2)+ijk[1]*self.nproc(axis=2)+ijk[2]
    def dset_from_key(self,key):
        assert self.has_key(key), "Snapshot does not contain requested key."
        if key[0]=='u': return 1
        elif key[0]=='v': return 2
        elif key[0]=='w': return 3
        elif key[0]=='p': return 4
        elif key[0]=='s': return int(key[1:])
        else: raise ValueError("Invalid value of 'key'.")
    def has_key(self,key):
        if key in ('u','v','w','p') and self.type=='uvwp':
            return True
        elif key[0]=='s' and self.type=='scal':
            return True
        else:
            return False
    def _initialize_nprocijk(self):
        self._initialize_buffer_proc()
        self._nprocijk = tuple([len(self._buffer_proc[2*ii]) for ii in range(0,3)])

class UCFAuxiliary(UCFTarFile):
    '''Handles a auxiliary.ucf.tar file.'''
    def __init__(self,file_tar,file_index=None):
        super(UCFAuxiliary,self).__init__(file_tar,file_index=file_index)
        # TBD: support statistics files
    def read_particles(self):
        data = self.handler.read('append_particles.bin')
        return read_particles(data,step=None,verbosity=self.verbose,debug=self.debug)
    def trajectories(self,select_col=None):
        from .particle import Particles,Trajectories
        pp,col,time = self.read_particles()
        part = Particles.from_array(pp[:,:,0],col,time[0],self.period(),select_col=select_col)
        traj = Trajectories(part,copy_particles=False)
        for ii in range(1,len(time)):
            traj += Particles.from_array(pp[:,:,ii],col,time[ii],self.period(),select_col=select_col)
        return traj

class Ustar:
    '''Minimalistic ustar implementation meant to be used with ucftar files'''
    def __init__(self,file_tar,file_index=None):
        self.path        = file_tar
        self.num_files   = 0
        self.file_name   = []
        self.file_size   = []
        self.file_offset = []
        if file_index:
            self.import_index_file(file_index)
        else:
            self.import_tar_file()
    def __del__(self):
        return
    def import_tar_file(self):
        '''Imports information on tar archive from scanning it'''
        from tarfile import TarFile,USTAR_FORMAT
        from struct import unpack
        with open(self.path,'rb') as f:
            tarinfos = TarFile(fileobj=f,format=USTAR_FORMAT).getmembers()
            self.num_files = 0
            for tarinfo in tarinfos:
                self.num_files += 1
                self.file_name.append(tarinfo.name)
                self.file_offset.append(tarinfo.offset_data)
                self.file_size.append(tarinfo.size)
            return
    def import_index_file(self,file_index):
        '''Imports information on tar archive from .taridx file'''
        from struct import unpack
        with open(file_index,'rb') as f:
            self.num_files = unpack('<q',f.read(8))[0]
            self.file_name = []
            self.file_size = []
            self.file_offset = []
            for ifile in range(0,self.num_files):
                self.file_name.append(f.read(256).decode().strip().rstrip('\0'))
                self.file_offset.append(unpack('<q',f.read(8))[0])
                self.file_size.append(unpack('<q',f.read(8))[0])
        return
    def read(self,file):
        '''Reads a file from the archive into memory. Data is returned as bytes.'''
        idx = self.file_name.index(file)
        with open(self.path,'rb') as f:
            f.seek(self.file_offset[idx])
            return f.read(self.file_size[idx])

def read_grid(file,verbosity=False,debug=False):
    obj = UCF(file=file,verbosity=verbosity,debug=debug)
    output = []
    for iset in range(0,obj.NumDataset):
        (data,params) = obj.readSet(step=1,dset=iset+1)
        nx = params[0]
        ny = params[1]
        nz = params[2]
        output.append((data[0:nx],data[nx:nx+ny],data[nx+ny:nx+ny+nz]))
    if obj.NumDataset<5:
        output.extend(output[-1:])
    obj.close()
    return output

def read_procgrid(file,verbosity=False,debug=False):
    obj = UCF(file=file,verbosity=verbosity,debug=debug)
    output = []
    for iset in range(0,obj.NumDataset):
        (data,params) = obj.readSet(step=1,dset=iset+1)
        nxp = params[0]
        nyp = params[1]
        nzp = params[2]
        output.append((
            data[0:nxp],                             # ibeg
            data[nxp:2*nxp],                         # iend
            data[2*nxp:2*nxp+nyp],                   # jbeg
            data[2*nxp+nyp:2*nxp+2*nyp],             # jend
            data[2*nxp+2*nyp:2*nxp+2*nyp+nzp],       # kbeg
            data[2*nxp+2*nyp+nzp:2*nxp+2*nyp*2*nzp]  # kend
            )) 
    #if obj.UCFVersion<2:
    if obj.NumDataset<5:
        output.extend(output[-1:])
    obj.close()
    return output

def read_chunk(file,step=1,dset=None,keep_ghost=True,verbosity=False,debug=False):
    obj = UCF(file=file,verbosity=verbosity,debug=debug)
    if dset is None:
        dset = range(1,obj.NumDataset+1) # fix that maybe later (this is maximum over all timesteps)
        returnList = True
    elif isinstance(dset,(list,tuple)):
        returnList = True
    else:
        dset = [dset]
        returnList = False
    output = []
    for ii in dset:
        tmp = dict()
        (data,params) = obj.readSet(step=step,dset=ii)
        tmp['ighost'] = params[0]
        tmp['ibeg']   = params[1]
        tmp['jbeg']   = params[2]
        tmp['kbeg']   = params[3]
        tmp['nxl']    = params[4]
        tmp['nyl']    = params[5]
        tmp['nzl']    = params[6]
        tmp['nx']     = params[7]
        tmp['ny']     = params[8]
        tmp['nz']     = params[9]
        tmp['data']   = data.reshape((tmp['nxl']+2*tmp['ighost'],
                                      tmp['nyl']+2*tmp['ighost'],
                                      tmp['nzl']+2*tmp['ighost']),
                                     order='F')
        tmp['rank']   = obj.IORank[0]
        tmp['rankijk']= obj.IORank[1:]
        if not keep_ghost and tmp['ighost']:
            tmp['data']  = tmp['data'][1:-1,1:-1,1:-1]
            tmp['ighost'] = 0
        output.append(tmp)
    obj.close()
    if not returnList:
        return output[0]
    return output

def read_particles(file,step=None,verbosity=False,debug=False):
    import numpy
    obj = UCF(file=file,verbosity=verbosity,debug=debug)
    if not isinstance(step,list):
        if step is None:
            step = range(1,obj.NumTimestep+1)
        else:
            step = [step]
    # The output will be the following: 
    # 1) numpy array with dimension (ncol,np,ntime) 
    # 2) dictionary which specifies the columns
    # We read the data step by step in a list, which is then converted to a 3D array
    pp = []
    time = []
    for ii in step:
      (data,params) = obj.readSet(step=ii,dset=1)
      npart       = params[0]
      ncol        = params[1]
      ncol_rank   = params[2]
      ncol_hybrid = params[3]
      ncol_dem    = params[4]
      ncol_scalar = params[5]
      nscal       = ncol_scalar//2
      pp.append(data.reshape((ncol,npart),order='F'))
      time.append(obj.getTime(ii))
    # Close UCF obeject
    obj.close()
    # Convert list of 2D arrays to 3D array
    pp = numpy.stack(pp,axis=2)
    time = numpy.array(time)
    # Create the dictionary
    col = colmap_from_flags(ncol_rank,ncol_hybrid,ncol_dem,nscal)
    # Return result
    return (pp,col,time)

def read_parameters(file):
    import configparser
    params = configparser.ConfigParser()
    if isinstance(file,str):
        with open(file,'r') as f:
            params.read(f)
    elif isinstance(file,bytes):
        params.read_string(file.decode('ascii'))
    return params
    
def colmap_from_flags(irank,ihybrid,idem,iscal):
    '''Creates a dictionary which specifies the columns of a particle array.'''
    col = {}
    ioffset = 0
    if irank>0:
      col['rank'] = ioffset; ioffset+=1
    if ihybrid>0:
      col['id'] = ioffset; ioffset+=1
      col['x']  = ioffset; ioffset+=1
      col['y']  = ioffset; ioffset+=1
      col['z']  = ioffset; ioffset+=1
      col['r']  = ioffset; ioffset+=1
      col['rho']= ioffset; ioffset+=1
      col['ax'] = ioffset; ioffset+=1
      col['ay'] = ioffset; ioffset+=1
      col['az'] = ioffset; ioffset+=1
      col['u']  = ioffset; ioffset+=1
      col['v']  = ioffset; ioffset+=1
      col['w']  = ioffset; ioffset+=1
      col['ox'] = ioffset; ioffset+=1
      col['oy'] = ioffset; ioffset+=1
      col['oz'] = ioffset; ioffset+=1
      col['fx'] = ioffset; ioffset+=1
      col['fy'] = ioffset; ioffset+=1
      col['fz'] = ioffset; ioffset+=1
      col['tx'] = ioffset; ioffset+=1
      col['ty'] = ioffset; ioffset+=1
      col['tz'] = ioffset; ioffset+=1
    if idem>0:
      col['fxc'] = ioffset; ioffset+=1
      col['fyc'] = ioffset; ioffset+=1
      col['fzc'] = ioffset; ioffset+=1
      col['txc'] = ioffset; ioffset+=1
      col['tyc'] = ioffset; ioffset+=1
      col['tzc'] = ioffset; ioffset+=1
    if iscal>0:
      for ii in range(iscal):
        col['s'+str(ii+1)] = ioffset; ioffset+=1
        col['q'+str(ii+1)] = ioffset; ioffset+=1
    return col

def grid_chunk(chunk,gridg):
    '''Returns the grid vectors for chunk (including ghost cells if they exist)'''
    import numpy
    xg,yg,zg = gridg
    # Shift indices so that they start from zero
    ib = chunk['ibeg']-1
    jb = chunk['jbeg']-1
    kb = chunk['kbeg']-1
    nxl = chunk['nxl']
    nyl = chunk['nyl']
    nzl = chunk['nzl']
    ighost = chunk['ighost']
    if ighost:
        nxg = len(xg)
        nyg = len(yg)
        nzg = len(zg)
        dx = xg[1]-xg[0]
        dy = yg[1]-yg[0]
        dz = zg[1]-zg[0]
        xl = numpy.zeros(nxl+2)
        yl = numpy.zeros(nyl+2)
        zl = numpy.zeros(nzl+2)
        xl[0]    = xg[ib]-dx
        xl[1:-1] = xg[ib:ib+nxl]
        xl[-1]   = xg[ib+nxl-1]+dx
        yl[0]    = yg[jb]-dy
        yl[1:-1] = yg[jb:jb+nyl]
        yl[-1]   = yg[jb+nyl-1]+dy
        zl[0]    = zg[kb]-dz
        zl[1:-1] = zg[kb:kb+nzl]
        zl[-1]   = zg[kb+nzl-1]+dz
    else:
        xl = xg[ibeg:ib+nxl-1]
        yl = yg[jbeg:jb+nyl-1]
        zl = zg[kbeg:kb+nzl-1]
    return (xl,yl,zl)

def grid_chunk_origin_spacing(chunk,gridg):
    '''Returns origin and spacing of regular grid of chunk data'''
    xg,yg,zg = gridg
    # Shift indices so that they start from zero
    ib = chunk['ibeg']-1
    jb = chunk['jbeg']-1
    kb = chunk['kbeg']-1
    dx = xg[1]-xg[0]
    dy = yg[1]-yg[0]
    dz = zg[1]-zg[0]
    xo = xg[ib]-chunk['ighost']*dx
    yo = yg[jb]-chunk['ighost']*dy
    zo = zg[kb]-chunk['ighost']*dz
    return (xo,yo,zo),(dx,dy,dz)