#!/usr/bin/env python
# Chooch.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.
#
# A wrapper for Chooch, for use in deciding what has happened during
# collection of experimental phasing data, and also for helping with
# MAD experiments.
#
# 18th December 2006
#
#

import os
import sys
import math
import string
import exceptions

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

if not os.environ.has_key('XIA2_ROOT'):
    raise RuntimeError, 'XIA2_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['XIA2_ROOT'] in sys.path:
    sys.path.append(os.environ['XIA2_ROOT'])

from Driver.DriverFactory import DriverFactory

# helper functions

def energy_to_wavelength(energy):
    h = 6.6260693e-34
    c = 2.9979246e8
    e = 1.6021765e-19

    return 1.0e10 * (h * c) / (e * energy)

def preprocess_scan(scan_file):
    '''Preprocess the scan file to a form that chooch will accept.'''

    try:
        i = int(open(scan_file, 'r').readlines()[1])
        return scan_file
    except:
        # assume that this is not in the friendly format...
        data = open(scan_file, 'r').readlines()
        more_data = []
        for d in data:
           if not '#' in d and d.strip():
               more_data.append(d)
        data = more_data
        count = len(data)
        out = open('xia2-chooch.raw', 'w')
        out.write('Chooch Scan File from xia2\n%d\n' % count)
        for d in data:
            out.write('%f %f\n' % tuple(map(float, d.split(',')[:2])))
        out.close()
        return 'xia2-chooch.raw'

def Chooch(DriverType = None):
    '''Factory for Chooch wrapper classes, with the specified
    Driver type.'''

    DriverInstance = DriverFactory.Driver(DriverType)

    class ChoochWrapper(DriverInstance.__class__):
        def __init__(self):

            DriverInstance.__class__.__init__(self)

            self.set_executable('chooch')
            self._scan = None
            self._edge_table = { }

            self._atom = 'se'

            self._data = []

        def set_scan(self, scan):
            '''Set a scan file for chooch to parse.'''

            self._scan = preprocess_scan(scan)
            return

        def set_atom(self, atom):
            '''Set the atom which should be in the scan.'''

            self._atom = atom
            return

        def scan(self):
            '''Run chooch.'''

            self.add_command_line('-e')
            self.add_command_line(self._atom)
            self.add_command_line(self._scan)

            self.start()
            self.close_wait()

            self.check_for_errors()

            self._data = []
            for o in open(os.path.join(self.get_working_directory(),
                                       'output.efs'), 'r').readlines():
                self._data.append(map(float, o.split()))

            output = self.get_all_output()
            collect = False

            for o in output:
                if collect:

                    if '-------' in o:
                        collect = False
                        continue

                    name, energy, fp, fpp = tuple(map(string.strip,
                                                      o.split('|')[1:5]))
                    self._edge_table[name] = {
                        'energy':float(energy),
                        'fp':float(fp),
                        'fpp':float(fpp),
                        'wave':energy_to_wavelength(float(energy))}

                if 'energy' in o and 'f\'' in o and 'f\'\'' in o:
                    collect = True

        def get_edges(self):
            return self._edge_table

        def get_fp_fpp(self, wave):
            '''Get the fp, fpp for a wavelength.'''

            if not self._data:
                raise RuntimeError, 'data array empty'

            fp = 0.0
            fpp = 0.0
            closest = 1.0e100

            for d in self._data:
                w = energy_to_wavelength(d[0])
                if math.fabs(w - wave) < closest:
                    fp = d[2]
                    fpp = d[1]
                    closest = math.fabs(w - wave)

            return fp, fpp

        def id_wavelength(self, wave):
            '''Try to identify a wavelength.'''

            min_wave_diff = 1.0e100
            name = None

            waves = []

            for edge in self._edge_table.keys():
                waves.append(self._edge_table[edge]['wave'])
                if math.fabs(self._edge_table[edge]['wave'] - wave) < \
                   min_wave_diff:
                    min_wave_diff = math.fabs(self._edge_table[edge]['wave'] -
                                              wave)
                    name = edge

            if min_wave_diff > 0.01:
                # assume that this is a remote...
                if wave > max(waves):
                    return 'LREM'
                else:
                    return 'HREM'

            return name.upper()

    return ChoochWrapper()

if __name__ == '__main__':

    if len(sys.argv) < 2:

        c = Chooch()
        c.set_scan(os.path.join(os.environ['XIA2_ROOT'], 'Data',
                                'Test', 'Scans',
                                'TM0486-9172-Se.raw'))
        c.set_atom('se')
        c.scan()

        edges = c.get_edges()

        for key in edges.keys():
            print '%s %6.2f %6.2f %8.6f' % (key,
                                            edges[key]['fp'],
                                            edges[key]['fpp'],
                                            edges[key]['wave'])

    else:
        for scan in sys.argv[1:]:

            print os.path.split(scan)[-1]

            try:

                c = Chooch()
                c.set_scan(scan)
                c.set_atom('se')
                c.scan()

                edges = c.get_edges()

                for key in edges.keys():
                    print '%s %6.2f %6.2f %8.6f' % (key,
                                                    edges[key]['fp'],
                                                    edges[key]['fpp'],
                                                    edges[key]['wave'])
            except exceptions.Exception, e:
                print 'failed (%s)' % str(e)


            print ''