#!/usr/bin/env python

"""
This module defines the classes relating to 3D lattices.
"""

from __future__ import division

__author__ = "Shyue Ping Ong, Michael Kocher"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__status__ = "Production"
__date__ = "Sep 23, 2011"

import numpy as np
import numpy.linalg as npl
from numpy import pi

from pymatgen.serializers.json_coders import MSONable


class Lattice(MSONable):
    '''
    A lattice object.  Essentially a matrix with conversion matrices. In general,
    it is assumed that length units are in Angstroms and angles are in degrees
    unless otherwise stated.
    '''

    def __init__(self, matrix):
        """
        Create a lattice from any sequence of 9 numbers. Note that the sequence
        is assumed to be read one row at a time. Each row represents one
        lattice vector.
        
        Args:
            matrix:
                Sequence of numbers in any form. Examples of acceptable
                input.
                i) An actual numpy array.
                ii) [[1, 0, 0],[0, 1, 0], [0, 0, 1]]
                iii) [1, 0, 0 , 0, 1, 0, 0, 0, 1]
                iv) (1, 0, 0, 0, 1, 0, 0, 0, 1)
        """

        self._matrix = np.array(matrix).reshape((3, 3))
        #Store these matrices for faster access
        self._md2c = np.transpose(self._matrix)
        self._mc2d = npl.inv(np.transpose(self._matrix))

    @property
    def md2c(self):
        '''Matrix for converting direct to cartesian coordinates'''
        return np.copy(self._md2c)

    @property
    def mc2d(self):
        '''Matrix for converting cartesian to direct coordinates'''
        return np.copy(self._mc2d)

    @property
    def matrix(self):
        '''Copy of matrix representing the Lattice'''
        return np.copy(self._matrix)

    def get_cartesian_coords(self, fractional_coords):
        """
        Returns the cartesian coordinates given fractional coordinates.
        
        Args:
            fractional_coords : Fractional coords.
            
        Returns:
            cartesian coordinates
        """
        return np.transpose(np.dot(self._md2c, np.transpose(fractional_coords)))

    def get_fractional_coords(self, cart_coords):
        """
        Returns the fractional coordinates given cartesian coordinates
        
        Args:
            cartesian_coords:
                Cartesian coords.
            
        Returns:
            fractional coordinates
        """
        return np.transpose(np.dot(self._mc2d, np.transpose(cart_coords)))

    @staticmethod
    def cubic(a):
        """
        Convenience constructor for a cubic lattice.
        
        Args:
            a:
                The *a* lattice parameter of the cubic cell.
            
        Returns:
            Cubic lattice of dimensions a x a x a.
        """
        return Lattice(np.array([[a, 0.0, 0.0], [0.0, a, 0.0], [0.0, 0.0, a]]))

    @staticmethod
    def tetragonal(a, c):
        """
        Convenience constructor for a tetragonal lattice.
        
        Args:
            a:
                The *a* lattice parameter of the tetragonal cell.
            c:
                The *c* lattice parameter of the tetragonal cell.
            
        Returns:
            Tetragonal lattice of dimensions a x a x c.
        """
        return Lattice.from_parameters(a, a, c, 90, 90, 90)

    @staticmethod
    def orthorhombic(a, b, c):
        """
        Convenience constructor for an orthorhombic lattice.
        
        Args:
            a:
                The *a* lattice parameter of the orthorhombic cell.
            b:
                The *b* lattice parameter of the orthorhombic cell.
            c:
                The *c* lattice parameter of the orthorhombic cell.
            
        Returns:
            Orthorhombic lattice of dimensions a x b x c.
        """
        return Lattice.from_parameters(a, b, c, 90, 90, 90)

    @staticmethod
    def monoclinic(a, b, c, alpha):
        """
        Convenience constructor for a monoclinic lattice.
        
        Args:
            a:
                The *a* lattice parameter of the monoclinc cell.
            b:
                The *b* lattice parameter of the monoclinc cell.
            c:
                The *c* lattice parameter of the monoclinc cell.
            alpha:
                The *alpha* angle between lattice vectors b and c.
            
        Returns:
            Monoclinic lattice of dimensions a x b x c with angle alpha between
            lattice vectors b and c.
        """
        return Lattice.from_parameters(a, b, c, alpha, 90, 90)

    @staticmethod
    def hexagonal(a, c):
        """
        Convenience constructor for a hexagonal lattice.
        
        Args:
            a:
                The *a* lattice parameter of the hexagonal cell.
            c:
                The *c* lattice parameter of the hexagonal cell.
            
        Returns:
            Hexagonal lattice of dimensions a x a x c.
        """
        return Lattice.from_parameters(a, a, c, 90.0, 90.0, 120.0)

    @staticmethod
    def rhombohedral(a, alpha):
        """
        Convenience constructor for a rhombohedral lattice.
        
        Args:
            a:
                The *a* lattice parameter of the rhombohedral cell.
            alpha:
                Angle for the rhombohedral lattice.
            
        Returns:
            Rhombohedral lattice of dimensions a x a x a.
        """
        return Lattice.from_parameters(a, a, a, alpha, alpha, alpha)

    @staticmethod
    def from_lengths_and_angles(abc, ang):
        '''
        Create a Lattice using unit cell lengths and angles (in degrees).
        
        Args:
            abc:
                lattice parameters, e.g. (4, 4, 5).
            ang: 
                lattice angles in degrees, e.g., (90,90,120).
                
        Returns:
            A Lattice with the specified lattice parameters.
        '''
        return Lattice.from_parameters(abc[0], abc[1], abc[2], ang[0], ang[1], ang[2])

    @staticmethod
    def from_parameters(a, b, c, alpha, beta, gamma):
        '''
        Create a Lattice using unit cell lengths and angles (in degrees).
        
        Args:
            a:
                The *a* lattice parameter of the monoclinc cell.
            b:
                The *b* lattice parameter of the monoclinc cell.
            c:
                The *c* lattice parameter of the monoclinc cell.
            alpha:
                The *alpha* angle.
            beta:
                The *beta* angle.
            gamma:
                The *gamma* angle.
                
        Returns:
            A Lattice with the specified lattice parameters.
        '''
        to_r = lambda degrees: np.radians(degrees)

        alpha_r = to_r(alpha)
        beta_r = to_r(beta)
        gamma_r = to_r(gamma)

        gamma_star = np.arccos((np.cos(alpha_r) * np.cos(beta_r)
                               - np.cos(gamma_r)) / (np.sin(alpha_r)
                               * np.sin(beta_r)))
        vector_a = [a * np.sin(to_r(beta)), 0.0, a * np.cos(to_r(beta))]
        vector_b = [-b * np.sin(to_r(alpha)) * np.cos(gamma_star), b
                    * np.sin(to_r(alpha)) * np.sin(gamma_star), b
                    * np.cos(to_r(alpha))]
        vector_c = np.array([0.0, 0.0, float(c)])
        return Lattice([vector_a, vector_b, vector_c])

    @staticmethod
    def from_dict(d):
        """
        Create a Lattice from a dictionary containing the a, b, c, alpha, beta, 
        and gamma parameters.
        """
        a = d["a"]
        b = d["b"]
        c = d["c"]
        alpha = d["alpha"]
        beta = d["beta"]
        gamma = d["gamma"]
        return Lattice.from_parameters(a, b, c, alpha, beta, gamma)

    @property
    def angles(self):
        """
        Returns the angles (alpha, beta, gamma) of the lattice.
        """
        return self.lengths_and_angles[1]

    @property
    def a(self):
        """
        *a* lattice parameter.
        """
        return self.abc[0]

    @property
    def b(self):
        """
        *b* lattice parameter.
        """
        return self.abc[1]

    @property
    def c(self):
        """
        *c* lattice parameter.
        """
        return self.abc[2]

    @property
    def abc(self):
        """
        Lengths of the lattice vectors, i.e. (a, b, c)
        """
        return self.lengths_and_angles[0]

    def _angle_between(self, x, y):
        """ 
        Internal method to calculate the angle between two vectors.
        """
        angle_between = 180.0 / np.pi * np.arccos(np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y)))
        return angle_between

    @property
    def alpha(self):
        """
        Angle alpha of lattice.
        """
        return self._angle_between(self._matrix[1], self._matrix[2])

    @property
    def beta(self):
        """
        Angle beta of lattice.
        """
        return self._angle_between(self._matrix[0], self._matrix[2])

    @property
    def gamma(self):
        """
        Angle gamma of lattice.
        """
        return self._angle_between(self._matrix[0], self._matrix[1])

    @property
    def volume(self):
        """
        Volume of the unit cell.
        """
        return npl.det(self._matrix)

    @property
    def lengths_and_angles(self):
        '''
        Returns (lattice lengths, lattice angles).
        '''
        prim = self._matrix
        lengths = np.sum(prim ** 2, axis=1) ** 0.5
        angles = np.zeros((3), float)
        angles[0] = np.arccos(np.dot(prim[1], prim[2]) / (lengths[1] * lengths[2])) * 180. / pi
        angles[1] = np.arccos(np.dot(prim[2], prim[0]) / (lengths[2] * lengths[0])) * 180. / pi
        angles[2] = np.arccos(np.dot(prim[0], prim[1]) / (lengths[0] * lengths[1])) * 180. / pi
        return lengths, np.around(angles, 9)

    @property
    def reciprocal_lattice(self):
        """
        Return the reciprocal lattice.
        """
        v = 2 * np.pi / self.volume
        k1 = np.cross(self._matrix[1], self._matrix[2]) * v
        k2 = np.cross(self._matrix[2], self._matrix[0]) * v
        k3 = np.cross(self._matrix[0], self._matrix[1]) * v
        return Lattice([k1, k2, k3])

    def __repr__(self):
        f = lambda x: '%0.6f' % x
        outs = []
        outs.append('Lattice')
        outs.append('    abc : ' + ' '.join(map(f, self.abc)))
        outs.append(' angles : ' + ' '.join(map(f, self.angles)))
        outs.append(' volume : %0.4f' % self.volume)
        outs.append('      A : ' + ' '.join(map(f, self._matrix[0])))
        outs.append('      B : ' + ' '.join(map(f, self._matrix[1])))
        outs.append('      C : ' + ' '.join(map(f, self._matrix[2])))
        return "\n".join(outs)

    def __eq__(self, other):
        """
        A lattice is considered to be equal to another if the internal matrix
        representation satisfies np.allclose(matrix1, matrix2) to be True.
        """
        if other == None:
            return False
        return np.allclose(self._matrix, other._matrix)

    def __hash__(self):
        return 7

    def __str__(self):
        return '\n'.join([' '.join(["%.6f" % i for i in row]) for row in self._matrix])

    @property
    def to_dict(self):
        '''
        Json-serialization dict representation of the Lattice.'''
        d = {
        'matrix': self._matrix.tolist(),
        'a': float(self.a),
        'b': float(self.b),
        'c': float(self.c),
        'alpha': float(self.alpha),
        'beta': float(self.beta),
        'gamma': float(self.gamma),
        'volume': float(self.volume),
        }
        d['module'] = self.__class__.__module__
        d['class'] = self.__class__.__name__
        return d

    def get_primitive_lattice(self, lattice_type):
        """
        Returns the primitive lattice of the lattice given the type.
        
        Args:
            lattice_type:
                An alphabet indicating whether the lattice type, P - primitive,
                R - rhombohedral, A - A-centered, B - B-centered, C - C-centered,
                I - body-centered, F - F-centered.
        """
        if lattice_type == 'P':
            return Lattice(self._matrix)
        conv_to_prim = {
            'R': np.array([[2 / 3, 1 / 3, 1 / 3], [-1 / 3, 1 / 3, 1 / 3], [-1 / 3, -2 / 3, 1 / 3]]),
            'A': np.array([[1, 0, 0], [0, 1 / 2, 1 / 2], [0, -1 / 2, 1 / 2]]),
            'B': np.array([[1 / 2, 0, 1 / 2], [0, 1, 0], [-1 / 2, 0, 1 / 2]]),
            'C': np.array([[1 / 2, 1 / 2, 0], [-1 / 2, 1 / 2, 0], [0, 0, 1]]),
            'I': np.array([[-1 / 2, 1 / 2, 1 / 2], [1 / 2, -1 / 2, 1 / 2], [1 / 2, 1 / 2, -1 / 2]]),
            'F': np.array([[1 / 2, 1 / 2, 0], [0, 1 / 2, 1 / 2], [1 / 2, 0, 1 / 2]])
            }
        return Lattice(np.dot(conv_to_prim[lattice_type], self._matrix))

    def get_most_compact_basis_on_lattice(self):
        """
        This method returns the alternative basis corresponding to the shortest
        3 linearly independent translational operations permitted.
        This tends to create larger angles for every elongated cells and is 
        beneficial for viewing crystal structure (especially when they are
        Niggli cells).
        """
        matrix = self.matrix
        a = matrix[0]
        b = matrix[1]
        c = matrix[2]
        while True:
            anychange = False
            """
            take care of c
            """
            if np.dot(a, b) > 0:
                diffvector = np.subtract(a, b)
            else:
                diffvector = np.add(a, b)
            if np.linalg.norm(diffvector) < np.linalg.norm(a) or np.linalg.norm(diffvector) < np.linalg.norm(b):
                if np.linalg.norm(a) < np.linalg.norm(b):
                    b = diffvector
                else:
                    a = diffvector
            anychange = True
            """
            take care of b
            """
            if np.dot(a, c) > 0:
                diffvector = np.subtract(a, c)
            else:
                diffvector = np.add(a, c)
            if np.linalg.norm(diffvector) < np.linalg.norm(a) or np.linalg.norm(diffvector) < np.linalg.norm(c):
                if(np.linalg.norm(a) < np.linalg.norm(c)):
                    c = diffvector
                else:
                    a = diffvector
            anychange = True
            """
            take care of a
            """
            if np.dot(c, b) > 0:
                diffvector = np.subtract(c, b)
            else:
                diffvector = np.add(c, b)
            if(np.linalg.norm(diffvector) < np.linalg.norm(c)
               or np.linalg.norm(diffvector) < np.linalg.norm(b)):
                if(np.linalg.norm(c) < np.linalg.norm(b)):
                    b = diffvector
                else:
                    c = diffvector
            anychange = True
            if anychange == True:
                break
        return Lattice([a, b, c])
