import logging
import math
import operator
import random
import time
from datetime import timedelta
from pickle import PickleError
from itertools import zip_longest
from collections import Counter

import numpy as np

import westpa
from .data_manager import weight_dtype
from .segment import Segment
from .states import InitialState
from . import extloader
from . import wm_ops


log = logging.getLogger(__name__)

EPS = np.finfo(weight_dtype).eps


def grouper(n, iterable, fillvalue=None):
    "Collect data into fixed-length chunks or blocks"
    # grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx
    args = [iter(iterable)] * n
    return zip_longest(fillvalue=fillvalue, *args)


class PropagationError(RuntimeError):
    pass


class WESimManager:
    def process_config(self):
        config = self.rc.config
        for entry, type_ in [('gen_istates', bool), ('block_size', int), ('save_transition_matrices', bool)]:
            config.require_type_if_present(['west', 'propagation', entry], type_)

        self.do_gen_istates = config.get(['west', 'propagation', 'gen_istates'], False)
        self.propagator_block_size = config.get(['west', 'propagation', 'block_size'], 1)
        self.save_transition_matrices = config.get(['west', 'propagation', 'save_transition_matrices'], False)
        self.max_run_walltime = config.get(['west', 'propagation', 'max_run_wallclock'], default=None)
        self.max_total_iterations = config.get(['west', 'propagation', 'max_total_iterations'], default=None)

    def __init__(self, rc=None):
        self.rc = rc or westpa.rc
        self.work_manager = self.rc.get_work_manager()
        self.data_manager = self.rc.get_data_manager()
        self.we_driver = self.rc.get_we_driver()
        self.system = self.rc.get_system_driver()

        # A table of function -> list of (priority, name, callback) tuples
        self._callback_table = {}
        self._valid_callbacks = set(
            (
                self.prepare_run,
                self.finalize_run,
                self.prepare_iteration,
                self.finalize_iteration,
                self.pre_propagation,
                self.post_propagation,
                self.pre_we,
                self.post_we,
                self.prepare_new_iteration,
            )
        )
        self._callbacks_by_name = {fn.__name__: fn for fn in self._valid_callbacks}
        self.n_propagated = 0

        # config items
        self.do_gen_istates = False
        self.propagator_block_size = 1
        self.save_transition_matrices = False
        self.max_run_walltime = None
        self.max_total_iterations = None
        self.process_config()

        # Per-iteration variables
        self.n_iter = None  # current iteration

        # Basis and initial states for this iteration, in case the propagator needs them
        self.current_iter_bstates = None  # BasisStates valid at this iteration
        self.current_iter_istates = None  # InitialStates used in this iteration

        # Basis states for next iteration
        self.next_iter_bstates = None  # BasisStates valid for the next iteration
        self.next_iter_bstate_cprobs = None  # Cumulative probabilities for basis states, used for selection

        # Tracking of this iteration's segments
        self.segments = None  # Mapping of seg_id to segment for all segments in this iteration
        self.completed_segments = None  # Mapping of seg_id to segment for all completed segments in this iteration
        self.incomplete_segments = None  # Mapping of seg_id to segment for all incomplete segments in this iteration

        # Tracking of binning
        self.bin_mapper_hash = None  # Hash of bin mapper from most recently-run WE, for use by post-WE analysis plugins

    def register_callback(self, hook, function, priority=0):
        '''Registers a callback to execute during the given ``hook`` into the simulation loop. The optional
        priority is used to order when the function is called relative to other registered callbacks.'''

        if hook not in self._valid_callbacks:
            try:
                hook = self._callbacks_by_name[hook]
            except KeyError:
                raise KeyError('invalid hook {!r}'.format(hook))

        # It's possible to register a callback that's a duplicate function, but at a different place in memory.
        #   For example, if you launch a run without clearing state from a previous run.
        #   More details on this are available in https://github.com/westpa/westpa/issues/182 but the below code
        #   handles specifically the problem that causes in plugin loading.
        try:
            # Before checking for set membership of (priority, function.__name__, function), just check
            #   function hash for collisions in this hook.
            hook_function_hash = [hash(callback[2]) for callback in self._callback_table[hook]]
        except KeyError:
            # If there's no entry in self._callback_table for this hook, then there definitely aren't any collisions
            #   because no plugins are registered to it yet in the first place.
            pass
        else:
            # If there are plugins registered to this hook, check for duplicate hash, which will definitely have the same name, module, function.
            try:
                if hash(function) in hook_function_hash:
                    log.info('{!r} has already been loaded, skipping'.format(function))
                    return
            except KeyError:
                pass

        try:
            self._callback_table[hook].add((priority, function.__name__, function))
        except KeyError:
            self._callback_table[hook] = set([(priority, function.__name__, function)])

        # Raise warning if there are multiple callback with same priority.
        for priority, count in Counter([callback[0] for callback in self._callback_table[hook]]).items():
            if count > 1:
                log.warning(
                    f'{count} callbacks in {hook} have identical priority {priority}. The order of callback execution is not guaranteed.'
                )
                log.warning(f'{hook}: {self._callback_table[hook]}')

        log.debug('registered callback {!r} for hook {!r}'.format(function, hook))

    def invoke_callbacks(self, hook, *args, **kwargs):
        callbacks = self._callback_table.get(hook, [])
        # Sort by priority, function name, then module name
        sorted_callbacks = sorted(callbacks, key=lambda x: (x[0], x[1], x[2].__module__))
        for priority, name, fn in sorted_callbacks:
            log.debug('invoking callback {!r} for hook {!r}'.format(fn, hook))
            fn(*args, **kwargs)

    def load_plugins(self, plugins=None):
        if plugins is None:
            plugins = []

        try:
            plugins_config = westpa.rc.config['west', 'plugins']
        except KeyError:
            plugins_config = []

        if plugins_config is None:
            plugins_config = []

        plugins += plugins_config

        for plugin_config in plugins:
            plugin_name = plugin_config['plugin']
            if plugin_config.get('enabled', True):
                log.info('loading plugin {!r}'.format(plugin_name))
                plugin = extloader.get_object(plugin_name)(self, plugin_config)
                log.debug('loaded plugin {!r}'.format(plugin))

    def report_bin_statistics(self, bins, target_states, save_summary=False):
        segments = list(self.segments.values())
        bin_counts = np.fromiter(map(len, bins), dtype=np.int_, count=len(bins))
        target_counts = self.we_driver.bin_target_counts

        # Do not include bins with target count zero (e.g. sinks, never-filled bins) in the (non)empty bins statistics
        n_active_bins = len(target_counts[target_counts != 0])

        if target_states:
            n_active_bins -= len(target_states)

        seg_probs = np.fromiter(map(operator.attrgetter('weight'), segments), dtype=weight_dtype, count=len(segments))
        bin_probs = np.fromiter(map(operator.attrgetter('weight'), bins), dtype=weight_dtype, count=len(bins))
        norm = seg_probs.sum()

        assert abs(1 - norm) < EPS * (len(segments) + n_active_bins)

        min_seg_prob = seg_probs[seg_probs != 0].min()
        max_seg_prob = seg_probs.max()
        seg_drange = math.log(max_seg_prob / min_seg_prob)
        min_bin_prob = bin_probs[bin_probs != 0].min()
        max_bin_prob = bin_probs.max()
        bin_drange = math.log(max_bin_prob / min_bin_prob)
        n_pop = len(bin_counts[bin_counts != 0])

        self.rc.pstatus('{:d} of {:d} ({:%}) active bins are populated'.format(n_pop, n_active_bins, n_pop / n_active_bins))
        self.rc.pstatus('per-bin minimum non-zero probability:       {:g}'.format(min_bin_prob))
        self.rc.pstatus('per-bin maximum probability:                {:g}'.format(max_bin_prob))
        self.rc.pstatus('per-bin probability dynamic range (kT):     {:g}'.format(bin_drange))
        self.rc.pstatus('per-segment minimum non-zero probability:   {:g}'.format(min_seg_prob))
        self.rc.pstatus('per-segment maximum non-zero probability:   {:g}'.format(max_seg_prob))
        self.rc.pstatus('per-segment probability dynamic range (kT): {:g}'.format(seg_drange))
        self.rc.pstatus('norm = {:g}, error in norm = {:g} ({:.2g}*epsilon)'.format(norm, (norm - 1), (norm - 1) / EPS))
        self.rc.pflush()

        if min_seg_prob < 1e-100:
            log.warning(
                '\nMinimum segment weight is < 1e-100 and might not be physically relevant. Please reconsider your progress coordinate or binning scheme.'
            )

        if save_summary:
            iter_summary = self.data_manager.get_iter_summary()
            iter_summary['n_particles'] = len(segments)
            iter_summary['norm'] = norm
            iter_summary['min_bin_prob'] = min_bin_prob
            iter_summary['max_bin_prob'] = max_bin_prob
            iter_summary['min_seg_prob'] = min_seg_prob
            iter_summary['max_seg_prob'] = max_seg_prob
            if np.isnan(iter_summary['cputime']):
                iter_summary['cputime'] = 0.0
            if np.isnan(iter_summary['walltime']):
                iter_summary['walltime'] = 0.0
            self.data_manager.update_iter_summary(iter_summary)

    def get_bstate_pcoords(self, basis_states, label='basis'):
        '''For each of the given ``basis_states``, calculate progress coordinate values
        as necessary.  The HDF5 file is not updated.'''

        self.rc.pstatus('Calculating progress coordinate values for {} states.'.format(label))
        futures = [self.work_manager.submit(wm_ops.get_pcoord, args=(basis_state,)) for basis_state in basis_states]
        fmap = {future: i for (i, future) in enumerate(futures)}
        for future in self.work_manager.as_completed(futures):
            basis_states[fmap[future]].pcoord = future.get_result().pcoord

    def report_basis_states(self, basis_states, label='basis'):
        pstatus = self.rc.pstatus
        pstatus('{:d} {} state(s) present'.format(len(basis_states), label), end='')
        if self.rc.verbose_mode:
            pstatus(':')
            pstatus(
                '{:6s}    {:12s}    {:20s}    {:20s}    {}'.format(
                    'ID', 'Label', 'Probability', 'Aux Reference', 'Progress Coordinate'
                )
            )
            for basis_state in basis_states:
                pstatus(
                    '{:<6d}    {:12s}    {:<20.14g}    {:20s}    {}'.format(
                        basis_state.state_id,
                        basis_state.label,
                        basis_state.probability,
                        basis_state.auxref or '',
                        ', '.join(map(str, basis_state.pcoord)),
                    )
                )
        pstatus()
        self.rc.pflush()

    def report_target_states(self, target_states):
        pstatus = self.rc.pstatus
        pstatus('{:d} target state(s) present'.format(len(target_states)), end='')
        if self.rc.verbose_mode and target_states:
            pstatus(':')
            pstatus('{:6s}    {:12s}    {}'.format('ID', 'Label', 'Progress Coordinate'))
            for target_state in target_states:
                pstatus(
                    '{:<6d}    {:12s}    {}'.format(
                        target_state.state_id, target_state.label, ','.join(map(str, target_state.pcoord))
                    )
                )
        pstatus()
        self.rc.pflush()

    def initialize_simulation(self, basis_states, target_states, start_states, segs_per_state=1, suppress_we=False):
        '''Initialize a new weighted ensemble simulation, taking ``segs_per_state`` initial
        states from each of the given ``basis_states``.

        ``w_init`` is the forward-facing version of this function'''

        data_manager = self.data_manager
        work_manager = self.work_manager
        pstatus = self.rc.pstatus
        system = self.system

        pstatus('Creating HDF5 file {!r}'.format(self.data_manager.we_h5filename))
        data_manager.prepare_backing()

        # Process target states
        data_manager.save_target_states(target_states)
        self.report_target_states(target_states)

        # Process basis states
        self.get_bstate_pcoords(basis_states)
        self.data_manager.create_ibstate_group(basis_states)
        self.data_manager.create_ibstate_iter_h5file(basis_states)
        self.report_basis_states(basis_states)

        # Process start states
        # Unlike the above, does not create an ibstate group.
        # TODO: Should it? I don't think so, if needed it can be traced back through basis_auxref

        # Here, we are trying to assign a state_id to the start state to be initialized, without actually
        # saving it to the ibstates records in any of the h5 files. It might actually be a problem
        # when tracing trajectories with westpa.analysis (especially with HDF5 framework) since it would
        # try to look for a basis state id > len(basis_state).
        # Since start states are only used while initializing and iteration 1, it's ok to not save it to save space. If necessary,
        # the structure can be traced directly to the parent file using the standard basis state logic referencing
        # west[iterations/iter_00000001/ibstates/istate_index/basis_auxref] of that istate.

        if len(start_states) > 0 and start_states[0].state_id is None:
            last_id = basis_states[-1].state_id
            for start_state in start_states:
                start_state.state_id = last_id + 1
                last_id += 1

        self.get_bstate_pcoords(start_states, label='start')
        self.report_basis_states(start_states, label='start')

        pstatus('Preparing initial states')
        initial_states = []
        weights = []
        if self.do_gen_istates:
            istate_type = InitialState.ISTATE_TYPE_GENERATED
        else:
            istate_type = InitialState.ISTATE_TYPE_BASIS

        for basis_state in basis_states:
            for _iseg in range(segs_per_state):
                initial_state = data_manager.create_initial_states(1, 1)[0]
                initial_state.basis_state_id = basis_state.state_id
                initial_state.basis_state = basis_state
                initial_state.istate_type = istate_type
                weights.append(basis_state.probability / segs_per_state)
                initial_states.append(initial_state)

        for start_state in start_states:
            for _iseg in range(segs_per_state):
                initial_state = data_manager.create_initial_states(1, 1)[0]
                initial_state.basis_state_id = start_state.state_id
                initial_state.basis_state = start_state
                initial_state.basis_auxref = start_state.auxref

                # Start states are assigned their own type, so they can be identified later
                initial_state.istate_type = InitialState.ISTATE_TYPE_START
                weights.append(start_state.probability / segs_per_state)
                initial_state.iter_used = 1
                initial_states.append(initial_state)

        if self.do_gen_istates:
            futures = [
                work_manager.submit(wm_ops.gen_istate, args=(initial_state.basis_state, initial_state))
                for initial_state in initial_states
            ]
            for future in work_manager.as_completed(futures):
                rbstate, ristate = future.get_result()
                initial_states[ristate.state_id].pcoord = ristate.pcoord
        else:
            for initial_state in initial_states:
                basis_state = initial_state.basis_state
                initial_state.pcoord = basis_state.pcoord
                initial_state.istate_status = InitialState.ISTATE_STATUS_PREPARED

        for initial_state in initial_states:
            log.debug('initial state created: {!r}'.format(initial_state))

        # save list of initial states just generated
        # some of these may not be used, depending on how WE shakes out
        data_manager.update_initial_states(initial_states, n_iter=1)

        if not suppress_we:
            self.we_driver.populate_initial(initial_states, weights, system)
            segments = list(self.we_driver.next_iter_segments)
            binning = self.we_driver.next_iter_binning
        else:
            segments = list(self.we_driver.current_iter_segments)
            binning = self.we_driver.final_binning

        bin_occupancies = np.fromiter(map(len, binning), dtype=np.uint, count=self.we_driver.bin_mapper.nbins)
        target_occupancies = np.require(self.we_driver.bin_target_counts, dtype=np.uint)

        # total_bins/replicas defined here to remove target state bin from "active" bins
        total_bins = len(bin_occupancies) - len(target_states)
        total_replicas = int(sum(target_occupancies)) - int(self.we_driver.bin_target_counts[-1]) * len(target_states)

        # Make sure we have
        for segment in segments:
            segment.n_iter = 1
            segment.status = Segment.SEG_STATUS_PREPARED
            assert segment.parent_id < 0
            assert initial_states[segment.initial_state_id].iter_used == 1

        data_manager.prepare_iteration(1, segments)
        data_manager.update_initial_states(initial_states, n_iter=1)

        if self.rc.verbose_mode:
            pstatus('\nSegments generated:')
            for segment in segments:
                pstatus('{!r}'.format(segment))

        pstatus(
            '''
        Total bins:            {total_bins:d}
        Initial replicas:      {init_replicas:d} in {occ_bins:d} bins, total weight = {weight:g}
        Total target replicas: {total_replicas:d}
        '''.format(
                total_bins=total_bins,
                init_replicas=int(sum(bin_occupancies)),
                occ_bins=len(bin_occupancies[bin_occupancies > 0]),
                weight=float(sum(segment.weight for segment in segments)),
                total_replicas=total_replicas,
            )
        )

        total_prob = float(sum(segment.weight for segment in segments))
        pstatus(f'1-prob: {1 - total_prob:.4e}')

        target_counts = self.we_driver.bin_target_counts
        # Do not include bins with target count zero (e.g. sinks, never-filled bins) in the (non)empty bins statistics
        n_active_bins = len(target_counts[target_counts != 0])
        seg_probs = np.fromiter(map(operator.attrgetter('weight'), segments), dtype=weight_dtype, count=len(segments))
        norm = seg_probs.sum()

        if not abs(1 - norm) < EPS * (len(segments) + n_active_bins):
            pstatus("Normalization check failed at w_init, explicitly renormalizing")
            for segment in segments:
                segment.weight /= norm

        # Send the segments over to the data manager to commit to disk
        data_manager.current_iteration = 1

        # Report statistics
        pstatus('Simulation prepared.')
        self.segments = {segment.seg_id: segment for segment in segments}
        self.report_bin_statistics(binning, target_states, save_summary=True)
        data_manager.flush_backing()
        data_manager.close_backing()

    def prepare_iteration(self):
        log.debug('beginning iteration {:d}'.format(self.n_iter))

        # the WE driver needs a list of all target states for this iteration
        # along with information about any new weights introduced (e.g. by recycling)
        target_states = self.data_manager.get_target_states(self.n_iter)
        new_weights = self.data_manager.get_new_weight_data(self.n_iter)

        self.we_driver.new_iteration(target_states=target_states, new_weights=new_weights)

        # Get basis states used in this iteration
        self.current_iter_bstates = self.data_manager.get_basis_states(self.n_iter)

        # Get the segments for this iteration and separate into complete and incomplete
        if self.segments is None:
            segments = self.segments = {segment.seg_id: segment for segment in self.data_manager.get_segments()}
            log.debug('loaded {:d} segments'.format(len(segments)))
        else:
            segments = self.segments
            log.debug('using {:d} pre-existing segments'.format(len(segments)))

        completed_segments = self.completed_segments = {}
        incomplete_segments = self.incomplete_segments = {}
        for segment in segments.values():
            if segment.status == Segment.SEG_STATUS_COMPLETE:
                completed_segments[segment.seg_id] = segment
            else:
                incomplete_segments[segment.seg_id] = segment
        log.debug('{:d} segments are complete; {:d} are incomplete'.format(len(completed_segments), len(incomplete_segments)))

        if len(incomplete_segments) == len(segments):
            # Starting a new iteration
            self.rc.pstatus('Beginning iteration {:d}'.format(self.n_iter))
        elif incomplete_segments:
            self.rc.pstatus('Continuing iteration {:d}'.format(self.n_iter))
        self.rc.pstatus(
            '{:d} segments remain in iteration {:d} ({:d} total)'.format(len(incomplete_segments), self.n_iter, len(segments))
        )

        # Get the initial states active for this iteration (so that the propagator has them if necessary)
        self.current_iter_istates = {
            state.state_id: state for state in self.data_manager.get_segment_initial_states(list(segments.values()))
        }
        log.debug('This iteration uses {:d} initial states'.format(len(self.current_iter_istates)))

        # Assign this iteration's segments' initial points to bins and report on bin population
        initial_pcoords = self.system.new_pcoord_array(len(segments))
        initial_binning = self.system.bin_mapper.construct_bins()
        for iseg, segment in enumerate(segments.values()):
            initial_pcoords[iseg] = segment.pcoord[0]
        initial_assignments = self.system.bin_mapper.assign(initial_pcoords)
        for segment, assignment in zip(iter(segments.values()), initial_assignments):
            initial_binning[assignment].add(segment)
        self.report_bin_statistics(initial_binning, [], save_summary=True)
        del initial_pcoords, initial_binning

        self.rc.pstatus('Waiting for segments to complete...')

        # Let the WE driver assign completed segments
        if completed_segments:
            self.we_driver.assign(list(completed_segments.values()))

        # load restart data
        self.data_manager.prepare_segment_restarts(
            incomplete_segments.values(), self.current_iter_bstates, self.current_iter_istates
        )

        # Get the basis states and initial states for the next iteration, necessary for doing on-the-fly recycling
        self.next_iter_bstates = self.data_manager.get_basis_states(self.n_iter + 1)
        self.next_iter_bstate_cprobs = np.add.accumulate([bstate.probability for bstate in self.next_iter_bstates])

        self.we_driver.avail_initial_states = {
            istate.state_id: istate for istate in self.data_manager.get_unused_initial_states(n_iter=self.n_iter + 1)
        }
        log.debug('{:d} unused initial states found'.format(len(self.we_driver.avail_initial_states)))

        # Invoke callbacks
        self.invoke_callbacks(self.prepare_iteration)

        # dispatch and immediately wait on result for prep_iter
        log.debug('dispatching propagator prep_iter to work manager')
        self.work_manager.submit(wm_ops.prep_iter, args=(self.n_iter, segments)).get_result()

    def finalize_iteration(self):
        '''Clean up after an iteration and prepare for the next.'''
        log.debug('finalizing iteration {:d}'.format(self.n_iter))

        self.invoke_callbacks(self.finalize_iteration)

        # dispatch and immediately wait on result for post_iter
        log.debug('dispatching propagator post_iter to work manager')
        self.work_manager.submit(wm_ops.post_iter, args=(self.n_iter, list(self.segments.values()))).get_result()

        # Move existing segments into place as new segments
        del self.segments
        self.segments = {segment.seg_id: segment for segment in self.we_driver.next_iter_segments}

        self.rc.pstatus("Iteration completed successfully")

    def get_istate_futures(self):
        '''Add ``n_states`` initial states to the internal list of initial states assigned to
        recycled particles.  Spare states are used if available, otherwise new states are created.
        If created new initial states requires generation, then a set of futures is returned
        representing work manager tasks corresponding to the necessary generation work.'''

        n_recycled = self.we_driver.n_recycled_segs
        n_istates_needed = self.we_driver.n_istates_needed

        log.debug('{:d} unused initial states available'.format(len(self.we_driver.avail_initial_states)))
        log.debug('{:d} new initial states required for recycling {:d} walkers'.format(n_istates_needed, n_recycled))

        futures = set()
        updated_states = []
        for _i in range(n_istates_needed):
            # Select a basis state according to its weight
            ibstate = np.digitize([random.random()], self.next_iter_bstate_cprobs)
            basis_state = self.next_iter_bstates[ibstate[0]]
            initial_state = self.data_manager.create_initial_states(1, n_iter=self.n_iter + 1)[0]
            initial_state.iter_created = self.n_iter
            initial_state.basis_state_id = basis_state.state_id
            initial_state.istate_status = InitialState.ISTATE_STATUS_PENDING

            if self.do_gen_istates:
                log.debug('generating new initial state from basis state {!r}'.format(basis_state))
                initial_state.istate_type = InitialState.ISTATE_TYPE_GENERATED
                futures.add(self.work_manager.submit(wm_ops.gen_istate, args=(basis_state, initial_state)))
            else:
                log.debug('using basis state {!r} directly'.format(basis_state))
                initial_state.istate_type = InitialState.ISTATE_TYPE_BASIS
                initial_state.pcoord = basis_state.pcoord.copy()
                initial_state.istate_status = InitialState.ISTATE_STATUS_PREPARED
                self.we_driver.avail_initial_states[initial_state.state_id] = initial_state
            updated_states.append(initial_state)
        self.data_manager.update_initial_states(updated_states, n_iter=self.n_iter + 1)
        return futures

    def propagate(self):
        segments = list(self.incomplete_segments.values())
        log.debug('iteration {:d}: propagating {:d} segments'.format(self.n_iter, len(segments)))

        # all futures dispatched for this iteration
        futures = set()
        segment_futures = set()

        # Immediately dispatch any necessary initial state generation
        istate_gen_futures = self.get_istate_futures()
        futures.update(istate_gen_futures)

        # Dispatch propagation tasks using work manager
        for segment_block in grouper(self.propagator_block_size, segments):
            segment_block = [_f for _f in segment_block if _f]
            pbstates, pistates = westpa.core.states.pare_basis_initial_states(
                self.current_iter_bstates, list(self.current_iter_istates.values()), segment_block
            )
            future = self.work_manager.submit(wm_ops.propagate, args=(pbstates, pistates, segment_block))
            futures.add(future)
            segment_futures.add(future)

        while futures:
            # TODO: add capacity for timeout or SIGINT here
            future = self.work_manager.wait_any(futures)
            futures.remove(future)

            if future in segment_futures:
                segment_futures.remove(future)
                incoming = future.get_result()
                self.n_propagated += 1

                self.segments.update({segment.seg_id: segment for segment in incoming})
                self.completed_segments.update({segment.seg_id: segment for segment in incoming})

                self.we_driver.assign(incoming)
                new_istate_futures = self.get_istate_futures()
                istate_gen_futures.update(new_istate_futures)
                futures.update(new_istate_futures)

                with self.data_manager.expiring_flushing_lock():
                    self.data_manager.update_segments(self.n_iter, incoming)

            elif future in istate_gen_futures:
                istate_gen_futures.remove(future)
                _basis_state, initial_state = future.get_result()
                log.debug('received newly-prepared initial state {!r}'.format(initial_state))
                initial_state.istate_status = InitialState.ISTATE_STATUS_PREPARED
                with self.data_manager.expiring_flushing_lock():
                    self.data_manager.update_initial_states([initial_state], n_iter=self.n_iter + 1)
                self.we_driver.avail_initial_states[initial_state.state_id] = initial_state
            else:
                log.error('unknown future {!r} received from work manager'.format(future))
                raise AssertionError('untracked future {!r}'.format(future))

        log.debug('done with propagation')
        self.save_bin_data()
        self.data_manager.flush_backing()

    def save_bin_data(self):
        '''Calculate and write flux and transition count matrices to HDF5. Population and rate matrices
        are likely useless at the single-tau level and are no longer written.'''
        # save_bin_data(self, populations, n_trans, fluxes, rates, n_iter=None)

        if self.save_transition_matrices:
            with self.data_manager.expiring_flushing_lock():
                iter_group = self.data_manager.get_iter_group(self.n_iter)
                for key in ['bin_ntrans', 'bin_fluxes']:
                    try:
                        del iter_group[key]
                    except KeyError:
                        pass
                iter_group['bin_ntrans'] = self.we_driver.transition_matrix
                iter_group['bin_fluxes'] = self.we_driver.flux_matrix

    def check_propagation(self):
        '''Check for failures in propagation or initial state generation, and raise an exception
        if any are found.'''

        failed_segments = [segment for segment in self.segments.values() if segment.status != Segment.SEG_STATUS_COMPLETE]

        if failed_segments:
            failed_ids = '  \n'.join(str(segment.seg_id) for segment in failed_segments)
            log.error('propagation failed for {:d} segment(s):\n{}'.format(len(failed_segments), failed_ids))
            raise PropagationError('propagation failed for {:d} segments'.format(len(failed_segments)))
        else:
            log.debug('propagation complete for iteration {:d}'.format(self.n_iter))

        failed_istates = [
            istate
            for istate in self.we_driver.used_initial_states.values()
            if istate.istate_status != InitialState.ISTATE_STATUS_PREPARED
        ]
        log.debug('{!r}'.format(failed_istates))
        if failed_istates:
            failed_ids = '  \n'.join(str(istate.state_id) for istate in failed_istates)
            log.error('initial state generation failed for {:d} states:\n{}'.format(len(failed_istates), failed_ids))
            raise PropagationError('initial state generation failed for {:d} states'.format(len(failed_istates)))
        else:
            log.debug('initial state generation complete for iteration {:d}'.format(self.n_iter))

    def run_we(self):
        '''Run the weighted ensemble algorithm based on the binning in self.final_bins and
        the recycled particles in self.to_recycle, creating and committing the next iteration's
        segments to storage as well.'''

        # The WE driver now does almost everything; we just have to record the
        # mapper used for binning this iteration, and update initial states
        # that have been used

        try:
            pickled, hashed = self.we_driver.bin_mapper.pickle_and_hash()
        except PickleError:
            pickled = hashed = ''

        self.bin_mapper_hash = hashed
        self.we_driver.construct_next()

        if self.we_driver.used_initial_states:
            for initial_state in self.we_driver.used_initial_states.values():
                initial_state.iter_used = self.n_iter + 1
            self.data_manager.update_initial_states(list(self.we_driver.used_initial_states.values()))

        self.data_manager.update_segments(self.n_iter, list(self.segments.values()))

        self.data_manager.require_iter_group(self.n_iter + 1)
        self.data_manager.save_iter_binning(self.n_iter + 1, hashed, pickled, self.we_driver.bin_target_counts)

        # Report on recycling
        recycling_events = {}
        for nw in self.we_driver.new_weights:
            try:
                recycling_events[nw.target_state_id].append(nw.weight)
            except KeyError:
                recycling_events[nw.target_state_id] = list([nw.weight])

        tstates_by_id = {state.state_id: state for state in self.we_driver.target_states.values()}

        for tstate_id, weights in recycling_events.items():
            tstate = tstates_by_id[tstate_id]
            self.rc.pstatus(
                'Recycled {:g} probability ({:d} walkers) from target state {!r}'.format(sum(weights), len(weights), tstate.label)
            )

    def prepare_new_iteration(self):
        '''Commit data for the coming iteration to the HDF5 file.'''
        self.invoke_callbacks(self.prepare_new_iteration)

        if self.rc.debug_mode:
            self.rc.pstatus('\nSegments generated:')
            for segment in self.we_driver.next_iter_segments:
                self.rc.pstatus('{!r} pcoord[0]={!r}'.format(segment, segment.pcoord[0]))

        self.data_manager.prepare_iteration(self.n_iter + 1, list(self.we_driver.next_iter_segments))
        self.data_manager.save_new_weight_data(self.n_iter + 1, self.we_driver.new_weights)

    def run(self):
        run_starttime = time.time()
        max_walltime = self.max_run_walltime
        if max_walltime:
            run_killtime = run_starttime + max_walltime
            self.rc.pstatus('Maximum wallclock time: %s' % timedelta(seconds=max_walltime or 0))
        else:
            run_killtime = None

        self.n_iter = self.data_manager.current_iteration
        max_iter = self.max_total_iterations or self.n_iter + 1

        iter_elapsed = 0
        while self.n_iter <= max_iter:
            if max_walltime and time.time() + 1.1 * iter_elapsed >= run_killtime:
                self.rc.pstatus('Iteration {:d} would require more than the allotted time. Ending run.'.format(self.n_iter))
                return

            try:
                iter_start_time = time.time()

                self.rc.pstatus('\n%s' % time.asctime())
                self.rc.pstatus('Iteration %d (%d requested)' % (self.n_iter, max_iter))

                self.prepare_iteration()
                self.rc.pflush()

                self.pre_propagation()
                self.propagate()
                self.rc.pflush()
                self.check_propagation()
                self.rc.pflush()
                self.post_propagation()

                cputime = sum(segment.cputime for segment in self.segments.values())

                self.rc.pflush()
                self.pre_we()
                self.run_we()
                self.post_we()
                self.rc.pflush()

                self.prepare_new_iteration()

                self.finalize_iteration()

                iter_elapsed = time.time() - iter_start_time
                iter_summary = self.data_manager.get_iter_summary()
                iter_summary['walltime'] += iter_elapsed
                iter_summary['cputime'] = cputime
                self.data_manager.update_iter_summary(iter_summary)

                self.n_iter += 1
                self.data_manager.current_iteration += 1

                try:
                    # This may give NaN if starting a truncated simulation
                    walltime = timedelta(seconds=float(iter_summary['walltime']))
                except ValueError:
                    walltime = 0.0

                try:
                    cputime = timedelta(seconds=float(iter_summary['cputime']))
                except ValueError:
                    cputime = 0.0

                self.rc.pstatus('Iteration wallclock: {0!s}, cputime: {1!s}\n'.format(walltime, cputime))
                self.rc.pflush()
            finally:
                self.data_manager.flush_backing()

        self.rc.pstatus('\n%s' % time.asctime())
        self.rc.pstatus('WEST run complete.')

    def prepare_run(self):
        '''Prepare a new run.'''
        self.data_manager.prepare_run()
        self.system.prepare_run()
        self.invoke_callbacks(self.prepare_run)

    def finalize_run(self):
        '''Perform cleanup at the normal end of a run'''
        self.invoke_callbacks(self.finalize_run)
        self.system.finalize_run()
        self.data_manager.finalize_run()

    def pre_propagation(self):
        self.invoke_callbacks(self.pre_propagation)

    def post_propagation(self):
        self.invoke_callbacks(self.post_propagation)

    def pre_we(self):
        self.invoke_callbacks(self.pre_we)

    def post_we(self):
        self.invoke_callbacks(self.post_we)
