import os
import shutil
from itertools import chain
from time import sleep
from glob import glob
import warnings
import pickle
import copy
import math

import joblib
from joblib.externals import cloudpickle

from ..utils import _get_fn_name, prod, progbar
from .prepare import _parse_combos, _parse_constants, _parse_attrs
from .combo_runner import _combo_runner, combo_runner_to_ds


BTCH_NM = "xyz-batch-{}.jbdmp"
RSLT_NM = "xyz-result-{}.jbdmp"
FNCT_NM = "xyz-function.clpkl"
INFO_NM = "xyz-settings.jbdmp"


class XYZError(Exception):
    pass


# --------------------------------- parsing --------------------------------- #

def parse_crop_details(fn, crop_name, crop_parent):
    """Work out how to structure the sowed data.

    Parameters
    ----------
    fn : callable, optional
        Function to infer name crop_name from, if not given.
    crop_name : str, optional
        Specific name to give this set of runs.
    crop_parent : str, optional
        Specific directory to put the ".xyz-{crop_name}/" folder in
        with all the cases and results.

    Returns
    -------
    crop_location : str
        Full path to the crop-folder.
    crop_name : str
        Name of the crop.
    crop_parent : str
        Parent folder of the crop.
    """
    if crop_name is None:
        if fn is None:
            raise ValueError("Either `fn` or `crop_name` must be give.")
        crop_name = _get_fn_name(fn)

    crop_parent = crop_parent if crop_parent is not None else os.getcwd()
    crop_location = os.path.join(crop_parent, ".xyz-{}".format(crop_name))

    return crop_location, crop_name, crop_parent


def parse_fn_runner_harvester(fn, runner, harvester):
    """
    """
    if harvester is not None:
        if (runner is not None) or (fn is not None):
            warnings.warn(
                "If `harvester` is set `runner` and `fn` are ignored.")
        harvester = harvester
        runner = harvester.runner
        fn = harvester.runner.fn
    elif runner is not None:
        if fn is not None:
            warnings.warn("If `runner` is set `fn` is ignored.")
        harvester = None
        runner = runner
        fn = runner.fn
    else:
        harvester = None
        runner = None
        fn = fn
    return fn, runner, harvester


