# Scalepack2Mtz.py
# Maintained by G.Winter
# 12th February 2007
#
# A module to carefully convert scalepack reflection files (merged or
# unmerged) into properly structured MTZ files. This is based on the last
# section of CCP4ScalerImplementation, as that does largely the same.
#
# This will:
#
# if (unmerged) combat -> scala for merging (else) convert to mtz
# truncate all wavelengths
# if (mad) compute cell, cad in cell, cad together
# add freeR column
# if (mad) look for inter radiation damage
#
# This assumes:
#
# Reflections are correctly indexed
# One reflection file per wavelength
# All files from same crystal and therefore project
#
# Should produce results similar to CCP4ScalerImplementation.
#
# Example data will come from xia2.
#
# Implementation fill follow line 1469++ of CCP4ScalerImplementation.
#
# Uses:
#
# scalepack2mtz wrapper for merged files
# combat, sortmtz, scala for unmerged files
# mtzdump, cad, freerflag to muddle with files
# truncate to compute F's from I's.
#

import os
import sys
import math

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.environ['XIA2_ROOT'])

# the wrappers that this will use - these are renamed so that the internal
# factory version can be used...
from Wrappers.CCP4.Scala import Scala as _Scala
from Wrappers.CCP4.Sortmtz import Sortmtz as _Sortmtz
from Wrappers.CCP4.Truncate import Truncate as _Truncate
from Wrappers.CCP4.Cad import Cad as _Cad
from Wrappers.CCP4.Freerflag import Freerflag as _Freerflag
from Wrappers.CCP4.Combat import Combat as _Combat
from Wrappers.CCP4.Scalepack2mtz import Scalepack2mtz as _Scalepack2mtz

from Handlers.Streams import Chatter
from Handlers.Files import FileHandler

from CCP4InterRadiationDamageDetector import CCP4InterRadiationDamageDetector

# jiffys

from lib.bits import auto_logfiler

