<commit_before>import ast
import collections
import copy
import inspect
import uuid
from operator import attrgetter

import astropy.units as u
import numpy as np
import os
import re
from yaml.reader import ReaderError

from astromodels.core.my_yaml import my_yaml
from astromodels.core.parameter import Parameter
from astromodels.core.tree import Node
from astromodels.utils.pretty_list import dict_to_list
from astromodels.utils.table import dict_to_table
from astromodels.core.memoization import memoize

__author__ = 'giacomov'


try:

    from IPython.display import display, HTML

except:

    has_ipython = False

else:

    has_ipython = True


class WarningNoTests(ImportWarning):
    pass


class FunctionDefinitionError(Exception):
    pass


class DesignViolation(Exception):
    pass


class ModelAssertionViolation(Exception):
    pass


class WrongDimensionality(Exception):
    pass


class TestSpecificationError(Exception):
    pass


class TestFailed(Exception):
    pass


class DocstringIsNotRaw(ValueError):
    pass


class UnknownFunction(ValueError):
    pass


class UnknownParameter(ValueError):
        pass


# Value to indicate that no latex formula has been given
NO_LATEX_FORMULA = '(no latex formula available)'



# This dictionary will contain the known function by name, so that the model_parser can instance
# them by looking into this dictionary. It will be filled by the FunctionMeta meta-class.

_known_functions = {}


