#!/usr/bin/env cctbx.python
# MtzFactory.py
#
#   Copyright (C) 2010 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 toolkit component to read MTZ format reflection files, wrapping the
# functionality in iotbx. This will return a data structure to represent
# merged and unmerged MTZ files.

import sys
import math
import os
import time

from iotbx import mtz

class mtz_dataset:
    '''A class to represent the MTZ dataset in the hierarchy. This will
    be instantiated in the mtz_crystal class below, and contain:

     - a list of columns

    Maybe more things will be added.'''

    def __init__(self, iotbx_dataset):
        self._name = iotbx_dataset.name()
        self._column_table = { }
        for column in iotbx_dataset.columns():
            self._column_table[column.label()] = column

        return

    def get_column_names(self):
        return list(self._column_table)

    def get_column(self, column_label):
        return self._column_table[column_label]

    def get_column_values(self, column_label, nan_value = 0.0):
        return self._column_table[column_label].extract_values(
            not_a_number_substitute = nan_value)

class mtz_crystal:
    '''A class to represent the MTZ crystal in the hierarchy. This will
    be instantiated by the factories below.'''

    def __init__(self, iotbx_crystal):
        self._name = iotbx_crystal.name()
        self._unit_cell = iotbx_crystal.unit_cell()

        self._dataset_table = { }
        for dataset in iotbx_crystal.datasets():
            self._dataset_table[dataset.name()] = mtz_dataset(dataset)

        self._column_table = { }

        for dname in self._dataset_table:
            dataset = self._dataset_table[dname]
            for column_name in dataset.get_column_names():
                assert(not column_name in self._column_table)
                self._column_table[column_name] = dataset.get_column(
                    column_name)

        return

    def get_dataset_names(self):
        return list(self._dataset_table)

    def get_dataset(self, dataset_name):
        return self._dataset_table[dataset_name]

    def get_unit_cell_parameters(self):
        return tuple(self._unit_cell.parameters())

    def get_unit_cell(self):
        return self._unit_cell

    def get_column_names(self):
        return list(self._column_table)

    def get_column(self, column_label):
        return self._column_table[column_label]

    def get_column_values(self, column_label, nan_value = 0.0):
        return self._column_table[column_label].extract_values(
            not_a_number_substitute = nan_value)

class mtz_file:
    '''A class to represent the full MTZ file in the hierarchy - this
    will have a list of one or more crystals contained within it each
    with its own unit cell and datasets.'''

    # FIXME need to keep in mind MTZ batch headers - can I access these?
    # yes - through stuff like this:
    #
    #    for batch in mtz_obj.batches():
    #        for token in dir(batch):
    #            print token
    #        print batch.num()
    #
    # but need to decide what I am looking for... will bake this in with
    # a future update (it would be interesting to look at how the UB
    # matrices behave.)

    def __init__(self, hklin):
        mtz_obj = mtz.object(hklin)

        self._miller_indices = mtz_obj.extract_miller_indices()
        self._resolution_range = mtz_obj.max_min_resolution()
        self._space_group = mtz_obj.space_group()

        self._crystal_table = { }

        for crystal in mtz_obj.crystals():
            self._crystal_table[crystal.name()] = mtz_crystal(crystal)

        self._column_table = { }

        for xname in self._crystal_table:
            crystal = self._crystal_table[xname]
            for dname in crystal.get_dataset_names():
                dataset = crystal.get_dataset(dname)
                for column_name in dataset.get_column_names():
                    assert(not column_name in self._column_table)
                    self._column_table[column_name] = dataset.get_column(
                        column_name)

        return

    def get_crystal_names(self):
        return list(self._crystal_table)

    def get_crystal(self, crystal_name):
        return self._crystal_table[crystal_name]

    def get_unit_cell(self):
        '''Get the unit cell object from HKL_base for other calculations.'''

        return self.get_crystal('HKL_base').get_unit_cell()

    def get_space_group(self):
        return self._space_group

    def get_resolution_range(self):
        return self._resolution_range

    def get_symmetry_operations(self):
        return [smx for smx in self._space_group.smx()]

    def get_centring_operations(self):
        return [ltr for ltr in self._space_group.ltr()]

    def get_miller_indices(self):
        return self._miller_indices

    def get_column_names(self):
        return list(self._column_table)

    def get_column(self, column_label):
        return self._column_table[column_label]

    def get_column_values(self, column_label, nan_value = 0.0):
        return self._column_table[column_label].extract_values(
            not_a_number_substitute = nan_value)

def mtz_dump(hklin):
    '''An implementation of mtzdump using the above classes.'''

    mtz = mtz_file(hklin)

    print 'Reading file: %s' % hklin
    print 'Spacegroup: %s' % mtz.get_space_group().type(
         ).universal_hermann_mauguin_symbol()

    print 'Centring operations:'
    for cenop in mtz.get_centring_operations():
        print cenop

    print 'Symmetry operations:'
    for symop in mtz.get_symmetry_operations():
        print symop

    for xname in mtz.get_crystal_names():
        crystal = mtz.get_crystal(xname)
        print 'Crystal: %s' % xname
        print 'Cell: %.3f %.3f %.3f %.3f %.3f %.3f' % \
              crystal.get_unit_cell_parameters()

        for dname in crystal.get_dataset_names():
            dataset = crystal.get_dataset(dname)
            print 'Dataset: %s' % dname
            print 'Columns (with min / max)'
            for column in dataset.get_column_names():
                values = dataset.get_column_values(column)
                print '%20s %.4e %.4e' % (column, min(values), max(values))

    print 'All columns:'
    for column in mtz.get_column_names():
        print column

if __name__ == '__main__':
    import sys

    mtz_dump(sys.argv[1])