#!/usr/bin/env python
# WedgeExpert.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.
#
# Some code to help figure out how the experiment was performed and tell
# Chef where we could consider cutting the data at... N.B. this will now
# take the digested wedges from the Chef wrapper.

import math
import sys
import pickle

def digest_wedges(wedges):
    '''Digest the wedges defined as a list of

    FIRST_DOSE FIRST_BATCH SIZE EXPOSURE DATASET

    to a set of potential "stop" points, in terms of the doses. This will
    return a list of DOSE values which should be treated as LIMITS (i.e.
    d:d < DOSE ok).'''

    # first digest to logical sweeps, keyed by the last image in the set
    # and the wavelength name, and containing the start and end dose.

    if False:
        return

    doses = { }

    belonging_wedges = { }

    for w in wedges:
        dose, batch, size, exposure, dataset = w
        if (dataset, batch) in doses:
            k_old = (dataset, batch)
            k_new = (dataset, batch + size)

            sweep = doses[k_old]
            doses[k_new] = (sweep[0], sweep[1] + exposure * size)

            belonging_wedges[k_new] = belonging_wedges[k_old]
            belonging_wedges[k_new].append(w)

            del(doses[k_old])
            del(belonging_wedges[k_old])

        else:
            end_dose = dose + exposure * (size - 1)

            doses[(dataset, batch + size)] = (dose, end_dose)
            belonging_wedges[(dataset, batch + size)] = [w]

    # now invert

    sweeps = { }

    for k in doses:
        sweeps[doses[k]] = k

    # now try to figure the sweeps which are overlapping - these will be
    # added to a list named "groups"

    groups = []

    last_time = sorted(sweeps)[0][0] - 1.0

    for s in sorted(sweeps):
        if s[0] > last_time:
            last_time = s[1]
            groups.append([s])
        else:
            last_time = s[1]
            groups[-1].append(s)

    dmaxes = []

    group_report = []

    for j, g in enumerate(groups):

        # check that the group structure is correct, i.e. all have the
        # uniform number of wedges...

        group_wedges = { }

        for s in g:
            all_wedges = belonging_wedges[sweeps[s]]
            dataset = all_wedges[0][-1]
            last_image = int(s[1])

            # Boom! this is wrong for inverse beam data sets

            group_wedges[(dataset, last_image)] = all_wedges

        size = 0

        datasets = list(set([k[0] for k in group_wedges]))

        if len(g) == 1:
            assert(len(belonging_wedges[sweeps[g[0]]]) == 1)
            group_report.append('Single wedge')
        elif len(datasets) == 1:
            assert(len(group_wedges) == 2)
            group_report.append('Inverse beam')
        else:
            group_report.append('%d-wedge data collection' % \
                                len(group_wedges))

        for dataset, last_image in group_wedges:
            if not size:
                size = len(group_wedges[(dataset, last_image)])

            assert(size == len(group_wedges[(dataset, last_image)]))

        for j in range(size):

            d_local = []

            for s in g:
                bw = belonging_wedges[sweeps[s]][j]
                d_local.append(bw[0] + bw[2] * bw[3])

            dmaxes.append(max(d_local))

    return dmaxes, group_report

if __name__ == '__main__':

    wedges = pickle.loads(open('test.pkl').read())

    dmaxes, group_report = digest_wedges(wedges)

    for gr in group_report:
        print gr