# The following is a metaclass for all the functions
class FunctionMeta(type):
    """
    A metaclass for the models, which takes care of setting up the parameters and the other attributes
    according to the definition given in the documentation of the function class.
    """

    def __new__(mcs, name, bases, dct):

        # We do the parsing of the parameters in the __new__ instead of the __init__ so this is the first thing
        # that runs when importing astromodels

        # Enforce the presence of the evaluate method

        if 'evaluate' not in dct:

            raise AttributeError("You have to implement the 'evaluate' method in %s" % name)

        # We also need the method _set_units

        if '_set_units' not in dct:
            raise AttributeError("You have to implement the '_set_units' method in %s" % name)

        # Now parse the documentation of the function which contains the parameter specification

        # The doc is a YAML document containing among other things the definition of the parameters

        # Parse it

        try:

            function_definition = my_yaml.load(dct['__doc__'])

        except ReaderError:  # pragma: no cover

            raise DocstringIsNotRaw("Docstring parsing has failed. "
                                    "Did you remember to specify the docstring of %s as raw? "
                                    "To do that, you have to put a r before the docstring, "
                                    '''like in \n\nr"""\n(docstring)\n"""\n\ninstead of just\n\n'''
                                    '''"""\ndocstring\n"""''' % name)

        else:

            # Store the function definition in the type

            dct['_function_definition'] = function_definition

        # Enforce the presence of a description and of a parameters dictionary

        assert "description" in function_definition.keys(), "You have to provide a 'description' token in the " \
                                                            "documentation of class %s" % name

        assert "parameters" in function_definition.keys(), "You have to provide a 'parameters' token in the " \
                                                           "documentation of class %s" % name

        # If there is a latex formula, store it in the type

        if 'latex' in function_definition:

            # First remove the escaping we did to overcome the limitation of the YAML parser

            latex_formula = function_definition['latex'].replace(r"\\", chr(92))

        else:

            latex_formula = NO_LATEX_FORMULA

        # Store latex formula in the type
        dct['_latex'] = latex_formula

        # Parse the parameters' dictionary
        assert isinstance(function_definition['parameters'], dict), "Wrong syntax in 'parameters' token. It must be " \
                                                                    "a dictionary. Refer to the documentation."

        # Add the parameters as attribute of the *type*. During the __call__ method below this dictionary will be used
        # to create a copy of each parameter which will be made available as child of the *instance*.

        dct['_parameters'] = collections.OrderedDict()

        for parameter_name, parameter_definition in function_definition['parameters'].iteritems():

            this_parameter = FunctionMeta.parse_parameter_definition(name, parameter_name, parameter_definition)

            dct['_parameters'][this_parameter.name] = this_parameter

        # Now perform a minimal check of the 'evaluate' function

        variables, parameters_in_calling_sequence = FunctionMeta.check_calling_sequence(name, 'evaluate',
                                                                                        dct['evaluate'],
                                                                                        ['x', 'y', 'z'])

        # Now check that all the parameters used in 'evaluate' are part of the documentation,
        # and that there are no unused parameters

        set1 = set(dct['_parameters'].keys())
        set2 = set(parameters_in_calling_sequence)

        if set1 != set2:

            # The parameters are different. Figure out who is missing and raise an exception accordingly

            if set1 > set2:

                missing = set1 - set2

                msg = "Parameters %s have init values but are not used in 'evaluate' in %s" % (",".join(missing), name)

            else:

                missing = set2 - set1

                msg = "Parameters %s are used in 'evaluate' but do not have init values in %s" % \
                      (",".join(missing), name)

            raise FunctionDefinitionError(msg)

        # Figure out the dimensionality of this function

        n_dim = len(variables)

        # Store the dimensionality in the *type*

        dct['_n_dim'] = n_dim

        # Now add the constructor to the class, if it does not provide one
        # You shouldn't usually provide a constructor, that's only for advance uses
        # like the TemplateModel

        if '_custom_init_' in dct:

            dct['__init__'] = dct['_custom_init_']

        else:

            dct['__init__'] = FunctionMeta.class_init

        # Finally, add the info() method to the type so that it can be called even without instancing the class

        def info():

            repr_dict = collections.OrderedDict()

            repr_dict['description'] = function_definition['description']

            if 'latex' in function_definition:
                repr_dict['formula'] = function_definition['latex']

            # Add the description of each parameter and their current value
            repr_dict['default parameters'] = collections.OrderedDict()

            for parameter_name in dct['_parameters'].keys():

                repr_dict['default parameters'][parameter_name] = dct['_parameters'][parameter_name].to_dict()

            if has_ipython:

                display(HTML(dict_to_list(repr_dict, html=True)))

            else:

                print(dict_to_list(repr_dict, html=False))

        dct['info'] = staticmethod(info)

        # Now call the __new__ of the "type" class (which then will call the __init__ of this metaclass)

        return super(FunctionMeta, mcs).__new__(mcs, name, bases, dct)

    def __init__(cls, name, bases, dct):

        # This is the MetaClass init, which is called after the __new__ is done

        # Store the name of the function in the type
        cls._name = cls.__name__

        # Add this as a known function

        _known_functions[name] = cls

        # Finally call the init of the type class

        super(FunctionMeta, cls).__init__(name, bases, dct)

    @staticmethod
    def class_init(instance, **kwargs):

        # This is what is going to be called as the __init__ of the class, every time a new instance
        # is created

        # Create a copy of the parameters dictionary which is in the type,
        # otherwise every instance would share the same dictionary

        copy_of_parameters = collections.OrderedDict()

        # Fill it by duplicating the parameters contained in the dictionary in the type

        for key, value in type(instance)._parameters.iteritems():

            copy_of_parameters[key] = value.duplicate()

            # If the user has specified a value in the constructor, update the
            # corresponding parameters value. This allow to use a constructor as:
            # my_powerlaw = powerlaw(logK=1.0, index=-2)

            if key in kwargs:

                copy_of_parameters[key].value = kwargs[key]

        # Now check that all the parameters specified in the kwargs are actually parameters of this function
        for key in kwargs.keys():

            try:

                copy_of_parameters[key]

            except KeyError:

                raise UnknownParameter("You specified an init value for %s, which is not a "
                                       "parameter of function %s" % (key, type(instance)._name))

        # Now call the init of the corresponding class
        n_dim = type(instance)._n_dim

        if n_dim == 1:

            Function1D.__init__(instance,
                                type(instance)._name,
                                type(instance)._function_definition,
                                copy_of_parameters)

        elif n_dim == 2:

            Function2D.__init__(instance,
                                type(instance)._name,
                                type(instance)._function_definition,
                                copy_of_parameters)

        elif n_dim == 3:

            Function3D.__init__(instance,
                                type(instance)._name,
                                type(instance)._function_definition,
                                copy_of_parameters)

        # Last, if the class provides a setup method, call it
        if hasattr(instance, "_setup"):

            instance._setup()

    @staticmethod
    def check_calling_sequence(name, function_name, function, possible_variables):
        """
        Check the calling sequence for the function looking for the variables specified.
        One or more of the variables can be in the calling sequence. Note that the
        order of the variables will be enforced.
        It will also enforce that the first parameter in the calling sequence is called 'self'.

        :param function: the function to check
        :param possible_variables: a list of variables to check, The order is important, and will be enforced
        :return: a tuple containing the list of found variables, and the name of the other parameters in the calling
        sequence
        """

        # Get calling sequence

        # If the function has been memoized, it will have a "input_object" member

        try:

            calling_sequence = inspect.getargspec(function.input_object).args

        except AttributeError:

            # This might happen if the function is with memoization

            calling_sequence = inspect.getargspec(function).args

        assert calling_sequence[0] == 'self', "Wrong syntax for 'evaluate' in %s. The first argument " \
                                              "should be called 'self'." % name

        # Figure out how many variables are used

        variables = filter(lambda var: var in possible_variables, calling_sequence)

        # Check that they actually make sense. They must be used in the same order
        # as specified in possible_variables

        assert len(variables) > 0, "The name of the variables for 'evaluate' in %s must be one or more " \
                                   "among %s, instead of %s" % (name, ','.join(possible_variables), ",".join(variables))

        if variables != possible_variables[:len(variables)]:
            raise AssertionError("The variables %s are out of order in '%s' of %s. Should be %s."
                                 % (",".join(variables), function_name, name, possible_variables[:len(variables)]))

        other_parameters = filter(lambda var: var not in variables and var != 'self', calling_sequence)

        return variables, other_parameters

    @staticmethod
    def parse_parameter_definition(func_name, par_name, definition):

        # Parse definition of parameter

        # Enforce the presence of attributes 'value' and 'desc'

        if 'initial value' not in definition:
            raise FunctionDefinitionError("Error for parameter %s of function %s: value for parameter must be"
                                          " specified" % (par_name, func_name))

        if 'desc' not in definition:
            raise FunctionDefinitionError("Error for parameter %s of function %s: desc for parameter must be"
                                          " specified" % (par_name, func_name))

        # Fetch attributes

        # Use unitless parameters when building the function, if no unit is specified, otherwise
        # use that unit
        if 'unit' not in definition or definition['unit'] is None or definition['unit'] == '':

            du = u.dimensionless_unscaled

        else:

            du = u.Unit(definition['unit'])

        def _parse_value(val):

            if isinstance(val, str):

                return eval(val)

            elif val is None:

                return None

            else:

                return float(val)

        value = _parse_value(definition['initial value'])
        desc = definition['desc']

        # Optional attributes are either None if not specified, or the value specified

        min_value = (None if 'min' not in definition else _parse_value(definition['min']))
        max_value = (None if 'max' not in definition else _parse_value(definition['max']))
        delta = (None if 'delta' not in definition else _parse_value(definition['delta']))
        unit = du

        # A parameter can be fixed by using fix=yes, otherwise it is free by default

        free = (True if 'fix' not in definition else not bool(definition['fix']))

        new_parameter = Parameter(par_name, value, min_value=min_value, max_value=max_value,
                                  delta=delta, desc=desc, free=free, unit=unit)

        return new_parameter


