"""
general_xrd.py - A set of python functions useful for doing calculations pertinent to XRD analaysis

Original author: Brendan Arnold, TCD - 2008

Version: 0.1 (some errors may exist!)

Do whatever you like with this code so long as credit for original basis is given in source.

The easiest way to import these functions into your code is to include it in
the same directory as your scripts and use 'from general_xrd import *'. For
more detailed instructions see http://docs.python.org/tut/node8.html
"""

from math import sin, cos, radians, degrees, sqrt, acos

COMMON_XRAY_WAVELENGTHS = {
  'CuKa1': 1.5405981E-10,
  'CuKa2': 1.54443E-10,
  'CuKb1': 1.39225E-10,
  'WLa1': 1.47642E-10,
  'WLa2': 1.48748E-10
}

def parse_miller_indices(p):
  """Pass in an interable of length 3 or 4. Tries to covnert each index to an integer and returns a tuple h, k, l

  Usage: parse_miller_indices([h,k,(j),l])"""
  try:
    n = len(p)
    if 5 < n < 2:
      raise
  except:
    raise ValueError, 'Argument must be iterable of 3 or 4 integers'
  if n == 3:
    h, k, l = p
  elif n == 4:
    # j is not used
    h, k, j, l = p
  try:
    h = int(h)
    k = int(k)
    l = int(l)
  except:
    raise ValueError, 'Values must be convertable to integers'
  return h, k, l

# The following functions come from Appendix 3 of Cullity 2nd ed. 'Elements of
# X-ray Diffraction' from Addison Wesley

def _S(a, b, c, alpha, beta, gamma):
  """Returns the 'S' convenience variables for triclinic calculations in this order:
  S11, S22, S33, S12, S23, S13

  Usage: S11, S22, S33, S12, S23, S13 = _S(a, b, c, alpha, beta, gamma)"""
  sa, sb, sg, ca, cb, cg = sin(radians(alpha)), sin(radians(beta)), sin(radians(gamma)), \
    cos(radians(alpha)), cos(radians(beta)), cos(radians(gamma))
  S11 = b**2*c**2*sa**2
  S22 = a**2*c**2*sb**2
  S33 = a**2*b**2*sg**2
  S12 = a*b*c**2*(ca*cb-cg)
  S23 = a**2*b*c*(cb*cg-ca)
  S13 = a*b**2*c*(cg*ca-cb)
  return S11, S22, S33, S12, S23, S13
  
def calc_volume_cubic(a=None):
  """Pass in a to get volume of unit cell

  Usage: calc_volume_cubic(a)"""
  return a**3

def calc_volume_tetragonal(a=None, c=None):
  """Pass in a, c to get volume of unit cell

  Usage: calc_volume_tetragonal(a, c)"""
  return c*a**2

def calc_volume_hexagonal(a=None, c=None):
  """Pass in a, c to get volume of unit cell

  Usage: calc_volume_hexagonal(a, c)"""
  return c*(a**2)*sqrt(3)/2

def calc_volume_rhombohedral(a=None, alpha=None):
  """Pass in a, alpha to get volume of unit cell

  Usage: calc_volume_rhombohedral(a, alpha)"""
  ca = cos(radians(alpha))
  return (a**3)*sqrt(1 - 3*ca**2 + 2*ca**3)

def calc_volume_orthorhombic(a=None, b=None, c=None):
  """Pass in a, b, c to get volume of unit cell

  Usage: calc_volume_orthorhombic(a, b, c)"""
  return a*b*c

def calc_volume_monoclinic(a=None, b=None, c=None, beta=None):
  """Pass in a, b, c, beta to get volume of unit cell

  Usage: calc_volume_monoclinic(a, b, c, beta)"""
  sb = sin(radians(beta))
  return a*b*c*sb

def calc_volume_triclinic(a=None, b=None, c=None, alpha=None, beta=None, gamma=None):
  """Pass in a, b, c, alpha, beta, gamma to get volume of unit cell

  Usage: calc_volume_triclinic(a, b, c, alpha, beta, gamma)"""
  ca, cb, cg = cos(radians(alpha)), cos(radians(beta)), cos(radians(gamma))
  return a*b*c*sqrt(1 - ca**2 - cb**2 - cg**2 + 2*ca*cb*cg)

