#!/usr/bin/env python
# MatrixExpert.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.
#
# 24th July 2007
#
# A small expert to handle orientation matrix calculations.
#

import os
import sys
import math

from cctbx import sgtbx
from cctbx import crystal
from cctbx import uctbx
from scitbx import matrix

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 Experts.SymmetryExpert import symop_to_mat, mat_to_symop
from Experts.LatticeExpert import l2s, s2l
from Wrappers.CCP4.Othercell import Othercell
from Wrappers.Phenix.LatticeSymmetry import LatticeSymmetry
from lib.SymmetryLib import lattice_to_spacegroup
from Handlers.Syminfo import Syminfo
from lib.bits import auto_logfiler
from Handlers.Streams import Debug

from ReferenceFrame import mosflm_to_xia2, xia2_to_mosflm

# jiffies to convert matrix format (messy)

def mat2vec(mat):
  return [[mat[0], mat[3], mat[6]],
          [mat[1], mat[4], mat[7]],
          [mat[2], mat[5], mat[8]]]

def vec2mat(vectors):
  return [vectors[0][0], vectors[1][0], vectors[2][0],
          vectors[0][1], vectors[1][1], vectors[2][1],
          vectors[0][2], vectors[1][2], vectors[2][2]]

# generic mathematical calculations for 3-vectors

# FIXME cite PRE as the source here for these rotns - N.B. these should be
# replaced with CCTBX code.

def rot_x(theta):
  '''Rotation matrix about Y of theta degrees.'''

  dtor = 180.0 / (4.0 * math.atan(1.0))

  c = math.cos(theta / dtor)
  s = math.sin(theta / dtor)

  return [1.0, 0.0, 0.0, 0.0, c, s, 0.0, -s, c]

def rot_y(theta):
  '''Rotation matrix about Y of theta degrees.'''

  dtor = 180.0 / (4.0 * math.atan(1.0))

  c = math.cos(theta / dtor)
  s = math.sin(theta / dtor)

  return [c, 0.0, -s, 0.0, 1.0, 0.0, s, 0.0, c]

def rot_z(theta):
  '''Rotation matrix about Y of theta degrees.'''

  dtor = 180.0 / (4.0 * math.atan(1.0))

  c = math.cos(theta / dtor)
  s = math.sin(theta / dtor)

  return [c, -s, 0.0, s, c, 0.0, 0.0, 0.0, 1.0]

def b_matrix(a, b, c, alpha, beta, gamma):
  '''Generate a B matric from a unit cell. Cite: Pflugrath in Methods
  Enzymology 276.'''

  dtor = 180.0 / (4.0 * math.atan(1.0))

  ca = math.cos(alpha / dtor)
  sa = math.sin(alpha / dtor)
  cb = math.cos(beta / dtor)
  sb = math.sin(beta / dtor)
  cg = math.cos(gamma / dtor)
  sg = math.sin(gamma / dtor)

  # invert the cell parameters
  # CITE: International Tables C Section 1.1

  V = a * b * c * math.sqrt(1 - ca * ca - cb * cb - cg * cg +
                            2 * ca * cb * cg)

  a_ = b * c * sa / V
  b_ = a * c * sb / V
  c_ = a * b * sg / V

  # NOTE well - these angles are in radians

  alpha_ = math.acos((cb * cg - ca) / (sb * sg))
  beta_ = math.acos((ca * cg - cb) / (sa * sg))
  gamma_ = math.acos((ca * cb - cg) / (sa * sb))

  ca_ = math.cos(alpha_)
  sa_ = math.sin(alpha_)
  cb_ = math.cos(beta_)
  sb_ = math.sin(beta_)
  cg_ = math.cos(gamma_)
  sg_ = math.sin(gamma_)

  # NEXT construct the B matrix - CITE Pflugrath in Methods E 276

  return [a_, b_ * cg_, c_ * cb_,
          0.0, b_ * sg_, - c_ * sb_ * ca_,
          0.0, 0.0, c_ * sb_ * sa_]

def dot(a, b):
  return sum([a[j] * b[j] for j in range(3)])

def cross(a, b):
  return [a[1] * b[2] - a[2] * b[1],
          a[2] * b[0] - a[0] * b[2],
          a[0] * b[1] - a[1] * b[0]]

def vecscl(vector, scale):
  return [vector[j] * scale for j in range(len(vector))]

def matscl(matrix, scale):
  return [matrix[j] * scale for j in range(len(matrix))]