class Function(Node):
    """
    Generic Function class. Will be subclassed in Function1D, Function2D and Function3D.

    """
    def __init__(self, name=None, function_definition=None, parameters=None):

        # I use default values only to avoid warnings from pycharm and other software about the
        # calling sequence of this contructor. We actually need to enforce its proper use,
        # with this assert

        assert name is not None and function_definition is not None and parameters is not None

        # Set up the node

        super(Function, self).__init__(name)

        # Store name, number of dimensions and the latex formula

        # Store also the function definition

        assert 'description' in function_definition,"Function definition must contain a description"

        if 'latex' not in function_definition:

            function_definition['latex'] = '$n.a.$'

        self._function_definition = function_definition

        # Add the parameters as children. Since the name of the key in the dictionary might
        # be different than the actual name of the parameter, use the .add_child method instead
        # of the add_children method

        self._parameters = collections.OrderedDict()

        for child_name, child in parameters.iteritems():

            self._parameters[child_name] = child

            # Add the parameter as a child of the function

            self._add_child(child)

        # Now generate a unique identifier (UUID) in a thread safe, multi-processing safe
        # way. This is used for example in the CompositeFunction class to keep track of the different
        # instances of the same function
        self._uuid = "{" + str(self._generate_uuid()) + "}"

        # Normal functions are able to change units, while some specific ones (such as the one from XSpec) are not.
        # In this second case, this variable will contain a tuple (x_unit, y_unit), but by default it should be
        # None

        self._fixed_units = None

    @property
    def n_dim(self):
        """
        :return: number of dimensions for this function (1, 2 or 3)
        """
        return type(self)._n_dim

    @property
    def free_parameters(self):
        """
        Returns a dictionary of free parameters for this function

        :return: dictionary of free parameters
        """

        free_parameters = collections.OrderedDict([(k,v) for k, v in self.parameters.iteritems() if v.free])

        return free_parameters

    @staticmethod
    def _generate_uuid():
        """
        Generate a unique identifier for this function.

        :return: the UUID
        """
        return uuid.UUID(bytes=os.urandom(16), version=4)

    def has_fixed_units(self):
        """
        Returns True if this function cannot change units, which is the case only for very specific functions (like
        models from foreign libraries like Xspec)

        :return: True or False
        """

        return not (self._fixed_units is None)

    @property
    def fixed_units(self):
        """
            Returns the fixed units if has_fixed_units is True (see has_fixed_units)

            :return: None, or a tuple (x_unit, y_unit)
        """

        return self._fixed_units

    @property
    def description(self):
        """
        Returns a description for this function
        """

        return self._function_definition['description']

    # Add a property returning the parameters dictionary
    @property
    def parameters(self):
        """
        Returns a dictionary of parameters
        """
        return self._parameters

    @property
    def latex(self):
        """
        Returns the LaTEX formula for this function
        """
        return self._function_definition['latex']

    # Define now all the operators which allow to combine functions. Each operator will return a new
    # instance of a CompositeFunction, which can then be used as a function on its own

    def of(self, another_function):
        """
        Compose this function with another as in this_function(another_function(x))
        :param another_function: another function to compose with the current one
        :return: a composite function instance
        """
        return CompositeFunction('of', self, another_function)

    def __neg__(self):

        return CompositeFunction('*-', self)

    def __abs__(self):

        return CompositeFunction('abs', self)

    def __pow__(self, other_instance):

        return CompositeFunction('**', self, other_instance)

    def __rpow__(self, other_instance):

        return CompositeFunction('**', other_instance, self)

    def __add__(self, other_instance):

        return CompositeFunction('+', self, other_instance)

    __radd__ = __add__

    def __sub__(self, other_instance):

        return CompositeFunction('-', self, other_instance)

    __rsub__ = __sub__

    def __mul__(self, other_instance):

        c = CompositeFunction('*', self, other_instance)

        # If the other instance is a function (and not a number), flag it so its units will be made dimensionless
        # in the set_units method of the composite function

        if isinstance(other_instance, Function):

            if self.has_fixed_units():

                # This is likely a XSpec model. The multiplication of two models with units will give the wrong
                # units to the results. So, depending on the type of the first and second model, we need to adjust
                # their units so that the result will keep the right units.

                if not self.fixed_units[1] == u.dimensionless_unscaled:

                    # This function has fixed unit and is not dimensionless (likely an additive XSpec model).
                    # We need to make the other function dimensionless so that the multiplication of them will keep
                    # the right units

                    other_instance._make_dimensionless = True

                else:

                    # This function has fixed unit, but it is dimensionless (likely a multiplicative XSpec model)
                    # The other function should keep its units, so we flag self instead

                    self._make_dimensionless = True

                    # However, if also the other function has fixed dimension and it is dimensionless
                    # (likely another XSpec multiplicative model) we need to flag that as well
                    if other_instance.has_fixed_units() and (other_instance.fixed_units[1] == u.dimensionless_unscaled):
                        other_instance._make_dimensionless = True

                        # Now we need to flag the composite function as fixed units and dimensionless
                        c._fixed_units = other_instance.fixed_units

            else:

                # We need to make the other instance dimensionless so that this function (which is not dimensionless)
                # multiplied by the other function (which we will make dimensionless) will give the right units as
                # the results

                other_instance._make_dimensionless = True

        return c

    __rmul__ = __mul__

    def __div__(self, other_instance):

        c = CompositeFunction('/', self, other_instance)

        # If the other instance is a function (and not a number), flag it so its units will be made dimensionless
        # in the set_units method of the composite function

        if isinstance(other_instance, Function):

            if self.has_fixed_units():

                # This is likely a XSpec model. Division is not supported.

                raise NotImplementedError("Division for XSpec models is not supported")

            else:

                # We need to make the other instance dimensionless

                other_instance._make_dimensionless = True

        return c

    def __rdiv__(self, other_instance):

        return CompositeFunction('/', other_instance, self)

    __truediv__ = __div__
    __rtruediv__ = __rdiv__

    def _repr__base(self, rich_output):

        repr_dict = collections.OrderedDict()

        repr_dict['description'] = self._function_definition['description']

        if 'latex' in self._function_definition:

            repr_dict['formula'] = self._function_definition['latex']

        # Add the description of each parameter and their current value
        repr_dict['parameters'] = collections.OrderedDict()

        for parameter in self._get_children():

            repr_dict['parameters'][parameter.name] = parameter.to_dict()

        return dict_to_list(repr_dict, rich_output)

    @property
    def uuid(self):
        """
        Returns the ID of the current function. The ID is used by the CompositeFunction class to keep track of the
        unique instances of each function. It should not be used by the user for any specific purpose.

        :return: (none)
        """
        return self._uuid

    def duplicate(self):
        """
        Create a copy of the current function with all the parameters equal to the current value

        :return: a new copy of the function
        """

        # Create a copy

        function_copy = copy.deepcopy(self)

        return function_copy

    def get_boundaries(self):  # pragma: no cover
        """
        Returns the boundaries of this function. By default there is no boundary, but subclasses can
        override this.

        :return: a tuple of tuples containing the boundaries for each coordinate (ra_min, ra_max), (dec_min, dec_max)
        """

        raise NotImplementedError("You have to implement this")

    def __call__(self, *args):  # pragma: no cover

        raise NotImplementedError("You have to implement this")

    def fast_call(self, *args):  # pragma: no cover

        raise NotImplementedError("You have to implement this")

    def evaluate_at(self, *args, **parameter_specification):  # pragma: no cover
        """
        Evaluate the function at the given x(,y,z) for the provided parameters, explicitly provided as part of the
        parameter_specification keywords.

        :param *args:
        :param **parameter_specification:
        :return:
        """

        # Set the parameters to the provided values
        for parameter in parameter_specification:

            self._get_child(parameter).value = parameter_specification[parameter]

        return self(*args)