class Crop(object):
    """Encapsulates all the details describing a single 'crop', that is,
    its location, name, and batch size/number. Also allows tracking of
    crop's progress, and experimentally, automatic submission of
    workers to grid engine to complete un-grown cases. Can also be instantiated
    directly from a :class:`~xyzpy.Runner` or :class:`~xyzpy.Harvester`
    instance.

    Parameters
    ----------
    fn : callable, optional
        Target function - Crop `name` will be inferred from this if
        not given explicitly. If given, `Sower` will also default
        to saving a version of `fn` to disk for `batch.grow` to use.
    name : str, optional
        Custom name for this set of runs - must be given if `fn`
        is not.
    parent_dir : str, optional
        If given, alternative directory to put the ".xyz-{name}/"
        folder in with all the cases and results.
    save_fn : bool, optional
        Whether to save the function to disk for `batch.grow` to use.
        Will default to True if `fn` is given.
    batchsize : int, optional
        How many cases to group into a single batch per worker.
        By default, batchsize=1. Cannot be specified if `num_batches`
        is.
    num_batches : int, optional
        How many total batches to aim for, cannot be specified if
        `batchsize` is.
    runner : xyzpy.Runner, optional
        A Runner instance, from which the `fn` can be inferred and
        which can also allow the Crop to reap itself straight to a
        dataset.
    harvester : xyzpy.Harvester, optional
        A Harvester instance, from which the `fn` can be inferred and
        which can also allow the Crop to reap itself straight to a
        on-disk dataset.
    autoload : bool, optional
        If True, check for the existence of a Crop written to disk
        with the same location, and if found, load it.

    See Also
    --------
    Runner.Crop, Harvester.Crop
    """

    def __init__(self, *,
                 fn=None,
                 name=None,
                 parent_dir=None,
                 save_fn=None,
                 batchsize=None,
                 num_batches=None,
                 runner=None,
                 harvester=None,
                 autoload=True):

        self._fn, self.runner, self.harvester = \
            parse_fn_runner_harvester(fn, runner, harvester)

        self.name = name
        self.parent_dir = parent_dir
        self.save_fn = save_fn
        self.batchsize = batchsize
        self.num_batches = num_batches
        self._batch_remainder = None

        # Work out the full directory for the crop
        self.location, self.name, self.parent_dir = \
            parse_crop_details(self._fn, self.name, self.parent_dir)

        # Save function so it can be automatically loaded with all deps?
        if (fn is None) and (save_fn is True):
            raise ValueError("Must specify a function for it to be saved!")
        self.save_fn = save_fn is not False

    # ------------------------------- methods ------------------------------- #

    def choose_batch_settings(self, combos):
        """Work out how to divide all cases into batches, i.e. ensure
        that ``batchsize * num_batches >= num_cases``.
        """
        n = prod(len(x) for _, x in combos)

        if (self.batchsize is not None) and (self.num_batches is not None):
            # Check that they are set correctly
            pos_tot = self.batchsize * self.num_batches
            if not (n <= pos_tot < n + self.batchsize):
                raise ValueError("`batchsize` and `num_batches` cannot both"
                                 "be specified if they do not not multiply"
                                 "to the correct number of total cases.")

        # Decide based on batchsize
        elif self.num_batches is None:
            if self.batchsize is None:
                self.batchsize = 1

            if not isinstance(self.batchsize, int):
                raise TypeError("`batchsize` must be an integer.")
            if self.batchsize < 1:
                raise ValueError("`batchsize` must be >= 1.")

            self.num_batches = math.ceil(n / self.batchsize)
            self._batch_remainder = 0

        # Decide based on num_batches:
        else:
            if not isinstance(self.num_batches, int):
                raise TypeError("`num_batches` must be an integer.")
            if self.num_batches < 1:
                raise ValueError("`num_batches` must be >= 1.")

            self.batchsize, self._batch_remainder = divmod(n, self.num_batches)

    def ensure_dirs_exists(self):
        """Make sure the directory structure for this crop exists.
        """
        os.makedirs(os.path.join(self.location, "batches"), exist_ok=True)
        os.makedirs(os.path.join(self.location, "results"), exist_ok=True)

    def save_info(self, combos):
        """Save information about the sowed cases.
        """
        # If saving Harvester or Runner, strip out function information so
        #   as just to use pickle.
        if self.harvester is not None:
            harvester_copy = copy.deepcopy(self.harvester)
            harvester_copy.runner.fn = None
            hrvstr_pkl = pickle.dumps(harvester_copy)
            runner_pkl = None
        elif self.runner is not None:
            hrvstr_pkl = None
            runner_copy = copy.deepcopy(self.runner)
            runner_copy.fn = None
            runner_pkl = pickle.dumps(runner_copy)
        else:
            hrvstr_pkl = None
            runner_pkl = None

        joblib.dump({
            'combos': combos,
            'batchsize': self.batchsize,
            'num_batches': self.num_batches,
            '_batch_remainder': self._batch_remainder,
            'harvester': hrvstr_pkl,
            'runner': runner_pkl,
        }, os.path.join(self.location, INFO_NM))

    def load_info(self):
        """Load the full settings from disk.
        """
        sfile = os.path.join(self.location, INFO_NM)

        if not os.path.isfile(sfile):
            raise XYZError("Settings can't be found at {}.".format(sfile))
        else:
            return joblib.load(sfile)

    def _sync_info_from_disk(self):
        """Load information about the saved cases.
        """
        settings = self.load_info()
        self.batchsize = settings['batchsize']
        self.num_batches = settings['num_batches']
        self._batch_remainder = settings['_batch_remainder']

        hrvstr_pkl = settings['harvester']
        harvester = None if hrvstr_pkl is None else pickle.loads(hrvstr_pkl)
        runner_pkl = settings['runner']
        runner = None if runner_pkl is None else pickle.loads(runner_pkl)

        self._fn, self.runner, self.harvester = \
            parse_fn_runner_harvester(None, runner, harvester)

    def save_function_to_disk(self):
        """Save the base function to disk using cloudpickle
        """
        joblib.dump(cloudpickle.dumps(self._fn),
                    os.path.join(self.location, FNCT_NM))

    def load_function(self):
        """Load the saved function from disk, and try to re-insert it back into
        Harvester or Runner if present.
        """
        self._fn = cloudpickle.loads(joblib.load(
            os.path.join(self.location, FNCT_NM)))

        if self.harvester is not None:
            if self.harvester.runner.fn is None:
                self.harvester.runner.fn = self._fn
            else:
                # TODO: check equality?
                raise XYZError("Trying to load this Crop's function, {}, from "
                               "disk but its Harvester already has a function "
                               "set: {}.".format(self._fn,
                                                 self.harvester.runner.fn))
        elif self.runner is not None:
            if self.runner.fn is None:
                self.runner.fn = self._fn
            else:
                # TODO: check equality?
                raise XYZError("Trying to load this Crop's function, {}, from "
                               "disk but its Runner already has a function "
                               "set: {}.".format(self._fn,
                                                 self.runner.fn))

    def prepare(self, combos):
        """Write information about this crop and the supplied combos to disk.
        Typically done at start of sow, not when Crop instantiated.
        """
        self.ensure_dirs_exists()
        if self.save_fn:
            self.save_function_to_disk()
        self.save_info(combos)

    def is_prepared(self):
        """Check whether this crop has been written to disk.
        """
        return os.path.exists(os.path.join(self.location, INFO_NM))

    def calc_progress(self):
        """Calculate how much progressed has been made in growing the cases.
        """
        if self.is_prepared():
            self._sync_info_from_disk()
            self._num_sown_batches = len(glob(
                os.path.join(self.location, "batches", BTCH_NM.format("*"))))
            self._num_results = len(glob(
                os.path.join(self.location, "results", RSLT_NM.format("*"))))
        else:
            self._num_sown_batches = -1
            self._num_results = -1

    def is_ready_to_reap(self):
        self.calc_progress()
        return (
            self._num_results > 0 and
            (self._num_results == self.num_sown_batches)
        )

    def missing_results(self):
        """
        """
        self.calc_progress()

        def no_result_exists(x):
            return not os.path.isfile(
                os.path.join(self.location, "results", RSLT_NM.format(x)))

        return tuple(filter(no_result_exists, range(1, self.num_batches + 1)))

    def delete_all(self):
        # delete everything
        shutil.rmtree(self.location)

    def __str__(self):
        # Location and name, underlined
        if not os.path.exists(self.location):
            return self.location + "\n * Not yet sown, or already reaped * \n"

        loc_len = len(self.location)
        name_len = len(self.name)

        self.calc_progress()
        percentage = 100 * self._num_results / self.num_batches

        # Progress bar
        total_bars = 20
        bars = int(percentage * total_bars / 100)

        return ("\n"
                "{location}\n"
                "{under_crop_dir}{under_crop_name}\n"
                "{num_results} / {total} batches of size {bsz} completed\n"
                "[{done_bars}{not_done_spaces}] : {percentage:.1f}%"
                "\n").format(
            location=self.location,
            under_crop_dir="-" * (loc_len - name_len),
            under_crop_name="=" * name_len,
            num_results=self._num_results,
            total=self.num_batches,
            bsz=self.batchsize,
            done_bars="#" * bars,
            not_done_spaces=" " * (total_bars - bars),
            percentage=percentage,
        )

    def __repr__(self):
        if not os.path.exists(self.location):
            progress = "*reaped or unsown*"
        else:
            self.calc_progress()
            progress = "{}/{}".format(self._num_results, self.num_batches)

        msg = "<Crop(name='{}', progress={}, batchsize={})>"
        return msg.format(self.name, progress, self.batchsize)

    def sow_combos(self, combos, constants=None, verbosity=1):
        """Sow to disk.
        """
        combos = _parse_combos(combos)
        constants = _parse_constants(constants)

        if self.runner is not None:
            constants = {**self.runner._constants, **constants}
            constants = {**self.runner._resources, **constants}

        # Sort to ensure order remains same for reaping results
        #   (don't want to hash kwargs)
        combos = sorted(combos, key=lambda x: x[0])

        with Sower(self, combos) as sow_fn:
            _combo_runner(fn=sow_fn, combos=combos, constants=constants,
                          verbosity=verbosity)

    def grow(self, batch_ids, **combo_runner_opts):
        """Grow specific batch numbers using this process.
        """
        if isinstance(batch_ids, int):
            batch_ids = (batch_ids,)

        _combo_runner(grow, combos=(('batch_number', batch_ids),),
                      constants={'verbosity': 0, 'crop': self},
                      **combo_runner_opts)

    def grow_missing(self, **combo_runner_opts):
        """Grow any missing results using this process.
        """
        self.grow(batch_ids=self.missing_results(), **combo_runner_opts)

    def reap_combos(self, wait=False, clean_up=True):
        """Reap already sown and grow results from specified crop.

        Parameters
        ----------
        wait : bool, optional
            Whether to wait for results to appear. If false (default) all
            results need to be in place before the reap.

        Returns
        -------
        results : nested tuple
            'N-dimensional' tuple containing the results.
        """
        if not (wait or self.is_ready_to_reap()):
            raise XYZError("This crop is not ready to reap yet - results are"
                           " missing.")

        # Load same combinations as cases saved with
        settings = joblib.load(os.path.join(self.location, INFO_NM))

        with Reaper(self, num_batches=settings['num_batches'],
                    wait=wait) as reap_fn:

            results = _combo_runner(fn=reap_fn,
                                    combos=settings['combos'],
                                    constants={})

        if clean_up:
            self.delete_all()

        return results

    def reap_combos_to_ds(self,
                          var_names=None,
                          var_dims=None,
                          var_coords=None,
                          constants=None,
                          attrs=None,
                          parse=True,
                          wait=False,
                          clean_up=True):
        """Reap a function over sowed combinations and output to a Dataset.

        Parameters
        ----------
        var_names : str, sequence of strings, or None
            Variable name(s) of the output(s) of `fn`, set to None if
            fn outputs data already labelled in a Dataset or DataArray.
        var_dims : sequence of either strings or string sequences, optional
            'Internal' names of dimensions for each variable, the values for
            each dimension should be contained as a mapping in either
            `var_coords` (not needed by `fn`) or `constants` (needed by `fn`).
        var_coords : mapping, optional
            Mapping of extra coords the output variables may depend on.
        constants : mapping, optional
            Arguments to `fn` which are not iterated over, these will be
            recorded either as attributes or coordinates if they are named
            in `var_dims`.
        resources : mapping, optional
            Like `constants` but they will not be recorded.
        attrs : mapping, optional
            Any extra attributes to store.
        wait : bool, optional
            Whether to wait for results to appear. If false (default) all
            results need to be in place before the reap.

        Returns
        -------
        xarray.Dataset
            Multidimensional labelled dataset contatining all the results.
        """
        if not (wait or self.is_ready_to_reap()):
            raise XYZError(
                "This crop is not ready to reap yet - results are missing.")

        # Load exact same combinations as cases saved with
        settings = joblib.load(os.path.join(self.location, INFO_NM))

        if parse:
            constants = _parse_constants(constants)
            attrs = _parse_attrs(attrs)

        with Reaper(self, num_batches=settings['num_batches'],
                    wait=wait) as reap_fn:

            # Move constants into attrs, so as not to pass them to the Reaper
            #   when if fact they were meant for the original function.
            ds = combo_runner_to_ds(fn=reap_fn,
                                    combos=settings['combos'],
                                    var_names=var_names,
                                    var_dims=var_dims,
                                    var_coords=var_coords,
                                    constants={},
                                    resources={},
                                    attrs={**attrs, **constants},
                                    parse=parse)

        if clean_up:
            self.delete_all()

        return ds

    def reap_runner(self, runner, wait=False, clean_up=True):
        """Reap a Crop over sowed combos and save to a dataset defined by a
        Runner.
        """
        # Can ignore `Runner.resources` as they play no part in desecribing the
        #   output, though they should be supplied to sow and thus grow.
        ds = self.reap_combos_to_ds(
            var_names=runner._var_names,
            var_dims=runner._var_dims,
            var_coords=runner._var_coords,
            constants=runner._constants,
            attrs=runner._attrs,
            parse=False,
            wait=wait,
            clean_up=clean_up)
        runner.last_ds = ds
        return ds

    def reap_harvest(self, harvester, wait=False, sync=True, overwrite=None,
                     clean_up=True):
        """
        """
        if harvester is None:
            raise ValueError("Cannot reap and harvest if no Harvester is set.")

        ds = self.reap_runner(harvester.runner, wait=wait, clean_up=clean_up)
        self.harvester.add_ds(ds, sync=sync, overwrite=overwrite)
        return ds

    def reap(self, wait=False, sync=True, overwrite=None, clean_up=True):
        """Reap sown and grown combos from disk. Return a dataset if a runner
        or harvester is set, otherwise, the raw nested tuple.

        Parameters
        ----------
        wait : bool, optional
            Whether to wait for results to appear. If false (default) all
            results need to be in place before the reap.
        sync : bool, optional
            Immediately sync the new dataset with the on-disk full dataset
            if a harvester is used.
        overwrite : bool, optional
            How to compare data when syncing to on-disk dataset.
            If ``None``, (default) merge as long as no conflicts.
            ``True``: overwrite with the new data. ``False``, discard any
            new conflicting data.
        clean_up : bool, optional
            Whether to delete all the batch files once the results have been
            gathered.

        Returns
        -------
        nested tuple or xarray.Dataset
        """
        if self.harvester is not None:
            return self.reap_harvest(self.harvester, clean_up=clean_up,
                                     wait=wait, sync=sync, overwrite=overwrite)
        elif self.runner is not None:
            return self.reap_runner(self.runner, wait=wait, clean_up=clean_up)
        else:
            return self.reap_combos(wait=wait, clean_up=clean_up)

    def check_bad(self, delete_bad=True):
        """Check that the result dumps are not bad -> sometimes length does not
        match the batch. Optionally delete these so that they can be re-grown.

        Parameters
        ----------
        delete_bad : bool
            Delete bad results as they are come across.

        Returns
        -------
        bad_ids : tuple
            The bad batch numbers.
        """
        # XXX: work out why this is needed sometimes on network filesystems.
        result_files = glob(
            os.path.join(self.location, "results", RSLT_NM.format("*")))

        bad_ids = []

        for result_file in result_files:
            # load corresponding batch file to check length.
            result_num = os.path.split(
                result_file)[-1].strip("xyz-result-").strip(".jbdmp")
            batch_file = os.path.join(
                self.location, "batches", BTCH_NM.format(result_num))

            batch = joblib.load(batch_file)

            try:
                result = joblib.load(result_file)
                unloadable = False
            except Exception as e:
                unloadable = True
                err = e

            if unloadable or (len(result) != len(batch)):
                msg = "result {} is bad".format(result_file)
                msg += "." if not delete_bad else " - deleting it."
                msg += " Error was: {}".format(err) if unloadable else ""
                print(msg)

                if delete_bad:
                    os.remove(result_file)

                bad_ids.append(result_num)

        return tuple(bad_ids)

    #  ----------------------------- properties ----------------------------- #

    def _get_fn(self):
        return self._fn

    def _set_fn(self, fn):
        if self.save_fn is None and fn is not None:
            self.save_fn = True
        self._fn = fn

    def _del_fn(self):
        self._fn = None
        self.save_fn = False

    fn = property(_get_fn, _set_fn, _del_fn,
                  "Function to save with the Crop for automatic loading and "
                  "running. Default crop name will be inferred from this if"
                  "not given explicitly as well.")

    @property
    def num_sown_batches(self):
        """Total number of batches to be run/grown.
        """
        self.calc_progress()
        return self._num_sown_batches

    @property
    def num_results(self):
        self.calc_progress()
        return self._num_results


