# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

"""
A module for NMR analysis
"""

from collections import namedtuple

import numpy as np

from pymatgen.core.periodic_table import Species
from pymatgen.core.structure import Site
from pymatgen.core.tensors import SquareTensor
from pymatgen.core.units import FloatWithUnit

__author__ = "Shyam Dwaraknath"
__copyright__ = "Copyright 2016, The Materials Project"
__version__ = "0.2"
__maintainer__ = "Shyam Dwaraknath"
__credits__ = "Xiaohui Qu"
__email__ = "shyamd@lbl.gov"
__date__ = "Mar 1, 2018"


class ChemicalShielding(SquareTensor):
    """
    This class extends the SquareTensor to perform extra analysis unique to
    NMR Chemical shielding tensors

    Three notations to describe chemical shielding tensor (RK Harris; Magn. Reson.
    Chem. 2008, 46, 582–598; DOI: 10.1002/mrc.2225) are supported.

    Authors: Shyam Dwaraknath, Xiaohui Qu
    """

    HaeberlenNotation = namedtuple("HaeberlenNotation", "sigma_iso, delta_sigma_iso, zeta, eta")
    MehringNotation = namedtuple("MehringNotation", "sigma_iso, sigma_11, sigma_22, sigma_33")
    MarylandNotation = namedtuple("MarylandNotation", "sigma_iso, omega, kappa")

    def __new__(cls, cs_matrix, vscale=None):
        """
        Create a Chemical Shielding tensor.
        Note that the constructor uses __new__
        rather than __init__ according to the standard method of
        subclassing numpy ndarrays.

        Args:
            cs_matrix (1x3 or 3x3 array-like): the 3x3 array-like
                representing the chemical shielding tensor
                or a 1x3 array of the primary sigma values corresponding
                to the principal axis system
            vscale (6x1 array-like): 6x1 array-like scaling the
                voigt-notation vector with the tensor entries
        """
        t_array = np.array(cs_matrix)

        if t_array.shape == (3,):
            return super().__new__(cls, np.diag(cs_matrix), vscale)
        if t_array.shape == (3, 3):
            return super().__new__(cls, cs_matrix, vscale)
        return None

    @property
    def principal_axis_system(self):
        """
        Returns a chemical shielding tensor aligned to the principle axis system
        so that only the 3 diagonal components are non-zero
        """
        return ChemicalShielding(np.diag(np.sort(np.linalg.eigvals(self.symmetrized))))

    @property
    def haeberlen_values(self):
        """
        Returns: the Chemical shielding tensor in Haeberlen Notation
        """
        pas = self.principal_axis_system
        sigma_iso = pas.trace() / 3
        sigmas = np.diag(pas)
        sigmas = sorted(sigmas, key=lambda x: np.abs(x - sigma_iso))
        sigma_yy, sigma_xx, sigma_zz = sigmas
        delta_sigma = sigma_zz - 0.5 * (sigma_xx + sigma_yy)
        zeta = sigma_zz - sigma_iso
        eta = (sigma_yy - sigma_xx) / zeta
        return self.HaeberlenNotation(sigma_iso, delta_sigma, zeta, eta)

    @property
    def mehring_values(self):
        """
        Returns: the Chemical shielding tensor in Mehring Notation
        """
        pas = self.principal_axis_system
        sigma_iso = pas.trace() / 3
        sigma_11, sigma_22, sigma_33 = np.diag(pas)
        return self.MehringNotation(sigma_iso, sigma_11, sigma_22, sigma_33)

    @property
    def maryland_values(self):
        """
        Returns: the Chemical shielding tensor in Maryland Notation
        """
        pas = self.principal_axis_system
        sigma_iso = pas.trace() / 3
        omega = np.diag(pas)[2] - np.diag(pas)[0]
        # There is a typo in equation 20 from Magn. Reson. Chem. 2008, 46, 582–598, the sign is wrong.
        # There correct order is presented in Solid State Nucl. Magn. Reson. 1993, 2, 285-288.
        kappa = 3.0 * (np.diag(pas)[1] - sigma_iso) / omega
        return self.MarylandNotation(sigma_iso, omega, kappa)

    @classmethod
    def from_maryland_notation(cls, sigma_iso, omega, kappa):
        """
        Initialize from Maryland notation.

        Args:
            sigma_iso ():
            omega ():
            kappa ():

        Returns:
            ChemicalShielding
        """
        sigma_22 = sigma_iso + kappa * omega / 3.0
        sigma_11 = (3.0 * sigma_iso - omega - sigma_22) / 2.0
        sigma_33 = 3.0 * sigma_iso - sigma_22 - sigma_11
        return cls(np.diag([sigma_11, sigma_22, sigma_33]))