class Function1D(Function):

    def __init__(self, name=None, function_definition=None, parameters=None):

        Function.__init__(self, name, function_definition, parameters)

        self._x_unit = None
        self._y_unit = None

    def evaluate(self, x, *args, **kwargs):  # pragma: no cover

        raise NotImplementedError("You have to re-implement this")

    def set_units(self, in_x_unit, in_y_unit):

        # Transform None in input to '', so that u.Unit() will generate a dimensionless unit

        in_x_unit = in_x_unit if in_x_unit is not None else ''
        in_y_unit = in_y_unit if in_y_unit is not None else ''

        # Get a Unit instance from the inputs

        try:

            in_x_unit = u.Unit(in_x_unit)
            in_y_unit = u.Unit(in_y_unit)

        except:

            raise TypeError("Could not get a Unit instance from provided units %s when setting units "
                            "for function %s" % ((in_x_unit, in_y_unit), self.name))

        # Now call the underlying method to set units, which is defined by each function
        new_units = self._set_units(in_x_unit, in_y_unit)

        # Store the units.
        # NOTE: the previous call to _set_units might return new units in special cases
        # (for example Xspec functions). So it is critical that we store them in the class' attributes
        # after the call to _set_units

        if new_units is not None:

            new_x_unit, new_y_unit = new_units

            self._x_unit = new_x_unit
            self._y_unit = new_y_unit

        else:

            self._x_unit = in_x_unit
            self._y_unit = in_y_unit

    def _set_units(self, x_unit, y_unit):  # pragma: no cover

        # This will be overridden by derived classes

        raise NotImplementedError("You have to implement the method _set_units for function %s" % self.name)

    @property
    def x_unit(self):
        """
        The unit of the independent variable
        :return: a astropy.Unit instance
        """
        return self._x_unit

    @property
    def y_unit(self):
        """
        The unit of the dependent variable
        :return: a astropy.Unit instance
        """
        return self._y_unit

    def __call__(self, x):

        # This method's code violates explicitly duck typing. The reason is that astropy.units introduce a very
        # significant overload on any computation. For this reason we treat differently the case with units from
        # the case without units, so that the latter case remains fast. Also, transforming an input
        # which is not an array into an array introduce a significant overload (10 microseconds or so), so we perform
        # this transformation only when strictly required

        # NOTE: for a single quantity such as q = (1.0 * u.keV), isinstance(q, np.ndarray) returns True

        if isinstance(x, np.ndarray):

            # We have an array as input (or a single quantity)

            if not isinstance(x, u.Quantity):

                # This is a normal array, let's use the fast call (without units)

                return self.fast_call(x)

            else:

                # This is an array with units or a single quantity, let's use the slow call which preserves units

                assert self.y_unit is not None, "In order to use units you need to use the function as a spectrum or " \
                                                "as something else, or you need to explicitly set the units."

                results = self._call_with_units(x)

                # Now convert to the expected y unit and return a astropy.Quantity by multiplying by the right unit
                return np.squeeze(results.to(self.y_unit).value) * self.y_unit

        else:

            # This is either a single number or a list

            # Transform the input to an array of floats. If x is a single number, this will be an array of size 1

            new_input = np.array(x, dtype=float, ndmin=1, copy=False)

            # Compute the function

            result = self.fast_call(new_input)

            # Now remove all dimensions of size 1. For example, an array of shape (1,) will become a single number,
            # so that if the input was a single number, also the output will be a single number

            return np.squeeze(result)

    def _call_with_units(self, x):

        # Gather the current parameters' values with units
        values = map(attrgetter("as_quantity"), self._get_children())

        try:

            results = self.evaluate(x, *values)

        except u.UnitsError:  # pragma: no cover<commit_msg>Using equivalencies so that it is now possible to use Angstrom and Hz in place of energies, for example<commit_after>import ast
import collections
import copy
import inspect
import uuid
from operator import attrgetter