class Scalepack2Mtz(object):
  '''A jiffy class to enable transformation between scalepack format
  (merged or unmerged) to a properly structured MTZ file.'''

  def __init__(self):
    # data that this will need to store
    # working directory stuff
    # pname/xname/dname/data file stuff

    self._working_directory = os.getcwd()

    # this will be keyed by dname

    self._hklin_files = { }
    self._pname = None
    self._xname = None
    self._dnames = []

    # cell and spacegroup information

    self._cell = None
    self._spacegroup = None

    return

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

  def get_working_directory(self):
    return self._working_directory

  # getter and setter methods

  def set_cell(self, cell):
    self._cell = cell
    return

  def set_spacegroup(self, spacegroup):
    self._spacegroup = spacegroup
    return

  def set_project_info(self, pname, xname):
    self._pname = pname
    self._xname = xname
    return

  def add_hklin(self, dname, hklin):
    if self._hklin_files.has_key(dname):
      raise RuntimeError, 'dataset name %s already exists' % dname

    self._hklin_files[dname] = hklin
    self._dnames.append(dname)

    return

  # factory methods...

  def Scala(self):
    '''Create a Scala wrapper from _Scala - set the working directory
    and log file stuff as a part of this...'''
    scala = _Scala()
    scala.set_working_directory(self.get_working_directory())
    auto_logfiler(scala)
    return scala

  def Scalepack2mtz(self):
    '''Create a Scalepack2mtz wrapper from _Scalepack2mtz - set the
    working directory and log file stuff as a part of this...'''
    scalepack2mtz = _Scalepack2mtz()
    scalepack2mtz.set_working_directory(self.get_working_directory())
    auto_logfiler(scalepack2mtz)
    return scalepack2mtz

  def Combat(self):
    '''Create a Combat wrapper from _Combat - set the working directory
    and log file stuff as a part of this...'''
    combat = _Combat()
    combat.set_working_directory(self.get_working_directory())
    auto_logfiler(combat)
    return combat

  def Sortmtz(self):
    '''Create a Sortmtz wrapper from _Sortmtz - set the working directory
    and log file stuff as a part of this...'''
    sortmtz = _Sortmtz()
    sortmtz.set_working_directory(self.get_working_directory())
    auto_logfiler(sortmtz)
    return sortmtz

  def Truncate(self):
    '''Create a Truncate wrapper from _Truncate - set the working directory
    and log file stuff as a part of this...'''
    truncate = _Truncate()
    truncate.set_working_directory(self.get_working_directory())
    auto_logfiler(truncate)
    return truncate

  def Mtz2various(self):
    '''Create a Mtz2various wrapper from _Mtz2various - set the working
    directory and log file stuff as a part of this...'''
    mtz2various = _Mtz2various()
    mtz2various.set_working_directory(self.get_working_directory())
    auto_logfiler(mtz2various)
    return mtz2various

  def Cad(self):
    '''Create a Cad wrapper from _Cad - set the working directory
    and log file stuff as a part of this...'''
    cad = _Cad()
    cad.set_working_directory(self.get_working_directory())
    auto_logfiler(cad)
    return cad

  def Freerflag(self):
    '''Create a Freerflag wrapper from _Freerflag - set the working
    directory and log file stuff as a part of this...'''
    freerflag = _Freerflag()
    freerflag.set_working_directory(self.get_working_directory())
    auto_logfiler(freerflag)
    return freerflag

  # helper functions

  def _merged_scalepack_to_mtz(self, dname):
    '''Convert a merged reflection file for dname from scalepack format
    to PNAME_XNAME_merged_tmp_DNAME.mtz.'''

    scalepack_file = self._hklin_files[dname]

    s2m = self.Scalepack2mtz()
    s2m.set_hklin(scalepack_file)
    s2m.set_hklout(os.path.join(self.get_working_directory(),
                                '%s_%s_merged_tmp_%s.mtz' % \
                                (self._pname, self._xname, dname)))
    s2m.set_spacegroup(self._spacegroup)
    s2m.set_cell(self._cell)
    s2m.set_project_info(self._pname, self._xname, dname)
    s2m.convert()

    return

  def _unmerged_scalepack_to_mtz(self, dname):
    '''Convert an unmerged reflection file for dname from scalepack format
    to PNAME_XNAME_merged_tmp_DNAME.mtz.'''

    scalepack_file = self._hklin_files[dname]

    # convert to multirecord MTZ

    hklout = os.path.join(self.get_working_directory(),
                          '%s_%s_unmerged_tmp_%s.mtz' % \
                          (self._pname, self._xname, dname))

    FileHandler.record_temporary_file(hklout)

    c = self.Combat()
    c.set_hklin(scalepack_file)
    c.set_hklout(hklout)
    c.set_spacegroup(self._spacegroup)
    c.set_cell(self._cell)
    c.set_project_info(self._pname, self._xname, dname)
    c.run()

    # sort

    hklin = hklout
    hklout = os.path.join(self.get_working_directory(),
                          '%s_%s_sorted_tmp_%s.mtz' % \
                          (self._pname, self._xname, dname))

    FileHandler.record_temporary_file(hklout)

    s = self.Sortmtz()
    s.set_hklin(hklin)
    s.set_hklout(hklout)
    s.sort()

    # merge - FIXME should this presume that the input is anomalous
    # data??

    hklin = hklout
    hklout = os.path.join(self.get_working_directory(),
                          '%s_%s_merged_tmp_%s.mtz' % \
                          (self._pname, self._xname, dname))

    FileHandler.record_temporary_file(hklout)

    # FIXME in here I need to know if we are wanting anomalous
    # pairs separated or merged, though it probably doesn't
    # matter too much...

    sc = self.Scala()
    sc.set_hklin(hklin)
    sc.set_hklout(hklout)
    sc.set_project_info(self._pname, self._xname, dname)
    sc.set_anomalous()
    sc.set_onlymerge()
    sc.merge()

    FileHandler.record_log_file('%s %s %s merge' % \
                                (self._pname, self._xname, dname),
                                sc.get_log_file())

    return

  def _decide_is_merged(self, file):
    '''Decide if the input file is merged, based on first 3 lines.'''

    f = open(file, 'r')
    start = [f.readline() for j in range(3)]
    f.close()

    # merged scalepack looks like number, number, cell

    if len(start[0].split()) == 1 and \
       len(start[1].split()) == 1 and \
       len(start[2].split()) > 6:
      return True

    # unmerged scalepack from scala looks like number spacegroup,
    # symop, symop, symop...

    if len(start[0].split()) > 2 and \
       len(start[1].split()) > 2 and \
       len(start[2].split()) > 2:
      return False

    # get to here is a problem

    raise RuntimeError, 'cannot decide file format'

  # merge method

  def scalepack_to_mtz(self, scalepack, hklout,
                       anomalous, spacegroup, cell, project_info = None):
    '''Convert a scalepack file to MTZ, maybe merging.'''

    # inspect to see if it is merged
    merged = self._decide_is_merged(scalepack)

    if merged:
      s2m = self.Scalepack2mtz()
      s2m.set_hklin(scalepack)
      s2m.set_hklout(hklout)
      s2m.set_spacegroup(spacegroup)
      s2m.set_cell(cell)
      if project_info:
        pname, xname, dname = project_info
        s2m.set_project_info(pname, xname, dname)
      s2m.convert()

      return

    else:

      hklout_c = os.path.join(
          self.get_working_directory(), 'combat-tmp.mtz')

      FileHandler.record_temporary_file(hklout_c)

      c = self.Combat()
      c.set_hklin(scalepack)
      c.set_hklout(hklout_c)
      c.set_spacegroup(spacegroup)
      c.set_cell(cell)
      if project_info:
        pname, xname, dname = project_info
        c.set_project_info(pname, xname, dname)
      c.run()

      hklin = hklout_c
      hklout_s = os.path.join(
          self.get_working_directory(), 'sortmtz-tmp.mtz')

      FileHandler.record_temporary_file(hklout_s)

      s = self.Sortmtz()
      s.set_hklin(hklin)
      s.set_hklout(hklout_s)
      s.sort()

      hklin = hklout_s

      sc = self.Scala()
      sc.set_hklin(hklin)
      sc.set_hklout(hklout)
      if project_info:
        pname, xname, dname = project_info
        sc.set_project_info(pname, xname, dname)
      sc.set_anomalous(anomalous)
      sc.set_onlymerge()
      sc.merge()

      result = {'symmary':sc.get_summary(),
                'loggraphs':sc.parse_ccp4_loggraph()}

      return result


  def convert(self):
    '''Transmogrify the input scalepack files to mtz.'''

    merged_files = []

    for dname in self._dnames:
      scalepack = self._hklin_files[dname]

      # inspect to see if it is merged
      merged = self._decide_is_merged(scalepack)

      # if (unmerged) merge with combat, sortmts, scala - this is
      # implemented in _unmerged_scalepack_to_mtz - takes dname

      hklout = os.path.join(self.get_working_directory(),
                            '%s_%s_merged_tmp_%s.mtz' % \
                            (self._pname, self._xname, dname))

      if merged:
        self._merged_scalepack_to_mtz(dname)
      else:
        self._unmerged_scalepack_to_mtz(dname)

      # truncate the reflections

      hklin = hklout

      hklout = os.path.join(self.get_working_directory(),
                            '%s_%s_truncated_tmp_%s.mtz' % \
                            (self._pname, self._xname, dname))

      FileHandler.record_temporary_file(hklout)

      t = self.Truncate()
      t.set_hklin(hklin)
      t.set_hklout(hklout)
      t.truncate()

      FileHandler.record_log_file('%s %s %s truncate' % \
                                  (self._pname,
                                   self._xname,
                                   dname),
                                  t.get_log_file())

      hklin = hklout
      hklout = os.path.join(self.get_working_directory(),
                            '%s_%s_truncated_cad_tmp_%s.mtz' % \
                            (self._pname, self._xname, dname))

      FileHandler.record_temporary_file(hklout)

      # need to assign the column labels here - this will probably
      # have to be a quick CAD run. This prevents duplicate column
      # names later on

      c = self.Cad()
      c.add_hklin(hklin)
      c.set_hklout(hklout)
      c.set_new_suffix(dname)
      c.update()

    # cad together

    c = self.Cad()
    for dname in self._dnames:
      hklin = os.path.join(self.get_working_directory(),
                           '%s_%s_truncated_cad_tmp_%s.mtz' % \
                           (self._pname, self._xname, dname))

      c.add_hklin(hklin)

    hklout = os.path.join(self.get_working_directory(),
                          '%s_%s_merged.mtz' % (self._pname,
                                                self._xname))

    Chatter.write('Merging all data sets to %s' % hklout)

    c.set_hklout(hklout)
    c.merge()

    # add FreeR column

    f = self.Freerflag()

    hklin = hklout
    hklout = os.path.join(self.get_working_directory(),
                          '%s_%s_merged_free.mtz' % (self._pname,
                                                     self._xname))

    f.set_hklin(hklin)
    f.set_hklout(hklout)
    f.add_free_flag()

    # if (mad) look for inter radiation damage

    return hklout

