#!/usr/bin/env python
# XDSScalerHelpers.py
#   Copyright (C) 2007 CCLRC, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# 5th July 2007
#
# Code to help the scaler along - this will basically be a bunch of jiffy
# functions...
#

import os
import sys

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'])

from Handlers.Streams import Debug
from lib.bits import auto_logfiler
from Wrappers.CCP4.Combat import Combat as _Combat
from Wrappers.CCP4.Pointless import Pointless as _Pointless

class XDSScalerHelper(object):
  '''A class which contains functions which will help the XDS Scaler
  with its work. This is implemented as a class to allow properties
  like working directories and so on to be maintained.'''

  def __init__(self):
    self._working_directory = os.getcwd()

    return

  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 Pointless(self):
    '''Create a Pointless wrapper from _Pointless - set working directory
    and log file stuff as a part of this...'''
    pointless = _Pointless()
    pointless.set_working_directory(self.get_working_directory())
    auto_logfiler(pointless)
    return pointless

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

  def get_working_directory(self):
    return self._working_directory

  def parse_xscale_ascii_header(self, xds_ascii_file):
    '''Parse out the input reflection files which contributed to this
    reflection file.'''

    file_map = { }

    for line in open(xds_ascii_file, 'r').readlines():
      if not line[0] == '!':
        break

      if 'ISET' in line and 'INPUT_FILE' in line:
        set = int(line.split()[2].strip())
        input_file = line.split('=')[2].strip()

        file_map[set] = input_file

        Debug.write('Set %d is from data %s' % (set, input_file))

    return file_map

  def parse_xscale_ascii_wavelength(self, xds_ascii_file):

    wavelength_dict = { }

    for line in open(xds_ascii_file, 'r').readlines():
      if not line[0] == '!':
        break

      if 'ISET' in line and 'X-RAY_WAVELENGTH' in line:
        set = int(line.split()[2].strip())
        wavelength = float(line.split('=')[2].split()[0])

        wavelength_dict[set] = wavelength

        Debug.write('Set %d wavelength %f' % (set, wavelength))

    if len(wavelength_dict.keys()) > 1:
      raise RuntimeError, 'more than one wavelength found'

    return wavelength_dict[wavelength_dict.keys()[0]]

  def split_xscale_ascii_file(self, xds_ascii_file, prefix):
    '''Split the output of XSCALE to separate reflection files for
    each run. The output files will be called ${prefix}${input_file}.'''

    file_map = self.parse_xscale_ascii_header(xds_ascii_file)

    files = { }
    return_map = { }

    keys = file_map.keys()

    for k in keys:
      files[k] = open(os.path.join(
          self.get_working_directory(),
          '%s%s' % (prefix, file_map[k])), 'w')

      return_map[file_map[k]] = '%s%s' % (prefix, file_map[k])

    # copy the header to all of the files

    for line in open(xds_ascii_file, 'r').readlines():
      if not line[0] == '!':
        break

      for k in keys:

        if 'ISET' in line and \
               int(line.split('ISET=')[1].split()[0]) != k:
          continue

        files[k].write(line)

    # next copy the appropriate reflections to each file

    for line in open(xds_ascii_file, 'r').readlines():
      if line[0] == '!':
        continue

      # FIXME this will not be correct if zero-dose correction
      # has been used as this applies an additional record at
      # the end... though it should always be #9
      k = int(line.split()[9])
      files[k].write(line)


    # then write the tailer

    for k in keys:
      files[k].write('!END_OF_DATA\n')
      files[k].close()


    return return_map

  def split_and_convert_xscale_output(self, input_file, prefix,
                                      project_info, scale_factor = 1.0):
    '''Split (as per method above) then convert files to MTZ
    format via combat. The latter step will add the
    pname / xname / dname things from the dictionary supplied.'''

    data_map = self.split_xscale_ascii_file(input_file, prefix)

    for token in data_map.keys():

      if not project_info.has_key(token):
        raise RuntimeError, 'project info for %s not available' % \
              token

      hklin = os.path.join(self.get_working_directory(),
                           data_map[token])
      hklout = os.path.join(self.get_working_directory(),
                            '%s.mtz' % hklin[:-4])

      wavelength = self.parse_xscale_ascii_wavelength(hklin)

      pname, xname, dname = project_info[token]

      # in here need to replace combat with pointless - which
      # will require a little juggling...

      if False:

        c = self.Combat()
        c.set_hklin(hklin)
        c.set_hklout(hklout)
        c.set_project_info(pname, xname, dname)
        if wavelength > 0.0:
          c.set_wavelength(wavelength)

        c.run()

      else:

        p = self.Pointless()
        p.set_xdsin(hklin)
        p.set_hklout(hklout)
        p.set_project_info(pname, xname, dname)
        p.set_scale_factor(scale_factor)
        p.xds_to_mtz()

      data_map[token] = hklout

    return data_map

if __name__ == '__main__':

  xsh = XDSScalerHelper()

  input_file = os.path.join(
      os.environ['X2TD_ROOT'], 'Test', 'UnitTest',
      'Modules', 'XDSScalerHelpers', '1VR9_NAT.HKL')
  project_info = {'NATIVE_NATIVE_HR.HKL':('JCSG', '1VR9', 'NATIVE'),
                  'NATIVE_NATIVE_LR.HKL':('JCSG', '1VR9', 'NATIVE')}
  print xsh.split_and_convert_xscale_output(input_file, 'SCALED_',
                                            project_info)