class Sower(object):
    """Class for sowing a 'crop' of batched combos to then 'grow' (on any
    number of workers sharing the filesystem) and then reap.
    """

    def __init__(self, crop, combos):
        """
        Parameters
        ----------
            crop : xyzpy.batch.Crop instance
                Description of where and how to store the cases and results.
            combos : mapping_like
                Description of combinations from which to sow cases from.

        """
        self.combos = combos
        self.crop = crop
        self.crop.choose_batch_settings(combos)
        # Internal:
        self._batch_cases = []  # collects cases to be written in single batch
        self._counter = 0  # counts how many cases are in batch so far
        self._batch_counter = 0  # counts how many batches have been written

    def save_batch(self):
        """Save the current batch of cases to disk using joblib.dump
         and start the next batch.
        """
        self._batch_counter += 1
        joblib.dump(self._batch_cases, os.path.join(
            self.crop.location, "batches", BTCH_NM.format(self._batch_counter))
        )
        self._batch_cases = []
        self._counter = 0

    # Context manager #

    def __enter__(self):
        self.crop.prepare(self.combos)
        return self

    def __call__(self, **kwargs):
        self._batch_cases.append(kwargs)
        self._counter += 1

        # when the number of cases doesn't divide the number of batches we
        #     distribute the remainder among the first crops.
        extra_batch = self._batch_counter < self.crop._batch_remainder

        if self._counter == self.crop.batchsize + int(extra_batch):
            self.save_batch()

    def __exit__(self, exception_type, exception_value, traceback):
        # Make sure any overfill also saved
        if self._batch_cases:
            self.save_batch()