def invert(matrix):
  vecs = mat2vec(matrix)
  scl = 1.0 / dot(vecs[0], cross(vecs[1], vecs[2]))

  return transpose(
      vec2mat([vecscl(cross(vecs[1], vecs[2]), scl),
               vecscl(cross(vecs[2], vecs[0]), scl),
               vecscl(cross(vecs[0], vecs[1]), scl)]))

def transpose(matrix):
  return [matrix[0], matrix[3], matrix[6],
          matrix[1], matrix[4], matrix[7],
          matrix[2], matrix[5], matrix[8]]

def det(matrix):
  vecs = mat2vec(matrix)
  return dot(vecs[0], cross(vecs[1], vecs[2]))

def matmul(b, a):
  avec = mat2vec(transpose(a))
  bvec = mat2vec(b)

  result = []
  for i in range(3):
    for j in range(3):
      result.append(dot(avec[i], bvec[j]))

  return result

def matvecmul(M, v):
  '''Multiply a vector v by a matrix M -> return M v.'''

  Mvec = mat2vec(transpose(M))
  result = []
  for i in range(3):
    result.append(dot(Mvec[i], v))
  return result

# things specific to mosflm matrix files...

def parse_matrix(matrix_text):
  '''Parse a matrix returning cell, a and u matrix.'''

  # this will need to be able to cope with the times
  # when the matrix includes columns merging together
  # (which sucks)

  # therefore parse this manually... or just add
  # a space before all '-'

  tokens = map(float, matrix_text.replace('-', ' -').split()[:30])

  cell = tokens[21:27]
  a = tokens[0:9]
  u = tokens[12:21]

  return cell, a, u

def format_matrix(cell, a, u):
  matrix_format = ' %11.8f %11.8f %11.8f\n' + \
                  ' %11.8f %11.8f %11.8f\n' + \
                  ' %11.8f %11.8f %11.8f\n'

  cell_format = ' %11.4f %11.4f %11.4f %11.4f %11.4f %11.4f\n'

  misset = '       0.000       0.000       0.000\n'

  return matrix_format % tuple(a) + misset + matrix_format % tuple(u) + \
         cell_format % tuple(cell) + misset

def transmogrify_matrix(lattice, matrix, target_lattice,
                        wavelength, wd = None):
  '''Transmogrify a matrix for lattice X into a matrix for lattice
  Y. This should work find for Mosflm... Will also return the new
  unit cell.'''

  cell, a, u = parse_matrix(matrix)

  o = _Othercell()

  if wd:
    o.set_working_directory(wd)
    auto_logfiler(o)

  o.set_cell(cell)
  o.set_lattice(lattice)
  o.generate()

  new_cell = o.get_cell(target_lattice)
  op = symop_to_mat(o.get_reindex_op(target_lattice))

  # HACK! again - bug # 3193

  if 'lattice_symmetry' in o.get_executable():
    op = transpose(op)

  # why is only one of the matrices inverted?!

  a = matmul(invert(op), a)
  u = matmul(op, u)

  # in here test that the given unit cell corresponds to the
  # one calculated from the A matrix.

  anew = matscl(a, 1.0 / wavelength)
  reala = transpose(invert(anew))
  _a, _b, _c = mat2vec(reala)

  la = math.sqrt(dot(_a, _a))
  lb = math.sqrt(dot(_b, _b))
  lc = math.sqrt(dot(_c, _c))

  if math.fabs(la - new_cell[0]) / new_cell[0] > 0.01:
    raise RuntimeError, 'cell check failed (wavelength != %f)' % \
          wavelength

  if math.fabs(lb - new_cell[1]) / new_cell[1] > 0.01:
    raise RuntimeError, 'cell check failed (wavelength != %f)' % \
          wavelength

  if math.fabs(lc - new_cell[2]) / new_cell[2] > 0.01:
    raise RuntimeError, 'cell check failed (wavelength != %f)' % \
          wavelength

  return format_matrix(new_cell, a, u)

def _Othercell():
  '''A factory to produce either a wrapper for LaticeSymmetry or
  OtherCell depending on what is available.'''

  o = None

  try:
    o = LatticeSymmetry()
    Debug.write('Using iotbx.lattice_symmetry')
  except:
    o = Othercell()
    Debug.write('Using othercell')

  return o