def calc_plane_spacing_cubic(p=None, a=None):
  """Pass in a plane [h,k,l] and a to get the plane spacing d

  Usage: calc_plane_spacing_cubic([h,k,l], a)"""
  h, k, l = parse_miller_indices(p)
  reciprocal_d2 = 1.0*(h**2 + k**2 + l**2)/a**2
  return sqrt(1.0/reciprocal_d2)

def calc_plane_spacing_tetragonal(p=None, a=None, c=None):
  """Pass in a plane [h,k,l], a and c to get the plane spacing d

  Usage: calc_plane_spacing_tetragonal([h,k,l], a, c)"""
  h, k, l = parse_miller_indices(p)
  reciprocal_d2 = 1.0*(h**2+k**2)/a**2 + 1.0*l**2/c**2
  return sqrt(1.0/reciprocal_d2)

def calc_plane_spacing_hexagonal(p=None, a=None, c=None):
  """Pass in a plane [h,k,l], a and c to get the plane spacing d

  Usage: calc_plane_spacing_hexagonal([h,k,l], a, c)"""
  h, k, l = parse_miller_indices(p)
  reciprocal_d2 = (4.0/3.0) * 1.0*(h**2+h*k+k**2)/a**2 + 1.0*l**2/c**2
  return sqrt(1.0/reciprocal_d2)
  
def calc_plane_spacing_rhombohedral(p=None, a=None, alpha=None):
  """Pass in a plane [h,k,l], a and alpha to get the plane spacing d

  Usage: calc_plane_spacing_rhombohedral([h,k,l], a, alpha)"""
  h, k, l = parse_miller_indices(p)
  sa, ca = sin(radians(alpha)), cos(radians(alpha))
  reciprocal_d2 = 1.0*( (h**2+k**2+l**2)*sa**2 + 2*(h*k+k*l+h*l)*(ca**2-ca) ) / \
    a**2*(1-3*ca**2+2*ca**3)
  return sqrt(1.0/reciprocal_d2)

def calc_plane_spacing_orthorhombic(p=None, a=None, b=None, c=None):
  """Pass in a plane [h,k,l], a, b and c to get the plane spacing d

  Usage: calc_plane_spacing_orthorhombic([h,k,l], a, b, c)"""
  h, k, l = parse_miller_indices(p)
  reciprocal_d2 = 1.0*h**2/a**2 + 1.0*k**2/b**2 + 1.0*l**2/c**2
  return sqrt(1.0/reciprocal_d2)

def calc_plane_spacing_monoclinic(p=None, a=None, b=None, c=None, beta=None):
  """Pass in a plane [h,k,l], a, b, c and beta to get the plane spacing d

  Usage: calc_plane_spacing_monoclinic([h,k,l], a, b, c, beta)"""
  h, k, l = parse_miller_indices(p)
  sb, cb = sin(radians(beta)), cos(radians(beta))
  reciprocal_d2 = (1.0/sb**2) * (1.0*h**2/a**2 + 1.0*k**2*sb**2/b**2 + 1.0*l**2/c**2 - 2.0*h*l*cb/(a*c))
  return sqrt(1.0/reciprocal_d2)

def calc_plane_spacing_triclinic(p, a, b, c, alpha, beta, gamma):
  """Pass in a plane [h,k,l], a, b, c, alpha, beta, gamma to get the plane spacing d

  Usage: calc_plane_spacing_triclinic([h,k,l], a, b, c, alpha, beta, gamma)"""
  h, k, l = parse_miller_indices(p)
  S11, S22, S33, S12, S23, S13 = _S(a, b, c, alpha, beta, gamma)
  V = calc_volume_triclinic(a, b, c, alpha, beta, gamma)
  reciprocal_d2 = 1.0/V**2 * (S11*h**2 + S22*k**2 + S33*l**2 + 2*S12*h*k + 2*S23*k*l + 2*S13*h*l) 
  return sqrt(1.0/reciprocal_d2)