def grow(batch_number, crop=None, fn=None, check_mpi=True,
         verbosity=2, debugging=False):
    """Automatically process a batch of cases into results. Should be run in an
    ".xyz-{fn_name}" folder.

    Parameters
    ----------
    batch_number : int
        Which batch to 'grow' into a set of results.
    crop : xyzpy.batch.Crop instance
        Description of where and how to store the cases and results.
    fn : callable, optional
        If specified, the function used to generate the results, otherwise
        the function will be loaded from disk.
    check_mpi : bool, optional
        Whether to check if the process is rank 0 and only save results if
        so - allows mpi functions to be simply used. Defaults to true,
        this should only be turned off if e.g. a pool of workers is being
        used to run different ``grow`` instances.
    verbosity : {0, 1, 2}, optional
        How much information to show.
    debugging : bool, optional
        Set logging level to DEBUG.
    """
    if debugging:
        import logging
        logger = logging.getLogger()
        logger.setLevel(logging.DEBUG)

    if crop is None:
        if os.path.relpath('.', '..')[:5] != ".xyz-":
            raise XYZError("`grow` should be run in a "
                           "\"{crop_parent}/.xyz-{crop_name}\" folder, else "
                           "`crop_parent` and `crop_name` (or `fn`) should be "
                           "specified.")
        crop_location = os.getcwd()
    else:
        crop_location = crop.location

    # load function
    if fn is None:
        fn = cloudpickle.loads(
            joblib.load(os.path.join(crop_location, FNCT_NM)))

    # load cases to evaluate
    cases = joblib.load(
        os.path.join(crop_location, "batches", BTCH_NM.format(batch_number)))

    if len(cases) == 0:
        raise ValueError("Something has gone wrong with the loading of "
                         "batch {} ".format(BTCH_NM.format(batch_number)) +
                         "for the crop at {}.".format(crop.location))

    # maybe want to run grow as mpiexec (i.e. `fn` itself in parallel),
    # so only save and delete on rank 0
    if check_mpi and 'OMPI_COMM_WORLD_RANK' in os.environ:  # pragma: no cover
        rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
    elif check_mpi and 'PMI_RANK' in os.environ:  # pragma: no cover
        rank = int(os.environ['PMI_RANK'])
    else:
        rank = 0

    if rank == 0:
        results = []

        descr = "Batch {}".format(batch_number)

        pbar = progbar(range(len(cases)), disable=verbosity <= 0, desc=descr)
        for i in pbar:
            if verbosity >= 2:
                pbar.set_description(descr + ": {}".format(cases[i]))

            # compute and store result!
            results.append(fn(**cases[i]))

        if len(results) != len(cases):
            raise ValueError("Something has gone wrong with processing "
                             "batch {} ".format(BTCH_NM.format(batch_number)) +
                             "for the crop at {}.".format(crop.location))

        # save to results
        joblib.dump(tuple(results), os.path.join(
            crop_location, "results", RSLT_NM.format(batch_number)))
    else:
        for case in cases:
            # worker: just help compute the result!
            fn(**case)