import astropy.units as u
import numpy as np
import os
import re
from yaml.reader import ReaderError

from astromodels.core.my_yaml import my_yaml
from astromodels.core.parameter import Parameter
from astromodels.core.tree import Node
from astromodels.utils.pretty_list import dict_to_list
from astromodels.utils.table import dict_to_table
from astromodels.core.memoization import memoize

__author__ = 'giacomov'


try:

    from IPython.display import display, HTML

except:

    has_ipython = False

else:

    has_ipython = True


class WarningNoTests(ImportWarning):
    pass


class FunctionDefinitionError(Exception):
    pass


class DesignViolation(Exception):
    pass


class ModelAssertionViolation(Exception):
    pass


class WrongDimensionality(Exception):
    pass


class TestSpecificationError(Exception):
    pass


class TestFailed(Exception):
    pass


class DocstringIsNotRaw(ValueError):
    pass


class UnknownFunction(ValueError):
    pass


class UnknownParameter(ValueError):
        pass


# Value to indicate that no latex formula has been given
NO_LATEX_FORMULA = '(no latex formula available)'



# This dictionary will contain the known function by name, so that the model_parser can instance
# them by looking into this dictionary. It will be filled by the FunctionMeta meta-class.

_known_functions = {}


# The following is a metaclass for all the functions
class FunctionMeta(type):
    """
    A metaclass for the models, which takes care of setting up the parameters and the other attributes
    according to the definition given in the documentation of the function class.
    """

    def __new__(mcs, name, bases, dct):

        # We do the parsing of the parameters in the __new__ instead of the __init__ so this is the first thing
        # that runs when importing astromodels

        # Enforce the presence of the evaluate method

        if 'evaluate' not in dct:

            raise AttributeError("You have to implement the 'evaluate' method in %s" % name)

        # We also need the method _set_units

        if '_set_units' not in dct:
            raise AttributeError("You have to implement the '_set_units' method in %s" % name)

        # Now parse the documentation of the function which contains the parameter specification

        # The doc is a YAML document containing among other things the definition of the parameters

        # Parse it

        try:

            function_definition = my_yaml.load(dct['__doc__'])

        except ReaderError:  # pragma: no cover

            raise DocstringIsNotRaw("Docstring parsing has failed. "
                                    "Did you remember to specify the docstring of %s as raw? "
                                    "To do that, you have to put a r before the docstring, "
                                    '''like in \n\nr"""\n(docstring)\n"""\n\ninstead of just\n\n'''
                                    '''"""\ndocstring\n"""''' % name)

        else:

            # Store the function definition in the type

            dct['_function_definition'] = function_definition

        # Enforce the presence of a description and of a parameters dictionary

        assert "description" in function_definition.keys(), "You have to provide a 'description' token in the " \
                                                            "documentation of class %s" % name

        assert "parameters" in function_definition.keys(), "You have to provide a 'parameters' token in the " \
                                                           "documentation of class %s" % name

        # If there is a latex formula, store it in the type

        if 'latex' in function_definition:

            # First remove the escaping we did to overcome the limitation of the YAML parser

            latex_formula = function_definition['latex'].replace(r"\\", chr(92))

        else:

            latex_formula = NO_LATEX_FORMULA

        # Store latex formula in the type
        dct['_latex'] = latex_formula

        # Parse the parameters' dictionary
        assert isinstance(function_definition['parameters'], dict), "Wrong syntax in 'parameters' token. It must be " \
                                                                    "a dictionary. Refer to the documentation."

        # Add the parameters as attribute of the *type*. During the __call__ method below this dictionary will be used
        # to create a copy of each parameter which will be made available as child of the *instance*.

        dct['_parameters'] = collections.OrderedDict()

        for parameter_name, parameter_definition in function_definition['parameters'].iteritems():

            this_parameter = FunctionMeta.parse_parameter_definition(name, parameter_name, parameter_definition)

            dct['_parameters'][this_parameter.name] = this_parameter

        # Now perform a minimal check of the 'evaluate' function

        variables, parameters_in_calling_sequence = FunctionMeta.check_calling_sequence(name, 'evaluate',
                                                                                        dct['evaluate'],
                                                                                        ['x', 'y', 'z'])

        # Now check that all the parameters used in 'evaluate' are part of the documentation,
        # and that there are no unused parameters

        set1 = set(dct['_parameters'].keys())
        set2 = set(parameters_in_calling_sequence)

        if set1 != set2:

            # The parameters are different. Figure out who is missing and raise an exception accordingly

            if set1 > set2:

                missing = set1 - set2

                msg = "Parameters %s have init values but are not used in 'evaluate' in %s" % (",".join(missing), name)

            else:

                missing = set2 - set1

                msg = "Parameters %s are used in 'evaluate' but do not have init values in %s" % \
                      (",".join(missing), name)

            raise FunctionDefinitionError(msg)

        # Figure out the dimensionality of this function

        n_dim = len(variables)

        # Store the dimensionality in the *type*

        dct['_n_dim'] = n_dim

        # Now add the constructor to the class, if it does not provide one
        # You shouldn't usually provide a constructor, that's only for advance uses
        # like the TemplateModel

        if '_custom_init_' in dct:

            dct['__init__'] = dct['_custom_init_']

        else:

            dct['__init__'] = FunctionMeta.class_init

        # Finally, add the info() method to the type so that it can be called even without instancing the class

        def info():

            repr_dict = collections.OrderedDict()

            repr_dict['description'] = function_definition['description']

            if 'latex' in function_definition:
                repr_dict['formula'] = function_definition['latex']

            # Add the description of each parameter and their current value
            repr_dict['default parameters'] = collections.OrderedDict()

            for parameter_name in dct['_parameters'].keys():

                repr_dict['default parameters'][parameter_name] = dct['_parameters'][parameter_name].to_dict()

            if has_ipython:

                display(HTML(dict_to_list(repr_dict, html=True)))

            else:

                print(dict_to_list(repr_dict, html=False))

        dct['info'] = staticmethod(info)

        # Now call the __new__ of the "type" class (which then will call the __init__ of this metaclass)

        return super(FunctionMeta, mcs).__new__(mcs, name, bases, dct)

    def __init__(cls, name, bases, dct):

        # This is the MetaClass init, which is called after the __new__ is done

        # Store the name of the function in the type
        cls._name = cls.__name__

        # Add this as a known function

        _known_functions[name] = cls

        # Finally call the init of the type class

        super(FunctionMeta, cls).__init__(name, bases, dct)

    @staticmethod
    def class_init(instance, **kwargs):

        # This is what is going to be called as the __init__ of the class, every time a new instance
        # is created

        # Create a copy of the parameters dictionary which is in the type,
        # otherwise every instance would share the same dictionary

        copy_of_parameters = collections.OrderedDict()

        # Fill it by duplicating the parameters contained in the dictionary in the type

        for key, value in type(instance)._parameters.iteritems():

            copy_of_parameters[key] = value.duplicate()

            # If the user has specified a value in the constructor, update the
            # corresponding parameters value. This allow to use a constructor as:
            # my_powerlaw = powerlaw(logK=1.0, index=-2)

            if key in kwargs:

                copy_of_parameters[key].value = kwargs[key]

        # Now check that all the parameters specified in the kwargs are actually parameters of this function
        for key in kwargs.keys():

            try:

                copy_of_parameters[key]

            except KeyError:

                raise UnknownParameter("You specified an init value for %s, which is not a "
                                       "parameter of function %s" % (key, type(instance)._name))

        # Now call the init of the corresponding class
        n_dim = type(instance)._n_dim

        if n_dim == 1:

            Function1D.__init__(instance,
                                type(instance)._name,
                                type(instance)._function_definition,
                                copy_of_parameters)

        elif n_dim == 2:

            Function2D.__init__(instance,
                                type(instance)._name,
                                type(instance)._function_definition,
                                copy_of_parameters)

        elif n_dim == 3:

            Function3D.__init__(instance,
                                type(instance)._name,
                                type(instance)._function_definition,
                                copy_of_parameters)

        # Last, if the class provides a setup method, call it
        if hasattr(instance, "_setup"):

            instance._setup()

    @staticmethod
    def check_calling_sequence(name, function_name, function, possible_variables):
        """
        Check the calling sequence for the function looking for the variables specified.
        One or more of the variables can be in the calling sequence. Note that the
        order of the variables will be enforced.
        It will also enforce that the first parameter in the calling sequence is called 'self'.

        :param function: the function to check
        :param possible_variables: a list of variables to check, The order is important, and will be enforced
        :return: a tuple containing the list of found variables, and the name of the other parameters in the calling
        sequence
        """

        # Get calling sequence

        # If the function has been memoized, it will have a "input_object" member

        try:

            calling_sequence = inspect.getargspec(function.input_object).args

        except AttributeError:

            # This might happen if the function is with memoization

            calling_sequence = inspect.getargspec(function).args

        assert calling_sequence[0] == 'self', "Wrong syntax for 'evaluate' in %s. The first argument " \
                                              "should be called 'self'." % name

        # Figure out how many variables are used

        variables = filter(lambda var: var in possible_variables, calling_sequence)

        # Check that they actually make sense. They must be used in the same order
        # as specified in possible_variables

        assert len(variables) > 0, "The name of the variables for 'evaluate' in %s must be one or more " \
                                   "among %s, instead of %s" % (name, ','.join(possible_variables), ",".join(variables))

        if variables != possible_variables[:len(variables)]:
            raise AssertionError("The variables %s are out of order in '%s' of %s. Should be %s."
                                 % (",".join(variables), function_name, name, possible_variables[:len(variables)]))

        other_parameters = filter(lambda var: var not in variables and var != 'self', calling_sequence)

        return variables, other_parameters

    @staticmethod
    def parse_parameter_definition(func_name, par_name, definition):

        # Parse definition of parameter

        # Enforce the presence of attributes 'value' and 'desc'

        if 'initial value' not in definition:
            raise FunctionDefinitionError("Error for parameter %s of function %s: value for parameter must be"
                                          " specified" % (par_name, func_name))

        if 'desc' not in definition:
            raise FunctionDefinitionError("Error for parameter %s of function %s: desc for parameter must be"
                                          " specified" % (par_name, func_name))

        # Fetch attributes

        # Use unitless parameters when building the function, if no unit is specified, otherwise
        # use that unit
        if 'unit' not in definition or definition['unit'] is None or definition['unit'] == '':

            du = u.dimensionless_unscaled

        else:

            du = u.Unit(definition['unit'])

        def _parse_value(val):

            if isinstance(val, str):

                return eval(val)

            elif val is None:

                return None

            else:

                return float(val)

        value = _parse_value(definition['initial value'])
        desc = definition['desc']

        # Optional attributes are either None if not specified, or the value specified

        min_value = (None if 'min' not in definition else _parse_value(definition['min']))
        max_value = (None if 'max' not in definition else _parse_value(definition['max']))
        delta = (None if 'delta' not in definition else _parse_value(definition['delta']))
        unit = du

        # A parameter can be fixed by using fix=yes, otherwise it is free by default

        free = (True if 'fix' not in definition else not bool(definition['fix']))

        new_parameter = Parameter(par_name, value, min_value=min_value, max_value=max_value,
                                  delta=delta, desc=desc, free=free, unit=unit)

        return new_parameter


