"""
NMR coupling values
===================

This is a compilation of NMR prefactors for calculation of relaxation times.
A note of caution is

"""

import abc
from math import log
from collections import OrderedDict

from ..utils.constants import gamma_full, hbar_joule, pi, gamma, mu0


__all__ = ['Quadrupolar', 'Czjzek', 'HeteroDipolar',
           'HomoDipolar', 'Constant', 'CSA']


class Coupling(metaclass=abc.ABCMeta):
    name = 'coupling'
    parameter = None
    choice = None
    unit = None
    equation = ''

    @staticmethod
    @abc.abstractmethod
    def relax(*args, **kwargs):
        raise NotImplementedError


class Quadrupolar(Coupling):
    """ Quadrupolar interaction
    """

    name = 'Quadrupolar'
    parameter = [r'\delta/C_{q}', r'\eta']
    unit = ['Hz', '']
    equation = r'3/50 \pi^{2} C_{q}^{2} (1+\eta^{2}/3) (2I+3)/(I^{2}(2I-1))'
    choice = [
        ('Spin', 'spin', OrderedDict([('1', 1), ('3/2', 1.5), ('2', 2), ('5/2', 2.5)])),
        ('Type', 'is_delta', {'Anisotropy': True, 'QCC': False})
    ]

    @staticmethod
    def relax(value: float, eta: float, spin: float = 1, is_delta: bool = True) -> float:
        r"""Calculates prefactor for quadrupolar relaxation

        .. math::
            \frac{3}{200} (2\pi C_q)^2 \left(1+\frac{\eta^{2}}{3}\right) \frac{2I+3}{I^{2}(2I-1)}

        If input is given as anisotropy :math:`\delta` then :math:`C_q` is calculated via

        .. math::
            C_q = \delta \frac{4I(2I-1)}{3(2m-1)}

        where `m` is 1 for integer spin and 3/2 else.

        Args:
            value (float): Quadrupolar coupling constant in Hz.
            eta (float): Asymmetry parameter
            spin (float): Spin `I`. Default is 1.
            is_delta (bool): If True, uses value as anisotropy :math:`\delta` else as QCC. Default is True.

        """
        m = 1 if (2*spin) % 2 == 0 else 1.5
        if is_delta:
            value *= 4*spin*(2*spin-1) / (3*(2*m-1))

        return 0.06 * (2*spin+3) * (pi*value)**2 * (1 + eta**2 / 3) / (spin**2 * (2*spin-1))

    @staticmethod
    def convert(in_value: float, to: str = 'aniso', spin: float = 1) -> float:
        r"""Convert between QCC and anisotropy.

        Relation between quadrupole coupling constant :math:`C_q` and anisotropy :math:`\delta` is given by

        .. math::
             \delta = \frac{3(2m-1)}{4I(2I-1)} C_q

        where `m` is 1 for integer spin and 3/2 else.

        Args:
            in_value (float): Input value
            to (str, {`aniso`, `qcc`}): Direction of conversion. Default is `aniso`
            spin (float): Spin number `I`. Default is 1.

        Returns:
            Converted value
        """
        if to not in ['aniso', 'qcc']:
            raise ValueError(f'Cannot convert to {to} use `aniso` or `qcc`')

        m = 1 if (2 * spin) % 2 == 0 else 1.5
        fac = 4 * spin * (2 * spin - 1) / (3 * (2 * m - 1))
        if to == 'delta':
            return in_value * fac
        else:
            return in_value / fac


class HomoDipolar(Coupling):
    """Homonuclear dipolar coupling"""

    name = 'Homonuclear Dipolar'
    parameter = ['r']
    unit = ['m']
    choice = [
        (r'\gamma', 'nucleus', {k: k for k in gamma}),
        ('Spin', 'spin', OrderedDict([('1/2', 0.5), ('1', 1), ('3/2', 1.5), ('2', 2), ('5/2', 2.5), ('3', 3)])),
    ]
    equation = r'2/5 [\mu_{0}/(4\pi) \hbar \gamma^{2}/r^{3}]^{2} I(I+1)'

    @staticmethod
    def relax(r: float, nucleus: str = '1H', spin: float = 0.5) -> float:
        """Calculate prefactor for homonuclear dipolar relaxation.

        Args:
            r (float): Distance in m.
            nucleus (str): Name of isotope. Default is `1H`
            spin (float): Spin `I` of isotope. Default is 1/2

        """

        if nucleus not in gamma_full:
            raise KeyError(f'Unknown nucleus {nucleus}')

        try:
            coupling = mu0 / (4*pi) * hbar_joule * gamma_full[nucleus]**2 / (r+1e-34)**3
        except ZeroDivisionError:
            return 1e318

        coupling **= 2
        coupling *= 2 * spin * (spin+1) / 5

        return coupling


