# ======================================================================
# matscipy - Python materials science tools
# https://github.com/libAtoms/matscipy
#
# Copyright (2014) James Kermode, King's College London
#                  Lars Pastewka, Karlsruhe Institute of Technology
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# ======================================================================

import itertools
import warnings

import numpy as np
from numpy.linalg import inv, norm
try:
    import scipy.stats as scipy_stats
except:
    scipy_stats = None

# import ase.units as units
# from ase.atoms import Atoms

# ##

# The indices of the full stiffness matrix of (orthorhombic) interest
Voigt_notation = [(0, 0), (1, 1), (2, 2), (1, 2), (0, 2), (0, 1)]

def full_3x3_to_Voigt_6_index(i, j):
    if i == j:
        return i
    return 6 - i - j

def Voigt_6_to_full_3x3_strain(strain_vector):
    """
    Form a 3x3 strain matrix from a 6 component vector in Voigt notation
    """
    e1, e2, e3, e4, e5, e6 = np.transpose(strain_vector)
    return np.transpose([[1.0 + e1, 0.5 * e6, 0.5 * e5],
                         [0.5 * e6, 1.0 + e2, 0.5 * e4],
                         [0.5 * e5, 0.5 * e4, 1.0 + e3]])


def Voigt_6_to_full_3x3_stress(stress_vector):
    """
    Form a 3x3 stress matrix from a 6 component vector in Voigt notation
    """
    s1, s2, s3, s4, s5, s6 = np.transpose(stress_vector)
    return np.transpose([[s1, s6, s5],
                         [s6, s2, s4],
                         [s5, s4, s3]])


def full_3x3_to_Voigt_6_strain(strain_matrix):
    """
    Form a 6 component strain vector in Voigt notation from a 3x3 matrix
    """
    strain_matrix = np.asarray(strain_matrix)
    return np.transpose([strain_matrix[..., 0, 0] - 1.0,
                         strain_matrix[..., 1, 1] - 1.0,
                         strain_matrix[..., 2, 2] - 1.0,
                         strain_matrix[..., 1, 2] + strain_matrix[..., 2, 1],
                         strain_matrix[..., 0, 2] + strain_matrix[..., 2, 0],
                         strain_matrix[..., 0, 1] + strain_matrix[..., 1, 0]])


def full_3x3_to_Voigt_6_stress(stress_matrix):
    """
    Form a 6 component stress vector in Voigt notation from a 3x3 matrix
    """
    stress_matrix = np.asarray(stress_matrix)
    return np.transpose([stress_matrix[..., 0, 0],
                         stress_matrix[..., 1, 1],
                         stress_matrix[..., 2, 2],
                         (stress_matrix[..., 1, 2] + stress_matrix[..., 1, 2]) / 2,
                         (stress_matrix[..., 0, 2] + stress_matrix[..., 0, 2]) / 2,
                         (stress_matrix[..., 0, 1] + stress_matrix[..., 0, 1]) / 2])


def Voigt_6x6_to_full_3x3x3x3(C):
    """
    Convert from the Voigt representation of the stiffness matrix to the full
    3x3x3x3 representation.

    Parameters
    ----------
    C : array_like
        6x6 stiffness matrix (Voigt notation).

    Returns
    -------
    C : array_like
        3x3x3x3 stiffness matrix.
    """

    C = np.asarray(C)
    C_out = np.zeros((3, 3, 3, 3), dtype=float)
    for i, j, k, l in itertools.product(list(range(3)), list(range(3)), list(range(3)), list(range(3))):
        Voigt_i = full_3x3_to_Voigt_6_index(i, j)
        Voigt_j = full_3x3_to_Voigt_6_index(k, l)
        C_out[i, j, k, l] = C[Voigt_i, Voigt_j]
    return C_out


def full_3x3x3x3_to_Voigt_6x6(C):
    """
    Convert from the full 3x3x3x3 representation of the stiffness matrix
    to the representation in Voigt notation. Checks symmetry in that process.
    """

    tol = 1e-3

    C = np.asarray(C)
    Voigt = np.zeros((6, 6))
    for i in range(6):
        for j in range(6):
            k, l = Voigt_notation[i]
            m, n = Voigt_notation[j]
            Voigt[i, j] = C[k, l, m, n]

            # print '---'
            # print k,l,m,n, C[k,l,m,n]
            # print m,n,k,l, C[m,n,k,l]
            # print l,k,m,n, C[l,k,m,n]
            # print k,l,n,m, C[k,l,n,m]
            # print m,n,l,k, C[m,n,l,k]
            # print n,m,k,l, C[n,m,k,l]
            # print l,k,n,m, C[l,k,n,m]
            # print n,m,l,k, C[n,m,l,k]
            # print '---'

            # Check symmetries
            assert abs(Voigt[i, j] - C[m, n, k, l]) < tol, \
                'Voigt[{},{}] = {}, C[{},{},{},{}] = {}' \
                .format(i, j, Voigt[i, j], m, n, k, l, C[m, n, k, l])
            assert abs(Voigt[i, j] - C[l, k, m, n]) < tol, \
                'Voigt[{},{}] = {}, C[{},{},{},{}] = {}' \
                .format(i, j, Voigt[i, j], k, l, m, n, C[l, k, m, n])
            assert abs(Voigt[i, j] - C[k, l, n, m]) < tol, \
                'Voigt[{},{}] = {}, C[{},{},{},{}] = {}' \
                .format(i, j, Voigt[i, j], k, l, n, m, C[k, l, n, m])
            assert abs(Voigt[i, j] - C[m, n, l, k]) < tol, \
                'Voigt[{},{}] = {}, C[{},{},{},{}] = {}' \
                .format(i, j, Voigt[i, j], m, n, l, k, C[m, n, l, k])
            assert abs(Voigt[i, j] - C[n, m, k, l]) < tol, \
                'Voigt[{},{}] = {}, C[{},{},{},{}] = {}' \
                .format(i, j, Voigt[i, j], n, m, k, l, C[n, m, k, l])
            assert abs(Voigt[i, j] - C[l, k, n, m]) < tol, \
                'Voigt[{},{}] = {}, C[{},{},{},{}] = {}' \
                .format(i, j, Voigt[i, j], l, k, n, m, C[l, k, n, m])
            assert abs(Voigt[i, j] - C[n, m, l, k]) < tol, \
                'Voigt[{},{}] = {}, C[{},{},{},{}] = {}' \
                .format(i, j, Voigt[i, j], n, m, l, k, C[n, m, l, k])

    return Voigt


