<commit_before># Copyright (C) 2017 <NAME> and <NAME>
#
# This file is part of WESTPA.
#
# WESTPA is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# WESTPA is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with WESTPA.  If not, see <http://www.gnu.org/licenses/>.

import numpy as np
import os, sys
import scipy.sparse as sp

from westtools import Plotter
import itertools

# A useful dataclass used as a wrapper for w_ipa to facilitate
# ease-of-use in ipython/jupyter notebooks/sessions.
# It basically just wraps up numpy arrays and dicts.

class WIPIDataset(object):
    def __init__(self, raw, key):
        self.__dict__ = {}
        self.raw = raw
        self.name = key
    def __repr__(self):
        if type(self.__dict__['raw']) == dict:
            return repr(self.__dir__())
        else:
            return repr(self.raw)
    def __getitem__(self, value):
        if type(value) != str:
            return self.__dict__['raw'][value]
        if value in self.__dict__['raw'].keys():
            return self.__dict__['raw'][value]
        elif value in self.__dict__.keys():
            return self.__dict__[value]
    def __setitem__(self, key, value):
        self.__dict__[key] = value
    def __getattr__(self, value):
        # Check if it's an attribute of the underlying datatype.
        # If not, just use the getitem function.
        if value in dir(self.__dict__['raw']):
            return getattr(self.__dict__['raw'], value)
        else:
            return self.__getitem__(value)
    def __setattr__(self, key, value):
        self.__dict__[key] = value
    def __dir__(self):
        dict_keys = self.__dict__.keys()
        remove = ['raw', 'name', '__dict__', 'plotter']
        for i in remove:
            try:
                dict_keys.remove(str(i))
            except:
                pass
        # We don't enforce that this is a dictionary.
        if type(self.__dict__['raw']) == dict:
            return sorted(set(list(self.raw.keys()) + dict_keys))
        else:
            return sorted(set(dict_keys))
    def keys(self):
        print(self.__dir__())
    # We want to override the basic math functions, now, so... this is only valid for numpy sets.
    def __add__(self, other):
        return self.__dict__['raw'] + other
    def __radd__(self, other):
        return other + self.__dict__['raw']
    def __sub__(self, other):
        return self.__dict__['raw'] - other
    def __rsub__(self, other):
        return other - self.__dict__['raw']
    def __mul__(self, other):
        return self.__dict__['raw'] * other
    def __rmul__(self, other):
        return other * self.__dict__['raw']
    def __div__(self, other):
        return self.__dict__['raw'] / other
    def __floordiv__(self, other):
        return self.__dict__['raw'] // other
    def __rdiv__(self, other):
        return other / self.__dict__['raw']
    def __mod__(self, other):
        return self.__dict__['raw'] % other
    def __pow__(self, other):
        return self.__dict__['raw'] ** other
    def __lshift__(self, other):
        return self.__dict__['raw'] << other
    def __rshift__(self, other):
        return self.__dict__['raw'] >> other
    def __and__(self, other):
        return self.__dict__['raw'] & other
    def __eq__(self, other):
        return self.__dict__['raw'] == other
    def __ne__(self, other):
        return self.__dict__['raw'] != other
    def __lt__(self, other):
        return self.__dict__['raw'] < other
    def __gt__(self, other):
        return self.__dict__['raw'] > other
    def __le__(self, other):
        return self.__dict__['raw'] <= other
    def __ge__(self, other):
        return self.__dict__['raw'] >= other
    def __xor__(self, other):
        return self.__dict__['raw'] ^ other
    def __or__(self, other):
        return self.__dict__['raw'] | other
    #def __iadd__(self, other):
    #    return self.__dict__['raw'] += other
    #def __isub__(self, other):
    #    return self.__dict__['raw'] -= other
    #def __imul__(self, other):
    #    return self.__dict__['raw'] *= other
    #def __idiv__(self, other):
    #    return self.__dict__['raw'] /= other
    #def __ifloordiv__(self, other):
    #    return self.__dict__['raw'] //= other
    #def __imod__(self, other):
    #    return self.__dict__['raw'] %= other
    #def __ipow__(self, other):
    #    return self.__dict__['raw'] **= other
    #def __ilshift__(self, other):
    #    return self.__dict__['raw'] <<= other
    #def __irshift__(self, other):
    #    return self.__dict__['raw'] >>= other
    #def __iand__(self, other):
    #    return self.__dict__['raw'] &= other
    #def __ixor__(self, other):
    #    return self.__dict__['raw'] ^= other
    #def __ior__(self, other):
    #    return self.__dict__['raw'] |= other