class ElectricFieldGradient(SquareTensor):
    """
    This class extends the SquareTensor to perform extra analysis unique to
    NMR Electric Field Gradient tensors in units of V/Angstrom^2

    Authors: Shyam Dwaraknath, Xiaohui Qu
    """

    def __new__(cls, efg_matrix, vscale=None):
        """
        Create a Chemical Shielding tensor.
        Note that the constructor uses __new__
        rather than __init__ according to the standard method of
        subclassing numpy ndarrays.

        Args:
            efg_matrix (1x3 or 3x3 array-like): the 3x3 array-like
                representing the electric field tensor
                or a 1x3 array of the primary values corresponding
                to the principal axis system
            vscale (6x1 array-like): 6x1 array-like scaling the
                voigt-notation vector with the tensor entries
        """
        t_array = np.array(efg_matrix)

        if t_array.shape == (3,):
            return super().__new__(cls, np.diag(efg_matrix), vscale)
        if t_array.shape == (3, 3):
            return super().__new__(cls, efg_matrix, vscale)
        return None

    @property
    def principal_axis_system(self):
        """
        Returns a electric field gradient tensor aligned to the principle axis system so that only the 3 diagonal
        components are non-zero
        """
        return ElectricFieldGradient(np.diag(np.sort(np.linalg.eigvals(self))))

    @property
    def V_xx(self):
        """
        Returns: First diagonal element
        """
        diags = np.diag(self.principal_axis_system)
        return sorted(diags, key=np.abs)[0]

    @property
    def V_yy(self):
        """
        Returns: Second diagonal element
        """
        diags = np.diag(self.principal_axis_system)
        return sorted(diags, key=np.abs)[1]

    @property
    def V_zz(self):
        """
        Returns: Third diagonal element
        """
        diags = np.diag(self.principal_axis_system)
        return sorted(diags, key=np.abs)[2]

    @property
    def asymmetry(self):
        """
        Asymmetry of the electric field tensor defined as:
            (V_yy - V_xx)/V_zz
        """
        diags = np.diag(self.principal_axis_system)
        V = sorted(diags, key=np.abs)
        return np.abs((V[1] - V[0]) / V[2])

    def coupling_constant(self, specie):
        """
        Computes the coupling constant C_q as defined in:
            Wasylishen R E, Ashbrook S E, Wimperis S. NMR of quadrupolar nuclei
            in solid materials[M]. John Wiley & Sons, 2012. (Chapter 3.2)

        C_q for a specific atom type for this electric field tensor:
                C_q=e*Q*V_zz/h
            h: Planck's constant
            Q: nuclear electric quadrupole moment in mb (millibarn
            e: elementary proton charge

        Args:
            specie: flexible input to specify the species at this site.
                    Can take a isotope or element string, Species object,
                    or Site object

        Return:

            the coupling constant as a FloatWithUnit in MHz
        """
        planks_constant = FloatWithUnit(6.62607004e-34, "m^2 kg s^-1")
        Vzz = FloatWithUnit(self.V_zz, "V ang^-2")
        e = FloatWithUnit(-1.60217662e-19, "C")

        # Convert from string to Species object
        if isinstance(specie, str):
            # isotope was provided in string format
            if len(specie.split("-")) > 1:
                isotope = str(specie)
                specie = Species(specie.split("-")[0])
                Q = specie.get_nmr_quadrupole_moment(isotope)
            else:
                specie = Species(specie)
                Q = specie.get_nmr_quadrupole_moment()
        elif isinstance(specie, Site):
            specie = specie.specie
            Q = specie.get_nmr_quadrupole_moment()
        elif isinstance(specie, Species):
            Q = specie.get_nmr_quadrupole_moment()
        else:
            raise ValueError("Invalid species provided for quadrupolar coupling constant calculations")

        return (e * Q * Vzz / planks_constant).to("MHz")
