#!/usr/bin/env python
# Truncate.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 Truncate, 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 Truncate(DriverType = None):
    '''A factory for TruncateWrapper classes.'''

    DriverInstance = DriverFactory.Driver(DriverType)
    CCP4DriverInstance = DecoratorFactory.Decorate(DriverInstance, 'ccp4')

    class TruncateWrapper(CCP4DriverInstance.__class__):
        '''A wrapper for Truncate, using the CCP4-ified Driver.'''

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

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

            self._anomalous = False
            self._nres = 0

            # should we do wilson scaling?
            self._wilson = True

            self._b_factor = 0.0
            self._moments = None

            self._wilson_fit_grad = 0.0
            self._wilson_fit_grad_sd = 0.0
            self._wilson_fit_m = 0.0
            self._wilson_fit_m_sd = 0.0
            self._wilson_fit_range = 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_anomalous(self, anomalous):
            self._anomalous = anomalous
            return

        def set_wilson(self, wilson):
            '''Set the use of Wilson scaling - if you set this to False
            Wilson scaling will be switched off...'''
            self._wilson = wilson

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

            self.check_hklin()
            self.check_hklout()

            self.start()

            # write the harvest files in the local directory, not
            # in $HARVESTHOME. Though we have set this for the project
            # so we should be fine to just plough ahead...
            # self.input('usecwd')

            if self._anomalous:
                self.input('anomalous yes')
            else:
                self.input('anomalous no')

            if self._nres:
                self.input('nres %d' % self._nres)

            if not self._wilson:
                self.input('scale 1')

            self.close_wait()

            try:
                self.check_for_errors()
                self.check_ccp4_errors()

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

                raise RuntimeError, 'truncate failure'

            # parse the output for interesting things, including the
            # numbers of reflections in and out (isn't that a standard CCP4
            # report?) and the number of absent reflections.

            self._nref_in, self._nref_out = self.read_nref_hklin_hklout(
                self.get_all_output())

            # FIXME guess I should be reading this properly...
            self._nabsent = self._nref_in - self._nref_out

            for line in self.get_all_output():
                if 'Least squares straight line gives' in line:
                    list = line.replace('=', ' ').split()
                    if not '***' in list[6]:
                        self._b_factor = float(list[6])
                    else:
                        Debug.write('no B factor available')

                if 'LSQ Line Gradient' in line:
                    self._wilson_fit_grad = float(line.split()[-1])
                    resol_width = max(self._wilson_fit_range) - \
                                  min(self._wilson_fit_range)
                    if self._wilson_fit_grad > 0 and resol_width > 1.0 \
                           and False:
                        raise RuntimeError, \
                              'wilson plot gradient positive: %.2f' % \
                              self._wilson_fit_grad
                    elif self._wilson_fit_grad > 0:
                        Debug.write(
                            'Positive gradient but not much wilson plot')


                if 'Uncertainty in Gradient' in line:
                    self._wilson_fit_grad_sd = float(line.split()[-1])
                if 'X Intercept' in line:
                    self._wilson_fit_m = float(line.split()[-1])
                if 'Uncertainty in Intercept' in line:
                    self._wilson_fit_m_sd = float(line.split()[-1])
                if 'Resolution range' in line:
                    self._wilson_fit_range = map(float, line.split()[-2:])


            results = self.parse_ccp4_loggraph()
            moments = transpose_loggraph(
                results['Acentric Moments of E for k = 1,3,4,6,8'])

            # keys we want in this are "Resln_Range" "1/resol^2" and
            # MomentZ2. The last of these should be around two, but is
            # likely to be a little different to this.
            self._moments = moments

            return

        def get_b_factor(self):
            return self._b_factor

        def get_wilson_fit(self):
            return self._wilson_fit_grad, self._wilson_fit_grad_sd, \
                   self._wilson_fit_m, self._wilson_fit_m_sd

        def get_wilson_fit_range(self):
            return self._wilson_fit_range

        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 read_nref_hklin_hklout(self, records):
            '''Look to see how many reflections came in through HKLIN, and
            how many went out again in HKLOUT.'''

            nref_in = 0
            nref_out = 0

            current_logical = None

            for record in records:
                if 'Logical Name' in record:
                    current_logical = record.split()[2]
                    assert(current_logical in ['HKLIN', 'HKLOUT', 'SYMINFO'])

                if 'Number of Reflections' in record:
                    if current_logical == 'HKLIN':
                        nref_in = int(record.split()[-1])
                    elif current_logical == 'HKLOUT':
                        nref_out = int(record.split()[-1])

            return nref_in, nref_out

    return TruncateWrapper()

if __name__ == '__main__':

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

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