def calc_plane_angle_cubic(p1=None, p2=None):
  """Supply two sets of miller indices, calculates the angle in degrees between each

  Usage: calc_plane_angle_cubic([h,k,l], [h,k,l])"""
  h1, k1, l1 = parse_miller_indices(p1)
  h2, k2, l2 = parse_miller_indices(p2)
  cos_phi = 1.0 * (h1*h2+k1*k2+l1*l2) / sqrt((h1**2+k1**2+l1**2)*(h2**2+k2**2+l2**2))
  # return cos_phi
  return degrees(acos(cos_phi))

def calc_plane_angle_tetragonal(p1=None, p2=None, a=None, c=None):
  """Supply two sets of miller indices, calculates the angle in degrees between each, as well as a and c magnitudes

  Usage: calc_plane_angle_tetragonal([h,k,l], [h,k,l], a, c)"""
  h1, k1, l1 = parse_miller_indices(p1)
  h2, k2, l2 = parse_miller_indices(p2)
  cos_phi = ( 1.0*(h1*h2+k1*k2)/a**2 + 1.0*l1*l2/c**2 ) / \
    sqrt( (1.0*(h1**2+k1**2)/a**2 + 1.0*l1**2/c**2)*(1.0*(h2**2+k2**2)/a**2 + 1.0*l2**2/c**2) )
  return degrees(acos(cos_phi))

def calc_plane_angle_hexagonal(p1=None, p2=None, a=None, c=None):
  """Supply two sets of miller indices, calculates the angle in degrees between each, as well as a and c magnitudes

  Usage: calc_plane_angle_hexagonal([h,k,l], [h,k,l], a, c)"""
  h1, k1, l1 = parse_miller_indices(p1)
  h2, k2, l2 = parse_miller_indices(p2)
  cos_phi = (h1*h2+k1*k2+0.5*h1*k2+0.5*h2*k1+(l1*l2*3.0*a**2)/(4.0*c**2)) / \
    sqrt((h1**2+k1**2+h1*k1+(3.0/4.0)*(1.0*a*l1/c)**2)*(h2**2+k2**2+h2*k2+(3.0/4.0)*(1.0*a*l2/c)**2))
  return degrees(acos(cos_phi))

def calc_plane_angle_rhombohedral(p1=None, p2=None, a=None, alpha=None):
  """Supply two sets of miller indices, a and alpha to calculate the angle in degrees between each

  Usage: calc_plane_angle_rhombohedral([h,k,l], [h,k,l], d1, d2, a, alpha)"""
  h1, k1, l1 = parse_miller_indices(p1)
  h2, k2, l2 = parse_miller_indices(p2)
  d1 = calc_plane_spacing_rhombohedral(p1, a, alpha)
  d2 = calc_plane_spacing_rhombohedral(p2, a, alpha)
  V = calc_volume_rhombohedral(a, alpha)
  sa, ca = sin(radians(alpha)), cos(radians(alpha))
  cos_phi = (1.0*(d1*d2*a**4)/V**2) * \
    ((h1*h2+k1*k2+l1*l2)*sa**2 + (ca**2 - ca)*(k1*l2+k2*l1+l1*h2+l2*h1+h1*k2+h2*k1))
  return degrees(acos(cos_phi))

def calc_plane_angle_orthorhombic(p1=None, p2=None, a=None, b=None, c=None):
  """Supply two sets of miller indices, a, b and c to calculate the angle in degrees between each

  Usage: calc_plane_angle_orthorhombic([h,k,l], [h,k,l], a, b, c)"""
  h1, k1, l1 = parse_miller_indices(p1)
  h2, k2, l2 = parse_miller_indices(p2)
  cos_phi = ( 1.0*h1*h2/a**2 + 1.0*k1*k2/b**2 + 1.0*l1*l2/c**2 ) / \
    sqrt((1.0*h1**2/a**2 + 1.0*k1**2/b**2 + 1.0*l1**2/c**2)*(1.0*h2**2/a**2 + 1.0*k2**2/b**2 + 1.0*l2**2/c**2))
  return degrees(acos(cos_phi))