# --------------------------------------------------------------------------- #
#                              Gathering results                              #
# --------------------------------------------------------------------------- #

class Reaper(object):
    """Class that acts as a stateful function to retrieve already sown and
    grow results.
    """

    def __init__(self, crop, num_batches, wait=False):
        """Class for retrieving the batched, flat, 'grown' results.

        Parameters
        ----------
            crop : xyzpy.batch.Crop instance
                Description of where and how to store the cases and results.
        """
        self.crop = crop

        files = (os.path.join(self.crop.location, "results", RSLT_NM.format(i))
                 for i in range(1, num_batches + 1))

        def _load(x):
            res = joblib.load(x)
            if (res is None) or len(res) == 0:
                raise ValueError("Something not right: result {} contains "
                                 "no data upon joblib.load".format(x))
            return res

        def wait_to_load(x):
            while not os.path.exists(x):
                sleep(0.2)

            if os.path.isfile(x):
                return _load(x)
            else:
                raise ValueError("{} is not a file.".format(x))

        self.results = chain.from_iterable(map(
            wait_to_load if wait else _load, files))

    def __enter__(self):
        return self

    def __call__(self, **kwargs):
        return next(self.results)

    def __exit__(self, exception_type, exception_value, traceback):
        # Check everything gone acccording to plan
        if tuple(self.results):
            raise XYZError("Not all results reaped!")


