# -*- coding: utf-8 -*-
"""Base Site class."""
from typing import ClassVar
from typing import Dict

from mrsimulator.utils.parseable import Parseable
from pydantic import validator

from .isotope import Isotope
from .tensors import AntisymmetricTensor
from .tensors import SymmetricTensor

__author__ = "Deepansh J. Srivastava"
__email__ = "deepansh2012@gmail.com"


class Site(Parseable):
    """
    Base class representing a single-site nuclear spin interaction tensor parameters.
    The single-site nuclear spin interaction tensors include the nuclear shielding
    and the electric quadrupolar tensor.

    .. rubric:: Attribute Documentation

    Attributes
    ----------

    isotope: str (optional).
        A string expressed as an atomic number followed by an isotope symbol, eg.,
        `'13C'`, `'17O'`. The default value is `'1H'`.

        Example
        -------

        >>> site = Site(isotope='2H')

    isotropic_chemical_shift: float (optional).
        The value is the isotropic chemical shift of the site in the unit of ppm. The
        default value is 0.

        Example
        -------

        >>> site.isotropic_chemical_shift = 43.3

    shielding_symmetric: :ref:`sy_api` or equivalent dict object (optional).
        The value of this attribute represents the irreducible second-rank traceless
        symmetric part of the nuclear shielding tensor. The default value is None.

        The allowed attributes of the :ref:`sy_api` class for `shielding_symmetric` are
        ``zeta``, ``eta``, ``alpha``, ``beta``, and ``gamma``, where ``zeta`` is the
        shielding anisotropy, in ppm, and ``eta`` is the shielding asymmetry parameter
        defined using the Haeberlen convention. The Euler angles ``alpha``, ``beta``,
        and ``gamma`` are in radians.

        Example
        -------

        >>> site.shielding_symmetric = {'zeta': 10, 'eta': 0.5}
        >>> # or equivalently
        >>> site.shielding_symmetric = SymmetricTensor(zeta=10, eta=0.5)

    shielding_antisymmetric: :ref:`asy_api` or equivalent dict object (optional).
        The value of this attribute represents the irreducible first-rank antisymmetric
        part of the nuclear shielding tensor. The default value is None.

        The allowed attributes of the :ref:`asy_api` class for
        `shielding_antisymmetric` are ``zeta``, ``alpha``, and ``beta``, where ``zeta``
        is the anisotropy parameter of the anti-symmetric first-rank tensor given in
        ppm. The angles ``alpha`` and ``beta`` are in radians.

        Example
        -------

        >>> site.shielding_antisymmetric = {'zeta': 20}
        >>> # or equivalently
        >>> site.shielding_antisymmetric = AntisymmetricTensor(zeta=20)

    quadrupolar: :ref:`sy_api` or equivalent dict object (optional).
        The value of this attribute represents the irreducible second-rank traceless
        symmetric part of the electric-field gradient tensor. The default value is
        None.

        The allowed attributes of the :ref:`sy_api` class for `quadrupolar` are ``Cq``,
        ``eta``, ``alpha``, ``beta``, and ``gamma``, where ``Cq`` is the quadrupolar
        coupling constant, in Hz, and ``eta`` is the quadrupolar asymmetry parameter.
        The Euler angles ``alpha``, ``beta``, and ``gamma`` are in radians.

        Example
        -------

        >>> site.quadrupolar = {'Cq': 3.2e6, 'eta': 0.52}
        >>> # or equivalently
        >>> site.quadrupolar = SymmetricTensor(Cq=3.2e6, eta=0.52)

    name: str (optional).
        The value is the name or id of the site. The default value is None.

        Example
        -------

        >>> site.name = '2H-0'
        >>> site.name
        '2H-0'

    label: str (optional).
        The value is a label for the site. The default value is None.

        Example
        -------

        >>> site.label = 'Quad site'
        >>> site.label
        'Quad site'

    description: str (optional).
        The value is a description of the site. The default value is None.

        Example
        -------

        >>> site.description = 'An example Quadrupolar site.'
        >>> site.description
        'An example Quadrupolar site.'

    Example
    -------

    The following are a few examples of setting the site object.

    >>> site1 = Site(
    ...     isotope='33S',
    ...     isotropic_chemical_shift=20, # in ppm
    ...     shielding_symmetric={
    ...         "zeta": 10, # in ppm
    ...         "eta": 0.5
    ...     },
    ...     quadrupolar={
    ...         "Cq": 5.1e6, # in Hz
    ...         "eta": 0.5
    ...     }
    ... )

    Using SymmetricTensor objects.

    >>> site1 = Site(
    ...     isotope='13C',
    ...     isotropic_chemical_shift=20, # in ppm
    ...     shielding_symmetric=SymmetricTensor(zeta=10, eta=0.5),
    ... )
    """

    name: str = None
    label: str = None
    description: str = None
    isotope: str = "1H"
    isotropic_chemical_shift: float = 0
    shielding_symmetric: SymmetricTensor = None
    shielding_antisymmetric: AntisymmetricTensor = None
    quadrupolar: SymmetricTensor = None

    property_unit_types: ClassVar = {"isotropic_chemical_shift": "dimensionless"}
    property_default_units: ClassVar = {"isotropic_chemical_shift": "ppm"}
    property_units: Dict = {"isotropic_chemical_shift": "ppm"}

    @validator("quadrupolar")
    def spin_must_be_at_least_one(cls, v, values):
        if v is None:
            return v
        isotope = values["isotope"]
        isotope = Isotope(**isotope) if isinstance(isotope, dict) else isotope
        spin_I = isotope.spin
        if spin_I >= 1:
            if "zeta" in v.property_units:
                v.property_units.pop("zeta")
            return v
        message = (
            f"with spin quantum number {spin_I} does not allow quadrupolar tensor."
        )
        raise ValueError(f"{isotope} {message}")

    @validator("shielding_symmetric")
    def shielding_symmetric_must_not_contain_Cq(cls, v, values):
        if "Cq" in v.property_units:
            v.property_units.pop("Cq")
        return v

    @validator("isotope", always=True)
    def validate_isotope(cls, v, *, values, **kwargs):
        return Isotope(symbol=v)

    class Config:
        validate_assignment = True

    @classmethod
    def parse_dict_with_units(cls, py_dict):
        """
        Parse the physical quantity from a dictionary representation of the Site
        object, where the physical quantity is expressed as a string with a number and
        a unit.

        Args:
            dict py_dict: A required python dict object.

        Returns:
            :ref:`site_api` object.

        Example
        -------

        >>> site_dict = {
        ...    "isotope": "13C",
        ...    "isotropic_chemical_shift": "20 ppm",
        ...    "shielding_symmetric": {"zeta": "10 ppm", "eta":0.5}
        ... }
        >>> site1 = Site.parse_dict_with_units(site_dict)
        """
        prop_mapping = {
            "shielding_symmetric": SymmetricTensor,
            "shielding_antisymmetric": AntisymmetricTensor,
            "quadrupolar": SymmetricTensor,
        }

        for k, v in prop_mapping.items():
            if k in py_dict:
                py_dict[k] = v.parse_dict_with_units(py_dict[k])

        return super().parse_dict_with_units(py_dict)

    def to_freq_dict(self, B0):
        """
        Serialize the Site object to a JSON compliant python dictionary object, where
        the attribute value is a number expressed in the attribute's default unit.
        The default unit for the attributes with respective dimensionalities is:

        - frequency: ``Hz``
        - angle: ``rad``

        Args:
            float B0: A required macroscopic magnetic flux density in units of T.

        Return:
            Python dict object.

        Example
        -------

        >>> pprint(site1.to_freq_dict(B0=9.4))
        {'description': None,
         'isotope': '13C',
         'isotropic_chemical_shift': -2013.1791999999998,
         'label': None,
         'name': None,
         'quadrupolar': None,
         'shielding_antisymmetric': None,
         'shielding_symmetric': {'alpha': None,
                                 'beta': None,
                                 'eta': 0.5,
                                 'gamma': None,
                                 'zeta': -1006.5895999999999}}
        """
        temp_dict = self.dict(exclude={"isotope"})
        temp_dict["isotope"] = self.isotope.symbol
        larmor_frequency = -self.isotope.gyromagnetic_ratio * B0  # in MHz
        for k in ["shielding_symmetric", "shielding_antisymmetric", "quadrupolar"]:
            if getattr(self, k):
                temp_dict[k] = getattr(self, k).to_freq_dict(larmor_frequency)
                if k == "shielding_symmetric":
                    temp_dict[k].pop("Cq")

        temp_dict["isotropic_chemical_shift"] *= larmor_frequency
        temp_dict.pop("property_units")
        return temp_dict