def Voigt_6x6_to_cubic(C):
    """
    Convert the Voigt 6x6 representation into the cubic elastic constants
    C11, C12 and C44.
    """

    tol = 1e-6

    C_check = np.zeros_like(C)
    C_check[np.diag_indices_from(C_check)] = C[np.diag_indices_from(C)]
    C_check[0:3, 0:3] = C[0:3, 0:3]
    if np.any(np.abs(C - C_check) > tol):
        raise ValueError('"C" does not have cubic symmetry.')

    C11s = np.array([C[0, 0], C[1, 1], C[2, 2]])
    C12s = np.array([C[1, 2], C[0, 2], C[0, 1]])
    C44s = np.array([C[3, 3], C[4, 4], C[5, 5]])

    C11 = np.mean(C11s)
    C12 = np.mean(C12s)
    C44 = np.mean(C44s)

    if np.any(np.abs(C11 - C11s) > tol) or np.any(np.abs(C12 - C12s) > tol) or \
            np.any(np.abs(C44 - C44s) > tol):
        raise ValueError('"C" does not have cubic symmetry.')

    return np.array([C11, C12, C44])


def cubic_to_Voigt_6x6(C11, C12, C44):
    return np.array([[C11, C12, C12, 0, 0, 0],
                     [C12, C11, C12, 0, 0, 0],
                     [C12, C12, C11, 0, 0, 0],
                     [0, 0, 0, C44, 0, 0],
                     [0, 0, 0, 0, C44, 0],
                     [0, 0, 0, 0, 0, C44]])

def _invariants(s, syy=None, szz=None, syz=None, sxz=None, sxy=None,
               full_3x3_to_Voigt_6=full_3x3_to_Voigt_6_stress):
    """
    Receives a list of stress tensors and returns the three tensor invariants.
    """
    if syy is None:
        s = np.asarray(s)
        if s.shape == (6,):
            s = s.reshape(1, -1)
        elif s.shape == (3, 3):
            s = full_3x3_to_Voigt_6(s)
        elif s.shape[-1] == 3 and s.shape[-2] == 3:
            s = full_3x3_to_Voigt_6(s)
    else:
        s = np.transpose([np.transpose(s),
                          np.transpose(syy),
                          np.transpose(szz),
                          np.transpose(syz),
                          np.transpose(sxz),
                          np.transpose(sxy)])
    I1 = s[..., 0] + s[..., 1] + s[..., 2]
    I2 = s[..., 0] * s[..., 1] + s[..., 1] * s[..., 2] + s[..., 2] * s[..., 0] - s[..., 3] ** 2 - \
        s[..., 4] ** 2 - s[..., 5] ** 2
    I3 = s[..., 0] * s[..., 1] * s[..., 2] + 2 * s[..., 3] * s[..., 4] * s[..., 5] - \
        s[..., 3] ** 2 * s[..., 2] - s[..., 4] ** 2 * s[..., 0] - s[..., 5] ** 2 * s[..., 1]

    return I1, I2, I3

def invariants(s, syy=None, szz=None, syz=None, sxz=None, sxy=None,
               full_3x3_to_Voigt_6=full_3x3_to_Voigt_6_stress):
    I1, I2, I3 = _invariants(s, syy=syy, szz=szz, syz=syz, sxz=sxz, sxy=sxy,
                             full_3x3_to_Voigt_6=full_3x3_to_Voigt_6)

    J2 = I1 ** 2 / 3 - I2
    J3 = 2 * I1 ** 3 / 27 - I1 * I2 / 3 + I3

    # Return hydrostatic pressure, octahedral shear stress and J3
    return -I1 / 3, np.sqrt(2 * J2 / 3), J3

# ##

def rotate_cubic_elastic_constants(C11, C12, C44, A, tol=1e-6):
    """
    Return rotated elastic moduli for a cubic crystal given the elastic
    constant in standard C11, C12, C44 notation.

    Parameters
    ----------
    C11, C12, C44 : float
        Cubic elastic moduli.
    A : array_like
        3x3 rotation matrix.

    Returns
    -------
    C : array
        6x6 matrix of rotated elastic constants (Voigt notation).
    """

    A = np.asarray(A)

    # Is this a rotation matrix?
    if np.sometrue(np.abs(np.dot(np.array(A), np.transpose(np.array(A))) -
                          np.eye(3, dtype=float)) > tol):
        raise RuntimeError('Matrix *A* does not describe a rotation.')

    # Invariant elastic constants
    la = C12
    mu = C44
    al = C11 - la - 2 * mu

    # Construct rotated C in Voigt notation
    C = []
    for i, j in Voigt_notation:
        for k, l in Voigt_notation:
            h = 0.0
            if i == j and k == l:
                h += la
            if i == k and j == l:
                h += mu
            if i == l and j == k:
                h += mu
            h += al * np.sum(A[i, :] * A[j, :] * A[k, :] * A[l, :])
            C += [h]

    C = np.array(C)
    C.shape = (6, 6)
    return C

# ##