# --------------------------------------------------------------------------- #
#                     Automatic Batch Submission Scripts                      #
# --------------------------------------------------------------------------- #

_BASE_QSUB_SCRIPT = """#!/bin/bash -l
#$ -S /bin/bash
#$ -l h_rt={hours}:{minutes}:{seconds},mem={gigabytes}G
#$ -l tmpfs={temp_gigabytes}G
{extra_resources}
#$ -N {name}
mkdir -p {output_directory}
#$ -wd {output_directory}
#$ -pe {pe} {num_procs}
#$ -t {run_start}-{run_stop}
cd {working_directory}
export OMP_NUM_THREADS={num_threads}
tmpfile=$(mktemp .xyzpy-qsub.XXXXXXXX)
cat <<EOF > $tmpfile
{setup}
from xyzpy.gen.batch import grow, Crop
crop = Crop(name='{name}')
"""

_QSUB_GROW_ALL_SCRIPT = """grow($SGE_TASK_ID, crop=crop, debugging={debugging})
"""

_QSUB_GROW_PARTIAL_SCRIPT = """batch_ids = {batch_ids}
grow(batch_ids[$SGE_TASK_ID - 1], crop=crop, debugging={debugging})
"""

_BASE_QSUB_SCRIPT_END = """EOF
{launcher} $tmpfile
rm $tmpfile
"""