def get_real_space_primitive_matrix(lattice, matrix, wd = None):
  '''Get the primitive real space vectors for the unit cell and
  lattice type. Note that the resulting matrix will need to be
  scaled by a factor equal to the wavelength in Angstroms.'''

  # parse the orientation matrix

  cell, a, u = parse_matrix(matrix)

  # generate other possibilities

  o = _Othercell()

  if wd:
    o.set_working_directory(wd)
    auto_logfiler(o)

  o.set_cell(cell)
  o.set_lattice(lattice)
  o.generate()

  # transform the possibly centred cell to the primitive setting

  new_cell = o.get_cell('aP')
  op = symop_to_mat(o.get_reindex_op('aP'))

  primitive_a = matmul(invert(op), a)

  # then convert to real space

  real_a = invert(primitive_a)

  return real_a[0:3], real_a[3:6], real_a[6:9]

def get_reciprocal_space_primitive_matrix(lattice, matrix, wd = None):
  '''Get the primitive reciprocal space vectors for this matrix.'''

  # parse the orientation matrix

  cell, a, u = parse_matrix(matrix)

  # generate other possibilities

  o = _Othercell()

  if wd:
    o.set_working_directory(wd)
    auto_logfiler(o)

  o.set_cell(cell)
  o.set_lattice(lattice)
  o.generate()

  # transform the possibly centred cell to the primitive setting

  new_cell = o.get_cell('aP')
  op = symop_to_mat(o.get_reindex_op('aP'))

  primitive_a = matmul(invert(op), a)

  return mat2vec(primitive_a)

def find_primitive_axes(lattice, matrix, wd = None):
  '''From an orientation matrix file, calculate the angles (phi) where
  the primitive cell axes a, b, c are in the plane of the detector
  (that is, orthogonal to the direct beam vector.'''

  a, b, c = get_real_space_primitive_matrix(lattice, matrix, wd)

  dtor = 180.0 / (4.0 * math.atan(1.0))

  return (dtor * math.atan( - a[2] / a[0]), \
          dtor * math.atan( - b[2] / b[0]), \
          dtor * math.atan( - c[2] / c[0]))

def find_primitive_reciprocal_axes(lattice, matrix, wd = None):
  '''From an orientation matrix file, calculate the angles (phi) where
  the primitive reciprical space cell axes a, b, c are in the plane of
  the detector (that is, orthogonal to the direct beam vector.'''

  a, b, c = get_reciprocal_space_primitive_matrix(lattice, matrix, wd)

  dtor = 180.0 / (4.0 * math.atan(1.0))

  return (dtor * math.atan( - a[2] / a[0]), \
          dtor * math.atan( - b[2] / b[0]), \
          dtor * math.atan( - c[2] / c[0]))

def mosflm_a_matrix_to_real_space(wavelength, lattice, matrix):
  '''Given a Mosflm A matrix and the associated spacegroup (think of this
  Bravais lattice (which will be converted to a spacegroup for the benefit
  of the CCTBX program lattice_symmetry) return the real space primative
  crystal lattice vectors in the xia2 reference frame. This reference frame
  corresponds to that defined for imgCIF.'''

  # convert the lattice to a spacegroup - not needed, see below
  # spacegroup_number = lattice_to_spacegroup(lattice)
  # spacegroup = Syminfo.spacegroup_name_to_number(spacegroup_number)

  # get the a, u, matrices and the unit cell
  cell, a, u = parse_matrix(matrix)

  # use iotbx.latice_symmetry to obtain the reindexing operator to
  # a primative triclinic lattice - this should not be specific to
  # using iotbx - othercell should be optionally supported too...
  ls = _Othercell()
  ls.set_cell(cell)
  ls.set_lattice(lattice)

  # ls.set_spacegroup(spacegroup)
  # cell, reindex = ls.generate_primative_reindex()
  ls.generate()

  cell = ls.get_cell('aP')
  reindex = ls.get_reindex_op('aP')
  reindex_matrix = symop_to_mat(reindex)

  # grim hack warning! - this is because for some reason (probably to
  # do with LS getting the inverse operation) othercell is returning
  # the transpose of the operator or something... FIXME bug # 3193

  if 'othercell' in ls.get_executable():
    reindex_matrix = transpose(reindex_matrix)

  # scale the a matrix
  a = matscl(a, 1.0 / wavelength)

  # convert to real space (invert) and apply this reindex operator to the a
  # matrix to get the primative real space triclinic cell axes
  real_a = matmul(reindex_matrix, transpose(invert(a)))

  # convert these to the xia2 reference frame
  a, b, c = mat2vec(real_a)

  ax = mosflm_to_xia2(a)
  bx = mosflm_to_xia2(b)
  cx = mosflm_to_xia2(c)

  # FIXME in here add check that the unit cell lengths are within 1% of
  # the correct value - if they are not, raise an exception as the wavelength
  # provided is probably wrong...

  la = math.sqrt(dot(ax, ax))
  lb = math.sqrt(dot(bx, bx))
  lc = math.sqrt(dot(cx, cx))

  # print out values - the cell may have been reindexed.

  Debug.write('Reindexed cell lengths: %.3f %.3f %.3f' % \
              (cell[0], cell[1], cell[2]))
  Debug.write('Calculated from matrix: %.3f %.3f %.3f' % \
              (la, lb, lc))

  if math.fabs(la - cell[0]) / cell[0] > 0.01:
    raise RuntimeError, 'cell check failed (wavelength != %f)' % \
          wavelength

  if math.fabs(lb - cell[1]) / cell[1] > 0.01:
    raise RuntimeError, 'cell check failed (wavelength != %f)' % \
          wavelength

  if math.fabs(lc - cell[2]) / cell[2] > 0.01:
    raise RuntimeError, 'cell check failed (wavelength != %f)' % \
          wavelength

  # return these vectors
  return ax, bx, cx