class HeteroDipolar(Coupling):
    """Heteronuclear dipolar coupling"""

    name = 'Heteronuclear Dipolar'
    parameter = ['r']
    unit = ['m']
    choice = [
        (r'\gamma_{1}', 'nucleus1', {k: k for k in gamma}),
        (r'\gamma_{2}', 'nucleus2', {k: k for k in gamma}),
        ('Spin', 'spin', dict([('1/2', 0.5), ('1', 1), ('3/2', 1.5), ('2', 2), ('5/2', 2.5), ('3', 3)]))
    ]
    equation = r'2/15 [\mu_{0}/(4\pi) \hbar \gamma^{2}/r^{3}]^{2} I(I+1)'

    @staticmethod
    def relax(r: float, nucleus1: str = '1H', nucleus2: str = '19F', spin: float = 0.5) -> float:
        r"""Calculate prefactor for feteronuclear dipolar relaxation.

        .. math::
            \frac{2}{15} \left(\frac{\mu_0}{4\pi} \hbar \frac{\gamma_1 \gamma_2}{r^3} \right)^2 S(S+1)

        Args:
            r (float): Distance in m.
            nucleus1 (str): Name of observed isotope. Default is `1H`
            nucleus2 (str): Name of coupled isotope. Default is `19F`
            spin (float): Spin `S` of the coupled isotope. Default is 1/2

        """
        if nucleus1 not in gamma_full:
            raise KeyError(f'Unknown nucleus {nucleus1}')

        if nucleus2 not in gamma_full:
            raise KeyError(f'Unknown nucleus {nucleus2}')

        try:
            coupling = mu0 / (4*pi) * hbar_joule * gamma_full[nucleus1]*gamma_full[nucleus2] / r**3
        except ZeroDivisionError:
            return 1e318

        coupling **= 2
        coupling *= 2 * spin * (spin+1) / 15

        return coupling


class Czjzek(Coupling):
    name = 'Czjzek'
    parameter = [r'\Sigma']
    unit = ['Hz']
    choice = [('Spin', 'spin', {'1': 1, '3/2': 1.5, '5/2': 2.5, '7/2': 3.5})]
    equation = r'3/[20log(2)] \pi^{2} \Sigma^{2} (2I+3)/(I^{2}(2I-1))'

    @staticmethod
    def relax(fwhm, spin=1):
        return 3*(2+spin+3) * pi**2 * fwhm**2 / 20 / log(2) / spin**2 / (2*spin-1)


class Constant(Coupling):
    name = 'Constant'
    parameter = ['C']
    equation = 'C'

    @staticmethod
    def relax(c):
        return c


class CSA(Coupling):
    """Chemical shift anisotropy"""

    name = 'CSA'
    equation = r'2/15 \Delta\sigma^{2} (1+\eta^{2}/2)'
    parameter = [r'\Delta\sigma', '\eta']
    unit = ['ppm', '']

    @staticmethod
    def relax(sigma: float, eta: float) -> float:
        r"""Relaxation prefactor for chemical shift anisotropy.

        .. math::
            C = \frac{2}{15} \Delta\sigma^2 \left(1+\frac{\eta^2}{3}\right)

        .. note::
            The influence of the external magnetic field is missing here but
            is multiplied in :func:`~nmreval.nmr.Relaxation.t1_csa`

        Args:
            sigma: Anisotropy (in ppm)
            eta: Asymmetry parameter

        Reference:
            Spiess, H.W.: Rotation of Molecules and Nuclear Spin Relaxation.
            In: NMR - Basic principles and progress, Vol. 15, (Springer, 1978)
            https://doi.org/10.1007/978-3-642-66961-3_2

        """
        return 2 * sigma**2 * (1+eta**2 / 3) / 15