def gen_qsub_script(crop, batch_ids=None, *,
                    hours=None,
                    minutes=None,
                    seconds=None,
                    gigabytes=2,
                    num_procs=1,
                    launcher='python',
                    setup="",
                    mpi=False,
                    temp_gigabytes=1,
                    output_directory=None,
                    extra_resources=None,
                    debugging=False):
    """Generate a qsub script to grow a Crop.

    Parameters
    ----------
    crop : Crop
        The crop to grow.
    batch_ids : int or tuple[int]
        Which batch numbers to grow, defaults to all missing batches.
    hours : int
        How many hours to request, default=0.
    minutes : int, optional
        How many minutes to request, default=20.
    seconds : int, optional
        How many seconds to request, default=0.
    gigabytes : int, optional
        How much memory to request, default: 2.
    num_procs : int, optional
        How many processes to request (threaded cores or MPI), default: 1.
    launcher : str, optional
        How to launch the script, default: ``'python'``. But could for example
        be ``'mpiexec python'`` for a MPI program.
    setup : str, optional
        Python script to run before growing, for things that shouldnt't be put
        in the crop function itself, e.g. one-time imports with side-effects
        like: ``"import tensorflow as tf; tf.enable_eager_execution()``".
    mpi : bool, optional
        Request MPI processes not threaded processes.
    temp_gigabytes : int, optional
        How much temporary on-disk memory.
    output_directory : str, optional
        What directory to write output to. Defaults to "$HOME/Scratch/output".
    extra_resources : str, optional
        Extra "#$ -l" resources, e.g. 'gpu=1'
    debugging : bool, optional
        Set the python log level to debugging.

    Returns
    -------
    str
    """
    if hours is minutes is seconds is None:
        hours, minutes, seconds = 1, 0, 0
    else:
        hours = 0 if hours is None else int(hours)
        minutes = 0 if minutes is None else int(minutes)
        seconds = 0 if seconds is None else int(seconds)

    if output_directory is None:
        from os.path import expanduser
        home = expanduser("~")
        output_directory = os.path.join(home, 'Scratch', 'output')

    crop.calc_progress()

    if extra_resources is None:
        extra_resources = ""
    else:
        extra_resources = "#$ -l {}".format(extra_resources)

    opts = {
        'hours': hours,
        'minutes': minutes,
        'seconds': seconds,
        'gigabytes': gigabytes,
        'name': crop.name,
        'num_procs': num_procs,
        'num_threads': 1 if mpi else num_procs,
        'run_start': 1,
        'launcher': launcher,
        'setup': setup,
        'pe': 'mpi' if mpi else 'smp',
        'temp_gigabytes': temp_gigabytes,
        'output_directory': output_directory,
        'working_directory': crop.parent_dir,
        'extra_resources': extra_resources,
        'debugging': debugging,
    }

    script = _BASE_QSUB_SCRIPT

    # grow specific ids
    if batch_ids is not None:
        script += _QSUB_GROW_PARTIAL_SCRIPT
        batch_ids = tuple(batch_ids)
        opts['run_stop'] = len(batch_ids)
        opts['batch_ids'] = batch_ids

    # grow all ids
    elif crop.num_results == 0:
        script += _QSUB_GROW_ALL_SCRIPT
        opts['run_stop'] = crop.num_batches

    # grow missing ids only
    else:
        script += _QSUB_GROW_PARTIAL_SCRIPT
        batch_ids = crop.missing_results()
        opts['run_stop'] = len(batch_ids)
        opts['batch_ids'] = batch_ids

    script += _BASE_QSUB_SCRIPT_END

    return script.format(**opts)


