"""
:class:`~pyUSID.processing.process.Process` - An abstract class for formulating scientific problems as computational
problems

Created on 7/17/16 10:08 AM

@author: Suhas Somnath, Chris Smith
"""

from __future__ import division, unicode_literals, print_function, \
    absolute_import
import numpy as np
import psutil
import time as tm
import h5py
from warnings import warn
from numbers import Number
from multiprocessing import cpu_count

from .comp_utils import parallel_compute, get_MPI, \
    group_ranks_by_socket, get_available_memory
from sidpy.base.num_utils import integers_to_slices
from sidpy.base.string_utils import validate_single_string_arg, format_time, \
    format_size
from sidpy.hdf.hdf_utils import write_simple_attrs, lazy_load_array

from ..io.hdf_utils import check_if_main, check_for_old
from ..io.usi_data import USIDataset

# TODO: internalize as many attributes as possible. Expose only those that will be required by the user


class Process(object):
    """
    An abstract class for formulating scientific problems as computational problems. This class handles the tedious,
    science-agnostic, file-operations, parallel-computations, and book-keeping operations such that children classes
    only need to specify application-relevant code for processing the data.
    """

    def __init__(self, h5_main, process_name, parms_dict=None, cores=None,
                 max_mem_mb=4*1024, mem_multiplier=1.0, lazy=False,
                 h5_target_group=None, verbose=False):
        """
        Parameters
        ----------
        h5_main : :class:`~pyUSID.io.usi_data.USIDataset`
            The USID main HDF5 dataset over which the analysis will be performed.
        process_name : str
            Name of the process
        cores : uint, optional
            How many cores to use for the computation. Default: all available cores - 2 if operating outside MPI context
        max_mem_mb : uint, optional
            How much memory to use for the computation.  Default 1024 Mb
        mem_multiplier : float, optional. Default = 1
            mem_multiplier is the number that will be multiplied with the
            (byte) size of a single position in the source dataset in order to
            better estimate the number of positions that can be processed at
            any given time (how many pixels of the source and results datasets
            can be retained in memory). The default value of 1.0 only accounts
            for the source dataset. A value greater than 1 would account for
            the size of results datasets as well. For example, if the result
            dataset is the same size and precision as the source dataset,
            the multiplier will be 2 (1 for source, 1 for result)
        lazy : bool, optional. Default = False
            If True, read_data_chunk and write_results_chunk will operate on
            dask arrays. If False - everything will be in numpy.
        h5_target_group : h5py.Group, optional. Default = None
            Location where to look for existing results and to place newly
            computed results. Use this kwarg if the results need to be written
            to a different HDF5 file. By default, this value is set to the
            parent group containing `h5_main`
        verbose : bool, Optional, default = False
            Whether or not to print debugging statements

        Attributes
        ----------
        self.h5_results_grp : :class:`h5py.Group`
            HDF5 group containing the HDF5 datasets that contain the results
            of the computation
        self.verbose : bool
            Whether or not to print debugging statements
        self.parms_dict : dict
            Dictionary of parameters for the computation
        self.duplicate_h5_groups : list
            List of :class:`h5py.Group` objects containing computational
            results that have been completely computed with the same
            set of parameters as those in self.parms_dict
        self.partial_h5_groups : list
            List of :class:`h5py.Group` objects containing computational
            results that have been partially computed with the same
            set of parameters as those in self.parms_dict
        self.process_name : str
            Name of the process. This is used for checking for existing
            completely and partially computed results as well as for naming
            the HDF5 group that will contain the results of the computation
        self._cores : uint
            Number of CPU cores to use for parallel computations.
            Ignored in the MPI context. Each rank gets 1 CPU core
        self._max_pos_per_read : uint
            Number of positions in the dataset to read per chunk
        self._status_dset_name : str
            Name of the HDF5 dataset that keeps track of the positions in the
            source dataset thave already been computed
        self._results : list
            List of objects returned as the result of computation performed by
            the self._map_function for each position in the current batch of
            positions that were processed
        self._h5_target_group : h5py.Group
            Location where existing / future results will be stored
        self.__resume_implemented : bool
            Whether or not this (child) class has implemented the
            self._get_existing_datasets() function
        self.__bytes_per_pos : uint
            Number of bytes used by one position of the source dataset
        self.mpi_comm : :class:`mpi4py.MPI.COMM_WORLD`
            MPI communicator. None if not running in an MPI context
        self.mpi_rank: uint
            MPI rank. Always 0 if not running in an MPI context
        self.mpi_size: uint
            Number of ranks in COMM_WORLD. 1 if not running in an MPI context
        self.__ranks_on_socket : uint
            Number of MPI ranks on a given CPU socket
        self.__socket_master_rank : uint
            Master MPI rank for a given CPU chip / socket
        self.__compute_jobs : array-like
            List of positions in the HDF5 dataset that need to be computed.
            This may not be a continuous list of numbers if multiple MPI
            workers had previously started computing and were interrupted.
        self.__start_pos : uint
            The index within self.__compute_jobs that a particular MPI rank /
            worker needs to start computing from.
        self.__rank_end_pos : uint
            The index within self.__compute_jobs that a particular MPI rank /
            worker needs to start computing till.
        self.__end_pos : uint
            The index within self.__compute_jobs that a particular MPI rank /
            worker needs to start computing till for the current batch of
            positions.
        self.__pixels_in_batch : array-like
            The positions being computed on by the current compute worker
        """
        MPI = get_MPI()

        # Ensure that the file is opened in the correct comm or something
        if MPI is not None and h5_main.file.driver != 'mpio':
            warn('Code was called in MPI context but HDF5 file was not opened '
                 'with the "mpio" driver. JobLib will be used instead of MPI '
                 'for parallel computation')
            MPI = None

        if MPI is not None:
            # If we came here then, the user has intentionally asked for multi-node computation
            comm = MPI.COMM_WORLD
            self.mpi_comm = comm
            self.mpi_rank = comm.Get_rank()
            self.mpi_size = comm.Get_size()

            if verbose:
                print("Rank {} of {} on {} sees {} logical cores on the socket".format(comm.Get_rank(), comm.Get_size(),
                                                                                       MPI.Get_processor_name(),
                                                                                       cpu_count()))

            # First, ensure that cores=logical cores in node. No point being economical / considerate
            cores = psutil.cpu_count()

            # It is sufficient if just one rank checks all this.
            if self.mpi_rank == 0:
                print('Working on {} ranks via MPI'.format(self.mpi_size))

            if verbose and self.mpi_rank == 0:
                print('Finished getting all necessary MPI information')

            """
            # Not sure how to check for this correctly
            messg = None
            try:
                if h5_main.file.comm != comm:
                    messg = 'The HDF5 file should have been opened with comm=MPI.COMM_WORLD. Currently comm={}'
                            ''.format(h5_main.file.comm)
            except AttributeError:
                messg = 'The HDF5 file should have been opened with comm=MPI.COMM_WORLD'
            if messg is not None:
                raise TypeError(messg)
            """

        else:
            if verbose:
                print('No mpi4py found or script was not called via mpixexec / mpirun. '
                      'Assuming single node computation')
            self.mpi_comm = None
            self.mpi_size = 1
            self.mpi_rank = 0

        # Checking if dataset is "Main"
        if not check_if_main(h5_main, verbose=verbose and self.mpi_rank == 0):
            raise ValueError('Provided dataset is not a "Main" dataset with necessary ancillary datasets')

        if h5_target_group is not None:
            if not isinstance(h5_target_group, (h5py.Group, h5py.File)):
                raise TypeError("'h5_target_group' must be a h5py.Group object")
        else:
            h5_target_group = h5_main.parent
        self._h5_target_group = h5_target_group

        if h5_target_group.file.mode == 'r':
            raise IOError('the file meant to contain the results '
                          '(h5_target_group) must not be in read-only mode to '
                          'write results to the file')

        process_name = validate_single_string_arg(process_name, 'process_name')

        if parms_dict is None:
            parms_dict = {}
        else:
            if not isinstance(parms_dict, dict):
                raise TypeError("Expected 'parms_dict' of type: dict")

        if MPI is not None:
            MPI.COMM_WORLD.barrier()
        # Not sure if we need a barrier here.

        if verbose and self.mpi_rank == 0:
            print('Rank {}: Upgrading from a regular h5py.Dataset to a USIDataset'.format(self.mpi_rank))

        # Generation of N-dimensional form would break things for some reason.
        self.h5_main = USIDataset(h5_main)

        if verbose and self.mpi_rank == 0:
            print('Rank {}: The HDF5 dataset is now a USIDataset'.format(self.mpi_rank))

        # Saving these as properties of the object:
        self.verbose = verbose
        self.__lazy = lazy
        self._cores = None
        self.__ranks_on_socket = 1
        self.__socket_master_rank = 0
        self._max_pos_per_read = None
        self.__bytes_per_pos = None

        # Now have to be careful here since the below properties are a function of the MPI rank
        self.__start_pos = None
        self.__rank_end_pos = None
        self.__end_pos = None
        self.__pixels_in_batch = None
        self.__compute_jobs = None

        # Determining the max size of the data that can be put into memory
        # all ranks go through this and they need to have this value any
        self._set_memory_and_cores(cores=cores, man_mem_limit=max_mem_mb,
                                   mem_multiplier=mem_multiplier)
        if verbose and self.mpi_rank == 0:
            print('Finished collecting info on memory and workers')
        self.duplicate_h5_groups = []
        self.partial_h5_groups = []
        self.process_name = process_name  # Reset this in the extended classes
        self.parms_dict = parms_dict

        """
        The name of the HDF5 dataset that should be present to signify which positions have already been computed
        This is NOT a fully private variable so that multiple processes can be run within a single group - Eg Fitter
        In the case of Fitter - this name can be changed from 'completed_guesses' to 'completed_fits'
        check_for_duplicates will be called by the Child class where they have the opportunity to change this
        variable before checking for duplicates
        """
        self._status_dset_name = 'completed_positions'

        self._results = None
        self.h5_results_grp = None

        # Check to see if the resuming feature has been implemented:
        self.__resume_implemented = False
        try:
            self._get_existing_datasets()
        except NotImplementedError:
            if verbose and self.mpi_rank == 0:
                print('It appears that this class may not be able to resume computations')
        except:
            # NameError for variables that don't exist
            # AttributeError for self.var_name that don't exist
            # TypeError (NoneType) etc.
            self.__resume_implemented = True

        if self.mpi_rank == 0:
            print('Consider calling test() to check results before calling compute() which computes on the entire'
                  ' dataset and writes results to the HDF5 file')

        self.duplicate_h5_groups, self.partial_h5_groups = self._check_for_duplicates()

    def __assign_job_indices(self):
        """
        Sets the start and end indices for each MPI rank
        """
        # First figure out what positions need to be computed
        self.__compute_jobs = np.where(self._h5_status_dset[()] == 0)[0]
        if self.verbose and self.mpi_rank == 0:
            if len(self.__compute_jobs) > 100:
                print('Among the {} positions in this dataset, {} positions '
                      'need to be computed'
                      '.'.format(self.h5_main.shape[0],
                                 len(self.__compute_jobs)))
            else:
                print('Among the {} positions in this dataset, the following '
                      'positions need to be computed: {}'
                      '.'.format(self.h5_main.shape[0], self.__compute_jobs))

        # integer division
        pos_per_rank = self.__compute_jobs.size // self.mpi_size
        if self.verbose and self.mpi_rank == 0:
            print('Each rank is required to work on {} of the {} (remaining) positions in this dataset'
                  '.'.format(pos_per_rank, self.__compute_jobs.size))

        # The start and end indices now correspond to the indices in the incomplete jobs rather than the h5 dataset
        self.__start_pos = self.mpi_rank * pos_per_rank
        self.__rank_end_pos = (self.mpi_rank + 1) * pos_per_rank
        self.__end_pos = int(min(self.__rank_end_pos, self.__start_pos + self._max_pos_per_read))
        if self.mpi_rank == self.mpi_size - 1:
            # Force the last rank to go to the end of the dataset
            self.__rank_end_pos = self.__compute_jobs.size

        if self.verbose:
            print('Rank {} will read positions {} to {} of {}'.format(self.mpi_rank, self.__start_pos,
                                                                      self.__rank_end_pos, self.h5_main.shape[0]))

    def _estimate_compute_time_per_pixel(self, *args, **kwargs):
        """
        Estimates how long it takes to compute an average pixel's worth of data. This information should be used by the
        user to limit the number of pixels that will be processed per batch to make best use of check-pointing. This
        function is exposed to the developer of the child classes. An approximate can be derived if it is simpler

        Returns
        -------

        """
        chosen_pos = np.random.randint(0, high=self.h5_main.shape[0]-1, size=5)
        t0 = tm.time()
        _ = parallel_compute(self.h5_main[chosen_pos, :], self._map_function, cores=1,
                             lengthy_computation=False, func_args=args, func_kwargs=kwargs, verbose=False)
        return (tm.time() - t0) / len(chosen_pos)

    def _get_pixels_in_current_batch(self):
        """
        Returns the indices of the pixels that will be processed in this batch.

        Returns
        -------
        pixels_in_batch : :class:`numpy.ndarray`
            1D array of unsigned integers denoting the pixels that will be read, processed, and written back to
        """
        return self.__pixels_in_batch

    def test(self, **kwargs):
        """
        Tests the process on a subset (for example a pixel) of the whole data. The class can be re-instantiated with
        improved parameters and tested repeatedly until the user is content, at which point the user can call
        :meth:`~pyUSID.processing.process.Process.compute` on the whole dataset.

        Notes
        -----
        This is not a function that is expected to be called in MPI

        Parameters
        ----------
        kwargs - dict, optional
            keyword arguments to test the process
        Returns
        -------
        """
        # All children classes should call super() OR ensure that they only work for self.mpi_rank == 0
        raise NotImplementedError('test_on_subset has not yet been implemented')

    def _check_for_duplicates(self):
        """
        Checks for instances where the process was applied to the same dataset with the same parameters

        Returns
        -------
        duplicate_h5_groups : list of h5py.Group objects
            List of groups satisfying the above conditions with completely computed results
        partial_h5_groups : list of h5py.Group objects
            List of groups satisfying the above conditions with partially computed results
        """
        if self.verbose and self.mpi_rank == 0:
            print('Checking for duplicates:')

        # This list will contain completed runs only
        existing = check_for_old(self.h5_main, self.process_name,
                                 new_parms=self.parms_dict,
                                 h5_parent_goup=self._h5_target_group,
                                 verbose=self.verbose and self.mpi_rank == 0)

        partial_h5_groups = []
        duplicate_h5_groups = []

        # First figure out which ones are partially completed:
        while len(existing) > 0:
            curr_group = existing.pop(0)
            """
            Earlier, we only checked the 'last_pixel' but to be rigorous we 
            should check self._status_dset_name
            The last_pixel attribute check may be deprecated in the future.
            Note that legacy computations did not have this dataset. We can add 
            to partially computed datasets
            """

            # Case 1: Modern book-keeping dataset available:
            if self._status_dset_name in curr_group.keys():

                status_dset = curr_group[self._status_dset_name]

                if not isinstance(status_dset, h5py.Dataset):
                    # We should not come here if things were implemented correctly
                    if self.mpi_rank == 0:
                        print('Results group: {} contained an object named: {} that should have been a dataset'
                              '.'.format(curr_group, self._status_dset_name))
                    continue

                if self.h5_main.shape[0] != status_dset.shape[0] or len(status_dset.shape) > 1 or \
                        status_dset.dtype != np.uint8:
                    if self.mpi_rank == 0:
                        print('Status dataset: {} was not of the expected shape or datatype'.format(status_dset))
                    continue

                # ##### ACTUAL COMPLETENESS TEST HERE #########

                completed_positions = np.sum(status_dset[()])

                if self.verbose and self.mpi_rank == 0:
                    print('{} has results that are {} % complete'
                          '.'.format(status_dset.name,
                                     int(100 * completed_positions / self.h5_main.shape[0])))

                # Case 1.A: Incomplete computation?
                if completed_positions < self.h5_main.shape[0]:
                    # If there are pixels uncompleted
                    # remove from duplicates and move to partial
                    if self.verbose and self.mpi_rank == 0:
                        print('moving {} to partial'.format(curr_group.name))
                    partial_h5_groups.append(curr_group)
                    # Let's write the legacy attribute for safety
                    curr_group.attrs['last_pixel'] = self.h5_main.shape[0]
                    # No further checks necessary
                    continue

                # Case 1.B: Complete computation:
                if self.verbose and self.mpi_rank == 0:
                    print('Moving {} to duplicate groups'.format(curr_group.name))
                duplicate_h5_groups.append(curr_group)
                continue

            # Case 2: Even the legacy book-keeping is absent:
            elif 'last_pixel' not in curr_group.attrs.keys():
                if self.mpi_rank == 0:
                    # Should not be coming here at all
                    print('Group: {} had neither the status HDF5 dataset or the legacy attribute: "last_pixel"'
                          '.'.format(curr_group))
                # Not sure what to do with such groups. Don't consider them
                continue

            # Case 3: Only the legacy book-keeping is available:
            else:
                last_pixel = curr_group.attrs['last_pixel']

                # Creating status dataset for forward compatibility:
                self._h5_status_dset = curr_group.create_dataset(
                    self._status_dset_name, dtype=np.uint8,
                    shape=(self.h5_main.shape[0],))
                if last_pixel > 0:
                    self._h5_status_dset[:last_pixel] = 1

                # Case 3.A: Partial
                if last_pixel < self.h5_main.shape[0]:
                    # move to partial
                    if self.verbose and self.mpi_rank == 0:
                        print('moving {} to partial since computation was {} % complete'
                              '.'.format(curr_group.name,
                                         int(100 * curr_group.attrs['last_pixel'] / self.h5_main.shape[0])))
                    partial_h5_groups.append(curr_group)
                    continue

                # Case 3.B: complete:
                else:
                    if self.verbose and self.mpi_rank == 0:
                        print('Moving {} to duplicate groups'.format(curr_group.name))
                    duplicate_h5_groups.append(curr_group)
                    continue

        if len(duplicate_h5_groups) > 0 and self.mpi_rank == 0:
            print('\nNote: ' + self.process_name + ' has already been performed with the same parameters before. '
                                                 'These results will be returned by compute() by default. '
                                                 'Set override to True to force fresh computation\n')
            print(duplicate_h5_groups)

        if len(partial_h5_groups) > 0 and self.mpi_rank == 0:
            print('\nNote: ' + self.process_name + ' has already been performed PARTIALLY with the same parameters. '
                                                 'compute() will resuming computation in the last group below. '
                                                 'To choose a different group call use_patial_computation()'
                                                 'Set override to True to force fresh computation or resume from a '
                                                 'data group besides the last in the list.\n')
            print(partial_h5_groups)

        return duplicate_h5_groups, partial_h5_groups

    def use_partial_computation(self, h5_partial_group=None):
        """
        Extracts the necessary parameters from the provided h5 group to resume computation

        Parameters
        ----------
        h5_partial_group : :class:`h5py.Group`
            Group containing partially computed results
        """
        # Attempt to automatically take partial results
        if h5_partial_group is None:
            if len(self.partial_h5_groups) < 1:
                raise ValueError('No group was found with partial results and no such group was provided')
            h5_partial_group = self.partial_h5_groups[-1]
        else:
            # Make sure that this group is among the legal ones already discovered:
            if h5_partial_group not in self.partial_h5_groups:
                raise ValueError('Provided group does not appear to be in the list of discovered groups')

        # Unnecessary since this will be defined at init
        # self.parms_dict = get_attributes(h5_partial_group)

        self.h5_results_grp = h5_partial_group

    def __set_cores(self, cores=None):
        """
        Checks number of CPU cores and sets the recommended number of cores to
        be used by analysis methods.
        This function can work with clusters with heterogeneous core counts
        (e.g. CADES SHPC Condo).

        Parameters
        ----------
        cores : uint, optional, Default = None (all or nearly all available)
            How many CPU cores to use for the computation.
        """
        if self.mpi_comm is None:
            min_free_cores = 1 + int(psutil.cpu_count() > 4)

            if cores is None:
                self._cores = max(1, psutil.cpu_count() - min_free_cores)
            else:
                if not isinstance(cores, int):
                    raise TypeError('cores should be an integer but got: {}'.format(cores))
                cores = int(abs(cores))
                self._cores = max(1, min(psutil.cpu_count(), cores))

            self.__socket_master_rank = 0
            self.__ranks_on_socket = 1
        else:
            # user-provided input cores will simply be ignored in an effort to use the entire CPU
            ranks_by_socket = group_ranks_by_socket(verbose=False)
            self.__socket_master_rank = ranks_by_socket[self.mpi_rank]
            # which ranks in this socket?
            ranks_on_this_socket = np.where(ranks_by_socket == self.__socket_master_rank)[0]
            # how many in this socket?
            self.__ranks_on_socket = ranks_on_this_socket.size
            # Force usage of all available memory
            man_mem_limit = None
            self._cores = 1
            # Disabling the following line since mpi4py and joblib didn't play well for Bayesian Inference
            # self._cores = self.__cores_per_rank = psutil.cpu_count() // self.__ranks_on_socket

    def _set_memory_and_cores(self, cores=None, man_mem_limit=None,
                              mem_multiplier=1.0):
        """
        Checks hardware limitations such as memory, number of CPU cores and sets the recommended data chunk sizes and
        the number of cores to be used by analysis methods. This function can work with clusters with heterogeneous
        memory sizes (e.g. CADES SHPC Condo).

        Parameters
        ----------
        cores : uint, optional, Default = 1
            How many cores to use for the computation.
        man_mem_limit : uint, optional, Default = None (all available memory)
            The amount a memory in Mb to use in the computation
        mem_multiplier : float, optional. Default = 1
            mem_multiplier is the number that will be multiplied with the
            (byte) size of a single position in the source dataset in order to
            better estimate the number of positions that can be processed at
            any given time (how many pixels of the source and results datasets
            can be retained in memory). The default value of 1.0 only accounts
            for the source dataset. A value greater than 1 would account for
            the size of results datasets as well. For example, if the result
            dataset is the same size and precision as the source dataset,
            the multiplier will be 2 (1 for source, 1 for result)
        """
        self.__set_cores(cores=cores)

        self.__set_memory(man_mem_limit=man_mem_limit,
                          mem_multiplier=mem_multiplier)

    def __set_memory(self, man_mem_limit=None, mem_multiplier=1.0):
        """
        Checks memory capabilities of each node and sets the recommended data
        chunk sizes to be used by analysis methods.
        This function can work with clusters with heterogeneous memory sizes
        (e.g. CADES SHPC Condo).

        Parameters
        ----------
        man_mem_limit : uint, optional, Default = None (all available memory)
            The amount a memory in Mb to use in the computation
        mem_multiplier : float, optional. Default = 1
            mem_multiplier is the number that will be multiplied with the
            (byte) size of a single position in the source dataset in order to
            better estimate the number of positions that can be processed at
            any given time (how many pixels of the source and results datasets
            can be retained in memory). The default value of 1.0 only accounts
            for the source dataset. A value greater than 1 would account for
            the size of results datasets as well. For example, if the result
            dataset is the same size and precision as the source dataset,
            the multiplier will be 2 (1 for source, 1 for result)
        """
        if not isinstance(mem_multiplier, float):
            raise TypeError('mem_multiplier must be a floating point number')
        mem_multiplier = abs(mem_multiplier)
        if mem_multiplier < 1:
            raise ValueError('mem_multiplier must be at least 1')

        avail_mem_bytes = get_available_memory()  # in bytes
        if self.verbose and self.mpi_rank == self.__socket_master_rank:
            # expected to be the same for all ranks so just use this.
            print('Rank {} - on socket with {} cores and {} avail. RAM shared '
                  'by {} ranks each given {} cores'
                  '.'.format(self.__socket_master_rank, psutil.cpu_count(),
                             format_size(avail_mem_bytes),
                             self.__ranks_on_socket, self._cores))

        if man_mem_limit is None:
            man_mem_limit = avail_mem_bytes
        else:
            if not isinstance(man_mem_limit, int):
                raise TypeError('man_mem_limit must be a whole number')
            # Note that man_mem_limit is specified in mega bytes
            man_mem_limit = abs(man_mem_limit) * 1024 ** 2  # in bytes
            if self.verbose and self.mpi_rank == 0:
                print('User has requested to use no more than {} of memory'
                      '.'.format(format_size(man_mem_limit)))

        max_mem_bytes = min(avail_mem_bytes, man_mem_limit)

        # Remember that multiple processes (either via MPI or joblib) will share this socket
        # This makes logical sense but there's always too much free memory and the
        # cores are starved.
        max_mem_per_worker = max_mem_bytes / (self._cores * self.__ranks_on_socket)
        if self.verbose and self.mpi_rank == self.__socket_master_rank:
            print('Rank {}: Each of the {} workers on this socket are allowed '
                  'to use {} of RAM'
                  '.'.format(self.mpi_rank,
                             self._cores * self.__ranks_on_socket,
                             format_size(max_mem_per_worker)))

        # Now calculate the number of positions OF RAW DATA ONLY that can be
        # stored in memory in one go PER worker
        self.__bytes_per_pos = self.h5_main.dtype.itemsize * self.h5_main.shape[1]
        if self.verbose and self.mpi_rank == 0:
            print('Each position in the SOURCE dataset is {} large'
                  '.'.format(format_size(self.__bytes_per_pos)))
        # Now multiply this with a factor that takes into account the expected
        # sizes of the results (Final and intermediate) datasets.
        self.__bytes_per_pos *= mem_multiplier
        if self.verbose and self.mpi_rank == 0 and mem_multiplier > 1:
            print('Each position of the source and results dataset(s) is {} '
                  'large.'.format(format_size(self.__bytes_per_pos)))

        self._max_pos_per_read = int(np.floor(max_mem_per_worker / self.__bytes_per_pos))

        if self.verbose and self.mpi_rank == self.__socket_master_rank:
            title = 'SOURCE dataset only'
            if mem_multiplier > 1:
                title = 'source and result(s) datasets'
            # expected to be the same for all ranks so just use this.
            print('Rank {}: Workers on this socket allowed to read {} '
                  'positions of the {} per chunk'
                  '.'.format(self.mpi_rank, self._max_pos_per_read, title))

    @staticmethod
    def _map_function(*args, **kwargs):
        """
        The function that manipulates the data on a single instance (position). This will be used by
        :meth:`~pyUSID.processing.process.Process._unit_computation` to process a chunk of data in parallel

        Parameters
        ----------
        args : list
            arguments to the function in the correct order
        kwargs : dict
            keyword arguments to the function
        Returns
        -------
        object
        """
        raise NotImplementedError('Please override the _unit_function specific to your process')

    def _read_data_chunk(self):
        """
        Reads a chunk of data for the intended computation into memory
        """
        if self.__start_pos < self.__rank_end_pos:
            self.__end_pos = int(min(self.__rank_end_pos, self.__start_pos + self._max_pos_per_read))

            # DON'T DIRECTLY apply the start and end indices anymore to the h5 dataset. Find out what it means first
            self.__pixels_in_batch = self.__compute_jobs[self.__start_pos: self.__end_pos]

            if self.verbose:
                print('Rank {} will read positions: {}'.format(self.mpi_rank, self.__pixels_in_batch))
                bytes_this_read = self.__bytes_per_pos * len(self.__pixels_in_batch)
                print('Rank {} will read {} of the SOURCE dataset'
                      '.'.format(self.mpi_rank, format_size(bytes_this_read)))
                if self.mpi_rank == self.__socket_master_rank:
                    tot_workers = self.__ranks_on_socket * self._cores
                    print('Rank: {} available memory: {}. '
                          '{} workers on this socket will in total read ~ {}'
                          '.'.format(self.mpi_rank,
                                     format_size(get_available_memory()),
                                     tot_workers,
                                     format_size(bytes_this_read * tot_workers)
                                     ))

            # Reading as Dask array to minimize memory copies when restructuring in child classes
            if self.__lazy:
                main_dset = lazy_load_array(self.h5_main)
            else:
                main_dset = self.h5_main

            self.data = main_dset[self.__pixels_in_batch, :]
            # DON'T update the start position

        else:
            if self.verbose:
                print('Rank {} - Finished reading all data!'.format(self.mpi_rank))
            self.data = None

    def _write_results_chunk(self):
        """
        Writes the computed results into appropriate datasets.
        This needs to be rewritten since the processed data is expected to be at least as large as the dataset
        """
        # Now update the start position
        self.__start_pos = self.__end_pos
        # This line can remain as is
        raise NotImplementedError('Please override the _set_results specific to your process')

    def _create_results_datasets(self):
        """
        Process specific call that will write the h5 group, guess dataset, corresponding spectroscopic datasets and also
        link the guess dataset to the spectroscopic datasets. It is recommended that the ancillary datasets be populated
        within this function.
        """
        raise NotImplementedError('Please override the _create_results_datasets specific to your process')

    def __create_compute_status_dataset(self):
        """
        Creates a dataset that keeps track of what pixels / rows have already been computed. Users are not expected to
        extend / modify this function.
        """
        # Check to make sure that such a group doesn't already exist
        if self._status_dset_name in self.h5_results_grp.keys():
            self._h5_status_dset = self.h5_results_grp[self._status_dset_name]
            if not isinstance(self._h5_status_dset, h5py.Dataset):
                raise ValueError('Provided results group: {} contains an expected object ({}) that is not a dataset'
                                 '.'.format(self.h5_results_grp, self._h5_status_dset))
            if self.h5_main.shape[0] != self._h5_status_dset.shape[0] or len(self._h5_status_dset.shape) > 1 or \
                    self._h5_status_dset.dtype != np.uint8:
                if self.mpi_rank == 0:
                    raise ValueError('Status dataset: {} was not of the expected shape or datatype'
                                     '.'.format(self._h5_status_dset))
        else:
            self._h5_status_dset = self.h5_results_grp.create_dataset(self._status_dset_name, dtype=np.uint8,
                                                                      shape=(self.h5_main.shape[0],))
            #  Could be fresh computation or resuming from a legacy computation
            if 'last_pixel' in self.h5_results_grp.attrs.keys():
                completed_pixels = self.h5_results_grp.attrs['last_pixel']
                if completed_pixels > 0:
                    self._h5_status_dset[:completed_pixels] = 1

    def _write_source_dset_provenance(self):
        """
        Writes path of HDF5 file and path of h5_main to the results group
        if results are being written to a new HDF5 file
        """
        if self.h5_main.file == self.h5_results_grp.file:
            return
        write_simple_attrs(self.h5_results_grp,
                           {'source_file_path': self.h5_main.file.filename,
                            'source_dataset_path': self.h5_main.name})

    def _get_existing_datasets(self):
        """
        The purpose of this function is to allow processes to resume from partly computed results
        Start with self.h5_results_grp
        """
        raise NotImplementedError('Please override the _get_existing_datasets specific to your process')

    def _unit_computation(self, *args, **kwargs):
        """
        The unit computation that is performed per data chunk. This allows room for any data pre / post-processing
        as well as multiple calls to parallel_compute if necessary
        """
        # TODO: Try to use the functools.partials to preconfigure the map function
        # cores = number of processes / rank here
        if self.verbose and self.mpi_rank == 0:
            print("Rank {} at Process class' default _unit_computation() that "
                  "will call parallel_compute()".format(self.mpi_rank))
        self._results = parallel_compute(self.data, self._map_function, cores=self._cores,
                                         lengthy_computation=False,
                                         func_args=args, func_kwargs=kwargs,
                                         verbose=self.verbose)

    def compute(self, override=False, *args, **kwargs):
        """
        Creates placeholders for the results, applies the :meth:`~pyUSID.processing.process.Process._unit_computation`
        to chunks of the dataset

        Parameters
        ----------
        override : bool, optional. default = False
            By default, compute will simply return duplicate results to avoid recomputing or resume computation on a
            group with partial results. Set to True to force fresh computation.
        args : list
            arguments to the mapped function in the correct order
        kwargs : dict
            keyword arguments to the mapped function

        Returns
        -------
        h5_results_grp : :class:`h5py.Group`
            Group containing all the results
        """

        class SimpleFIFO(object):
            """
            Simple class that maintains a moving average of some numbers.
            """

            def __init__(self, length=5):
                """
                Create a SimpleFIFO object

                Parameters
                ----------
                length : unsigned integer
                    Number of values that need to be maintained for the moving average
                """
                self.__queue = list()
                if not isinstance(length, int):
                    raise TypeError('length must be a positive integer')
                if length <= 0:
                    raise ValueError('length must be a positive integer')
                self.__max_length = length
                self.__count = 0

            def put(self, item):
                """
                Adds the item to the internal queue. If the size of the queue exceeds its capacity, the oldest
                item is removed.

                Parameters
                ----------
                item : float or int
                    Any real valued number
                """
                if (not isinstance(item, Number)) or isinstance(item, complex):
                    raise TypeError('Provided item: {} is not a Number'.format(item))
                self.__queue.append(item)
                self.__count += 1
                if len(self.__queue) > self.__max_length:
                    _ = self.__queue.pop(0)

            def get_mean(self):
                """
                Returns the average of the elements within the queue

                Returns
                -------
                avg : number.Number
                    Mean of all elements within the queue
                """
                return np.mean(self.__queue)

            def get_cycles(self):
                """
                Returns the number of items that have been added to the queue in total

                Returns
                -------
                count : int
                    number of items that have been added to the queue in total
                """
                return self.__count

        if not override:
            if len(self.duplicate_h5_groups) > 0:
                if self.mpi_rank == 0:
                    print('Returned previously computed results at ' + self.duplicate_h5_groups[-1].name)
                self.h5_results_grp = self.duplicate_h5_groups[-1]
                return self.duplicate_h5_groups[-1]
            elif len(self.partial_h5_groups) > 0 and self.h5_results_grp is None:
                if self.mpi_rank == 0:
                    print('Resuming computation in group: ' + self.partial_h5_groups[-1].name)
                self.use_partial_computation()

        resuming = False
        if self.h5_results_grp is None:
            # starting fresh
            if self.verbose and self.mpi_rank == 0:
                print('Creating HDF5 group and datasets to hold results')
            self._create_results_datasets()
            self._write_source_dset_provenance()
        else:
            # resuming from previous checkpoint
            resuming = True
            self._get_existing_datasets()

        self.__create_compute_status_dataset()

        if resuming and self.mpi_rank == 0:
            percent_complete = int(100 * len(np.where(self._h5_status_dset[()] == 1)[0]) /
                                   self._h5_status_dset.shape[0])
            print('Resuming computation. {}% completed already'.format(percent_complete))

        self.__assign_job_indices()

        # Not sure if this is necessary but I don't think it would hurt either
        if self.mpi_comm is not None:
            self.mpi_comm.barrier()

        compute_times = SimpleFIFO(5)
        write_times = SimpleFIFO(5)
        orig_rank_start = self.__start_pos

        if self.mpi_rank == 0 and self.mpi_size == 1:
            if self.__resume_implemented:
                print('\tThis class (likely) supports interruption and resuming of computations!\n'
                      '\tIf you are operating in a python console, press Ctrl+C or Cmd+C to abort\n'
                      '\tIf you are in a Jupyter notebook, click on "Kernel">>"Interrupt"\n'
                      '\tIf you are operating on a cluster and your job gets killed, re-run the job to resume\n')
            else:
                print('\tThis class does NOT support interruption and resuming of computations.\n'
                      '\tIn order to enable this feature, simply implement the _get_existing_datasets() function')

        if self.verbose and self.mpi_rank == self.__socket_master_rank:
            print('Rank: {} - with nothing loaded has {} free memory'
                  ''.format(self.mpi_rank, format_size(get_available_memory())))

        self._read_data_chunk()

        if self.mpi_comm is not None:
            self.mpi_comm.barrier()

        if self.verbose and self.mpi_rank == self.__socket_master_rank:
            print('Rank: {} - with only raw data loaded has {} free memory'
                  ''.format(self.mpi_rank, format_size(get_available_memory())))

        while self.data is not None:

            num_jobs_in_batch = self.__end_pos - self.__start_pos

            t_start_1 = tm.time()

            self._unit_computation(*args, **kwargs)

            comp_time = np.round(tm.time() - t_start_1, decimals=2)  # in seconds
            time_per_pix = comp_time / num_jobs_in_batch
            compute_times.put(time_per_pix)

            if self.verbose:
                print('Rank {} - computed chunk in {} or {} per pixel. Average: {} per pixel'
                      '.'.format(self.mpi_rank, format_time(comp_time), format_time(time_per_pix),
                                 format_time(compute_times.get_mean())))

            # Ranks can become memory starved. Check memory usage - raw data + results in memory at this point
            if self.verbose and self.mpi_rank == self.__socket_master_rank:
                print('Rank: {} - now holding onto raw data + results has {} free memory'
                      ''.format(self.mpi_rank, format_size(get_available_memory())))

            t_start_2 = tm.time()
            self._write_results_chunk()

            # NOW, update the positions. Users are NOT allowed to touch start and end pos
            self.__start_pos = self.__end_pos
            # Leaving in this provision that will allow restarting of processes
            if self.mpi_size == 1:
                self.h5_results_grp.attrs['last_pixel'] = self.__end_pos
            # Child classes don't even have to worry about flushing. Process will do it.
            self.h5_main.file.flush()

            dump_time = np.round(tm.time() - t_start_2, decimals=2)
            write_times.put(dump_time / num_jobs_in_batch)

            if self.verbose:
                print('Rank {} - wrote its {} pixel chunk in {}'.format(self.mpi_rank,
                                                                        num_jobs_in_batch,
                                                                        format_time(dump_time)))

            time_remaining = (self.__rank_end_pos - self.__end_pos) * \
                             (compute_times.get_mean() + write_times.get_mean())

            if self.verbose or self.mpi_rank == 0:
                percent_complete = int(100 * (self.__end_pos - orig_rank_start) /
                                       (self.__rank_end_pos - orig_rank_start))
                print('Rank {} - {}% complete. Time remaining: {}'.format(self.mpi_rank, percent_complete,
                                                                          format_time(time_remaining)))

            # All ranks should mark the pixels for this batch as completed. 'last_pixel' attribute will be updated later
            # Setting each section to 1 independently
            for curr_slice in integers_to_slices(self.__pixels_in_batch):
                self._h5_status_dset[curr_slice] = 1

            self._read_data_chunk()

        if self.verbose:
            print('Rank {} - Finished computing all jobs!'.format(self.mpi_rank))

        if self.mpi_comm is not None:
            self.mpi_comm.barrier()

        if self.mpi_rank == 0:
            print('Finished processing the entire dataset!')

        # Update the legacy 'last_pixel' attribute here:
        if self.mpi_rank == 0:
            self.h5_results_grp.attrs['last_pixel'] = self.h5_main.shape[0]

        return self.h5_results_grp