def calc_plane_angle_monoclinic(p1=None, p2=None, a=None, b=None, c=None, beta=None):
  """Supply two sets of miller indices, a, b, c and beta to calculate the angle in degrees between each

  Usage: calc_plane_angle_monoclinic([h,k,l], [h,k,l], d1, d2, a, b, c, beta)"""
  h1, k1, l1 = parse_miller_indices(p1)
  h2, k2, l2 = parse_miller_indices(p2)
  d1, d2 = calc_plane_spacing_monoclinic(p1, a, b, c, beta), \
    calc_plane_spacing_monoclinic(p2, a, b, c, beta)
  sb, cb = sin(radians(beta)), cos(radians(beta))
  cos_phi = (1.0*d1*d2/sb**2) * (1.0*h1*h2/a**2 + 1.0*k1*k2*sb**2/b**2 + 1.0*l1*l2/c**2 - 1.0*(l1*h2+l2*h1)*cb/(a*c))
  return degrees(acos(cos_phi))

def calc_plane_angle_triclinic(p1=None, p2=None, a=None, b=None, c=None, alpha=None, beta=None, gamma=None):
  """Supply two sets of miller indices, a, b, c, alpha, beta and gamma to calculate the angle in degrees between each

  Usage: calc_plane_angle_triclinic([h,k,l], [h,k,l], a, b, c, alpha, beta, gamma)"""
  h1, k1, l1 = parse_miller_indices(p1)
  h2, k2, l2 = parse_miller_indices(p2)
  S11, S22, S33, S12, S23, S13 = _S(a, b, c, alpha, beta, gamma)
  d1, d2 = calc_plane_spacing_triclinic(p1, a, b, c, alpha, beta, gamma), \
    calc_plane_spacing_triclinic(p2, a, b, c, alpha, beta, gamma)
  V = calc_volume_triclinic(a, b, c, alpha, beta, gamma)
  cos_phi = (1.0*d1*d2/V**2) * (S11*h1*h2 + S22*k1*k2 + S33*l1*l2 + S23*(k1*l2+k2*l1) + S13*(l1*h2+l2*h1) + S12*(h1*k2+h2*k1))
  return degrees(acos(cos_phi))

# Class based functionality

class _Lattice(object):
  """Parent class for lattice objects"""
  def set_lattice_param(self, val, param):
    param = str(param)
    if hasattr(self, param):
      try:
        val = float(val)
      except:
        raise ValueError, 'Lattice parameters must be convertable to a float'
      setattr(self, val, param)
    else:
      raise ValueError, 'Lattice of type %s does not require lattice parameter %s' % (self.type, param)
  def get_lattice_param(self, param):
    param = str(param)
    if hasattr(self, param):
      return getattr(self, param)
    else:
      raise ValueError, 'Lattice of type %s does not have lattice parameter %s' % (self.type, param)

class CubicLattice(_Lattice):
  """Create cubic lattice object, pass a on creating an instance"""
  def __init__(self, a=None):
    self.a = float(a)
  def get_description(self):
    return '%s (a=%.3f)' % (self.type, self.a)
  description = property(get_description)
  def get_type(self):
    return 'cubic'
  type = property(get_type)
  def get_cell_volume(self):
    return calc_volume_cubic(self.a)
  cell_volume = property(get_cell_volume)
  def get_angle_between_planes(self, p1, p2):
    return calc_plane_angle_cubic(p1, p2)

class TetragonalLattice(_Lattice):
  """Create tetragonal lattice object, pass a and c on creating an instance"""
  def __init__(self, a=None, c=None):
    self.a = float(a)
    self.c = float(c)
  def get_description(self):
    return '%s (a=%.3f, c=%.3f)' % (self.type, self.a, self.c)
  description = property(get_description)
  def get_type(self):
    return 'tetragonal'
  type = property(get_type)
  def get_cell_volume(self):
    return calc_volume_tetragonal(self.a, self.c)
  cell_volume = property(get_cell_volume)
  def get_angle_between_planes(self, p1, p2):
    return calc_plane_angle_tetragonal(p1, p2, self.a, self.c)

