#!/usr/bin/env cctbx.python
# PyChef.py
#
#   Copyright (C) 2009 Diamond Light Source, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# A reimplementation of the FORTRAN program CHEF (Winter, PhD thesis)
# into Python, using the CCTBX library. The idea is to analyse scaled
# intensity measurements as a function of cumulative dose, to assess the
# impact of radiation damage.
#

import sys
import math
import os
import time
import threading

from PyChefHelpers import get_mtz_column_list, compute_unique_reflections

from iotbx import mtz

class PyChef:
    '''The main PyChef class.'''

    def __init__(self):

        # assuming that we will be using BATCH by default... unless there is a
        # column named DOSE

        self._base_column = 'BATCH'

        self._hklin_list = []

        self._reflections = { }
        self._unit_cells = { }
        self._space_groups = { }

        self._range_min = None
        self._range_max = None
        self._range_width = None

        # Dose / batch information for reporting purposes

        self._dose_information = { }

        self._resolution_high = None
        self._resolution_low = None

        self._resolution_bins = 8

        self._anomalous = False

        self._title = None

        self._ncpu = 1

        return

    def set_base_column(self, base_column):
        '''Set the baseline for analysis: note well that this should
        be called before add_hklin, so that add_hklin can check that
        this column is present.'''

        self._base_column = base_column

        return

    def add_hklin(self, hklin):
        '''Add a reflection file to the list for analysis... '''

        columns = get_mtz_column_list(hklin)

        assert(self._base_column in columns)

        self._hklin_list.append(hklin)

        return

    def set_range(self, range_min, range_max, range_width):
        '''Set the range of e.g. dose to consider for analysis, with
        binning width set.'''

        assert(range_max > range_min)
        assert(range_max - range_min > range_width)

        self._range_min = range_min
        self._range_max = range_max
        self._range_width = range_width

        return

    def set_resolution(self, resolution_high, resolution_low = None):
        '''Set the resolution range for analysis.'''

        self._resolution_high = resolution_high

        if resolution_low:
            assert(resolution_low > resolution_high)
            self._resolution_low = resolution_low

        return

    def set_anomalous(self, anomalous):
        '''Set the separation of anomalous pairs on or off.'''

        self._anomalous = anomalous
        return

    def set_title(self, title):
        self._title = title
        return

    def set_ncpu(self, ncpu):
        '''Set the number of cpus to be used for parallel processing.'''
        self._ncpu = ncpu
        return

    def init(self):
        '''Initialise the program - this will read all of the reflections
        and so on into memory and set up things like the unit cell objects.'''

        symmetry = None

        overall_dmin = None
        overall_dmax = None

        overall_range_min = None
        overall_range_max = None
        overall_range_width = None

        for hklin in self._hklin_list:

            mtz_obj = mtz.object(hklin)

            mi = mtz_obj.extract_miller_indices()
            dmax, dmin = mtz_obj.max_min_resolution()

            if overall_dmin is None:
                overall_dmin = dmin
            else:
                if dmin > overall_dmin:
                    overall_dmin = dmin

            if overall_dmax is None:
                overall_dmax = dmax
            else:
                if dmax < overall_dmax:
                    overall_dmax = dmax

            crystal_name = None
            dataset_name = None
            nref = 0
            uc = None

            # chef does not care about systematic absences from e.g.
            # screw axes => patterson group not space group. No,
            # patterson group is always centric?!

            sg = mtz_obj.space_group()

            # .build_derived_patterson_group()

            if not symmetry:
                symmetry = sg
            else:
                assert(symmetry == sg)

            # now have a rummage through to get the columns out that I want

            base_column = None
            misym_column = None
            i_column = None
            sigi_column = None

            batch_column = None
            dose_column = None

            for crystal in mtz_obj.crystals():

                for dataset in crystal.datasets():
                    if dataset.name() != 'HKL_base':
                        dataset_name = dataset.name()

                if crystal.name() != 'HKL_base':
                    crystal_name = crystal.name()

                uc = crystal.unit_cell()

                for dataset in crystal.datasets():

                    for column in dataset.columns():

                        # for recording purposes for the BATCH / DOSE
                        # mapping...

                        if column.label() == 'BATCH':
                            batch_column = column
                        elif column.label() == 'DOSE':
                            dose_column = column

                        if column.label() == self._base_column:
                            base_column = column
                        if column.label() == 'M_ISYM':
                            misym_column = column
                        if column.label() == 'I':
                            i_column = column
                        if column.label() == 'SIGI':
                            sigi_column = column

            assert(base_column != None)
            assert(misym_column != None)
            assert(i_column != None)
            assert(sigi_column != None)

            if batch_column and dose_column:

                # read, accumulate the dose information - assume that these
                # are correctly structured as per xia2 handling...

                batch_values = batch_column.extract_values(
                    not_a_number_substitute = 0.0)
                dose_values = dose_column.extract_values(
                    not_a_number_substitute = 0.0)

                assert(len(batch_values) == len(dose_values))

                if min(dose_values) != max(dose_values):

                    for j in range(len(batch_values)):
                        batch = int(round(batch_values[j]))
                        dose = dose_values[j]

                        if not batch in self._dose_information:
                            self._dose_information[batch] = dose, dataset_name

            print 'Reading in data from %s/%s' % (crystal_name, dataset_name)
            print 'Cell: %6.2f %6.2f %6.2f %6.2f %6.2f %6.2f' % \
                  tuple(uc.parameters())
            print 'Spacegroup: %s' % sg.type(
                ).universal_hermann_mauguin_symbol()
            print 'Using: %s/%s/%s' % \
                  (i_column.label(), sigi_column.label(), base_column.label())

            base_values = base_column.extract_values(
                not_a_number_substitute = 0.0)

            min_base = min(base_values)
            max_base = max(base_values)

            if not self._range_width:
                bases = sorted(set(base_values))
                mean_shift = sum([bases[j + 1] - bases[j] for j in \
                                  range(len(bases) - 1)]) / (len(bases) - 1)

                if overall_range_width is None:
                    overall_range_width = mean_shift
                if mean_shift < overall_range_width:
                    overall_range_width = mean_shift

            if overall_range_min == None:
                overall_range_min = min_base
            else:
                if min_base < overall_range_min:
                    overall_range_min = min_base

            if overall_range_max == None:
                overall_range_max = max_base
            else:
                if max_base > overall_range_max:
                    overall_range_max = max_base

            misym_values = misym_column.extract_values(
                not_a_number_substitute = 0.0)

            i_values = i_column.extract_values(
                not_a_number_substitute = 0.0)
            i_values_valid = i_column.selection_valid()

            sigi_values = sigi_column.extract_values(
                not_a_number_substitute = 0.0)
            sigi_values_valid = sigi_column.selection_valid()

            reflections = { }

            for j in range(mi.size()):

                if not i_values_valid[j]:
                    continue

                if not sigi_values_valid[j]:
                    continue

                h, k, l = mi[j]
                base = base_values[j]
                misym = misym_values[j]
                i = i_values[j]
                sigi = sigi_values[j]

                # test whether this is within the range of interest

                if self._range_min != None:
                    if base < self._range_min:
                        continue

                if self._range_max != None:
                    if base > self._range_max:
                        continue

                if self._resolution_high:
                    if uc.d([h, k, l]) < self._resolution_high:
                        continue

                if self._resolution_low:
                    if uc.d([h, k, l]) > self._resolution_low:
                        continue

                if self._anomalous:
                    pm = int(round(misym)) % 2
                else:
                    pm = 0

                if not (h, k, l) in reflections:
                    reflections[(h, k, l)] = []

                reflections[(h, k, l)].append((pm, base, i, sigi))

            # ok, copy the reflections to the class for future analysis

            self._reflections[(crystal_name, dataset_name)] = reflections
            self._unit_cells[(crystal_name, dataset_name)] = uc
            self._space_groups[(crystal_name, dataset_name)] = sg

        if not self._resolution_low:
            print 'Assigning low resolution limit: %.2f' % overall_dmax
            self._resolution_low = overall_dmax

        if not self._resolution_high:
            print 'Assigning high resolution limit: %.2f' % overall_dmin
            self._resolution_high = overall_dmin

        if not self._range_width:
            print 'Assigning baseline width:   %.2f' % \
                  overall_range_width
            self._range_width = overall_range_width

        if not self._range_min:
            print 'Assigning baseline minimum: %.2f' % \
                  (overall_range_min - self._range_width)
            self._range_min = overall_range_min - self._range_width

        if not self._range_max:
            print 'Assigning baseline maximum: %.2f' % \
                  overall_range_max
            self._range_max = overall_range_max

        # if > 1 data set assume that this is anomalous data...

        if len(self._reflections) > 1:
            print 'More than one data set: assume anomalous'
            self._anomalous = True

        # FIXME add warning if measurements don't reach the edge...

        return

    def print_completeness_vs_dose(self):
        '''Print the completeness vs. dose for each input reflection file.
        This will need to read the MTZ file, get the dataset cell constants
        and space group, compute the expected number of reflections, then
        as a function of dose compute what fraction of these are measured.
        For native data this is going to simply write out:

        dose, fraction

        but for MAD data it will need to print out:

        dose, fractionI+, fractionI-, fractionI, fractionI+andI-

        would be nice to have this described a little tidier.'''

        for crystal_name, dataset_name in sorted(self._reflections):

            reflections = self._reflections[(crystal_name, dataset_name)]
            uc = self._unit_cells[(crystal_name, dataset_name)]
            sg = self._space_groups[(crystal_name, dataset_name)]

            nref = len(compute_unique_reflections(uc, sg, self._anomalous,
                                                  self._resolution_high,
                                                  self._resolution_low))

            nref_n = len(compute_unique_reflections(uc, sg, False,
                                                    self._resolution_high,
                                                    self._resolution_low))

            print 'Cumulative completeness analysis for %s/%s' % \
                  (crystal_name, dataset_name)
            print 'Expecting %d reflections' % nref

            # Right, in here I need to get the lowest dose for a given
            # h, k, l and then add this reflection for completeness to all
            # high dose bins, as it is measured already. So, if it is
            #
            # centric and anomalous, add to I+ and I-
            # centric and not anomalous, add to I
            # acentric and anomalous, add to I+ or I-
            # acentric and not anomalous, add to I
            #
            # To do this the best thing to do is to read through all of the
            # reflections and keep a dictionary of the lowest dose at which
            # a given reflection was recorded. Then iterate through this
            # list to see how many we have as a function of dose...
            #
            # Ok, so after trying to implement this cleanly I think that the
            # only way to do this is to actually read all of the reflections
            # in and store them in e.g. a dictionary. Could be expensive
            # for large data sets, but worry about that ... later. Can at
            # least store this in-memory representation. Accordingly will also
            # need to read in the 'I' column...

            # this will be a dictionary indexed by the Miller indices and
            # containing anomalous flag (1: I+, 0:I- or native) baseline
            # intensity values and the error estimate..

            # now construct the completeness tables for I or I+ & I-,
            # then populate from this list of lowest doses...

            if self._anomalous:

                print '$TABLE : Completeness vs. %s, %s/%s:' % \
                      (self._base_column, crystal_name, dataset_name)
                print '$GRAPHS: Completeness:N:1,2,3,4,5: $$'
                print '%8s %5s %5s %5s %5s $$ $$' % \
                      (self._base_column, 'I+', 'I-', 'I', 'dI')

                iplus_count = []
                iminus_count = []
                ieither_count = []
                iboth_count = []

                nsteps = 1 + int(
                    (self._range_max - self._range_min) / self._range_width)

                for j in range(nsteps):
                    iplus_count.append(0)
                    iminus_count.append(0)
                    ieither_count.append(0)
                    iboth_count.append(0)

                for h, k, l in reflections:
                    base_min_iplus = self._range_max + self._range_width
                    base_min_iminus = self._range_max + self._range_width

                    for pm, base, i, sigi in reflections[(h, k, l)]:
                        if sg.is_centric((h, k, l)):
                            if base < base_min_iplus:
                                base_min_iplus = base
                            if base < base_min_iminus:
                                base_min_iminus = base
                        elif pm:
                            if base < base_min_iplus:
                                base_min_iplus = base
                        else:
                            if base < base_min_iminus:
                                base_min_iminus = base

                    start_iplus = int((base_min_iplus - self._range_min)
                                      / self._range_width)

                    start_iminus = int((base_min_iminus - self._range_min)
                                      / self._range_width)

                    if start_iplus < nsteps:
                        iplus_count[start_iplus] += 1
                    if start_iminus < nsteps:
                        iminus_count[start_iminus] += 1
                    if min(start_iplus, start_iminus) < nsteps:
                        ieither_count[min(start_iplus, start_iminus)] += 1
                    if max(start_iplus, start_iminus) < nsteps:
                        iboth_count[max(start_iplus, start_iminus)] += 1

                # now sum up

                for j in range(1, nsteps):
                    iplus_count[j] += iplus_count[j - 1]
                    iminus_count[j] += iminus_count[j - 1]
                    ieither_count[j] += ieither_count[j - 1]
                    iboth_count[j] += iboth_count[j - 1]

                for j in range(nsteps):
                    iplus = iplus_count[j] / float(nref_n)
                    iminus = iminus_count[j] / float(nref_n)
                    ieither = ieither_count[j] / float(nref_n)
                    iboth = iboth_count[j] / float(nref_n)

                    print '%8.1f %5.3f %5.3f %5.3f %5.3f' % \
                          (self._range_min + j * self._range_width,
                           iplus, iminus, ieither, iboth)

                print '$$'

            else:

                print '$TABLE : Completeness vs. %s, %s/%s:' % \
                      (self._base_column, crystal_name, dataset_name)
                print '$GRAPHS: Completeness:N:1, 2: $$'
                print '%8s %5s $$ $$' % (self._base_column, 'I')

                i_count = []

                nsteps = 1 + int(
                    (self._range_max - self._range_min) / self._range_width)

                for j in range(nsteps):
                    i_count.append(0)

                for h, k, l in reflections:
                    base_min = self._range_max + self._range_width

                    for pm, base, i, sigi in reflections[(h, k, l)]:
                        if base < base_min:
                            base_min = base

                    start = int((base_min - self._range_min)
                                / self._range_width)

                    # for j in range(start, nsteps):
                    i_count[start] += 1

                for j in range(1, nsteps):
                    i_count[j] += i_count[j - 1]

                for j in range(nsteps):
                    i = i_count[j] / float(nref)

                    print '%8.1f %5.3f' % \
                          (self._range_min + j * self._range_width, i)

                print '$$'

        return

    def rd(self):
        '''Calculate Rd (Diederichs, 2006) for the data in the individual
        wavelengths. This will obviously result in one plot for each of the
        input data sets.'''

        nsteps = 1 + int(
            (self._range_max - self._range_min) / self._range_width)

        for crystal_name, dataset_name in sorted(self._reflections):

            rd_top = [0.0 for j in range(nsteps)]
            rd_bottom = [0.0 for j in range(nsteps)]

            reflections = self._reflections[(crystal_name, dataset_name)]

            if self._anomalous:

                for hkl in reflections:

                    observations = self._reflections[
                        (crystal_name, dataset_name)][hkl]

                    iplus = []
                    iminus = []

                    for pm, base, i, sigi in observations:
                        if pm:
                            iplus.append((base, i, sigi))
                        else:
                            iminus.append((base, i, sigi))

                    for n, (base, i, sigi) in enumerate(iplus):
                        for _base, _i, _sigi in iplus[n + 1:]:
                            d = int(round(math.fabs(base - _base) /
                                          self._range_width))
                            rd_top[d] += math.fabs(i - _i)
                            rd_bottom[d] += 0.5 * (i + _i)

                    for n, (base, i, sigi) in enumerate(iminus):
                        for _base, _i, _sigi in iminus[n + 1:]:
                            d = int(round(math.fabs(base - _base) /
                                          self._range_width))
                            rd_top[d] += math.fabs(i - _i)
                            rd_bottom[d] += 0.5 * (i + _i)

            else:

                for hkl in reflections:

                    observations = self._reflections[
                        (crystal_name, dataset_name)][hkl]


                    for n, (pm, base, i, sigi) in enumerate(observations):
                        for _pm, _base, _i, _sigi in observations[n + 1:]:
                            d = int(round(math.fabs(base - _base) /
                                          self._range_width))
                            rd_top[d] += math.fabs(i - _i)
                            rd_bottom[d] += 0.5 * (i + _i)

            # print the report...

            print '$TABLE : R vs. %s difference, %s/%s:' % \
                  (self._base_column, crystal_name, dataset_name)
            print '$GRAPHS: Rd:N:1, 2: $$'
            print '%8s %5s $$ $$' % (self._base_column, 'Rd')

            for j in range(nsteps):
                d = self._range_width * j
                if rd_bottom[j]:
                    r = rd_top[j] / rd_bottom[j]
                else:
                    r = 0.0
                print '%8.1f %5.3f' % (d, r)

            print '$$'

        return

    def scp(self):
        '''Perform the scp = rcp / ercp calculation as a function of
        assumulated dose across a number of resolution bins, from
        measurements already cached in memory.'''

        rcp_top = { }
        rcp_bottom = { }
        isigma = { }
        count = { }

        if self._resolution_low:
            smin = 1.0 / (self._resolution_low * self._resolution_low)
        else:
            smin = 0.0

        smax = 1.0 / (self._resolution_high * self._resolution_high)

        nsteps = 1 + int(
            (self._range_max - self._range_min) / self._range_width)

        # lay out the storage

        for j in range(self._resolution_bins + 1):

            rcp_top[j] = []
            rcp_bottom[j] = []
            isigma[j] = []
            count[j] = []

            for k in range(nsteps):
                rcp_top[j].append(0.0)
                rcp_bottom[j].append(0.0)
                isigma[j].append(0.0)
                count[j].append(0)

        # then populate

        for xname, dname in sorted(self._reflections):

            print 'Accumulating from %s %s' % (xname, dname)

            for h, k, l in self._reflections[(xname, dname)]:

                d = self._unit_cells[(xname, dname)].d([h, k, l])

                s = 1.0 / (d * d)

                bin = int(self._resolution_bins * (s - smin) / (smax - smin))

                observations = self._reflections[(xname, dname)][(h, k, l)]

                iplus = []
                iminus = []

                for pm, base, i, sigi in observations:
                    if pm:
                        iplus.append((base, i, sigi))
                    else:
                        iminus.append((base, i, sigi))

                # compute contributions

                for n, (base, i, sigi) in enumerate(iplus):

                    for _base, _i, _sigi in iplus[n + 1:]:
                        start = int((max(base, _base) - self._range_min) /
                                    self._range_width)

                        ra = math.fabs(i - _i)
                        rb = 0.5 * (i + _i)

                        rcp_top[bin][start] += ra
                        rcp_bottom[bin][start] += rb

                        isigma[bin][start] += (i / sigi) + (_i / _sigi)
                        count[bin][start] += 2

                for n, (base, i, sigi) in enumerate(iminus):

                    for _base, _i, _sigi in iminus[n + 1:]:
                        start = int((max(base, _base) - self._range_min) /
                                    self._range_width)

                        ra = math.fabs(i - _i)
                        rb = 0.5 * (i + _i)

                        rcp_top[bin][start] += ra
                        rcp_bottom[bin][start] += rb

                        isigma[bin][start] += (i / sigi) + (_i / _sigi)
                        count[bin][start] += 2

        # now accumulate as a funtion of time...

        for k in range(self._resolution_bins):
            for j in range(1, nsteps):
                rcp_top[k][j] += rcp_top[k][j - 1]
                rcp_bottom[k][j] += rcp_bottom[k][j - 1]
                isigma[k][j] += isigma[k][j - 1]
                count[k][j] += count[k][j - 1]

        # now digest the results - as a function of dose and resolution...

        if self._title:
            print '$TABLE : Normalised radiation damage analysis (%s):' % \
                  self._title

        else:
            print '$TABLE : Normalised radiation damage analysis:'

        print '$GRAPHS: Scp(d):N:1,%d: $$' % (self._resolution_bins + 2)

        columns = ''
        for j in range(self._resolution_bins):
            columns += ' S%d' % j

        print '%s %s Scp(d) $$ $$' % (self._base_column, columns)
        format = '%8.1f %6.4f'
        for k in range(self._resolution_bins):
            format += ' %6.4f'

        for j in range(nsteps):
            base = j * self._range_width + self._range_min
            values = [base]

            for k in range(self._resolution_bins):

                if rcp_bottom[k][j] and count[k][j] > 100:
                    rcp = rcp_top[k][j] / rcp_bottom[k][j]
                    isig = isigma[k][j] / count[k][j]
                    scp = rcp / (1.1284 / isig)
                else:
                    scp = 0.0
                    rcp = 0.0
                    isig = 0.0

                values.append(scp)

            values.append((sum(values[1:]) / self._resolution_bins))

            print format % tuple(values)

        print '$$'

        return

    def rcp(self):
        '''Perform the rcp calculation as a function of
        assumulated dose across a number of resolution bins, from
        measurements already cached in memory.'''

        rcp_top = { }
        rcp_bottom = { }

        if self._resolution_low:
            smin = 1.0 / (self._resolution_low * self._resolution_low)
        else:
            smin = 0.0

        smax = 1.0 / (self._resolution_high * self._resolution_high)

        nsteps = 1 + int(
            (self._range_max - self._range_min) / self._range_width)

        # lay out the storage

        for j in range(self._resolution_bins + 1):

            rcp_top[j] = []
            rcp_bottom[j] = []

            for k in range(nsteps):
                rcp_top[j].append(0.0)
                rcp_bottom[j].append(0.0)

        # then populate

        for xname, dname in sorted(self._reflections):

            print 'Accumulating from %s %s' % (xname, dname)

            for h, k, l in self._reflections[(xname, dname)]:

                d = self._unit_cells[(xname, dname)].d([h, k, l])

                s = 1.0 / (d * d)

                bin = int(self._resolution_bins * (s - smin) / (smax - smin))

                observations = self._reflections[(xname, dname)][(h, k, l)]

                iplus = []
                iminus = []

                for pm, base, i, sigi in observations:
                    if pm:
                        iplus.append((base, i, sigi))
                    else:
                        iminus.append((base, i, sigi))

                # compute contributions

                for n, (base, i, sigi) in enumerate(iplus):

                    for _base, _i, _sigi in iplus[n + 1:]:
                        start = int((max(base, _base) - self._range_min) /
                                    self._range_width)

                        ra = math.fabs(i - _i)
                        rb = 0.5 * (i + _i)

                        rcp_top[bin][start] += ra
                        rcp_bottom[bin][start] += rb

                for n, (base, i, sigi) in enumerate(iminus):

                    for _base, _i, _sigi in iminus[n + 1:]:
                        start = int((max(base, _base) - self._range_min) /
                                    self._range_width)

                        ra = math.fabs(i - _i)
                        rb = 0.5 * (i + _i)

                        rcp_top[bin][start] += ra
                        rcp_bottom[bin][start] += rb

        # now accumulate as a funtion of time...

        for k in range(self._resolution_bins):
            for j in range(1, nsteps):
                rcp_top[k][j] += rcp_top[k][j - 1]
                rcp_bottom[k][j] += rcp_bottom[k][j - 1]

        # now digest the results - as a function of dose and resolution...

        if self._title:
            print '$TABLE : Cumulative radiation damage analysis (%s):' % \
                  self._title

        else:
            print '$TABLE : Cumulative radiation damage analysis:'

        print '$GRAPHS: Rcp(d):N:1,%d: $$' % (self._resolution_bins + 2)

        columns = ''
        for j in range(self._resolution_bins):
            columns += ' S%d' % j

        print '%s %s Rcp(d) $$ $$' % (self._base_column, columns)
        format = '%8.1f %6.4f'
        for k in range(self._resolution_bins):
            format += ' %6.4f'

        for j in range(nsteps):
            base = j * self._range_width + self._range_min
            values = [base]

            for k in range(self._resolution_bins):

                if rcp_bottom[k][j]:
                    rcp = rcp_top[k][j] / rcp_bottom[k][j]
                else:
                    rcp = 0.0

                values.append(rcp)

            klist = [k in range(self._resolution_bins)]

            ot = sum([rcp_top[k][j] for k in klist])
            ob = sum([rcp_bottom[k][j] for k in klist])

            if ob:
                overall = ot / ob
            else:
                overall = 0.0

            values.append(overall)

            print format % tuple(values)

        print '$$'

        return

    def print_dose_profile(self):

        if not self._dose_information:
            return

        print '$TABLE: Dose vs. BATCH:'
        print '$GRAPHS: Dose Profile:N:1,2: $$'
        print 'BATCH DOSE DATASET $$ $$'

        for j in sorted(self._dose_information):
            d, n = self._dose_information[j]
            print '%d %f %s' % (j, d, n)

        print '$$'

        return

    # threaded functions to improve speed of above calculations

    def calculate_completeness_vs_dose_anomalous(self, reflections, hkls):
        '''Compute arrays of number vs. dose for anomalous data for
        these named reflections.'''

        iplus_count = []
        iminus_count = []
        ieither_count = []
        iboth_count = []

        nsteps = 1 + int(
            (self._range_max - self._range_min) / self._range_width)

        for j in range(nsteps):
            iplus_count.append(0)
            iminus_count.append(0)
            ieither_count.append(0)
            iboth_count.append(0)

        for hkl in hkls:
            base_min_iplus = self._range_max + self._range_width
            base_min_iminus = self._range_max + self._range_width

            for pm, base, i, sigi in reflections[hkl]:
                if sg.is_centric(hkl):
                    if base < base_min_iplus:
                        base_min_iplus = base
                    if base < base_min_iminus:
                        base_min_iminus = base
                elif pm:
                    if base < base_min_iplus:
                        base_min_iplus = base
                else:
                    if base < base_min_iminus:
                        base_min_iminus = base

            start_iplus = int((base_min_iplus - self._range_min)
                              / self._range_width)

            start_iminus = int((base_min_iminus - self._range_min)
                               / self._range_width)

            if start_iplus < nsteps:
                iplus_count[start_iplus] += 1
            if start_iminus < nsteps:
                iminus_count[start_iminus] += 1
            if min(start_iplus, start_iminus) < nsteps:
                ieither_count[min(start_iplus, start_iminus)] += 1
            if max(start_iplus, start_iminus) < nsteps:
                iboth_count[max(start_iplus, start_iminus)] += 1

        for j in range(1, nsteps):
            iplus_count[j] += iplus_count[j - 1]
            iminus_count[j] += iminus_count[j - 1]
            ieither_count[j] += ieither_count[j - 1]
            iboth_count[j] += iboth_count[j - 1]

        return iplus_count, iminus_count, iboth_count, ieither_count

    def parallel_calculate_completeness_vs_dose_anomalous(self, reflections):
        '''Compute arrays of completeness vs. dose for anomalous data for
        these reflections. This version will split the reflections into
        different chunks and accumulate afterwards.'''

        uc = self._unit_cells[(crystal_name, dataset_name)]
        sg = self._space_groups[(crystal_name, dataset_name)]

        nref = len(compute_unique_reflections(uc, sg, self._anomalous,
                                              self._resolution_high,
                                              self._resolution_low))

        nref_n = len(compute_unique_reflections(uc, sg, False,
                                                self._resolution_high,
                                                self._resolution_low))

        iplus_count = []
        iminus_count = []
        ieither_count = []
        iboth_count = []

        nsteps = 1 + int(
            (self._range_max - self._range_min) / self._range_width)

        for j in range(nsteps):
            iplus_count.append(0)
            iminus_count.append(0)
            ieither_count.append(0)
            iboth_count.append(0)

        hkls = list(reflections)

        chunk_size = len(hkls) / self._ncpu

        chunks = [hkls[j: j + chunk_size] \
                  for j in range(0, len(hkls), chunk_size)]

        # here need to spawn parallel threads

        for hkl in reflections:
            base_min_iplus = self._range_max + self._range_width
            base_min_iminus = self._range_max + self._range_width

            for pm, base, i, sigi in reflections[hkl]:
                if sg.is_centric(hkl):
                    if base < base_min_iplus:
                        base_min_iplus = base
                    if base < base_min_iminus:
                        base_min_iminus = base
                elif pm:
                    if base < base_min_iplus:
                        base_min_iplus = base
                else:
                    if base < base_min_iminus:
                        base_min_iminus = base

            start_iplus = int((base_min_iplus - self._range_min)
                              / self._range_width)

            start_iminus = int((base_min_iminus - self._range_min)
                               / self._range_width)

            if start_iplus < nsteps:
                iplus_count[start_iplus] += 1
            if start_iminus < nsteps:
                iminus_count[start_iminus] += 1
            if min(start_iplus, start_iminus) < nsteps:
                ieither_count[min(start_iplus, start_iminus)] += 1
            if max(start_iplus, start_iminus) < nsteps:
                iboth_count[max(start_iplus, start_iminus)] += 1

        # now sum up

        for j in range(1, nsteps):
            iplus_count[j] += iplus_count[j - 1]
            iminus_count[j] += iminus_count[j - 1]
            ieither_count[j] += ieither_count[j - 1]
            iboth_count[j] += iboth_count[j - 1]

        # now compute this as fractions

        comp_iplus = [ip / float(nref_n) for ip in iplus_count]
        comp_iminus = [im / float(nref_n) for im in iminus_count]
        comp_ieither = [ie / float(nref_n) for ie in iplus_count]
        comp_iboth = [ib / float(nref_n) for ib in iplus_count]

        # and return

        return comp_iplus, comp_iminus, comp_ieither, comp_iboth

    def help_calculate_completeness_vs_dose_anomalous(self, reflections, hkls):

        iplus_count = []
        iminus_count = []
        ieither_count = []
        iboth_count = []

        nsteps = 1 + int(
            (self._range_max - self._range_min) / self._range_width)

        for j in range(nsteps):
            iplus_count.append(0)
            iminus_count.append(0)
            ieither_count.append(0)
            iboth_count.append(0)

        for hkl in hkls:
            base_min_iplus = self._range_max + self._range_width
            base_min_iminus = self._range_max + self._range_width

            for pm, base, i, sigi in reflections[hkl]:
                if sg.is_centric(hkl):
                    if base < base_min_iplus:
                        base_min_iplus = base
                    if base < base_min_iminus:
                        base_min_iminus = base
                elif pm:
                    if base < base_min_iplus:
                        base_min_iplus = base
                else:
                    if base < base_min_iminus:
                        base_min_iminus = base

            start_iplus = int((base_min_iplus - self._range_min)
                              / self._range_width)

            start_iminus = int((base_min_iminus - self._range_min)
                               / self._range_width)

            if start_iplus < nsteps:
                iplus_count[start_iplus] += 1
            if start_iminus < nsteps:
                iminus_count[start_iminus] += 1
            if min(start_iplus, start_iminus) < nsteps:
                ieither_count[min(start_iplus, start_iminus)] += 1
            if max(start_iplus, start_iminus) < nsteps:
                iboth_count[max(start_iplus, start_iminus)] += 1

        return iplus_count, iminus_count, iboth_count, ieither_count