#!/usr/bin/env python
# Polarrfn.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.
#
# 20th June 2007
#

import sys
import os
import math

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
from Decorators.DecoratorFactory import DecoratorFactory
from Handlers.Streams import Chatter

from Experts.SymmetryExpert import gen_rot_mat_euler, symop_to_mat
from Handlers.Syminfo import Syminfo

def matrix_diff(a, b):
    return sum([(a[i] - b[i]) * (a[i] - b[i]) for i in range(9)])

def Polarrfn(DriverType = None):
    '''A factory for PolarrfnWrapper classes.'''

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

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

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

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

            self._labin_f = 'F'
            self._labin_sigf = 'SIGF'
            self._height = 50

            return

        def set_height(self, height):
            self._height = height

        def set_labin(self, labin_f, labin_sigf):
            self._labin_f = labin_f
            self._labin_sigf = labin_sigf
            return

        def polarrfn(self):

            self.check_hklin()

            self.start()

            self.input('self 20 4')
            self.input('resolution 15 3')
            self.input('crystal file 1')
            self.input('labin file 1 F=%s SIGF=%s' % (self._labin_f,
                                                      self._labin_sigf))
            self.input('noprint')

            self.input('find %d 100' % self._height)

            self.close_wait()

            # should check for errors etc here

            # next try and read the various peaks...

            output = self.get_all_output()

            j = 0

            peaks = { }

            symops = { }

            while j < len(self.get_all_output()):

                line = output[j]

                if 'Space group =' in line:
                    operations = Syminfo.get_symops(
                        int(line.replace(')', '').split()[-1]))
                    for op in operations:
                        symops[tuple(symop_to_mat(op))] = op

                current_peak = 0

                if 'Alpha   Beta   Gamma      Peak' in line:

                    while not 'The rotation given' in line:
                        line = output[j]
                        if ['Peak'] == line.split()[:1]:
                            current_peak = int(line.split()[-1])
                            peaks[current_peak] = []
                        else:
                            try:
                                a = int(line.split()[0])
                                b = int(line.split()[1])
                                if a == 1 and b == 1:
                                    peak = map(float, line.split()[2:5])
                                    peak.append(float(line.split()[5]))
                                    peaks[current_peak].append(peak)
                            except:
                                pass


                        j += 1

                j += 1

            numbers = peaks.keys()
            numbers.sort()
            if 0 in numbers:
                numbers.remove(0)

            unique = { }
            keys = []

            for n in numbers:
                for p in peaks[n]:
                    key = p[0], p[1], p[2]
                    if unique.has_key(key):
                        if unique[key] > p[3]:
                            continue

                    unique[key] = p[3]
                    if not key in keys:
                        keys.append(key)

            Chatter.write('... Alpha Beta  Gamma Height (matching symop)')

            for k in keys:

                # in here compute a matrix for this rotation
                # compare it to the matrices derived from the
                # symmetry operations - if it matches, record
                # the symop, else don't.

                matrix = gen_rot_mat_euler(k[2], k[1], k[0])

                matstr = '[%.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f]' % \
                         tuple(matrix)

                symop = ''

                for mat in symops.keys():
                    diff = matrix_diff(list(mat), matrix)
                    if diff < 0.1:
                        symop = symops[mat]


                if symop:
                    Chatter.write('... %5.1f %5.1f %5.1f %5.1f  (%s)' % \
                                  (k[0], k[1], k[2], unique[k], symop))
                else:
                    Chatter.write('... %5.1f %5.1f %5.1f %5.1f' % \
                                  (k[0], k[1], k[2], unique[k]))

            return


    return PolarrfnWrapper()


if __name__ == '__main__':

    from Ecalc import Ecalc

    if len(sys.argv) < 2:
        raise RuntimeError, '%s hklin [F SIGF]' % sys.argv[0]

    hklin = sys.argv[1]

    if len(sys.argv) == 4:
        labin_f = sys.argv[2]
        labin_sigf = sys.argv[3]
    else:
        labin_f = 'F'
        labin_sigf = 'SIGF'

    ecalc = Ecalc()
    ecalc.set_hklin(hklin)
    ecalc.set_hklout('temp.mtz')
    ecalc.set_labin(labin_f, labin_sigf)
    ecalc.ecalc()

    polar = Polarrfn()
    polar.set_hklin('temp.mtz')
    polar.set_labin('E', 'SIGE')
    polar.polarrfn()