def qsub_grow(crop, batch_ids=None, *,
              hours=None,
              minutes=None,
              seconds=None,
              gigabytes=2,
              num_procs=1,
              launcher='python',
              setup="",
              mpi=False,
              temp_gigabytes=1,
              output_directory=None,
              extra_resources=None,
              debugging=False):  # pragma: no cover
    """Automagically submit SGE jobs to grow all missing results.

    Parameters
    ----------
    crop : Crop
        The crop to grow.
    batch_ids : int or tuple[int]
        Which batch numbers to grow, defaults to all missing batches.
    hours : int
        How many hours to request, default=0.
    minutes : int, optional
        How many minutes to request, default=20.
    seconds : int, optional
        How many seconds to request, default=0.
    gigabytes : int, optional
        How much memory to request, default: 2.
    num_procs : int, optional
        How many processes to request (threaded cores or MPI), default: 1.
    launcher : str, optional
        How to launch the script, default: ``'python'``. But could for example
        be ``'mpiexec python'`` for a MPI program.
    setup : str, optional
        Python script to run before growing, for things that shouldnt't be put
        in the crop function itself, e.g. one-time imports with side-effects
        like: ``"import tensorflow as tf; tf.enable_eager_execution()``".
    mpi : bool, optional
        Request MPI processes not threaded processes.
    temp_gigabytes : int, optional
        How much temporary on-disk memory.
    output_directory : str, optional
        What directory to write output to. Defaults to "$HOME/Scratch/output".
    extra_resources : str, optional
        Extra "#$ -l" resources, e.g. 'gpu=1'
    debugging : bool, optional
        Set the python log level to debugging.
    """
    if crop.is_ready_to_reap():
        print("Crop ready to reap: nothing to submit.")
        return

    import subprocess

    script = gen_qsub_script(
        crop,
        batch_ids=batch_ids,
        hours=hours,
        minutes=minutes,
        seconds=seconds,
        gigabytes=gigabytes,
        temp_gigabytes=temp_gigabytes,
        output_directory=output_directory,
        num_procs=num_procs,
        launcher=launcher,
        setup=setup,
        mpi=mpi,
        extra_resources=extra_resources,
        debugging=debugging,
    )

    script_file = os.path.join(crop.location, "__qsub_script__.sh")

    with open(script_file, mode='w') as f:
        f.write(script)

    subprocess.run(['qsub', script_file])

    os.remove(script_file)


Crop.gen_qsub_script = gen_qsub_script
Crop.qsub_grow = qsub_grow
