#!/usr/bin/env python
# PyChefHelpers.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.
#
# Functions to help with the implementation of PyChef - little jiffy
# bits which don't really need to be embedded in the main program.
#

import sys
import math
import os

from iotbx.mtz import object as mtz_factory

from cctbx.miller import build_set
from cctbx.crystal import symmetry as crystal_symmetry

def get_mtz_column_list(hklin):

    mtz_obj = mtz_factory(file_name = hklin)

    # construct a list of columns in the file

    cnames = []

    batch_column = None
    batch_dataset = None

    for crystal in mtz_obj.crystals():
        for dataset in crystal.datasets():
            for column in dataset.columns():
                cnames.append(column.label())

    return cnames

def compute_unique_reflections(unit_cell,
                               space_group,
                               anomalous,
                               high_resolution_limit,
                               low_resolution_limit = None):
    '''Compute the list of unique reflections from the unit cell and space
    group.'''

    cs = crystal_symmetry(unit_cell = unit_cell,
                          space_group = space_group)

    return [hkl for hkl in build_set(cs, anomalous,
                                     d_min = high_resolution_limit,
                                     d_max = low_resolution_limit).indices()]

if __name__ == '__main__':

    for hklin in sys.argv[1:]:

        mtz_obj = mtz_factory(file_name = hklin)
        sg = mtz_obj.space_group().build_derived_patterson_group()

        for crystal in mtz_obj.crystals():
            uc = crystal.unit_cell()

            for dataset in crystal.datasets():

                print crystal.name(), dataset.name()

        dmax, dmin = mtz_obj.max_min_resolution()

        print len(compute_unique_reflections(uc, sg, True,
                                             dmin, dmax))

        ms = set()

        miller_indices = mtz_obj.extract_miller_indices()

        for hkl in miller_indices:
            ms.add(hkl)

        print len(ms)