"""
Collection of methods and classes definition useful for constraints computation
"""
# standard libraries imports
from __future__ import print_function

# external libraries imports
import numpy as np

# fullrmc library imports
from ..Globals import LOGGER, FLOAT_TYPE, PI
from ..Globals import str, long, unicode, bytes, basestring, range, xrange, maxint
from ..Core.Collection import is_number, get_caller_frames

class ShapeFunction(object):
    """
    Shape function used to correct for particle shape. The shape
    function is subtracted from the total G(r) of g(r). It must
    be used when non-periodic boundary conditions are used to take
    into account the atomic density drop and to correct for the
    :math:`\\rho_{0}` approximation.

    :Parameters:
        #. engine (Engine): The fitting engine.
        #. weighting (string): The elements weighting.
        #. qmin (number): The minimum reciprocal distance q in
           :math:`\\AA^{-1}` considered to compute the shape function.
        #. qmax (number): The maximum reciprocal distance q in
           :math:`\\AA^{-1}` considered to compute the shape function.
        #. dq (number): The reciprocal distance bin size in
           :math:`\\AA^{-1}` considered to compute the shape function.
        #. rmin (number): The minimum distance in :math:`\\AA` considered
           upon building the histogram prior to computing the shape function.
        #. rmax (number): The maximum distance in :math:`\\AA`
           considered upon building the histogram prior to computing the shape
           function.
        #. dr (number): The bin size in :math:`\\AA` considered upon
           building the histogram prior to computing the shape function.

    **N.B: tweak qmax as small as possible to reduce the wriggles ...**
    """
    def __init__(self, engine, weighting="atomicNumber",
                       qmin=0.001, qmax=1, dq=0.005,
                       rmin=0.00, rmax=100, dr=1):
        # get qmin
        assert is_number(qmin), LOGGER.error("qmin must be a number")
        qmin = FLOAT_TYPE(qmin)
        assert qmin>0, LOGGER.error("qmin '%s' must be bigger than 0"%qmin)
        # get qmin
        assert is_number(qmax), LOGGER.error("qmax must be a number")
        qmax = FLOAT_TYPE(qmax)
        assert qmax>qmin, LOGGER.error("qmax '%s' must be bigger than qmin '%s'"%(qmin,qmax))
        # get dq
        assert is_number(dq), LOGGER.error("dq must be a number")
        dq = FLOAT_TYPE(dq)
        assert dq>0, LOGGER.error("dq '%s' must be bigger than 0"%dq)
        # import StructureFactorConstraint
        from fullrmc.Constraints.StructureFactorConstraints import StructureFactorConstraint
        # overload constraint
        class _ShapeFunctionStructureFactor(StructureFactorConstraint):
            """This overloading is needed to avoid the constraint being
            saved upon setting the different properties"""
            def _get_repository(self, *args, **kwargs):
                return None
            def _dump_to_repository(self, *args, **kwargs):
                return
        # create StructureFactorConstraint
        Q = np.arange(qmin, qmax, dq)
        D = np.transpose([Q, np.zeros(len(Q))]).astype(FLOAT_TYPE)
        #self._SFC = StructureFactorConstraint(rmin=rmin, rmax=rmax, dr=dr, experimentalData=D, weighting="atomicNumber")
        self._SFC = _ShapeFunctionStructureFactor(rmin=rmin, rmax=rmax, dr=dr, experimentalData=D, weighting="atomicNumber")
        self._SFC._set_engine(engine)
        self._SFC.listen(message="engine set")
        # set parameters
        self._rmin = FLOAT_TYPE(rmin)
        self._rmax = FLOAT_TYPE(rmax)
        self._dr   = FLOAT_TYPE(dr)
        self._qmin = FLOAT_TYPE(qmin)
        self._qmax = FLOAT_TYPE(qmax)
        self._dq   = FLOAT_TYPE(dq)
        self._weighting = weighting

    def __get_Gr_from_Sq(self, qValues, rValues, Sq):
        Gr    = np.zeros(len(rValues), dtype=FLOAT_TYPE)
        sq_1  = Sq-1
        qsq_1 = qValues*sq_1
        dq = qValues[1]-qValues[0]
        for ridx, r in enumerate(rValues):
            sinqr_dq = dq * np.sin(qValues*r)
            Gr[ridx] = (2./PI) * np.sum( qsq_1 * sinqr_dq )
        return Gr

    def get_Gr_shape_function(self, rValues, compute=True):
        """
        Get shape function of G(r) used in a PairDistributionConstraint.

        :Parameters:
            #. rValues (numpy.ndarray): The r values array.
            #. compute (boolean): whether to recompute shape
               function reciprocal data.

        :Returns:
            #. shapeFunction (numpy.ndarray): The compute shape function.
        """
        # compute data
        if compute: self._SFC.compute_data()
        # get shape function
        return self.__get_Gr_from_Sq( qValues=self._SFC.experimentalQValues,
                                      rValues=rValues,
                                      Sq=self._SFC.get_constraint_value()['total'])

    def get_gr_shape_function(self, rValues, compute=True):
        """
        Get shape function of g(r) used in a PairCorrelationConstraint.

        :Parameters:
            #. rValues (numpy.ndarray): The r values array.
            #. compute (boolean): whether to recompute shape
               function reciprocal data.

        :Returns:
            #. shapeFunction (numpy.ndarray): The compute shape function.
        """
        sFunc = self.get_Gr_shape_function(rValues=rValues, compute=compute)
        rho0  = self.engine.numberDensity #(self._SFC.engine.numberOfAtoms/self._SFC.engine.volume).astype(FLOAT_TYPE)
        return sFunc / (FLOAT_TYPE(4.)*PI*rho0*rValues)