# add a unit test

if __name__ == '__main__':
  # UNMERGED

  data_directory = os.path.join(os.environ['X2TD_ROOT'],
                                'Test', 'UnitTest', 'Interfaces',
                                'Scaler', 'Unmerged')

  s2m = Scalepack2Mtz()

  s2m.add_hklin('INFL', os.path.join(data_directory,
                                     'TS00_13185_unmerged_INFL.sca'))
  s2m.add_hklin('LREM', os.path.join(data_directory,
                                     'TS00_13185_unmerged_LREM.sca'))
  s2m.add_hklin('PEAK', os.path.join(data_directory,
                                     'TS00_13185_unmerged_PEAK.sca'))

  s2m.set_cell((57.73, 76.93, 86.57, 90.00, 90.00, 90.00))
  s2m.set_spacegroup('P212121')
  s2m.set_project_info('TS00', '13185')

  print s2m.convert()

  # MERGED

  data_directory = os.path.join(os.environ['X2TD_ROOT'],
                                'Test', 'UnitTest', 'Interfaces',
                                'Scaler', 'Merged')

  s2m = Scalepack2Mtz()

  s2m.add_hklin('INFL', os.path.join(data_directory,
                                     'TS00_13185_scaled_INFL.sca'))
  s2m.add_hklin('LREM', os.path.join(data_directory,
                                     'TS00_13185_scaled_LREM.sca'))
  s2m.add_hklin('PEAK', os.path.join(data_directory,
                                     'TS00_13185_scaled_PEAK.sca'))

  s2m.set_cell((57.73, 76.93, 86.57, 90.00, 90.00, 90.00))
  s2m.set_spacegroup('P212121')
  s2m.set_project_info('TS00A', '13185')

  print s2m.convert()