import json
from importlib.resources import files
from typing import List, Dict
from warnings import warn

from numpy import exp, abs, pi
from pandas import DataFrame, concat
from .weights import Weights


class XRayFormFactors(Weights):
    r"""This class generates sample weights corresponding to X-ray form factors,
    specifically the non-dispersion corrected parametrized form factors.
    In general, the form factor for may be written as

    .. math::
        f(q, \omega) = f_0(q) + f'(q, \omega) + if''(q, \omega)

    where :math:`q` is a scalar, corresponding to the norm
    of the desired reciprocal lattice point.

    The weights generated by this class corresponds to :math:`f_0(q)`.
    There are two possible parametrizations of :math:`f_0(q)`
    to choose from, both based on a sum of exponentials of the form

    .. math::

        f_0(q) = \sum_{i=1}^k a_i \exp(-b_i q^2) + c.

    Two parametrizations are available:

    * ``'waasmaier-1995'`` corresponds to Table 1 from D. Waasmaier, A. Kirfel,
      Acta Crystallographica Section A **51**, 416 (1995);
      `doi: 10.1107/S0108767394013292 <https://doi.org/10.1107/S0108767394013292>`_.
      This parametrization uses five exponentials (:math:`k=5`) and extends up to
      :math:`6.0\,\mathrm{Å}^{-1}`.
    * ``'international-iv-1974'`` corresponds to Table 2.2B from
      *International Tables for X-ray Crystallography, Vol. IV*,
      The Kynoch Press: Birmingham, England, 1974. This parametrization uses
      four exponentials (:math:`k=4`) and extends up to
      :math:`2.0\,\mathrm{Å}^{-1}`.

    In practice differences are expected to be insignificant. It is unlikely that
    you have to deviate from the default, which is ``'waasmaier-1995'``.

    Parameters
    ----------
    atom_types
        List of atomic species for which to retrieve scattering lengths.
    source
        Source to use for parametrization of the form factors :math:`f_0(q)`.
        Allowed values are ``'waasmaier-1995'`` and ``'international-iv-1974'``
        (see above).
    """

    def __init__(
            self,
            atom_types: List[str],
            source: str = 'waasmaier-1995'
    ):
        self._source = source
        form_factors = self._read_form_factors(source)
        # Select the relevant species
        form_factors = form_factors[form_factors.index.isin(atom_types)]  # Atom species is index
        self._form_factors = form_factors

        # Check if any of the fetched form factors are missing,
        # indicating that it is missing in the experimental database.
        for s in atom_types:
            row = form_factors[form_factors.index == s]
            if row.empty:
                if s == 'H':
                    # Manually insert values for H such that the form factor is
                    # zero and raise a warning, since it is such a common element.
                    warn('No parametrization for H. Setting form factor for H to zero.')
                    values = [['DUMMY', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]
                    form_factors = concat([DataFrame(data=values,
                                                     index=['H'],
                                                     columns=form_factors.columns),
                                           form_factors])
                    continue
                raise ValueError('Missing tabulated values '
                                 f'for requested species {s}.')

        weights_coh = form_factors.to_dict(orient='index')
        supports_currents = False
        super().__init__(weights_coh, None, supports_currents=supports_currents)

    def get_weight_coh(self, atom_type, q_norm):
        """Get the coherent weight for a given atom type and q-vector norm."""
        return self._compute_f0(self._weights_coh[atom_type], q_norm, self._source)

    def _compute_f0(
            self,
            coefficients: Dict,
            q_norm: float,
            source: str):
        r"""Compute f_0(q) based on the chosen parametrization.
        There are two possible parametrizations of f_0(q) to choose from,
        both based on a sum of exponentials of the form

        .. math::

            f_0(q) = \sum_{i=1}^k a_i * exp(-b_i * q**2) + c

        Parameters
        ----------
        coefficients
            Parametrization parameters, read from the corresponding source file.
        q_norm
            The |q|-value at which to evaluate the form factor.
        source
            Either 'waasmaier-1995' or 'international-iv-1974',
            corresponding to the two available sources for the
            parametrization of the :math:`f_0(q)` term of the form factors.
        """
        if source == 'waasmaier-1995':
            s = q_norm / (4 * pi)  # q in dynasor is q = 4 pi sin(theta) / lambda.
            s_squared = s*s        # s = sin(theta) / lambda in the Waasmaier paper.
            if abs(s) > 6.0:
                warn('Waasmaier parametrization is not reliable for q '
                     'above 75.398 rad/Å (corresponding to s=6.0 1/Å)')
            # Very verbose and repetitive, but I think it is more readable
            # than a reduction over the Series object.
            # Suggestions for a nice & readable solution are welcome.
            f0 = coefficients['c'] + \
                coefficients['a1'] * exp(- coefficients['b1'] * s_squared) + \
                coefficients['a2'] * exp(- coefficients['b2'] * s_squared) + \
                coefficients['a3'] * exp(- coefficients['b3'] * s_squared) + \
                coefficients['a4'] * exp(- coefficients['b4'] * s_squared) + \
                coefficients['a5'] * exp(- coefficients['b5'] * s_squared)
            return f0
        elif source == 'international-iv-1974':
            raise NotImplementedError('The ITC 1974 parametrization has not been implemented yet.')
        else:
            raise ValueError(f'Unknown source {source}')

    def _read_form_factors(self, source: str) -> DataFrame:
        r"""
        Extracts the parametrization for the form factors :math:`f_0(q)`,
        based on either of two sources.

        Parameters
        ----------
        source
            Either 'waasmaier-1995' or 'international-iv-1974',
            corresponding to the two available sources for the
            parametrization of the :math:`f_0(q)` term of the form factors.
        """
        if source == 'waasmaier-1995':
            data_file = files(__package__) / \
                'form-factors/x-ray-parameters-waasmaier-kirfel-1995.json'
            with open(data_file) as fp:
                coefficients = json.load(fp)

            form_factors = DataFrame.from_dict(coefficients)
            form_factors.index.names = ['species']
        elif source == 'international-iv-1974':
            raise NotImplementedError('The ITC 1974 parametrization has not been implemented yet.')
        else:
            raise ValueError(f'Unknown source {source}')

        return form_factors
