#!/usr/bin/env python
# Ctruncate.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.
#
# 26th October 2006
#
# A wrapper for the CCP4 program Ctruncate, which calculates F's from
# I's and gives a few useful statistics about the data set.
#
# FIXME 26/OCT/06 this needs to be able to take into account the solvent
#                 content of the crystal (at the moment it will be assumed
#                 to be 50%.)
#
# FIXME 16/NOV/06 need to be able to get the estimates B factor from the
#                 Wilson plot and also second moment stuff, perhaps?
#
# FIXME 02/FEB/11 read the number of reflections excluded from the data set
#                 as systematic absences and verify that this is equal to
#                 the number of reflections going in - number coming out.

import os
import sys

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

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

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

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

from Driver.DriverFactory import DriverFactory
from Decorators.DecoratorFactory import DecoratorFactory
from lib.bits import transpose_loggraph
from Handlers.Streams import Chatter, Debug

def Ctruncate(DriverType = None):
  '''A factory for CtruncateWrapper classes.'''

  DriverInstance = DriverFactory.Driver(DriverType)

  class CtruncateWrapper(DriverInstance.__class__):
    '''A wrapper for Ctruncate, using the regular Driver.'''

    def __init__(self):
      # generic things
      DriverInstance.__class__.__init__(self)

      self.set_executable(os.path.join(
          os.environ.get('CBIN', ''), 'ctruncate'))

      self._anomalous = False
      self._nres = 0

      self._b_factor = 0.0
      self._moments = None

      # numbers of reflections in and out, and number of absences
      # counted

      self._nref_in = 0
      self._nref_out = 0
      self._nabsent = 0

      return

    def set_hklin(self, hklin):
      self._hklin = hklin
      return

    def set_hklout(self, hklout):
      self._hklout = hklout
      return

    def set_nres(self, nres):
      self._nres = nres
      return

    def set_anomalous(self, anomalous):
      self._anomalous = anomalous
      return

    def truncate(self):
      '''Actually perform the truncation procedure.'''

      if not self._hklin:
        raise RuntimeError, 'hklin not defined'
      if not self._hklout:
        raise RuntimeError, 'hklout not defined'

      self.add_command_line('-hklin')
      self.add_command_line(self._hklin)
      self.add_command_line('-hklout')
      self.add_command_line(self._hklout)

      if self._nres:
        self.add_command_line('-nres')
        self.add_command_line('%d' % self._nres)

      if self._anomalous:
        self.add_command_line('-colano')
        self.add_command_line('/*/*/[I(+),SIGI(+),I(-),SIGI(-)]')

      self.add_command_line('-colin')
      self.add_command_line('/*/*/[IMEAN,SIGIMEAN]')

      self.start()
      self.close_wait()

      try:
        self.check_for_errors()

      except RuntimeError, e:
        try:
          os.remove(self._hklout)
        except:
          pass

        raise RuntimeError, 'ctruncate failure'

      nref = 0

      for record in self.get_all_output():
        if 'Number of reflections:' in record:
          nref = int(record.split()[-1])

        if 'B =' in record and 'intercept' in record:
          self._b_factor = float(record.split()[2])

      self._nref_in, self._nref_out = nref, nref
      self._nabsent = 0

      moments = None

      results = self.parse_ccp4_loggraph()

      if 'Acentric moments of E using Truncate method' in results:
        moments = transpose_loggraph(
            results['Acentric moments of E using Truncate method'])
      elif 'Acentric moments of I' in results:
        moments = transpose_loggraph(
            results['Acentric moments of I'])
      elif 'Acentric moments of E' in results:
        moments = transpose_loggraph(
            results['Acentric moments of E'])
      else:
        Debug.write('Acentric moments of E/I not found')

      self._moments = moments

      return

    def get_b_factor(self):
      return self._b_factor

    def get_moments(self):
      return self._moments

    def get_nref_in(self):
      return self._nref_in

    def get_nref_out(self):
      return self._nref_out

    def get_nabsent(self):
      return self._nabsent

    def parse_ccp4_loggraph(self):
      '''Look through the standard output of the program for
      CCP4 loggraph text. When this is found store it in a
      local dictionary to allow exploration.'''

      # reset the loggraph store
      self._loggraph = { }

      output = self.get_all_output()

      for i in range(len(output)):
        line = output[i]
        if '$TABLE' in line:

          n_dollar = line.count('$$')

          current = line.split(':')[1].replace('>',
                                                        '').strip()
          self._loggraph[current] = { }
          self._loggraph[current]['columns'] = []
          self._loggraph[current]['data'] = []

          loggraph_info = ''

          while n_dollar < 4:
            n_dollar += line.count('$$')
            loggraph_info += line

            if n_dollar == 4:
              break

            i += 1
            line = output[i]

          tokens = loggraph_info.split('$$')
          self._loggraph[current]['columns'] = tokens[1].split()

          if len(tokens) < 4:
            raise RuntimeError, 'loggraph "%s" broken' % current

          data = tokens[3].split('\n')

          columns = len(self._loggraph[current]['columns'])

          for j in range(len(data)):
            record = data[j].split()
            if len(record) == columns:
              self._loggraph[current]['data'].append(record)

      return self._loggraph

  return CtruncateWrapper()

if __name__ == '__main__':

  ctruncate = Ctruncate()
  ctruncate.set_hklin(sys.argv[1])
  ctruncate.set_hklout(sys.argv[2])
  ctruncate.truncate()

  print ctruncate.get_nref_in(), ctruncate.get_nref_out(), \
        ctruncate.get_nabsent()