#!/usr/bin/env python
# SymmetryExpert.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.
#
# 21st May 2007
#
# A small expert to handle symmetry calculations.
#

import os
import sys
import math
from scitbx import matrix
from cctbx import sgtbx

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 Wrappers.XIA.Mat2symop import Mat2symop
# from Wrappers.XIA.Symop2mat import Symop2mat
from Handlers.Syminfo import Syminfo

def gen_rot_mat_euler(alpha, beta, gamma):
    '''Compute a rotation matrix (stored as e11 e12 e13 e22 e23...)
    as product R(x, gamma).R(y, beta).R(z, alpha).'''

    from MatrixExpert import rot_x, rot_y, rot_z

    rz = rot_z(alpha)
    ry = rot_y(beta)
    rx = rot_x(gamma)

    r = _multiply_symmetry_matrix(ry, rz)
    return _multiply_symmetry_matrix(rx, r)

def _multiply_symmetry_matrix(a, b):
    '''compute a * b, for e.g. h_ = a * b * h, e.g. apply b before a.'''

    if True:
        return (matrix.sqr(a) * matrix.sqr(b)).elems

    result = []

    result.append(a[0] * b[0] +
                  a[1] * b[3] +
                  a[2] * b[6])

    result.append(a[0] * b[1] +
                  a[1] * b[4] +
                  a[2] * b[7])

    result.append(a[0] * b[2] +
                  a[1] * b[5] +
                  a[2] * b[8])


    result.append(a[3] * b[0] +
                  a[4] * b[3] +
                  a[5] * b[6])

    result.append(a[3] * b[1] +
                  a[4] * b[4] +
                  a[5] * b[7])

    result.append(a[3] * b[2] +
                  a[4] * b[5] +
                  a[5] * b[8])


    result.append(a[6] * b[0] +
                  a[7] * b[3] +
                  a[8] * b[6])

    result.append(a[6] * b[1] +
                  a[7] * b[4] +
                  a[8] * b[7])

    result.append(a[6] * b[2] +
                  a[7] * b[5] +
                  a[8] * b[8])

    return result

def r_to_rt(r):
    '''Convert R matrix to RT, assuming T=0.'''

    result = []

    for i in range(3):
        for j in range(3):
            result.append(r[i * 3 + j])
        result.append(0)

    return result

def rt_to_r(rt):
    '''Convert RT matrix to R, removing T.'''

    result = []
    for i in range(3):
        for j in range(3):
            result.append(rt[4 * i + j])

    return result

def compose_matrices_rt(mat_a, mat_b):
    '''Compose symmetry matrix files for XDS. These are 12 element
    matrices...'''

    mat_c = _multiply_symmetry_matrix(rt_to_r(mat_a),
                                      rt_to_r(mat_b))

    return r_to_rt(mat_c)

def compose_matrices_r(mat_a, mat_b):
    '''Compose symmetry matrix applying b then a.'''

    mat_c = _multiply_symmetry_matrix(mat_a,
                                      mat_b)

    return mat_c

def compose_symops(a, b):
    '''Compose operation c, which is applying b then a.'''

    # symop2mat = Symop2mat()
    mat_a = symop_to_mat(a)
    mat_b = symop_to_mat(b)

    mat_c = _multiply_symmetry_matrix(mat_a, mat_b)

    # mat2symop = Mat2symop()
    result = mat_to_symop(mat_c).strip()

    return result

def symop_to_mat(symop):
    symop = symop.replace('h', 'x').replace('k', 'y').replace('l', 'z')
    return matrix.sqr(sgtbx.change_of_basis_op(
        symop).c().as_double_array()[:9]).elems

def mat_to_symop(mat):
    return sgtbx.change_of_basis_op(sgtbx.rt_mx(
        matrix.sqr(mat), (0, 0, 0),
        r_den = 12, t_den = 144)).as_xyz()

def lattice_to_spacegroup_number(lattice):
    '''Return the spacegroup number corresponding to the lowest symmetry
    possible for a given Bravais lattice.'''

    _lattice_to_spacegroup_number = {'aP':1,
                                     'mP':3,
                                     'mC':5,
                                     'oP':16,
                                     'oC':20,
                                     'oF':22,
                                     'oI':23,
                                     'tP':75,
                                     'tI':79,
                                     'hP':143,
                                     'hR':146,
                                     'cP':195,
                                     'cF':196,
                                     'cI':197}

    if not lattice in _lattice_to_spacegroup_number.keys():
        raise RuntimeError, 'lattice %s unknown' % lattice

    return _lattice_to_spacegroup_number[lattice]

def modulo(m, x):
    '''Return x modulo m for floating values m, x.'''

    while x < 0:
        x += m
    while x > m:
        x -= m

    return x

if __name__ == '__main__':

    s = compose_symops('-y,x+y,z', 'h,k,l')
    m = symop_to_mat(s)

    assert(s == mat_to_symop(m))

    s = '-h,-k,h+l'

    m = symop_to_mat(s)

    for i in map(int, m):
        print i,