#!/usr/bin/env python
# Scaler.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.
#
# An interface for programs which do scaling - this will handle all of the
# input and output, delegating the actual implementation to a wrapper which
# implements this interface via inheritance.
#
# This interface is designed to work with "high level" scaling, that
# is the case where all of the expertise about the scaling is delegated
# to the wrapper.
#
# The following cases need to be handled:
#
# (1) multiple sweeps contributing to one wavelength of data (e.g.
#     1VR9 native data.)
# (2) multiple sweeps at different wavelengths for MAD data collection
#     (e.g. 1VR9 SeMet MAD data.)
#
# These cases need to be handled implicitly, which means that the collection
# order will have to be handled. A mechanism for separating out the data
# for different crystals will be needed, to allow the radiation damage
# handling stuff to do it's job.
#
# The overall data model will correspond to the CCP4 MTZ hierarchy, that
# is project/crystal/dataset. In this interface it is assumed that all
# data will correspond to a single project, since anything else is simply
# perverse!
#
# Input data will take the form of handles to Integrater implementations,
# which can provide the required data as and when it is asked for. At some
# point I will need to think about how to handle the issue that XSCALE does
# the best job of scaling data from XDS...
#
# At least, I need to implement a mechanism for handling this. More effort
# is needed in the design of the Factories... Since this applies also for
# things like the Indexer in the xia2process implementation I should
# probably resolve this first.
#
# Update 08/SEP/06
# ----------------
#
# Factory for Integrater takes Indexer as argument. By analogy, a Scaler
# Factory will take one or more Integraters as input. These will then
# allow the best scaler to be selected.
#
# The scaling step should implicitly include scaling & reindexing to the
# "standard" setting. This should raise an exception if the diferent sweeps
# have been integrated with different programs.
#
# In XModel terms, this will be available from:
#
# XSweep - for characterization of a single sweep, looking for in sweep
#          radiation damage, resolution limits & so on.
# XWavelength - for scaling together multiple passes which belong to the
#               same wavelength, e.g. low and high resolution pass.
#               Also for looking for radiation damage.
# XCrystal - for scaling together all of the data measured for a given
#            crystal, e.g. multiwavelength, multi passes as for XWavelength,
#            looking for multi-set radiation damage. This is assumed to
#            provide the final reflection output.
#
# Note well: The XCrystal level scaling will also be responsible for producing
# the specialised data sets, e.g. for phasing & refinement. The former should
# optimise the "signal", while the latter should optimise the "resolution"
# and "quality" (this is to be assessed, for the moment think in terms of
# limiting radiation damage.)
#
# This will be most complicated, and will depend to a critical extent on the
# way in which the scaling is managed.
#
# The scaling process should allow the following to be obtained:
#
# merged reflections for phasing
# merged reflections for refinement
# unmerged reflections for phasing
# unmerged reflections for refinement
# r_merge
# r_pim
# resolution
# "anomalous signal/difference"
# "anomalous dispersion"
# twinning information (implies that Truncate will be included)
#
# probably other things but they will have to wait.
#
# As input, the following will be appropriate:
#
# A managed list of Integrater implementations. These should be grouped
# into wavelengths &c. All must be integrated with a compatible unit cell.
# A resolution limit.
# Anomalous true, false.
# "standard" unit cell (optional)
#
# Output formats:
#
# The output will be available as MTZ, providing that the proper mtz hierarchy
# is available (need to make sure there is a way of providing this as input.)
# Scalepack format files will be provided for unmerged. These will be named
# according to something which matches the MTZ hierarchy, e.g.
# crystal_dataset.sca.
#
# Refinement data will be merged native.
#
# Organisation 21/SEP/06
#
# Ok, this is reasonably complicated, because I don't want to tie this
# directly into the .xinfo hierarchy, so I will need to be able to express
# the input structure of the integraters in some relatively clever way.
#
# In here, think in terms of XDS/XSCALE & Mosflm->Scala. For Mosflm->Scala
# this will need to:
#
# (1) check that the unit cells are compatible.
# (2) sort the reflection files together.
# (3) organise the batches, runs - these need to be put in as
#     information about the project/crystal/dataset which correspond
#     to each run.
# (4) actually perform the scaling.
# (5) after scaling will have multiple reflection files - these need to have
#     the unit cells "standardised" and then be merged again via CAD.
#     Could alternatively just keep them as separate reflection files, just
#     record the fact in the .xinfo output.
#
# From a scala example (documented in Ph.D./Chapter 4)
#
# run 1 batch N to M
# name run 1 project foo crystal bar dataset peak
# base dataset [define remote here] # defines what the dispersive differences
#                                   # are relative to in the analysis
#
# This means that the project, crystal, dataset information needs to come
# in, along with the sweeps (reflection files) and the epochs of data
# collection for the radiation damage analysis. Latter may be NULL, in which
# case process the reflection files in alphabetical order.
#
# Implementation
# --------------
#
# Ok, in terms of the implementation this could be more complicated. This
# is not going to be particularly easy to implement by a single program
# wrapper, so perhaps I will have to actually implement CCP4Scaler,
# XDSScaler &c., which will be a composite class which performs the operation,
# using wrapper classes for the different programs...

