#!/usr/bin/env python
# XScaleHelpers.py
#   Copyright (C) 2006 CCLRC, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# Helpers for the wrapper for XSCALE, the XDS Scaling program.
#

import math
import sys

def _generate_resolution_shells(low, high):
    '''Generate 9 evenly spaced in reciprocal space resolution
    shells from low to high resolution, e.g. in 1/d^2.'''

    dmin = (1.0 / high) * (1.0 / high)
    dmax = (1.0 / low) * (1.0 / low)
    diff = (dmin -  dmax) / 8.0

    shells = [1.0 / math.sqrt(dmax)]

    for j in range(8):
        shells.append(1.0 / math.sqrt(dmax + diff * (j + 1)))

    return shells

def generate_resolution_shells_str(low, high):
    '''Generate a string of 8 evenly spaced in reciprocal space resolution
    shells from low to high resolution, e.g. in 1/d^2.'''

    result = ''
    shells = _generate_resolution_shells(low, high)

    for s in shells:
        result += ' %.2f' % s

    return result

def get_correlation_coefficients_and_group(xscale_lp):
    '''Get and group correlation coefficients between data sets from the
    xscale log file. Also access the reflection file names to show which ones
    should be scaled together.'''

    ccs = { }

    file_names = { }

    xmax = 0

    records = open(xscale_lp).readlines()

    # first scan through to get the file names...

    for j, record in enumerate(records):
        if 'NUMBER OF UNIQUE REFLECTIONS' in record:

            k = j + 5

            while len(records[k].split()) == 5:
                values = records[k].split()
                file_names[int(values[0])] = values[-1]

                k += 1

            break

    if not file_names:
        for j, record in enumerate(records):
            if 'SET# INTENSITY  ACCEPTED REJECTED' in record:

                k = j + 1

                while len(records[k].split()) == 5:
                    values = records[k].split()
                    file_names[int(values[0])] = values[-1]

                    k += 1

                break


    for j, record in enumerate(records):

        if 'CORRELATIONS BETWEEN INPUT DATA SETS' in record:

            k = j + 5

            while len(records[k].split()) == 6:
                values = records[k].split()

                _i = int(values[0])
                _j = int(values[1])
                _n = int(values[2])
                _cc = float(values[3])

                ccs[(_i, _j)] = (_n, _cc)
                ccs[(_j, _i)] = (_n, _cc)

                xmax = _i + 1

                k += 1

            break

    used = []
    groups = { }

    for j in range(xmax):
        test_file = file_names[j + 1]
        if test_file in used:
            continue
        used.append(test_file)
        groups[test_file] = [test_file]
        for k in range(j + 1, xmax):
            if ccs[(j + 1, k + 1)][1] > 0.9:
                groups[test_file].append(file_names[k + 1])
                used.append(file_names[k + 1])

    return groups

if __name__ == '__main__':
    get_correlation_coefficients_and_group(sys.argv[1])