# Similar to the above, but slightly expanded to contain information from analysis files.
class KineticsIteration(object):
    def __init__(self, kin_h5file, index, assign, iteration=-1):
        self.__dict__ = {}
        self.h5file = kin_h5file
        # Keys:
        _2D_h5keys = [ 'conditional_flux_evolution', 'rate_evolution' ]
        _1D_h5keys = [ 'state_pop_evolution', 'color_prob_evolution' ]
        for key in _2D_h5keys:
            self.__dict__[key] = self.__2D_with_error__(key, index, assign)
        for key in _1D_h5keys:
            self.__dict__[key] = self.__1D_with_error__(key, index, assign)
        try:
            self.__dict__['total_fluxes'] = WIPIDataset(raw=self.h5file['total_fluxes'], key=None)
        except:
            pass

    def __repr__(self):
        return repr(self.__dir__())
    def __getitem__(self, value):
        if value in self.__dict__.keys():
            return self.__dict__[value]
    def __setitem__(self, key, value):
        self.__dict__[key] = value
    def __getattr__(self, value):
        if value in self.__dict__.keys():
            return self.__dict__[value]
    def __setattr__(self, key, value):
        self.__dict__[key] = value
    def __dir__(self):
        dict_keys = self.__dict__.keys()
        # We don't want to show the plotter class; just the plot function
        remove = [ 'h5file', '__dict__']
        for i in remove:
            try:
                dict_keys.remove(str(i))
            except:
                pass
        return sorted(set(dict_keys))
    def keys(self):
        print(self.__dir__())

    # We seriously need to rename this.
    # It's similar to the global WIPDataset, but has some nice pretty print functions.
    class __custom_dataset__(object):
        # This is just allow it to be indexed via properties.
        # Not a huge thing, but whatever.
        def __init__(self, raw, assign, key):
            self.__dict__ = {}
            self.raw = raw
            self.name = key
            self.assign = assign
            self.nstates = assign.attrs['nstates']
            self.dim = len(raw.shape)
        def __repr__(self):
            return repr(self.__dir__())
        def __getitem__(self, value):
            if value in self.__dict__['raw'].dtype.names:
                return self.__dict__['raw'][value]
            elif value in self.__dict__.keys():
                return self.__dict__[value]
        def __setitem__(self, key, value):
            self.__dict__[key] = value
        def __getattr__(self, value):
            if value in self.__dict__['raw'].dtype.names:
                return self.__dict__['raw'][value]
            elif value in self.__dict__.keys():
                return self.__dict__[value]
        def __setattr__(self, key, value):
            self.__dict__[key] = value
        def __dir__(self):
            dict_keys = self.__dict__.keys()
            # We don't want to show the plotter class; just the plot function
            remove = ['assign', 'dim', 'nstates', 'plotter', '__dict__']
            for i in remove:
                try:
                    dict_keys.remove(str(i))
                except:
                    pass
            return sorted(set(list(self.raw.dtype.names) + dict_keys))
        def keys(self):
            print(self.__dir__())
        def _repr_pretty_(self, p, cycle):
            if self.dim == 1:
                return self._1D_repr_pretty_(p, cycle)
            if self.dim == 2:
                return self._2D_repr_pretty_(p, cycle)
        def _1D_repr_pretty_(self, p, cycle):
           # We're just using this as a way to print things in a pretty way.  They can still be indexed appropriately.
           # Stolen shamelessly from westtools/kinetics_tool.py
            maxlabellen = max(map(len,self.assign['state_labels']))
            p.text('')
            p.text('{name} data:\n'.format(name=self.name))
            for istate in xrange(self.nstates):
                p.text('{:{maxlabellen}s}: mean={:21.15e} CI=({:21.15e}, {:21.15e}) * tau^-1\n'
                    .format(self.assign['state_labels'][istate],
                    self.raw['expected'][istate],
                    self.raw['ci_lbound'][istate],
                    self.raw['ci_ubound'][istate],
                    maxlabellen=maxlabellen))
            p.text('To access data, index via the following names:\n')
            p.text(str(self.__dir__()))
            return " "
        def _2D_repr_pretty_(self, p, cycle):
            # We're just using this as a way to print things in a pretty way.  They can still be indexed appropriately.
            # Stolen shamelessly from westtools/kinetics_tool.py
            maxlabellen = max(map(len,self.assign['state_labels']))
            p.text('')
            p.text('{name} data:\n'.format(name=self.name))
            for istate in xrange(self.nstates):
                for jstate in xrange(self.nstates):
                    if istate == jstate: continue
                    p.text('{:{maxlabellen}s} -> {:{maxlabellen}s}: mean={:21.15e} CI=({:21.15e}, {:21.15e}) * tau^-1\n'
                        .format(self.assign['state_labels'][istate], self.assign['state_labels'][jstate],
                        self.raw['expected'][istate, jstate],
                        self.raw['ci_lbound'][istate, jstate],
                        self.raw['ci_ubound'][istate, jstate],
                        maxlabellen=maxlabellen))
            p.text('To access data, index via the following names:\n')
            p.text(str(self.__dir__()))
            return " "


    def __2D_with_error__(self, h5key, index, assign):
        # Check the start and stop, calculate the block size, and index appropriately.
        # While we could try and automatically generate this above, it's a little more consistent to try it here.
        # This should show the first block for which the current iteration has contributed data.
        self.step_iter = (self.h5file[h5key]['iter_stop'][0] - self.h5file[h5key]['iter_start'][0])[1,0]
        value = ((index-2) // self.step_iter)
        if value < 0:
            value = 0
        raw = self.h5file[h5key][value, :, :]
        error = (raw['ci_ubound'] - raw['ci_lbound']) / (2*raw['expected'])
        expected = raw['expected']
        raw = self.__custom_dataset__(raw, assign, h5key)
        raw.error = error
        raw.plotter = Plotter(self.h5file, h5key, iteration=value, interface='text')
        raw.plot = raw.plotter.plot
        return raw
    def __1D_with_error__(self, h5key, index, assign):
        self.step_iter = (self.h5file[h5key]['iter_stop'][0] - self.h5file[h5key]['iter_start'][0])[1]
        value = ((index-1) // self.step_iter)
        if value < 0:
            value = 0
        raw = self.h5file[h5key][value, :]
        error = (raw['ci_ubound'] - raw['ci_lbound']) / (2*raw['expected'])
        expected = raw['expected']
        raw = self.__custom_dataset__(raw, assign, h5key)
        raw.error = error
        raw.plotter = Plotter(self.h5file, h5key, iteration=value, interface='text')
        raw.plot = raw.plotter.plot
        return raw


class __get_data_for_iteration__(object):
    '''
    All interesting data from an iteration (current/past).  Whenever you change the scheme or iteration,
    this dictionary is automatically updated.  For the current iteration, it's keyed to the current seg_id.
    For the past iteration, it's keyed to the seg_id in the CURRENT iteration such that:

        w.current[X] & w.past[X]

    returns information about seg_id X in the current iteration and information on seg_ID X's PARENT in the
    preceding iteration.

    Can be indexed via a seg_id, or like a dictionary with the following keys:

        kinavg, weights, pcoord, auxdata (optional), parents, summary, seg_id, walkers, states, bins

    kinavg, states, and bins refer to the output from w_kinavg and w_assign for this iteration
    and analysis scheme.  They are NOT dynamics bins, but the bins defined in west.cfg.  
    
    Has the following properties:

        .minweight, .maxweight

    which return all properties of the segment that matches those criteria in the selected iteration.

    If you change the analysis scheme, so, too, will the important values.
    '''

    def __init__(self, parent, value, seg_ids = None):
        '''
        Initializes and sets the correct data.
        '''
        # We've classed this so that we can override some of the normal functions and allow indexing via seg_id
        self.__dict__ = {}
        # Is this function thread safe?
        iter_group = parent.data_reader.get_iter_group(value)
        #iter_group = parent.west['iterations/iter_{num:08d}'.format(num=value)]
        self.parent = parent
        current = {}
        current['iteration'] = value
        if seg_ids == None:
            seg_ids = xrange(0, iter_group['seg_index']['weight'].shape[0])
        # Just make these easier to access.
        current['weights'] = iter_group['seg_index']['weight'][seg_ids]
        current['pcoord'] = iter_group['pcoord'][...][seg_ids, :, :]
        try:
            current['auxdata'] = {}
            for key in iter_group['auxdata'].keys():
                current['auxdata'][key] = iter_group['auxdata'][key][...][seg_ids, :]
        except:
            pass
        current['parents'] = iter_group['seg_index']['parent_id'][seg_ids]
        current['summary'] = parent.data_reader.data_manager.get_iter_summary(int(value))
        current['seg_id'] = np.array(range(0, iter_group['seg_index'].shape[0]))[seg_ids]
        current['walkers'] = current['summary']['n_particles']
        current['states'] = parent.assign['trajlabels'][value-1, :current['walkers'], :][seg_ids]
        current['bins'] = parent.assign['assignments'][value-1, :current['walkers'], :][seg_ids]
        # Calculates the bin population for this iteration.
        nbins = parent.assign['state_map'].shape[0]
        # We have to take the 'unknown' state into account
        nstates = parent.assign['state_labels'].shape[0] + 1
        # Temporarily disabled while I sort out the fact that we shouldn't be using data from w_assign for state populations.
        #current['plot'] = Plotter(parent.direct, parent.reweight, parent.iteration, parent.assign['bin_labels'], parent.assign['state_labels'], current['populations'].states, current['populations'].bins, parent.interface)
        # Now we'll load up the results of the kinetics analysis.
        current['direct'] = KineticsIteration(parent.direct, value, parent.assign, value)
        evolution_datasets = [ 'rate_evolution', 'conditional_flux_evolution', 'state_pop_evolution', 'color_prob_evolution' ]
        # We want to load these up as... oh, who knows, I suppose?
        try:
            current['reweight'] = KineticsIteration(parent.reweight, value, parent.assign, value)
            # We'll make this not a sparse matrix...
            matrix = parent.reweight['iterations/iter_{:08d}'.format(value)]
            # Assume color.
            current['instant_matrix'] = sp.coo_matrix((matrix['flux'][...], (matrix['rows'][...], matrix['cols'][...])), shape=((nbins-1)*2, (nbins-1)*2)).todense()
            reweighting = True
        except:<commit_msg>added in global.<commit_after># Copyright (C) 2017 <NAME> and <NAME>
#
# This file is part of WESTPA.
#
# WESTPA is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# WESTPA is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with WESTPA.  If not, see <http://www.gnu.org/licenses/>.

import numpy as np
import os, sys
import scipy.sparse as sp

from westtools import Plotter
import itertools

# A useful dataclass used as a wrapper for w_ipa to facilitate
# ease-of-use in ipython/jupyter notebooks/sessions.
# It basically just wraps up numpy arrays and dicts.

class WIPIDataset(object):
    def __init__(self, raw, key):
        self.__dict__ = {}
        self.raw = raw
        self.name = key
    def __repr__(self):
        if type(self.__dict__['raw']) == dict:
            return repr(self.__dir__())
        else:
            return repr(self.raw)
    def __getitem__(self, value):
        if type(value) != str:
            return self.__dict__['raw'][value]
        if value in self.__dict__['raw'].keys():
            return self.__dict__['raw'][value]
        elif value in self.__dict__.keys():
            return self.__dict__[value]
    def __setitem__(self, key, value):
        self.__dict__[key] = value
    def __getattr__(self, value):
        # Check if it's an attribute of the underlying datatype.
        # If not, just use the getitem function.
        if value in dir(self.__dict__['raw']):
            return getattr(self.__dict__['raw'], value)
        else:
            return self.__getitem__(value)
    def __setattr__(self, key, value):
        self.__dict__[key] = value
    def __dir__(self):
        dict_keys = self.__dict__.keys()
        remove = ['raw', 'name', '__dict__', 'plotter']
        for i in remove:
            try:
                dict_keys.remove(str(i))
            except:
                pass
        # We don't enforce that this is a dictionary.
        if type(self.__dict__['raw']) == dict:
            return sorted(set(list(self.raw.keys()) + dict_keys))
        else:
            return sorted(set(dict_keys))
    def keys(self):
        print(self.__dir__())
    # We want to override the basic math functions, now, so... this is only valid for numpy sets.
    def __add__(self, other):
        return self.__dict__['raw'] + other
    def __radd__(self, other):
        return other + self.__dict__['raw']
    def __sub__(self, other):
        return self.__dict__['raw'] - other
    def __rsub__(self, other):
        return other - self.__dict__['raw']
    def __mul__(self, other):
        return self.__dict__['raw'] * other
    def __rmul__(self, other):
        return other * self.__dict__['raw']
    def __div__(self, other):
        return self.__dict__['raw'] / other
    def __floordiv__(self, other):
        return self.__dict__['raw'] // other
    def __rdiv__(self, other):
        return other / self.__dict__['raw']
    def __mod__(self, other):
        return self.__dict__['raw'] % other
    def __pow__(self, other):
        return self.__dict__['raw'] ** other
    def __lshift__(self, other):
        return self.__dict__['raw'] << other
    def __rshift__(self, other):
        return self.__dict__['raw'] >> other
    def __and__(self, other):
        return self.__dict__['raw'] & other
    def __eq__(self, other):
        return self.__dict__['raw'] == other
    def __ne__(self, other):
        return self.__dict__['raw'] != other
    def __lt__(self, other):
        return self.__dict__['raw'] < other
    def __gt__(self, other):
        return self.__dict__['raw'] > other
    def __le__(self, other):
        return self.__dict__['raw'] <= other
    def __ge__(self, other):
        return self.__dict__['raw'] >= other
    def __xor__(self, other):
        return self.__dict__['raw'] ^ other
    def __or__(self, other):
        return self.__dict__['raw'] | other
    #def __iadd__(self, other):
    #    return self.__dict__['raw'] += other
    #def __isub__(self, other):
    #    return self.__dict__['raw'] -= other
    #def __imul__(self, other):
    #    return self.__dict__['raw'] *= other
    #def __idiv__(self, other):
    #    return self.__dict__['raw'] /= other
    #def __ifloordiv__(self, other):
    #    return self.__dict__['raw'] //= other
    #def __imod__(self, other):
    #    return self.__dict__['raw'] %= other
    #def __ipow__(self, other):
    #    return self.__dict__['raw'] **= other
    #def __ilshift__(self, other):
    #    return self.__dict__['raw'] <<= other
    #def __irshift__(self, other):
    #    return self.__dict__['raw'] >>= other
    #def __iand__(self, other):
    #    return self.__dict__['raw'] &= other
    #def __ixor__(self, other):
    #    return self.__dict__['raw'] ^= other
    #def __ior__(self, other):
    #    return self.__dict__['raw'] |= other

# Similar to the above, but slightly expanded to contain information from analysis files.
class KineticsIteration(object):
    def __init__(self, kin_h5file, index, assign, iteration=-1):
        self.__dict__ = {}
        self.h5file = kin_h5file
        # Keys:
        _2D_h5keys = [ 'conditional_flux_evolution', 'rate_evolution' ]
        _1D_h5keys = [ 'state_pop_evolution', 'color_prob_evolution' ]
        for key in _2D_h5keys:
            self.__dict__[key] = self.__2D_with_error__(key, index, assign)
        for key in _1D_h5keys:
            self.__dict__[key] = self.__1D_with_error__(key, index, assign)
        try:
            self.__dict__['total_fluxes'] = WIPIDataset(raw=self.h5file['total_fluxes'], key=None)
        except:
            pass

    def __repr__(self):
        return repr(self.__dir__())
    def __getitem__(self, value):
        if value in self.__dict__.keys():
            return self.__dict__[value]
    def __setitem__(self, key, value):
        self.__dict__[key] = value
    def __getattr__(self, value):
        if value in self.__dict__.keys():
            return self.__dict__[value]
    def __setattr__(self, key, value):
        self.__dict__[key] = value
    def __dir__(self):
        dict_keys = self.__dict__.keys()
        # We don't want to show the plotter class; just the plot function
        remove = [ 'h5file', '__dict__']
        for i in remove:
            try:
                dict_keys.remove(str(i))
            except:
                pass
        return sorted(set(dict_keys))
    def keys(self):
        print(self.__dir__())

    # We seriously need to rename this.
    # It's similar to the global WIPDataset, but has some nice pretty print functions.
    class __custom_dataset__(object):
        # This is just allow it to be indexed via properties.
        # Not a huge thing, but whatever.
        def __init__(self, raw, assign, key):
            self.__dict__ = {}
            self.raw = raw
            self.name = key
            self.assign = assign
            self.nstates = assign.attrs['nstates']
            self.dim = len(raw.shape)
        def __repr__(self):
            return repr(self.__dir__())
        def __getitem__(self, value):
            if value in self.__dict__['raw'].dtype.names:
                return self.__dict__['raw'][value]
            elif value in self.__dict__.keys():
                return self.__dict__[value]
        def __setitem__(self, key, value):
            self.__dict__[key] = value
        def __getattr__(self, value):
            if value in self.__dict__['raw'].dtype.names:
                return self.__dict__['raw'][value]
            elif value in self.__dict__.keys():
                return self.__dict__[value]
        def __setattr__(self, key, value):
            self.__dict__[key] = value
        def __dir__(self):
            dict_keys = self.__dict__.keys()
            # We don't want to show the plotter class; just the plot function
            remove = ['assign', 'dim', 'nstates', 'plotter', '__dict__']
            for i in remove:
                try:
                    dict_keys.remove(str(i))
                except:
                    pass
            return sorted(set(list(self.raw.dtype.names) + dict_keys))
        def keys(self):
            print(self.__dir__())
        def _repr_pretty_(self, p, cycle):
            if self.dim == 1:
                return self._1D_repr_pretty_(p, cycle)
            if self.dim == 2:
                return self._2D_repr_pretty_(p, cycle)
        def _1D_repr_pretty_(self, p, cycle):
           # We're just using this as a way to print things in a pretty way.  They can still be indexed appropriately.
           # Stolen shamelessly from westtools/kinetics_tool.py
            maxlabellen = max(map(len,self.assign['state_labels']))
            p.text('')
            p.text('{name} data:\n'.format(name=self.name))
            for istate in xrange(self.nstates):
                p.text('{:{maxlabellen}s}: mean={:21.15e} CI=({:21.15e}, {:21.15e}) * tau^-1\n'
                    .format(self.assign['state_labels'][istate],
                    self.raw['expected'][istate],
                    self.raw['ci_lbound'][istate],
                    self.raw['ci_ubound'][istate],
                    maxlabellen=maxlabellen))
            p.text('To access data, index via the following names:\n')
            p.text(str(self.__dir__()))
            return " "
        def _2D_repr_pretty_(self, p, cycle):
            # We're just using this as a way to print things in a pretty way.  They can still be indexed appropriately.
            # Stolen shamelessly from westtools/kinetics_tool.py
            maxlabellen = max(map(len,self.assign['state_labels']))
            p.text('')
            p.text('{name} data:\n'.format(name=self.name))
            for istate in xrange(self.nstates):
                for jstate in xrange(self.nstates):
                    if istate == jstate: continue
                    p.text('{:{maxlabellen}s} -> {:{maxlabellen}s}: mean={:21.15e} CI=({:21.15e}, {:21.15e}) * tau^-1\n'
                        .format(self.assign['state_labels'][istate], self.assign['state_labels'][jstate],
                        self.raw['expected'][istate, jstate],
                        self.raw['ci_lbound'][istate, jstate],
                        self.raw['ci_ubound'][istate, jstate],
                        maxlabellen=maxlabellen))
            p.text('To access data, index via the following names:\n')
            p.text(str(self.__dir__()))
            return " "


    def __2D_with_error__(self, h5key, index, assign):
        # Check the start and stop, calculate the block size, and index appropriately.
        # While we could try and automatically generate this above, it's a little more consistent to try it here.
        # This should show the first block for which the current iteration has contributed data.
        self.step_iter = (self.h5file[h5key]['iter_stop'][0] - self.h5file[h5key]['iter_start'][0])[1,0]
        value = ((index-2) // self.step_iter)
        if value < 0:
            value = 0
        raw = self.h5file[h5key][value, :, :]
        error = (raw['ci_ubound'] - raw['ci_lbound']) / (2*raw['expected'])
        expected = raw['expected']
        raw = self.__custom_dataset__(raw, assign, h5key)
        raw.error = error
        raw.plotter = Plotter(self.h5file, h5key, iteration=value, interface='text')
        raw.plot = raw.plotter.plot
        return raw
    def __1D_with_error__(self, h5key, index, assign):
        self.step_iter = (self.h5file[h5key]['iter_stop'][0] - self.h5file[h5key]['iter_start'][0])[1]
        value = ((index-1) // self.step_iter)
        if value < 0:
            value = 0
        raw = self.h5file[h5key][value, :]
        error = (raw['ci_ubound'] - raw['ci_lbound']) / (2*raw['expected'])
        expected = raw['expected']
        raw = self.__custom_dataset__(raw, assign, h5key)
        raw.error = error
        raw.plotter = Plotter(self.h5file, h5key, iteration=value, interface='text')
        raw.plot = raw.plotter.plot
        return raw


class __get_data_for_iteration__(object):
    '''
    All interesting data from an iteration (current/past).  Whenever you change the scheme or iteration,
    this dictionary is automatically updated.  For the current iteration, it's keyed to the current seg_id.
    For the past iteration, it's keyed to the seg_id in the CURRENT iteration such that:

        w.current[X] & w.past[X]

    returns information about seg_id X in the current iteration and information on seg_ID X's PARENT in the
    preceding iteration.

    Can be indexed via a seg_id, or like a dictionary with the following keys:

        kinavg, weights, pcoord, auxdata (optional), parents, summary, seg_id, walkers, states, bins

    kinavg, states, and bins refer to the output from w_kinavg and w_assign for this iteration
    and analysis scheme.  They are NOT dynamics bins, but the bins defined in west.cfg.  
    
    Has the following properties:

        .minweight, .maxweight

    which return all properties of the segment that matches those criteria in the selected iteration.

    If you change the analysis scheme, so, too, will the important values.
    '''

    def __init__(self, parent, value, seg_ids = None):
        '''
        Initializes and sets the correct data.
        '''
        # We've classed this so that we can override some of the normal functions and allow indexing via seg_id
        self.__dict__ = {}
        # Is this function thread safe?
        iter_group = parent.data_reader.get_iter_group(value)
        #iter_group = parent.west['iterations/iter_{num:08d}'.format(num=value)]
        self.parent = parent
        current = {}
        current['iteration'] = value
        if seg_ids == None:
            seg_ids = xrange(0, iter_group['seg_index']['weight'].shape[0])
        # Just make these easier to access.
        current['weights'] = iter_group['seg_index']['weight'][seg_ids]
        current['pcoord'] = iter_group['pcoord'][...][seg_ids, :, :]
        try:
            current['auxdata'] = {}
            for key in iter_group['auxdata'].keys():
                current['auxdata'][key] = iter_group['auxdata'][key][...][seg_ids, :]
        except:
            pass
        current['parents'] = iter_group['seg_index']['parent_id'][seg_ids]
        current['summary'] = parent.data_reader.data_manager.get_iter_summary(int(value))
        current['seg_id'] = np.array(range(0, iter_group['seg_index'].shape[0]))[seg_ids]
        current['walkers'] = current['summary']['n_particles']
        current['states'] = parent.assign['trajlabels'][value-1, :current['walkers'], :][seg_ids]
        current['bins'] = parent.assign['assignments'][value-1, :current['walkers'], :][seg_ids]
        # Calculates the bin population for this iteration.
        nbins = parent.assign['state_map'].shape[0]
        # We have to take the 'unknown' state into account
        nstates = parent.assign['state_labels'].shape[0] + 1
        # Temporarily disabled while I sort out the fact that we shouldn't be using data from w_assign for state populations.
        #current['plot'] = Plotter(parent.direct, parent.reweight, parent.iteration, parent.assign['bin_labels'], parent.assign['state_labels'], current['populations'].states, current['populations'].bins, parent.interface)
        # Now we'll load up the results of the kinetics analysis.
        current['direct'] = KineticsIteration(parent.direct, value, parent.assign, value)
        evolution_datasets = [ 'rate_evolution', 'conditional_flux_evolution', 'state_pop_evolution', 'color_prob_evolution' , 'total_fluxes']
        # We want to load these up as... oh, who knows, I suppose?
        try:
            current['reweight'] = KineticsIteration(parent.reweight, value, parent.assign, value)
            # We'll make this not a sparse matrix...
            matrix = parent.reweight['iterations/iter_{:08d}'.format(value)]
            # Assume color.
            current['instant_matrix'] = sp.coo_matrix((matrix['flux'][...], (matrix['rows'][...], matrix['cols'][...])), shape=((nbins-1)*2, (nbins-1)*2)).todense()
            reweighting = True
        except:
{'fraction_non_alphanumeric': 0.12216742807006928, 'fraction_numerical': 0.0064744501255559255, 'mean_word_length': 2.440616217341522, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3691573', 'n_tokens_mistral': 9884, 'n_tokens_neox': 9646, 'n_words': 2718}