import os
import sys
import inspect

if not os.environ.has_key('XIA2_ROOT'):
  raise RuntimeError, 'XIA2_ROOT not defined'

if not os.environ['XIA2_ROOT'] in sys.path:
  sys.path.append(os.path.join(os.environ['XIA2_ROOT']))

from lib.bits import inherits_from
from lib.SymmetryLib import lauegroup_to_lattice, sort_lattices

from Handlers.Streams import Chatter, Debug

# file conversion (and merging) jiffies

from Modules.Scalepack2Mtz import Scalepack2Mtz
from Modules.Mtz2Scalepack import Mtz2Scalepack

class Scaler(object):
  '''An interface to present scaling functionality in a similar way to the
  integrater interface.'''

  def __init__(self):
    # set up a framework for storing all of the input information...
    # this should really only consist of integraters...

    # key this by the epoch, if available, else will need to
    # do something different.
    self._scalr_integraters = { }

    # the corrections to apply - see trac #162
    self._scalr_corrections = False
    self._scalr_correct_decay = None
    self._scalr_correct_modulation = None
    self._scalr_correct_absorption = None
    self._scalr_correct_partiality = None

    # integraters have the following methods for pulling interesting
    # information out:
    #
    # get_integrater_project_info() - pname, xname, dname
    # get_integrater_epoch() - measurement of first frame

    self.scaler_reset()

    self._scalr_reference_reflection_file = None
    self._scalr_freer_file = None

    # user input to guide spacegroup choices
    self._scalr_input_spacegroup = None
    self._scalr_input_pointgroup = None

    # places to hold the output

    # this should be a dictionary keyed by datset / format, or
    # format / dataset
    self._scalr_scaled_reflection_files = None

    # this needs to be a dictionary keyed by dataset etc, e.g.
    # key = pname, xname, dname
    self._scalr_statistics = None

    # and also the following keys:
    self._scalr_statistics_keys = [
        'High resolution limit', 'Low resolution limit',
        'Completeness', 'Multiplicity',
        'I/sigma', 'Rmerge',
        'Rmeas(I)', 'Rmeas(I+/-)',
        'Rpim(I)', 'Rpim(I+/-)',
        'Wilson B factor', 'Partial bias',
        'Anomalous completeness', 'Anomalous multiplicity',
        'Anomalous correlation', 'Anomalous slope',
        'dF/F', 'dI/s(dI)',
        'Total observations', 'Total unique']

    # information for returning "interesting facts" about the data
    self._scalr_highest_resolution = 0.0
    self._scalr_cell = None
    self._scalr_likely_spacegroups = []
    self._scalr_unlikely_spacegroups = []

    # do we want anomalous pairs separating?
    self._scalr_anomalous = False

    # admin junk
    self._working_directory = os.getcwd()
    self._scalr_pname = None
    self._scalr_xname = None

    # link to parent xcrystal
    self._scalr_xcrystal = None

    return

  # FIXME x1698 these not currently used yet

  def _scale_list_likely_pointgroups(self, integrater):
    raise RuntimeError, 'overload me'

  def _scale_reindex_to_reference(self, reference, integrater):
    raise RuntimeError, 'overload me'

  # FIXME to here

  def _scale_prepare(self):
    raise RuntimeError, 'overload me'

  def _scale(self):
    raise RuntimeError, 'overload me'

  def _scale_finish(self):
    pass

  def set_working_directory(self, working_directory):
    self._working_directory = working_directory
    return

  def get_working_directory(self):
    return self._working_directory

  def set_scaler_input_spacegroup(self, spacegroup):
    self._scalr_input_spacegroup = spacegroup
    return

  def set_scaler_input_pointgroup(self, pointgroup):
    self._scalr_input_pointgroup = pointgroup
    return

  def get_scaler_input_spacegroup(self):
    return self._scaler_input_spacegroup

  def get_scaler_input_pointgroup(self):
    return self._scaler_input_pointgroup

  def set_scaler_xcrystal(self, xcrystal):
    self._scalr_xcrystal = xcrystal
    return

  def get_scaler_xcrystal(self):
    return self._scalr_xcrystal

  def set_scaler_project_info(self, pname, xname):
    '''Set the project and crystal this scaler is working with.'''

    self._scalr_pname = pname
    self._scalr_xname = xname

    return

  def get_scaler_project_info(self):
    '''Get the scaler project and crystal.'''

    return self._scalr_pname, self._scalr_xname

  def set_scaler_reference_reflection_file(self, reference_reflection_file):
    self._scalr_reference_reflection_file = reference_reflection_file
    return

  def get_scaler_reference_reflection_file(self):
    return self._scalr_reference_reflection_file

  def set_scaler_freer_file(self, freer_file):
    self._scalr_freer_file = freer_file
    return

  def get_scaler_freer_file(self):
    return self._scalr_freer_file

  def set_scaler_prepare_done(self, done = True):

    frm = inspect.stack()[1]
    mod = inspect.getmodule(frm[0])
    Debug.write('Called scaler prepare done from %s %d (%s)' %
                (mod.__name__, frm[0].f_lineno, done))

    self._scalr_prepare_done = done
    return

  def set_scaler_done(self, done = True):

    frm = inspect.stack()[1]
    mod = inspect.getmodule(frm[0])
    Debug.write('Called scaler done from %s %d (%s)' %
                (mod.__name__, frm[0].f_lineno, done))

    self._scalr_done = done
    return

  def set_scaler_finish_done(self, done = True):

    frm = inspect.stack()[1]
    mod = inspect.getmodule(frm[0])
    Debug.write('Called scaler finish done from %s %d (%s)' %
                (mod.__name__, frm[0].f_lineno, done))

    self._scalr_finish_done = done
    return

  def set_scaler_anomalous(self, anomalous):
    self._scalr_anomalous = anomalous
    return

  def get_scaler_anomalous(self):
    return self._scalr_anomalous

  def scaler_reset(self):

    Debug.write('Scaler reset')

    self._scalr_done = False
    self._scalr_prepare_done = False
    self._scalr_finish_done = False
    self._scalr_result = None
    return

  # getters for the scaling model which was used - first see that the
  # corrections were applied, then the individual getters for the
  # separate corrections

  def get_scaler_corrections(self):
    return self._scalr_corrections

  def get_scaler_correct_decay(self):
    return self._scalr_correct_decay

  def get_scaler_correct_modulation(self):
    return self._scalr_correct_modulation

  def get_scaler_correct_absorption(self):
    return self._scalr_correct_absorption

  def get_scaler_correct_partiality(self):
    return self._scalr_correct_partiality

  # getters of the status - note how the gets cascade to ensure that
  # everything is up-to-date...

  def get_scaler_prepare_done(self):
    return self._scalr_prepare_done

  def get_scaler_done(self):
    if not self.get_scaler_prepare_done():
      Debug.write('Resetting Scaler done as prepare not done')
      self.set_scaler_done(False)
    return self._scalr_done

  def get_scaler_finish_done(self):
    if not self.get_scaler_done():
      Debug.write(
          'Resetting scaler finish done as scaling not done')
      self.set_scaler_finish_done(False)
    return self._scalr_finish_done

  def add_scaler_integrater(self, integrater):
    '''Add an integrater to this scaler, to provide the input.'''

    epoch = integrater.get_integrater_epoch()

    if epoch == 0 and self._scalr_integraters:
      raise RuntimeError, 'multi-sweep integrater has epoch 0'

    if epoch in self._scalr_integraters.keys():
      raise RuntimeError, 'integrater with epoch %d already exists' % \
            epoch

    self._scalr_integraters[epoch] = integrater

    self.scaler_reset()

    return

  # FIXME x1698 these not currently used yet

  def _scale_setup_integrater(self, integrater):
    '''Check that the pointgroup for a data set is consistent with
    the lattice used for integration, then determine the pointgroup for
    the data.'''

    # FIXME will have to handle gracefully user provided pointgroup

    pointgroups = self._scale_list_likely_pointgroups(integrater)
    indexer = integrater.get_integrater_indexer()
    lattices = [lauegroup_to_lattice(p) for p in pointgroups]

    correct_lattice = None

    for lattice in lattices:
      state = indexer.set_indexer_asserted_lattice(lattice)

      if state == indexer.LATTICE_CORRECT:
        correct_lattice = lattice
        break

      elif state == indexer.LATTICE_IMPOSSIBLE:
        continue

      elif state == indexer.LATTICE_POSSIBLE:
        currect_lattice = lattice
        break

    assert(correct_lattice)

    # run this analysis again, which may respond in different conclusions
    # if it triggers the reprocessing of the data with a new lattice

    pointgroups = self._scale_list_likely_pointgroups(integrater)
    lattices = [lauegroup_to_lattice(p) for p in pointgroups]

    return pointgroups[lattices.index(correct_lattice)]

  def _scale_setup(self):
    '''Set things up for scaling, in particular mediate pointgroup /
    lattice with the indexers.'''

    assert(self._scalr_integraters)

    epochs = sorted(self._scalr_integraters)
    integraters = [self._scalr_integraters[e] for e in epochs]

    pointgroups = [self._scale_setup_integrater(i) for i in integraters]
    lattices = [lauegroup_to_lattice(p) for p in pointgroups]

    unique_lattices = list(set(lattices))

    # consider the situation that they arrived at more than one conclusion

    if len(unique_lattices) > 1:
      consensus_lattice = sort_lattices(unique_lattices)[0]

      for integrater in integraters:
        indexer = integrater.get_integrater_indexer()
        state = indexer.set_indexer_asserted_lattice(consensus_lattice)

        assert(state != indexer.LATTICE_IMPOSSIBLE)

    # then decide on the consensus pointgroup

    pointgroups = set([])

    for integrater in integraters:
      pointgroups = self._scale_list_likely_pointgroups(integrater)
      lattices = [lauegroup_to_lattice(p) for p in pointgroups]
      poingroups.add(pointgroups[lattices.index(consensus_lattice)])

    # FIXME will need to handle twinned cases more gracefully sometime
    # FIXME also need to "mend" the integrater set spacegroup API

    assert(len(pointgroups) == 1)

    for integrater in integraters:
      integrater.set_integrater_spacegroup_number(pointgroup)

    # now reindex to the correct setting

    reference = integraters[0]

    for integrater in integraters[1:]:
      self._scale_reindex_to_reference(reference, integrater)

    return pointgroups[0]

  # FIXME to here

  def scale(self):
    '''Actually perform the scaling - this is delegated to the
    implementation.'''

    if self._scalr_integraters == { }:
      raise RuntimeError, \
            'no Integrater implementations assigned for scaling'

    xname = self._scalr_xcrystal.get_name()

    while not self.get_scaler_finish_done():
      while not self.get_scaler_done():
        while not self.get_scaler_prepare_done():

          Chatter.banner('Preparing %s' % xname)

          self._scalr_prepare_done = True
          self._scale_prepare()

        Chatter.banner('Scaling %s' % xname)

        self._scalr_done = True
        self._scalr_result = self._scale()

      self._scalr_finish_done = True
      self._scale_finish()

    return self._scalr_result

  def get_scaled_reflections(self, format):
    '''Get a specific format of scaled reflection files. This may
    trigger transmogrification of files.'''

    if not format in ['mtz', 'sca', 'sca_unmerged']:
      raise RuntimeError, 'format %s unknown' % format

    self.scale()

    if format in self._scalr_scaled_reflection_files.keys():
      return self._scalr_scaled_reflection_files[format]

    # specific code to handle generation cases

    if format == 'sca_unmerged':
      raise RuntimeError, 'cannot generate unmerged reflections'

    if format == 'mtz':
      # generate and return an mtz file - ideally this will be from
      # merged reflections, though unmerged will do - though this
      # will need the unit cell information etc. Aha! this is
      # available from this interface...

      cell = self.get_scaler_cell()
      spacegroup = self.get_scaler_likely_spacegroups()[0]

      wavelengths = self._scalr_scaled_reflection_files['sca'].keys()
      project, crystal = self.get_scaler_project_info()

      s2m = Scalepack2Mtz()

      for w in wavelengths:
        s2m.add_hklin(w, self._scalr_scaled_reflection_files[
            'sca'][w])
      s2m.set_cell(cell)
      s2m.set_spacegroup(spacegroup)
      s2m.set_project_info(project, crystal)

      self._scalr_scaled_reflection_files['mtz'] = s2m.convert()
      return self._scalr_scaled_reflection_files['mtz']

    if format == 'sca':
      # generate merged scalepack format reflections from the mtz
      # format

      m2s = Mtz2Scalepack()

      m2s.set_hklin(self._scalr_scaled_reflection_files['mtz'])
      self._scalr_scaled_reflection_files['sca'] = m2s.convert()
      return self._scalr_scaled_reflection_files['sca']

    raise RuntimeError, 'cannot possibly reach this point'


  def get_scaled_merged_reflections(self):
    '''Return the reflection files and so on.'''

    self.scale()
    return self._scalr_scaled_reflection_files

  def get_scaler_statistics(self):
    '''Return the overall scaling statistics.'''

    self.scale()

    return self._scalr_statistics

  def get_scaler_cell(self):
    '''Return the final unit cell from scaling.'''

    self.scale()
    return self._scalr_cell

  def get_scaler_likely_spacegroups(self):
    '''Return a list of likely spacegroups - you should try using
    the first in this list first.'''

    self.scale()
    return self._scalr_likely_spacegroups

  def get_scaler_unlikely_spacegroups(self):
    '''Return a list of unlikely spacegroups - you should try using
    the likely ones first. These are spacegroups in the correct
    pointgroup but with systematic absences which dont match up.'''

    self.scale()
    return self._scalr_unlikely_spacegroups

  def get_scaler_highest_resolution(self):
    '''Get the highest resolution achieved by the crystal.'''

    self.scale()
    return self._scalr_highest_resolution