class HexagonalLattice(_Lattice):
  """Create hexagonal lattice object, pass a and c on creating an instance"""
  def __init__(self, a=None, c=None):
    self.a = float(a)
    self.c = float(c)
  def get_description(self):
    return '%s (a=%.3f, c=%.3f)' % (self.type, self.a, self.c)
  description = property(get_description)
  def get_type(self):
    return 'hexagonal'
  type = property(get_type)
  def get_cell_volume(self):
    return calc_volume_hexagonal(self.a, self.c)
  cell_volume = property(get_cell_volume)
  def get_angle_between_planes(self, p1, p2):
    return calc_plane_angle_hexagonal(p1, p2, self.a, self.c)

class RhombohedralLattice(_Lattice):
  """Create rhombohedral lattice object, pass a and alpha on creating an instance"""
  def __init__(self, a=None, alpha=None):
    self.a = float(a)
    self.alpha = float(alpha)
  def get_description(self):
    return '%s (a=%.3f, alpha=%.3f)' % (self.type, self.a, self.alpha)
  description = property(get_description)
  def get_type(self):
    return 'rhombohedral'
  type = property(get_type)
  def get_cell_volume(self):
    return calc_volume_rhombohedral(self.a, self.alpha)
  cell_volume = property(get_cell_volume)
  def get_angle_between_planes(self, p1, p2):
    return calc_plane_angle_rhombohedral(p1, p2, self.a, self.alpha)

class OrthorhombicLattice(_Lattice):
  """Create orthorhombic lattice object, pass a, b and c on creating an instance"""
  def __init__(self, a=None, b=None, c=None):
    self.a = float(a)
    self.b = float(b)
    self.c = float(c)
  def get_description(self):
    return '%s (a=%.3f, b=%.3f, c=%.3f)' % (self.type, self.a, self.b, self.c)
  description = property(get_description)
  def get_type(self):
    return 'orthorhombic'
  type = property(get_type)
  def get_cell_volume(self):
    return calc_volume_orthorhombic(self.a, self.b, self.c)
  cell_volume = property(get_cell_volume)
  def get_angle_between_planes(self, p1, p2):
    return calc_plane_angle_orthorhombic(p1, p2, self.a, self.b, self.c)

class MonoclinicLattice(_Lattice):
  """Create monoclinic lattice object, pass a, b, c and beta on creating an instance"""
  def __init__(self, a=None, b=None, c=None, beta=None):
    self.a = float(a)
    self.b = float(b)
    self.c = float(c)
    self.beta = float(beta)
  def get_description(self):
    return '%s (a=%.3f, b=%.3f, c=%.3f, beta=%.3f)' % (self.type, self.a, self.b, self.c, self.beta)
  description = property(get_description)
  def get_type(self):
    return 'monoclinic'
  type = property(get_type)
  def get_cell_volume(self):
    return calc_volume_monoclinic(self.a, self.b, self.c, self.beta)
  cell_volume = property(get_cell_volume)
  def get_angle_between_planes(self, p1, p2):
    return calc_plane_angle_monoclinic(p1, p2, self.a, self.b, self.c, self.beta)

class TriclinicLattice(_Lattice):
  """Create triclinic lattice object, pass a, b, c, alpha, beta and gamma on creating an instance"""
  def __init__(self, a=None, b=None, c=None, alpha=None, beta=None, gamma=None):
    self.a = float(a)
    self.b = float(b)
    self.c = float(c)
    self.beta = float(beta)
    self.alpha = float(alpha)
    self.gamma = float(gamma)
  def get_description(self):
    return '%s (a=%.3f, b=%.3f, c=%.3f, alpha=%.3f, beta=%.3f, gamma=%.3f)' % (self.type, self.a, self.b, self.c, self.alpha, self.beta, self.gamma)
  description = property(get_description)
  def get_type(self):
    return 'triclinic'
  type = property(get_type)
  def get_cell_volume(self):
    return calc_volume_triclinic(self.a, self.b, self.c, self.alpha, self.beta, self.gamma)
  cell_volume = property(get_cell_volume)
  def get_angle_between_planes(self, p1, p2):
    return calc_plane_angle_triclinic(p1, p2, self.a, self.b, self.c, self.alpha, self.beta, self.gamma)