class Function(Node):
    """
    Generic Function class. Will be subclassed in Function1D, Function2D and Function3D.

    """
    def __init__(self, name=None, function_definition=None, parameters=None):

        # I use default values only to avoid warnings from pycharm and other software about the
        # calling sequence of this contructor. We actually need to enforce its proper use,
        # with this assert

        assert name is not None and function_definition is not None and parameters is not None

        # Set up the node

        super(Function, self).__init__(name)

        # Store name, number of dimensions and the latex formula

        # Store also the function definition

        assert 'description' in function_definition,"Function definition must contain a description"

        if 'latex' not in function_definition:

            function_definition['latex'] = '$n.a.$'

        self._function_definition = function_definition

        # Add the parameters as children. Since the name of the key in the dictionary might
        # be different than the actual name of the parameter, use the .add_child method instead
        # of the add_children method

        self._parameters = collections.OrderedDict()

        for child_name, child in parameters.iteritems():

            self._parameters[child_name] = child

            # Add the parameter as a child of the function

            self._add_child(child)

        # Now generate a unique identifier (UUID) in a thread safe, multi-processing safe
        # way. This is used for example in the CompositeFunction class to keep track of the different
        # instances of the same function
        self._uuid = "{" + str(self._generate_uuid()) + "}"

        # Normal functions are able to change units, while some specific ones (such as the one from XSpec) are not.
        # In this second case, this variable will contain a tuple (x_unit, y_unit), but by default it should be
        # None

        self._fixed_units = None

    @property
    def n_dim(self):
        """
        :return: number of dimensions for this function (1, 2 or 3)
        """
        return type(self)._n_dim

    @property
    def free_parameters(self):
        """
        Returns a dictionary of free parameters for this function

        :return: dictionary of free parameters
        """

        free_parameters = collections.OrderedDict([(k,v) for k, v in self.parameters.iteritems() if v.free])

        return free_parameters

    @staticmethod
    def _generate_uuid():
        """
        Generate a unique identifier for this function.

        :return: the UUID
        """
        return uuid.UUID(bytes=os.urandom(16), version=4)

    def has_fixed_units(self):
        """
        Returns True if this function cannot change units, which is the case only for very specific functions (like
        models from foreign libraries like Xspec)

        :return: True or False
        """

        return not (self._fixed_units is None)

    @property
    def fixed_units(self):
        """
            Returns the fixed units if has_fixed_units is True (see has_fixed_units)

            :return: None, or a tuple (x_unit, y_unit)
        """

        return self._fixed_units

    @property
    def description(self):
        """
        Returns a description for this function
        """

        return self._function_definition['description']

    # Add a property returning the parameters dictionary
    @property
    def parameters(self):
        """
        Returns a dictionary of parameters
        """
        return self._parameters

    @property
    def latex(self):
        """
        Returns the LaTEX formula for this function
        """
        return self._function_definition['latex']

    # Define now all the operators which allow to combine functions. Each operator will return a new
    # instance of a CompositeFunction, which can then be used as a function on its own

    def of(self, another_function):
        """
        Compose this function with another as in this_function(another_function(x))
        :param another_function: another function to compose with the current one
        :return: a composite function instance
        """
        return CompositeFunction('of', self, another_function)

    def __neg__(self):

        return CompositeFunction('*-', self)

    def __abs__(self):

        return CompositeFunction('abs', self)

    def __pow__(self, other_instance):

        return CompositeFunction('**', self, other_instance)

    def __rpow__(self, other_instance):

        return CompositeFunction('**', other_instance, self)

    def __add__(self, other_instance):

        return CompositeFunction('+', self, other_instance)

    __radd__ = __add__

    def __sub__(self, other_instance):

        return CompositeFunction('-', self, other_instance)

    __rsub__ = __sub__

    def __mul__(self, other_instance):

        c = CompositeFunction('*', self, other_instance)

        # If the other instance is a function (and not a number), flag it so its units will be made dimensionless
        # in the set_units method of the composite function

        if isinstance(other_instance, Function):

            if self.has_fixed_units():

                # This is likely a XSpec model. The multiplication of two models with units will give the wrong
                # units to the results. So, depending on the type of the first and second model, we need to adjust
                # their units so that the result will keep the right units.

                if not self.fixed_units[1] == u.dimensionless_unscaled:

                    # This function has fixed unit and is not dimensionless (likely an additive XSpec model).
                    # We need to make the other function dimensionless so that the multiplication of them will keep
                    # the right units

                    other_instance._make_dimensionless = True

                else:

                    # This function has fixed unit, but it is dimensionless (likely a multiplicative XSpec model)
                    # The other function should keep its units, so we flag self instead

                    self._make_dimensionless = True

                    # However, if also the other function has fixed dimension and it is dimensionless
                    # (likely another XSpec multiplicative model) we need to flag that as well
                    if other_instance.has_fixed_units() and (other_instance.fixed_units[1] == u.dimensionless_unscaled):
                        other_instance._make_dimensionless = True

                        # Now we need to flag the composite function as fixed units and dimensionless
                        c._fixed_units = other_instance.fixed_units

            else:

                # We need to make the other instance dimensionless so that this function (which is not dimensionless)
                # multiplied by the other function (which we will make dimensionless) will give the right units as
                # the results

                other_instance._make_dimensionless = True

        return c

    __rmul__ = __mul__

    def __div__(self, other_instance):

        c = CompositeFunction('/', self, other_instance)

        # If the other instance is a function (and not a number), flag it so its units will be made dimensionless
        # in the set_units method of the composite function

        if isinstance(other_instance, Function):

            if self.has_fixed_units():

                # This is likely a XSpec model. Division is not supported.

                raise NotImplementedError("Division for XSpec models is not supported")

            else:

                # We need to make the other instance dimensionless

                other_instance._make_dimensionless = True

        return c

    def __rdiv__(self, other_instance):

        return CompositeFunction('/', other_instance, self)

    __truediv__ = __div__
    __rtruediv__ = __rdiv__

    def _repr__base(self, rich_output):

        repr_dict = collections.OrderedDict()

        repr_dict['description'] = self._function_definition['description']

        if 'latex' in self._function_definition:

            repr_dict['formula'] = self._function_definition['latex']

        # Add the description of each parameter and their current value
        repr_dict['parameters'] = collections.OrderedDict()

        for parameter in self._get_children():

            repr_dict['parameters'][parameter.name] = parameter.to_dict()

        return dict_to_list(repr_dict, rich_output)

    @property
    def uuid(self):
        """
        Returns the ID of the current function. The ID is used by the CompositeFunction class to keep track of the
        unique instances of each function. It should not be used by the user for any specific purpose.

        :return: (none)
        """
        return self._uuid

    def duplicate(self):
        """
        Create a copy of the current function with all the parameters equal to the current value

        :return: a new copy of the function
        """

        # Create a copy

        function_copy = copy.deepcopy(self)

        return function_copy

    def get_boundaries(self):  # pragma: no cover
        """
        Returns the boundaries of this function. By default there is no boundary, but subclasses can
        override this.

        :return: a tuple of tuples containing the boundaries for each coordinate (ra_min, ra_max), (dec_min, dec_max)
        """

        raise NotImplementedError("You have to implement this")

    def __call__(self, *args):  # pragma: no cover

        raise NotImplementedError("You have to implement this")

    def fast_call(self, *args):  # pragma: no cover

        raise NotImplementedError("You have to implement this")

    def evaluate_at(self, *args, **parameter_specification):  # pragma: no cover
        """
        Evaluate the function at the given x(,y,z) for the provided parameters, explicitly provided as part of the
        parameter_specification keywords.

        :param *args:
        :param **parameter_specification:
        :return:
        """

        # Set the parameters to the provided values
        for parameter in parameter_specification:

            self._get_child(parameter).value = parameter_specification[parameter]

        return self(*args)