def rotate_elastic_constants(C, A, tol=1e-6):
    """
    Return rotated elastic moduli for a general crystal given the elastic
    constant in Voigt notation.

    Parameters
    ----------
    C : array_like
        6x6 matrix of elastic constants (Voigt notation).
    A : array_like
        3x3 rotation matrix.

    Returns
    -------
    C : array
        6x6 matrix of rotated elastic constants (Voigt notation).
    """

    A = np.asarray(A)

    # Is this a rotation matrix?
    if np.sometrue(np.abs(np.dot(np.array(A), np.transpose(np.array(A))) -
                          np.eye(3, dtype=float)) > tol):
        raise RuntimeError('Matrix *A* does not describe a rotation.')

    # Rotate
    return full_3x3x3x3_to_Voigt_6x6(np.einsum('ia,jb,kc,ld,abcd->ijkl',
                                               A, A, A, A,
                                               Voigt_6x6_to_full_3x3x3x3(C)))
# 
# # ##
# 
# class CubicElasticModuli:
#     tol = 1e-6
# 
#     def __init__(self, C11, C12, C44):
#         """
#         Initialize a cubic system with elastic constants C11, C12, C44
#         """
# 
#         warnings.warn('CubicElasticModuli is deprecated. Use '
#                       'rotate_elastic_constants function instead.')
# 
#         # la, mu, al are the three invariant elastic constants
#         self.la = C12
#         self.mu = C44
#         self.al = C11 - self.la - 2 * self.mu
# 
#         A = np.eye(3, dtype=float)
# 
#         # Compute initial stiffness matrix
#         self.rotate(A)
# 
# 
#     def rotate(self, A):
#         """
#         Compute the rotated stiffness matrix
#         """
# 
#         A = np.asarray(A)
# 
#         # Is this a rotation matrix?
#         if np.sometrue(np.abs(np.dot(np.array(A), np.transpose(np.array(A))) - 
#                               np.eye(3, dtype=float)) > self.tol):
#             raise RuntimeError('Matrix *A* does not describe a rotation.')
# 
#         C = [ ]
#         for i, j in Voigt_notation:
#             for k, l in Voigt_notation:
#                 h = 0.0
#                 if i == j and k == l:
#                     h += self.la
#                 if i == k and j == l:
#                     h += self.mu
#                 if i == l and j == k:
#                     h += self.mu
#                 h += self.al * np.sum(A[i, :] * A[j, :] * A[k, :] * A[l, :])
#                 C += [ h ]
# 
#         self.C = np.asarray(C)
#         self.C.shape = (6, 6)
#         return self.C
# 
# 
#     def _rotate_explicit(self, A):
#         """
#         Compute the rotated stiffness matrix by applying the rotation to the
#         full stiffness matrix. This function is for debugging purposes only.
#         """
# 
#         A = np.asarray(A)
# 
#         # Is this a rotation matrix?
#         if np.sometrue(np.abs(np.dot(np.array(A), np.transpose(np.array(A))) - 
#                               np.eye(3, dtype=float)) > self.tol):
#             raise RuntimeError('Matrix *A* does not describe a rotation.')
# 
#         C = np.zeros((3, 3, 3, 3), dtype=float)
# 
#         # Construct unrotated stiffness matrix
#         for i in range(3):
#             for j in range(3):
#                 for k in range(3):
#                     for m in range(3):
#                         h = 0.0
#                         if i == j and k == m:
#                             h += self.la
#                         if i == k and j == m:
#                             h += self.mu
#                         if i == m and j == k:
#                             h += self.mu
#                         if i == j and j == k and k == m:
#                             h += self.al
#                         C[i, j, k, m] = h
# 
#         # Rotate
#         C = np.einsum('ia,jb,kc,ld,abcd->ijkl', A, A, A, A, C)
# 
#         self.C = full_3x3x3x3_to_Voigt_6x6(C)
#         return self.C
# 
# 
#     def stiffness(self):
#         """
#         Return the elastic constants
#         """
# 
#         return self.C
# 
# 
#     def compliance(self):
#         """
#         Return the compliance coefficients
#         """
# 
#         return inv(self.C)
# 
# # ##
# 
# def measure_triclinic_elastic_constants(a, delta=0.001, optimizer=None,
#                                         logfile=None, **kwargs):
#     """
#     Brute-force measurement of elastic constants for a triclinic (general)
#     unit cell.
# 
#     Parameters
#     ----------
#     a : ase.Atoms
#         Atomic configuration.
#     optimizer : ase.optimizer.*
#         Optimizer to use for atomic position. Does not optimize atomic
#         position if set to None.
#     delta : float
#         Strain increment for analytical derivatives of stresses.
# 
#     Returns
#     -------
#     C : array_like
#         6x6 matrix of the elastic constants in Voigt notation.
#     """
# 
#     if optimizer is not None:
#         optimizer(a, logfile=logfile).run(**kwargs)
# 
#     r0 = a.positions.copy()
# 
#     cell = a.cell.copy()
#     volume = a.get_volume()
# 
#     C = np.zeros((3, 3, 3, 3), dtype=float)
#     for i in range(3):
#         for j in range(3):
#             a.set_cell(cell, scale_atoms=True)
#             a.set_positions(r0)
#         
#             D = np.eye(3)
#             D[i, j] += 0.5 * delta
#             D[j, i] += 0.5 * delta
#             a.set_cell(np.dot(D, cell.T).T, scale_atoms=True)
#             if optimizer is not None:
#                 optimizer(a, logfile=logfile).run(**kwargs)
#             sp = Voigt_6_to_full_3x3_stress(a.get_stress() * a.get_volume())
# 
#             D = np.eye(3)
#             D[i, j] -= 0.5 * delta
#             D[j, i] -= 0.5 * delta
#             a.set_cell(np.dot(D, cell.T).T, scale_atoms=True)
#             if optimizer is not None:
#                 optimizer(a, logfile=logfile).run(**kwargs)
#             sm = Voigt_6_to_full_3x3_stress(a.get_stress() * a.get_volume())
# 
#             C[:, :, i, j] = (sp - sm) / (2 * delta * volume)
# 
#     a.set_cell(cell, scale_atoms=True)
#     a.set_positions(r0)
# 
#     return full_3x3x3x3_to_Voigt_6x6(C)
# 
# 
# 
# Cij_symmetry = {
#    'cubic':           np.array([[1, 7, 7, 0, 0, 0],
#                                 [7, 1, 7, 0, 0, 0],
#                                 [7, 7, 1, 0, 0, 0],
#                                 [0, 0, 0, 4, 0, 0],
#                                 [0, 0, 0, 0, 4, 0],
#                                 [0, 0, 0, 0, 0, 4]]),
# 
#    'trigonal_high':   np.array([[1, 7, 8, 9, 10, 0],
#                                 [7, 1, 8, 0, -9, 0],
#                                 [8, 8, 3, 0, 0, 0],
#                                 [9, -9, 0, 4, 0, 0],
#                                 [10, 0, 0, 0, 4, 0],
#                                 [0, 0, 0, 0, 0, 6]]),
# 
#    'trigonal_low':    np.array([[1, 7, 8, 9, 10, 0 ],
#                                 [7, 1, 8, -9, -10, 0 ],
#                                 [8, 8, 3, 0, 0, 0 ],
#                                 [9, -9, 0, 4, 0, -10],
#                                 [10, -10, 0, 0, 4, 9 ],
#                                 [0, 0, 0, -10 , 9, 6 ]]),
# 
#    'tetragonal_high': np.array([[1, 7, 8, 0, 0, 0],
#                                 [7, 1, 8, 0, 0, 0],
#                                 [8, 8, 3, 0, 0, 0],
#                                 [0, 0, 0, 4, 0, 0],
#                                 [0, 0, 0, 0, 4, 0],
#                                 [0, 0, 0, 0, 0, 6]]),
# 
#    'tetragonal_low':  np.array([[1, 7, 8, 0, 0, 11],
#                                 [7, 1, 8, 0, 0, -11],
#                                 [8, 8, 3, 0, 0, 0],
#                                 [0, 0, 0, 4, 0, 0],
#                                 [0, 0, 0, 0, 4, 0],
#                                 [11, -11, 0, 0, 0, 6]]),
# 
#    'orthorhombic':    np.array([[ 1, 7, 8, 0, 0, 0],
#                                 [ 7, 2, 12, 0, 0, 0],
#                                 [ 8, 12, 3, 0, 0, 0],
#                                 [ 0, 0, 0, 4, 0, 0],
#                                 [ 0, 0, 0, 0, 5, 0],
#                                 [ 0, 0, 0, 0, 0, 6]]),
# 
#    'monoclinic':      np.array([[ 1, 7, 8, 0, 10, 0],
#                                 [ 7, 2, 12, 0, 14, 0],
#                                 [ 8, 12, 3, 0, 17, 0],
#                                 [ 0, 0, 0, 4, 0, 20],
#                                 [10, 14, 17, 0, 5, 0],
#                                 [ 0, 0, 0, 20, 0, 6]]),
#     
#     'triclinic':       np.array([[ 1, 7, 8, 9, 10, 11],
#                                  [ 7, 2, 12, 13, 14, 15],
#                                  [ 8, 12, 3, 16, 17, 18],
#                                  [ 9, 13, 16, 4, 19, 20],
#                                  [10, 14, 17, 19, 5, 21],
#                                  [11, 15, 18, 20, 21, 6 ]]),
#    }
# 
# def _dec(pattern):
#     # Decrease C_ij indices patterns by one.
#     # Used to make strain_patterns definitions below more readable.
#     return [(i - 1, j - 1) for (i, j) in pattern]
# 
# strain_patterns = {
# 
#    'cubic': [
#       # strain pattern e1+e4, yields C11, C21, C31 and C44, then C12 is average of C21 and C31
#       [ np.array([1, 0, 0, 1, 0, 0]), _dec([(1, 1), (2, 1), (3, 1), (4, 4)])]
#    ],
# 
#    'trigonal_high': [
#       # strain pattern e3 yield C13, C23 and C33
#       [ np.array([0, 0, 1, 0, 0, 0]), _dec([(1, 3), (2, 3), (3, 3)])],
# 
#       # strain pattern e1+e4 yields C11 C21 C31 and C44
#       [ np.array([1, 0, 0, 1, 0, 0]), _dec([(1, 1), (2, 1), (3, 1), (4, 4)])],
# 
#       # strain pattern e1 yields C11 C21 C31 C41 C51
#       [ np.array([1, 0, 0, 0, 0, 0]), _dec([(1, 1), (2, 1), (3, 1), (4, 1), (5, 1)])],
# 
#       # strain pattern e3+e4
#       [ np.array([0, 0, 1, 1, 0, 0]), _dec([(3, 3), (4, 4)])]
# 
#    ],
# 
#    'trigonal_low': [
#      # strain pattern e1, yields C11, C21, C31, C41, C51
#      [ np.array([1, 0, 0, 0, 0, 0]), _dec([(1, 1), (2, 1), (3, 1), (4, 1), (5, 1)])],
# 
#      # strain pattern e3 + e4, yields C33, C44
#      [ np.array([0, 0, 1, 1, 0, 0]), _dec([(3, 3), (4, 4)]) ],
# 
#      [ np.array([0, 0, 0, 0, 0, 1]), _dec([(6, 6)]) ]
#    ],
# 
#    'tetragonal': [
#      # strain pattern e1+e4
#      [ np.array([1, 0, 0, 1, 0, 0]), _dec([(1, 1), (2, 1), (3, 1), (6, 1), (4, 4)]) ],
# 
#      # strain pattern e3+e6
#      [ np.array([0, 0, 1, 0, 0, 1]), _dec([(3, 3), (6, 6)]) ]
#    ],
# 
#    'orthorhombic': [
#       # strain pattern e1+e4
#       [ np.array([1, 0, 0, 1, 0, 0]), _dec([(1, 1), (2, 1), (3, 1), (4, 4)]) ],
# 
#       # strain pattern e2+e5
#       [ np.array([0, 1, 0, 0, 1, 0]), _dec([(1, 2), (2, 2), (3, 2), (5, 5)]) ],
# 
#       # strain pattern e3+e6
#       [ np.array([0, 0, 1, 0, 0, 1]), _dec([(1, 3), (2, 3), (3, 3), (6, 6)]) ]
#    ],
# 
#    'monoclinic': [
#       # strain pattern e1+e4
#       [ np.array([1, 0, 0, 1, 0, 0]), _dec([(1, 1), (2, 1), (3, 1), (4, 4), (5, 1), (6, 4)]) ],
# 
#       # strain pattern e3+e6
#       [ np.array([0, 0, 1, 0, 0, 1]), _dec([(1, 3), (2, 3), (3, 3), (5, 3), (4, 6), (6, 6)]) ],
# 
#       # strain pattern e2
#       [ np.array([0, 1, 0, 0, 0, 0]), _dec([(1, 2), (2, 2), (3, 2), (5, 2)]) ],
# 
#       # strain pattern e5
#       [ np.array([0, 0, 0, 0, 1, 0]), _dec([(1, 5), (2, 5), (3, 5), (5, 5)]) ]
#    ],
# 
#    'triclinic': [
#       [ np.array([1, 0, 0, 0, 0, 0]), _dec([(1, 1), (2, 1), (3, 1), (4, 1), (5, 1), (6, 1)])],
#       [ np.array([0, 1, 0, 0, 0, 0]), _dec([(1, 2), (2, 2), (3, 2), (4, 2), (5, 2), (6, 2)])],
#       [ np.array([0, 0, 1, 0, 0, 0]), _dec([(1, 3), (2, 3), (3, 3), (4, 3), (5, 3), (6, 3)])],
#       [ np.array([0, 0, 0, 1, 0, 0]), _dec([(1, 4), (2, 4), (3, 4), (4, 4), (5, 4), (6, 4)])],
#       [ np.array([0, 0, 0, 0, 1, 0]), _dec([(1, 5), (2, 5), (3, 5), (4, 5), (5, 5), (6, 5)])],
#       [ np.array([0, 0, 0, 0, 0, 1]), _dec([(1, 6), (2, 6), (3, 6), (4, 6), (5, 6), (6, 6)])],
#    ]
# 
#    }
# 
# Cij_symmetry['hexagonal'] = Cij_symmetry['trigonal_high']
# Cij_symmetry[None] = Cij_symmetry['triclinic']
# 
# strain_patterns['hexagonal'] = strain_patterns['trigonal_high']
# strain_patterns['tetragonal_high'] = strain_patterns['tetragonal_low'] = strain_patterns['tetragonal']
# strain_patterns[None] = strain_patterns['triclinic']
# 
# 
# def generate_strained_configs(at0, symmetry='triclinic', N_steps=5, delta=1e-2):
#     """
#     Generate a sequence of strained configurations
# 
#     Parameters
#     ----------
#     a : ase.Atoms
#         Bulk crystal configuration - both unit cell and atomic positions
#         should be relaxed before calling this routine.
#     symmetry : string
#         Symmetry to use to determine which strain patterns are necessary.
#         Default is 'triclininc', i.e. no symmetry.
#     N_steps : int
#         Number of atomic configurations to generate for each strain pattern.
#         Default is 5. Absolute strain values range from -delta*N_steps/2
#         to +delta*N_steps/2.
#     delta : float
#         Strain increment for analytical derivatives of stresses. Default 1e-2
# 
#     Returns
#     -------
#     Generator which yields a sequence of ase.Atoms instances corresponding
#     to the minima set of strained conifugurations required to evaluate the
#     full 6x6 C_ij matrix under the assumed symmetry.
#     """
# 
#     if not symmetry in strain_patterns:
#         raise ValueError('Unknown symmetry %s. Valid options are %s' % (symmetry, strain_patterns.keys()))
# 
#     for pindex, (pattern, fit_pairs) in enumerate(strain_patterns[symmetry]):
#         for step in range(N_steps):
#             strain = np.where(pattern == 1, delta * (step + 1 - (N_steps + 1) / 2.0), 0.0)
#             at = at0.copy()
#             if at0.get_calculator() is not None:
#                 at.set_calculator(at0.get_calculator())
#             T = Voigt_6_to_full_3x3_strain(strain)
#             at.set_cell(np.dot(T, at.cell.T).T, scale_atoms=False)
#             at.positions[:] = np.dot(T, at.positions.T).T
#             at.info['strain'] = T
#             yield at
# 
# 
# # Elastic constant calculation.
# 
# # Code adapted from elastics.py script, available from
# # http://github.com/djw/elastic-constants
# #
# # Copyright (c) 2008, Dan Wilson
# # All rights reserved.
# #
# # Redistribution and use in source and binary forms, with or without
# # modification, are permitted provided that the following conditions are met:
# #     * Redistributions of source code must retain the above copyright
# #       notice, this list of conditions and the following disclaimer.
# #     * Redistributions in binary form must reproduce the above copyright
# #       notice, this list of conditions and the following disclaimer in the
# #       documentation and/or other materials provided with the distribution.
# #     * Neither the name of the copyright holder nor the
# #       names of its contributors may be used to endorse or promote products
# #       derived from this software without specific prior written permission.
# #
# # THIS SOFTWARE IS PROVIDED BY DAN WILSON ''AS IS'' AND ANY
# # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# # DISCLAIMED. IN NO EVENT SHALL DAN WILSON BE LIABLE FOR ANY
# # DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# # (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# # LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# # ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# # SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# 
# def fit_elastic_constants(a, symmetry='triclinic', N_steps=5, delta=1e-2, optimizer=None,
#                           verbose=True, graphics=False, logfile=None, **kwargs):
#     """
#     Compute elastic constants by linear regression of stress vs. strain
# 
#     Parameters
#     ----------
#     a : ase.Atoms or list of ase.Atoms
#         Either a single atomic configuration or a list of strained
#         configurations. If a single configuration is given, it is
#         passed :func:`generate_strained_configs` along with `symmetry`, `N_steps`,
#         and `delta` to generate the set of strained configurations.
#     symmetry : string
#         Symmetry to use to determine which strain patterns are necessary.
#         Default is 'triclininc', i.e. no symmetry.
#     N_steps : int
#         Number of atomic configurations to generate for each strain pattern.
#         Default is 5. Absolute strain values range from -delta*N_steps/2
#         to +delta*N_steps/2.    
#     delta : float
#         Strain increment for analytical derivatives of stresses.
#         Default is 1e-2.
#     optimizer : ase.optimizer.*
#         Optimizer to use for atomic positions (internal degrees of freedom)
#         for each applied strain. Initial config `a` is not optimised, and should
#         already have relaxed cell and atomic positions. Does not optimize atomic
#         positions if set to None.
#     verbose : bool
#         If True, print additional infomation about the quality of fit
#         and summarise results of C_ij and estimated errors. Default True.
#     graphics : bool
#         If True, use :mod:`matplotlib.pyplot` to plot the stress vs. strain
#         curve for each C_ij component fitted. Default True. 
#     logfile : bool
#         Log file to write optimizer output to. Default None (i.e. suppress
#         output).
#     **kwargs : dict
#         Additional arguments to pass to `optimizer.run()` method e.g. `fmax`.
# 
#     Returns
#     -------
#     C : array_like
#         6x6 matrix of the elastic constants in Voigt notation.
#     C_err : array_like
#         If scipy.stats module is available then error estimates for each C_ij
#         component are obtained from the accuracy of the linear regression.
#         Otherwise an array of np.zeros((6,6)) is returned.
#     
#     Notes
#     -----
# 
#     Code originally adapted from elastics.py script, available from
#     http://github.com/djw/elastic-constants
#     """
# 
#     if isinstance(a, Atoms):
#         # we've been passed a single Atoms object: use it to generate
#         # set of strained configurations according to symmetry
#         strained_configs = generate_strained_configs(a, symmetry, N_steps, delta)
#     else:
#         # assume we've been passed a list of strained configs
#         strained_configs = a
# 
#     if graphics:
#         import matplotlib.pyplot as plt
# 
#     def do_fit(index1, index2, stress, strain, patt):
#         if verbose:
#             print('Fitting C_%d%d' % (index1 + 1, index2 + 1))
#             print('Strain %r' % strain[:, index2])
#             print('Stress %r GPa' % (stress[:, index1] / units.GPa))
# 
#         if scipy_stats is not None:
#             cijFitted, intercept, r, tt, stderr = scipy_stats.linregress(strain[:, index2], stress[:, index1])
#         else:
#             cijFitted, intercept = np.polyfit(strain[:, index2], stress[:, index1], 1)
#             r, tt, stderr = 0.0, None, 0.0
# 
#         if verbose:
#             # print info about the fit
#             print('Cij (gradient) / GPa    :    ', cijFitted / units.GPa)
#             if scipy_stats is not None:
#                 print('Error in Cij / GPa      :    ', stderr / units.GPa)
#                 if abs(r) > 0.9:
#                     print('Correlation coefficient :    ', r)
#                 else:
#                     print('Correlation coefficient :    ', r, '     <----- WARNING')
#             else:
#                 print('scipy.stats not available, no error estimation performed')
# 
#         if graphics:
#             # position this plot in a 6x6 grid
#             sp = plt.subplot(6, 6, 6 * index1 + (index2 + 1))
#             sp.set_axis_on()
# 
#             # change the labels on the axes
#             xlabels = sp.get_xticklabels()
#             plt.setp(xlabels, 'rotation', 90, fontsize=7)
#             ylabels = sp.get_yticklabels()
#             plt.setp(ylabels, fontsize=7)
# 
#             # colour the plot depending on the strain pattern
#             colourDict = {0: '#BAD0EF', 1:'#FFCECE', 2:'#BDF4CB', 3:'#EEF093', 4:'#FFA4FF', 5:'#75ECFD'}
#             sp.set_axis_bgcolor(colourDict[patt])
# 
#             # plot the data
#             plt.plot([strain[0, index2], strain[-1, index2]],
#                      [(cijFitted * strain[0, index2] + intercept) / units.GPa,
#                       (cijFitted * strain[-1, index2] + intercept) / units.GPa])
#             plt.plot(strain[:, index2], stress[:, index1] / units.GPa, 'ro')
#             plt.xticks(strain[:, index2])
# 
#         return cijFitted, stderr
# 
#     if not symmetry in strain_patterns:
#         raise ValueError('Unknown symmetry %s. Valid options are %s' % (symmetry, strain_patterns.keys()))
# 
#     # There are 21 independent elastic constants
#     Cijs = {}
#     Cij_err = {}
# 
#     # Construct mapping from (i,j) to index into Cijs in range 1..21
#     # (upper triangle only to start with)
#     Cij_map = {}
#     Cij_map_sym = {}
#     for i in range(6):
#         for j in range(i, 6):
#             Cij_map[(i, j)] = Cij_symmetry[None][i, j]
#             Cij_map_sym[(i, j)] = Cij_symmetry[symmetry][i, j]
# 
#     # Reverse mapping, index 1..21 -> tuple (i,j) with i, j in range 0..5
#     Cij_rev_map = dict(zip(Cij_map.values(), Cij_map.keys()))
# 
#     # Add the lower triangle to Cij_map, e.g. C21 = C12
#     for (i1, i2) in Cij_map.copy().keys():
#         Cij_map[(i2, i1)] = Cij_map[(i1, i2)]
#         Cij_map_sym[(i2, i1)] = Cij_map_sym[(i1, i2)]
# 
# 
#     N_pattern = len(strain_patterns[symmetry])
#     configs = iter(strained_configs)
# 
#     strain = np.zeros((N_pattern, N_steps, 6))
#     stress = np.zeros((N_pattern, N_steps, 6))
# 
#     if graphics:
#         fig = plt.figure(num=1, figsize=(9.5, 8), facecolor='white')
#         fig.clear()
#         fig.subplots_adjust(left=0.07, right=0.97, top=0.97, bottom=0.07, wspace=0.5, hspace=0.5)
# 
#         for index1 in range(6):
#             for index2 in range(6):
#                 # position this plot in a 6x6 grid
#                 sp = plt.subplot(6, 6, 6 * (index1) + index2 + 1)
#                 sp.set_axis_off()
#                 plt.text(0.4, 0.4, "n/a")
# 
#     # Fill in strain and stress arrays from config Atoms list
#     for pattern_index, (pattern, fit_pairs) in enumerate(strain_patterns[symmetry]):
#         for step in range(N_steps):
#             at = next(configs)
#             if optimizer is not None:
#                 optimizer(at, logfile=logfile).run(**kwargs)
#             strain[pattern_index, step, :] = full_3x3_to_Voigt_6_strain(at.info['strain'])
#             stress[pattern_index, step, :] = at.get_stress()
# 
#     # Do the linear regression
#     for pattern_index, (pattern, fit_pairs) in enumerate(strain_patterns[symmetry]):
#         for (index1, index2) in fit_pairs:
#             fitted, err = do_fit(index1, index2,
#                                  stress[pattern_index, :, :],
#                                  strain[pattern_index, :, :],
#                                  pattern_index)
# 
#             index = abs(Cij_map_sym[(index1, index2)])
# 
#             if not index in Cijs:
#                 if verbose:
#                     print('Setting C%d%d (%d) to %f +/- %f' % 
#                           (index1 + 1, index2 + 1, index, fitted, err))
#                 Cijs[index] = [fitted]
#                 Cij_err[index] = [err]
#             else:
#                 if verbose:
#                     print('Updating C%d%d (%d) with value %f +/- %f' % 
#                           (index1 + 1, index2 + 1, index, fitted, err))
#                 Cijs[index].append(fitted)
#                 Cij_err[index].append(err)
#             if verbose:
#                 print('\n')
# 
# 
#     C = np.zeros((6, 6))
#     C_err = np.zeros((6, 6))
#     C_labels = np.zeros((6, 6), dtype='S4')
#     C_labels[:] = '    '
# 
#     # Convert lists to mean
#     for k in Cijs:
#         Cijs[k] = np.mean(Cijs[k])
# 
#     # Combine statistical errors
#     for k, v in Cij_err.items():
#         Cij_err[k] = np.sqrt(np.sum(np.array(v) ** 2)) / np.sqrt(len(v))
# 
#     if symmetry.startswith('trigonal'):
#         # Special case for trigonal lattice: C66 = (C11 - C12)/2
#         Cijs[Cij_map[(5, 5)]] = 0.5 * (Cijs[Cij_map[(0, 0)]] - Cijs[Cij_map[(0, 1)]])
#         Cij_err[Cij_map[(5, 5)]] = np.sqrt(Cij_err[Cij_map[(0, 0)]] ** 2 + Cij_err[Cij_map[(0, 1)]] ** 2)
# 
#     # Generate the 6x6 matrix of elastic constants
#     # - negative values signify a symmetry relation
#     for i in range(6):
#         for j in range(6):
#             index = Cij_symmetry[symmetry][i, j]
#             if index > 0:
#                 C[i, j] = Cijs[index]
#                 C_err[i, j] = Cij_err[index]
#                 ii, jj = Cij_rev_map[index]
#                 C_labels[i, j] = ' C%d%d' % (ii + 1, jj + 1)
#                 C_err[i, j] = Cij_err[index]
#             elif index < 0:
#                 C[i, j] = -Cijs[-index]
#                 C_err[i, j] = Cij_err[-index]
#                 ii, jj = Cij_rev_map[-index]
#                 C_labels[i, j] = '-C%d%d' % (ii + 1, jj + 1)
# 
#     if verbose:
#         print(np.array2string(C_labels).replace("'", ""))
#         print('\n = \n')
#         print(np.array2string(C / units.GPa,
#                               suppress_small=True,
#                               precision=2))
#         print
# 
#         # Summarise the independent components of C_ij matrix
#         printed = {}
#         for i in range(6):
#             for j in range(6):
#                 index = Cij_symmetry[symmetry][i, j]
#                 if index <= 0 or index in printed: continue
#                 print('C_%d%d = %-4.2f +/- %-4.2f GPa' % 
#                       (i + 1, j + 1, C[i, j] / units.GPa, C_err[i, j] / units.GPa))
#                 printed[index] = 1
# 
#     return C, C_err
# 
# 
# def youngs_modulus(C, l):
#     """
#     Calculate approximate Youngs modulus E_l from 6x6 elastic constants matrix C_ij
# 
#     This is the modulus for loading in the l direction. For the exact answer, taking
#     into account elastic anisotropuy, rotate the C_ij matrix to the correct frame,
#     compute the compliance matrix::
# 
#        C = ...  # 6x6 C_ij matrix in crystal frame
#        A = ...  # rotation matrix
#        Cr = rotate_elastic_constants(C, A)
#        S = np.inv(Cr)
#        E_x = 1/S[0, 0]  # Young's modulus for a pull in x direction
#        E_y = 1/S[1, 1]  # Young's modulus for a pull in y direction
#        E_z = 1/S[0, 0]  # Young's modulus for a pull in z direction
# 
#     Notes
#     -----
# 
#     Formula is from W. Brantley, Calculated elastic constants for stress problems associated 
#     with semiconductor devices. J. Appl. Phys., 44, 534 (1973).
#     """  
# 
#     S = inv(C)  # Compliance matrix
#     lhat = l / norm(l)  # Normalise directions
# 
#     # Youngs modulus in direction l, ratio of stress sigma_l 
#     # to strain response epsilon_l
#     E = 1.0 / (S[0, 0] - 2.0 * (S[0, 0] - S[0, 1] - 0.5 * S[3, 3]) * (lhat[0] * lhat[0] * lhat[1] * lhat[1] + 
#          lhat[1] * lhat[1] * lhat[2] * lhat[2] + 
#          lhat[0] * lhat[0] * lhat[2] * lhat[2]))
#     return E
# 
# 
# def poisson_ratio(C, l, m):
#     """
#     Calculate approximate Poisson ratio \nu_{lm} from 6x6 elastic constant matrix C_{ij}
# 
#     This is the response in `m` direction to pulling in `l` direction. Result is dimensionless.
# 
#     Notes
#     -----
# 
#     Formula is from W. Brantley, Calculated elastic constants for stress problems associated 
#     with semiconductor devices. J. Appl. Phys., 44, 534 (1973).
#     """
#     
#     S = inv(C)  # Compliance matrix
#     lhat = l / norm(l)  # Normalise directions
#     mhat = m / norm(m)
# 
#     # Poisson ratio v_lm: response in m direction to strain in 
#     # l direction, v_lm = - epsilon_m/epsilon_l
#     v = -((S[0, 1] + (S[0, 0] - S[0, 1] - 0.5 * S[3, 3]) * (lhat[0] * lhat[0] * mhat[0] * mhat[0] + 
#          lhat[1] * lhat[1] * mhat[1] * mhat[1] + 
#          lhat[2] * lhat[2] * mhat[2] * mhat[2])) / 
#          (S[0, 0] - 2.0 * (S[0, 0] - S[0, 1] - 0.5 * S[3, 3]) * (lhat[0] * lhat[0] * lhat[1] * lhat[1] + 
#          lhat[1] * lhat[1] * lhat[2] * lhat[2] + 
#          lhat[0] * lhat[0] * lhat[2] * lhat[2])))
#     return v
# 
# 
# def elastic_moduli(C, l=np.array([1, 0, 0]), R=None, tol=1e-6):
#     """
#     Calculate elastic moduli from 6x6 elastic constant matrix C_{ij}.
#     
#     The elastic moduli calculated are: Young's muduli, Poisson's ratios,
#     shear moduli, bulk mudulus and bulk mudulus tensor.
#     
#     If a direction l is specified, the system is rotated to have it as its
#     x direction (see Notes for details). If R is specified the system is
#     rotated according to it.
#     
#     Parameters
#     ----------
#     C : array_like
#         6x6 stiffness matrix (Voigt notation).
#     l : array_like, optional
#         3D direction vector for pull (the default is the x direction
#         of the original system)
#     R : array_like, optional
#         3x3 rotation matrix.
#     
#     Returns
#     -------
#     E : array_like
#         Young's modulus for a stress in each of the three directions
#         of the rotated system.
#     nu : array_like
#          3x3 matrix with Poisson's ratios.
#     Gm : array_like
#          3x3 matrix with shear moduli.
#     B : float
#         Bulk modulus.
#     K : array_like
#         3x3 matrix with bulk modulus tensor.
#     
#     Other Parameters
#     ----------------
#     tol : float, optional
#         tolerance for checking validity of rotation and comparison
#         of vectors.
#     
#     Notes
#     ---
#     It works by rotating the elastic constant tensor to the desired
#     direction, so it should be valid for arbitrary crystal structures.
#     If only l is specified there is an infinite number of possible
#     rotations. The chosen one is a rotation along the axis orthogonal
#     to the plane defined by the vectors (1, 0, 0) and l.
#     
#     Bulk modulus tensor as defined in
#     O. Rand and V. Rovenski, "Analytical Methods in Anisotropic
#     Elasticity", Birkh\"auser (2005), pp. 71.
#     
#     """
#     
#     if R is not None:
#         R = np.asarray(R)
# 
#         # Is this a rotation matrix?
#         if np.sometrue(np.abs(np.dot(np.array(R), np.transpose(np.array(R))) - 
#                               np.eye(3, dtype=float)) > tol):
#             raise RuntimeError('Matrix *R* does not describe a rotation.')
#     else:
#         u_a = np.array([1, 0, 0])
#         u_b = l / norm(l)  # Normalise directions
#         R = np.eye(3)
# 
#         if not np.allclose(l, u_a, rtol=tol, atol=tol):
#             u_v = np.cross(u_a, u_b)
#             u_v_mat = np.array([[ 0, -u_v[2], u_v[1]],
#                                 [ u_v[2], 0, -u_v[0]],
#                                 [-u_v[1], u_v[0], 0]])
# 
#             R = R + u_v_mat + \
#                 np.dot(u_v_mat, u_v_mat) * (1 - np.dot(u_a, u_b)) / \
#                 np.linalg.norm(u_v) ** 2
# 
#     Cr = rotate_elastic_constants(C, R)
#     S = np.linalg.inv(Cr)
# 
#     # Young's modulus for a stress in \alpha direction; \alpha = x, y, z
#     E = np.zeros(3)
#     E[0] = 1 / S[0, 0]
#     E[1] = 1 / S[1, 1]
#     E[2] = 1 / S[2, 2]
# 
#     # Poisson's ratio ($\alpha$, $\beta$); $\alpha$, $\beta$ = x, y, z
#     nu = np.array([
#             [-1, -S[1, 0] / S[0, 0], -S[2, 0] / S[0, 0]],
#             [-S[0, 1] / S[1, 1], -1, -S[2, 1] / S[1, 1]],
#             [-S[0, 2] / S[2, 2], -S[1, 2] / S[2, 2], -1]
#             ])
# 
#     # Shear modulus ($\alpha$, $\beta$); $\alpha$, $\beta$ = x, y, z
#     G = np.zeros(3)
#     G[0] = 1 / S[3, 3]  # Shear modulus yz
#     G[1] = 1 / S[4, 4]  # Shear modulus zx
#     G[2] = 1 / S[5, 5]  # Shear modulus xy
#     Gm = np.array([
#             [E[0] / 4, G[2], G[1]],
#             [G[2], E[1] / 4, G[0]],
#             [G[1], G[0], E[2] / 4]
#             ])
# 
#     # Bulk modulus
#     B = 1 / np.sum(S[0:3, 0:3])
#     
#     # Bulk modulus tensor
#     Crt = Voigt_6x6_to_full_3x3x3x3(Cr)
#     K = np.einsum('ijkk', Crt)
# 
#     return E, nu, Gm, B, K