def reindex_sym_related(A, A_ref):
  '''Calculate a reindexing matrix to move the indices referred to in
  A to the reference frame in A_ref: both are orientation matrices from
  Mosflm.'''

  Amat = matrix.sqr(parse_matrix(A)[1])
  Amat_ref = matrix.sqr(parse_matrix(A_ref)[1])
  R = Amat_ref.inverse() * Amat
  Debug.write('%5.3f %5.3f %5.3f %5.3f %5.3f %5.3f %5.3f %5.3f %5.3f' % \
              R.elems)

  reindex = mat_to_symop(R)
  return reindex

def mosflm_matrix_centred_to_primitive(lattice, mosflm_a_matrix):
  '''Convert a mosflm orientation matrix from a centred setting to a
  primitive one (i.e. same lattice, but without the centering operations,
  which therefore corresponds to a different basis).'''

  space_group_number = l2s(lattice)
  spacegroup = sgtbx.space_group_symbols(space_group_number).hall()
  sg = sgtbx.space_group(spacegroup)

  rtod = 180.0 / math.pi

  if not (sg.n_ltr() - 1):
    return mosflm_a_matrix

  cell, amat, umat = parse_matrix(mosflm_a_matrix)

  # first derive the wavelength

  mi = matrix.sqr(amat)
  m = mi.inverse()

  A = matrix.col(m.elems[0:3])
  B = matrix.col(m.elems[3:6])
  C = matrix.col(m.elems[6:9])

  a = math.sqrt(A.dot())
  b = math.sqrt(B.dot())
  c = math.sqrt(C.dot())

  alpha = rtod * B.angle(C)
  beta = rtod * C.angle(A)
  gamma = rtod * A.angle(B)

  wavelength = ((cell[0] / a) + (cell[1] / b) + (cell[2] / c)) / 3.0

  # then use this to rescale the A matrix

  mi = matrix.sqr([a / wavelength for a in amat])
  m = mi.inverse()

  sgp = sg.build_derived_group(True, False)
  lattice_p = s2l(sgp.type().number())
  symm = crystal.symmetry(unit_cell = cell,
                          space_group = sgp)

  rdx = symm.change_of_basis_op_to_best_cell()
  symm_new = symm.change_basis(rdx)

  # now apply this to the reciprocal-space orientation matrix mi

  cb_op = rdx
  R = cb_op.c_inv().r().as_rational().as_float().transpose().inverse()
  mi_r = mi * R

  # now re-derive the cell constants, just to be sure

  m_r = mi_r.inverse()
  Ar = matrix.col(m_r.elems[0:3])
  Br = matrix.col(m_r.elems[3:6])
  Cr = matrix.col(m_r.elems[6:9])

  a = math.sqrt(Ar.dot())
  b = math.sqrt(Br.dot())
  c = math.sqrt(Cr.dot())

  alpha = rtod * Br.angle(Cr)
  beta = rtod * Cr.angle(Ar)
  gamma = rtod * Ar.angle(Br)

  cell = uctbx.unit_cell((a, b, c, alpha, beta, gamma))

  amat = [wavelength * e for e in mi_r.elems]
  bmat = matrix.sqr(cell.fractionalization_matrix())
  umat = mi_r * bmat.inverse()

  new_matrix = ['%s\n' % r for r in \
                format_matrix((a, b, c, alpha, beta, gamma),
                              amat, umat.elems).split('\n')]

  return new_matrix

if __name__ == '__main__':

  lattice = sys.argv[1]
  mosflm_a_matrix = open(sys.argv[2]).read()

  result = mosflm_matrix_centred_to_primitive(lattice, mosflm_a_matrix)

  for r in result:
    print r[:-1]