class Function1D(Function):

    def __init__(self, name=None, function_definition=None, parameters=None):

        Function.__init__(self, name, function_definition, parameters)

        self._x_unit = None
        self._y_unit = None

    def evaluate(self, x, *args, **kwargs):  # pragma: no cover

        raise NotImplementedError("You have to re-implement this")

    def set_units(self, in_x_unit, in_y_unit):

        # Transform None in input to '', so that u.Unit() will generate a dimensionless unit

        in_x_unit = in_x_unit if in_x_unit is not None else ''
        in_y_unit = in_y_unit if in_y_unit is not None else ''

        # Get a Unit instance from the inputs

        try:

            in_x_unit = u.Unit(in_x_unit)
            in_y_unit = u.Unit(in_y_unit)

        except:

            raise TypeError("Could not get a Unit instance from provided units %s when setting units "
                            "for function %s" % ((in_x_unit, in_y_unit), self.name))

        # Now call the underlying method to set units, which is defined by each function
        new_units = self._set_units(in_x_unit, in_y_unit)

        # Store the units.
        # NOTE: the previous call to _set_units might return new units in special cases
        # (for example Xspec functions). So it is critical that we store them in the class' attributes
        # after the call to _set_units

        if new_units is not None:

            new_x_unit, new_y_unit = new_units

            self._x_unit = new_x_unit
            self._y_unit = new_y_unit

        else:

            self._x_unit = in_x_unit
            self._y_unit = in_y_unit

    def _set_units(self, x_unit, y_unit):  # pragma: no cover

        # This will be overridden by derived classes

        raise NotImplementedError("You have to implement the method _set_units for function %s" % self.name)

    @property
    def x_unit(self):
        """
        The unit of the independent variable
        :return: a astropy.Unit instance
        """
        return self._x_unit

    @property
    def y_unit(self):
        """
        The unit of the dependent variable
        :return: a astropy.Unit instance
        """
        return self._y_unit

    def __call__(self, x):

        # This method's code violates explicitly duck typing. The reason is that astropy.units introduce a very
        # significant overload on any computation. For this reason we treat differently the case with units from
        # the case without units, so that the latter case remains fast. Also, transforming an input
        # which is not an array into an array introduce a significant overload (10 microseconds or so), so we perform
        # this transformation only when strictly required

        # NOTE: for a single quantity such as q = (1.0 * u.keV), isinstance(q, np.ndarray) returns True

        if isinstance(x, np.ndarray):

            # We have an array as input (or a single quantity)

            if not isinstance(x, u.Quantity):

                # This is a normal array, let's use the fast call (without units)

                return self.fast_call(x)

            else:

                # This is an array with units or a single quantity, let's use the slow call which preserves units

                assert self.y_unit is not None, "In order to use units you need to use the function as a spectrum or " \
                                                "as something else, or you need to explicitly set the units."

                results = self._call_with_units(x)

                # Now convert to the expected y unit and return a astropy.Quantity by multiplying by the right unit
                return np.squeeze(results.to(self.y_unit).value) * self.y_unit

        else:

            # This is either a single number or a list

            # Transform the input to an array of floats. If x is a single number, this will be an array of size 1

            new_input = np.array(x, dtype=float, ndmin=1, copy=False)

            # Compute the function

            result = self.fast_call(new_input)

            # Now remove all dimensions of size 1. For example, an array of shape (1,) will become a single number,
            # so that if the input was a single number, also the output will be a single number

            return np.squeeze(result)

    def _call_with_units(self, x):

        # Gather the current parameters' values with units
        values = map(attrgetter("as_quantity"), self._get_children())

        try:

            results = self.evaluate(x.to(self.x_unit, equivalencies=u.spectral()), *values)

        except u.UnitsError:  # pragma: no cover
{'fraction_non_alphanumeric': 0.06726652394969226, 'fraction_numerical': 0.0014383195076264382, 'mean_word_length': 2.1949238578680204, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 64, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4130856', 'n_tokens_mistral': 15328, 'n_tokens_neox': 14373, 'n_words': 5928}