"""Implementation of :class:`RealMesh` and :class:`ComplexMesh`, along with FFT engines."""

import os
import ctypes
import numbers

import numpy as np
from numpy import ctypeslib
from numpy.lib.mixins import NDArrayOperatorsMixin as NDArrayLike

from . import utils
from .utils import BaseClass, BaseMetaClass

class MeshError(Exception):

    """Exception raised when issue with mesh."""

class SetterProperty(object):
    """
    Attribute setter, runs ``func`` when setting a class attribute.
    Taken from https://stackoverflow.com/questions/17576009/python-class-property-use-setter-but-evade-getter
    """
    def __init__(self, func, doc=None):
        self.func = func
        self.__doc__ = doc if doc is not None else func.__doc__

    def __set__(self, obj, value):
        return self.func(obj, value)

class BaseMesh(NDArrayLike,BaseClass,metaclass=BaseMetaClass):
    """
    Base implementation for mesh.
    What follows are just methods to make :class:`BaseMesh` behave like a numpy array.
    numpy functions can be applied directly to any instance ``mesh`` through e.g.::

        np.sum(mesh)

    Note
    ----
    To get a deep copy of the mesh (including :attr:`value`), use :meth:`deepcopy`.
    :meth:`copy` will return a shallow copy.

    Attributes
    ----------
    value : array
        Numpy array holding mesh values, or ``None`` when unset.
        Can be set any time using ``mesh.value = newvalue``.

    info : MeshInfo
        Mesh information (boxsize, boxcenter, nmesh, etc.).

    attrs : dict
        Dictionary of other attributes.

    boxsize : array
        See :class:`MeshInfo`.

    boxcenter : array
        See :class:`MeshInfo`.

    nmesh : array
        See :class:`MeshInfo`.

    cellsize : array
        See :class:`MeshInfo`.

    ndim : array
        See :class:`MeshInfo`.
    """
    _attrs = ['info', 'nthreads', 'attrs']
    _HANDLED_TYPES = (np.ndarray, numbers.Number)

    _path_lib = os.path.join(utils.lib_dir, 'mesh_{}.so')

    def __init__(self, value=None, info=None, nthreads=None, attrs=None, **kwargs):
        """
        Initalize :class:`BaseMesh`.

        Parameters
        ----------
        value : array, default=None
            Numpy array holding mesh values, or ``None`` (can set later through ``mesh.value = value``.

        info : MeshInfo, default=None
            Mesh information (boxsize, boxcenter, nmesh, etc.),
            copied and updated with ``kwargs``.

        nthreads : int, default=None
            Number of threads to use in mesh calculations; defaults to OpenMP's default.

        attrs : dict
            Dictionary of other attributes.

        kwargs : dict
            Arguments for :class:`MeshInfo`.
        """
        if info is None:
            self.info = MeshInfo(value=value, **kwargs)
        else:
            self.info = info.clone(value=value, **kwargs)
        self.value = None
        self.dtype = self.info.dtype
        self.value = value
        self.set_num_threads(nthreads)
        self.attrs = attrs or {}
        self.fft_engine = None

    @property
    def shape(self):
        return tuple(self.nmesh)

    @property
    def size(self):
        return np.prod(self.shape)

    @property
    def _type_float_mesh(self):
        # Return ctypes-type for numpy array
        return ctypeslib.ndpointer(dtype=self._type_float, shape=self.size, flags='C')

    @property
    def dtype(self):
        return self.info.dtype

    @dtype.setter
    def dtype(self, dtype):
        """Called when setting :attr:`dtype`, loading the relevant C-library."""
        self.info.dtype = dtype
        self.value = self.value
        self._lib = ctypes.CDLL(self._path_lib.format(self._precision), mode=ctypes.RTLD_LOCAL)

    def set_num_threads(self, nthreads=None):
        """Set number of OpenMP threads used in mesh calculations."""
        if nthreads is not None:
            func = self._lib.set_num_threads
            func.argtypes = (ctypes.c_int,)
            func(nthreads)

    @property
    def nthreads(self):
        """Number of OpenMP threads."""
        func = self._lib.get_num_threads
        func.restype = ctypes.c_int
        return func()

    def __mul__(self, other):
        r = self.deepcopy(copy_value=False)
        r.value = r.value * other
        return r

    def __imul__(self, other):
        self.value *= other
        return self

    def __div__(self, other):
        r = self.deepcopy(copy_value=False)
        r.value = r.value / other
        return r

    __truediv__ = __div__

    def __rdiv__(self, other):
        r = self.deepcopy()
        r.value = other / r.value
        return r

    __rtruediv__ = __rdiv__

    def __idiv__(self, other):
        self.value /= other
        return self

    __itruediv__ = __idiv__

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        # Taken from https://numpy.org/doc/stable/reference/generated/numpy.lib.mixins.NDArrayOperatorsMixin.html
        # See also https://github.com/rainwoodman/pmesh/blob/master/pmesh/pm.py
        out = kwargs.get('out', ())
        for x in inputs + out:
            # Only support operations with instances of _HANDLED_TYPES.
            # Use BaseMesh instead of type(self) for isinstance to
            # allow subclasses that don't override __array_ufunc__ to
            # handle BaseMesh objects.
            if not isinstance(x, self._HANDLED_TYPES + (BaseMesh,)):
                return NotImplemented

        # Defer to the implementation of the ufunc on unwrapped values.
        inputs = tuple(x.value if isinstance(x, BaseMesh) else x
                       for x in inputs)
        if out:
            kwargs['out'] = tuple(
                x.value if isinstance(x, BaseMesh) else x
                for x in out)
        result = getattr(ufunc, method)(*inputs, **kwargs)

        def cast(result):
            # booleans, cannot be reasonable BaseMesh objects
            # just return the ndarray
            if result.dtype == '?':
                return result
            # different shape, cannot be reasonable Field objects
            # just return the ndarray
            if result.shape != self.shape:
                return result
            # really only cast when we are using simple +-* **, etc.
            new = self.deepcopy(copy_value=False)
            new.value = result
            return new

        if type(result) is tuple:
            # multiple return values
            return tuple(cast(x) for x in result)
        elif method == 'at':
            # no return value
            return None
        else:
            # one return value
            return cast(result)

    def __getitem__(self, index):
        return self.value.__getitem__(index)

    def __setitem__(self, index, value):
        return self.value.__setitem__(index, value)

    def __array__(self, dtype=None):
        return self.value

    @SetterProperty
    def value(self, value):
        if value is not None:
            if isinstance(value, np.ndarray) and value.size == self.size:
                value.shape = self.shape
                value = value.astype(self.dtype, copy=False, order='C')
            else:
                value_ = value
                value = np.empty(shape=self.shape, dtype=self.dtype, order='C')
                value[...] = value_
        self.__dict__['value'] = value

    def zeros_like(self):
        new = self.deepcopy(copy_value=False)
        new.value = np.zeros(shape=self.shape, dtype=self.dtype, order='C')
        return new

    def empty_like(self, *args, **kwargs):
        new = self.deepcopy(copy_value=False)
        new.value = np.empty(shape=self.shape, dtype=self.dtype, order='C')
        return new

    def __repr__(self):
        info = ['{}={}'.format(name,getattr(self.info, name)) for name in self.info._attrs]
        return '{}({})'.format(self.__class__.__name__,', '.join(info))

    def _copy_value(self, out=None):
        if out is None: out = np.empty_like(self.value)
        func = self._lib.copy
        func.argtypes = (self._type_float_mesh, self._type_float_mesh, ctypes.c_size_t)
        func.restype = ctypes.c_int
        self.value.shape = out.shape = -1
        flag = func(self.value, out, self.size)
        if (flag != 0):
            raise MeshError('Issue with _copy_value')
        self.value.shape = out.shape = self.shape
        return out

    def deepcopy(self, copy_value=True):
        kwargs = {name:getattr(self,name) for name in self._attrs}
        kwargs['info'] = kwargs['info'].deepcopy()
        new = self.__class__(self._copy_value() if copy_value and self.value is not None else self.value,**kwargs)
        new.fft_engine = self.fft_engine
        return new

    def get_fft_engine(self, engine='numpy', **kwargs):
        """
        Return engine for fast Fourier transform.

        Parameters
        ----------
        engine : string, BaseFFTEngine, default='numpy'
            If string, use 'numpy' or 'fftw' (package pyfftw must be installed);
            else a FFT engine.

        kwargs : dict
            Options for the FFT engines, used if ``engine`` is a FFT engine name (string).
            See :class:`NumpyFFTEngine` and :class:`FFTWEngine`.

        Returns
        -------
        engine : BaseFFTEngine
            FFT engine.
        """
        kwargs.setdefault('nthreads',self.nthreads)
        return get_fft_engine(engine,shape=self.shape,type_real=self.dtype,**kwargs)

    def set_fft_engine(self, engine='numpy', **kwargs):
        """
        Set engine for fast Fourier transform.
        See :meth:`get_fft_engine`.
        """
        self.fft_engine = self.get_fft_engine(engine=engine, **kwargs)

def _make_property(name):

    @property
    def func(self):
        return getattr(self.info, name)

    return func

for name in ['boxsize', 'boxcenter', 'nmesh', 'offset', 'cellsize', 'ndim', '_precision', '_type_float']:
    setattr(BaseMesh, name, _make_property(name))

class MeshInfo(BaseClass):
    """
    Class holding mesh information.

    Attributes
    ----------
    dtype : np.dtype
        Type for mesh array.

    nmesh : array
        Mesh size, i.e. number of mesh nodes along each axis.

    boxsize : array
        Physical size of the box.

    boxcenter : array
        Box center.
    """
    _attrs = ['dtype', 'nmesh', 'boxsize', 'boxcenter']

    def __init__(self, nmesh=None, boxsize=None, boxcenter=None, cellsize=None, value=None, positions=None, boxpad=1.5, dtype=None):
        """
        Initalize :class:`MeshInfo`.

        Parameters
        ----------
        nmesh : array, int, default=None
            Mesh size, i.e. number of mesh nodes along each axis.
            If not provided, see ``value``.

        boxsize : array, float, default=None
            Physical size of the box.
            If not provided, see ``positions``.

        boxcenter : array, float, default=None
            Box center.
            If not provided, see ``positions``.

        cellsize : array, float, default=None
            Physical size of mesh cells.
            If not ``None``, and mesh size ``nmesh`` is not ``None``, used to set ``boxsize`` as ``nmesh * cellsize``.
            If ``nmesh`` is ``None``, it is set as (the nearest integer(s) to) ``boxsize/cellsize``.

        value : array, default=None
            Only used to get mesh size.

        positions : array of shape (N,3), default=None
            If ``boxsize`` and / or ``boxcenter`` is ``None``, use these positions
            to determine ``boxsize`` and / or ``boxcenter``.

        boxpad : float, default=1.5
            When ``boxsize`` is determined from ``positions``, take ``boxpad`` times the smallest box enclosing ``positions`` as ``boxsize``.

        dtype : string, np.dtype, defaut=None
            Type for :attr:`value` array.
            If ``None``, defaults to ``np.asarray(value).dtype`` if ``value`` is not ``None``, else 'f8'.
        """
        if value is not None:
            value = np.asarray(value)
            if dtype is None: dtype = value.dtype
            if nmesh is None: nmesh = value.shape
        dtype = np.dtype(dtype if dtype is not None else 'f8')

        if boxsize is None or boxcenter is None:
            if positions is None:
                raise MeshError('boxsize and boxcenter must be specified if positions are not provided')
            pos_min, pos_max = positions.min(axis=0), positions.max(axis=0)
            delta = np.abs(pos_max - pos_min)
            if boxcenter is None: boxcenter = 0.5 * (pos_min + pos_max)
            if boxsize is None:
                if cellsize is not None and nmesh is not None:
                    boxsize = nmesh * cellsize
                else:
                    boxsize = delta.max() * boxpad
            if (boxsize < delta).any(): raise MeshError('boxsize too small to contain all data')

        if nmesh is None:
            if cellsize is not None:
                nmesh = np.rint(boxsize/cellsize).astype(int)
            else:
                raise MeshError('nmesh (or cellsize) must be specified')

        self.__dict__['dtype'] = np.dtype(dtype)
        self.boxsize = boxsize
        self.boxcenter = boxcenter
        self.nmesh = nmesh

    def clone(self, **kwargs):
        """Clone current :class:`MeshInfo` instance, optionally updating attributes with ``kwargs``."""
        for name in self._attrs: kwargs.setdefault(name, getattr(self, name))
        return self.__class__(**kwargs)

    @SetterProperty
    def dtype(self, dtype):
        self.__dict__['dtype'] = np.dtype(dtype)
        self.boxsize = self.boxsize # set correct type
        self.boxcenter = self.boxcenter # set correct type

    @property
    def _precision(self):
        # Return float if float32, double if float64
        return self._type_float.__name__[len('c_'):]

    @property
    def _type_float(self):
        # Return ctypes-type corresponding to numpy-dtype
        # Take care of complex type
        dtype = np.empty(0, dtype=self.dtype).real.dtype
        return ctypeslib.as_ctypes_type(dtype)

    @SetterProperty
    def boxsize(self, boxsize):
        # Called when setting :attr:`boxsize`, enforcing array of shape (3,).
        _boxsize = np.empty(self.ndim, dtype=self._type_float, order='C')
        _boxsize[:] = boxsize
        self.__dict__['boxsize'] = _boxsize

    @SetterProperty
    def boxcenter(self, boxcenter):
        # Called when setting :attr:`boxcenter`, enforcing array of shape (3,).
        _boxcenter = np.empty(self.ndim, dtype=self._type_float, order='C')
        _boxcenter[:] = boxcenter
        self.__dict__['boxcenter'] = _boxcenter

    @SetterProperty
    def nmesh(self, nmesh):
        # Called when setting :attr:`nmesh`, enforcing array of shape (3,).
        _nmesh = np.empty(self.ndim, dtype=ctypes.c_int, order='C')
        _nmesh[:] = nmesh
        self.__dict__['nmesh'] = _nmesh

    @property
    def offset(self):
        """Coordinates of the (0,0,0) corner of the box."""
        return self.boxcenter - self.boxsize/2.

    @property
    def cellsize(self):
        "Physical size of mesh cells."
        return self.boxsize/self.nmesh

    @property
    def ndim(self):
        """Number of dimensions: 3."""
        return 3

    def wrap(self, positions):
        """Wrap input positions."""
        return ((positions - self.offset) % self.boxsize) + self.offset

    def deepcopy(self):
        import copy
        return copy.deepcopy(self)

class RealMesh(BaseMesh):

    """Class holding a 3D real mesh."""

    _path_lib = os.path.join(utils.lib_dir,'mesh_{}.so')

    def __init__(self, value=None, dtype=None, info=None, nthreads=None, attrs=None, **kwargs):
        """
        Initalize :class:`RealMesh`.

        Parameters
        ----------
        value : array, default=None
            Numpy array holding mesh values, or ``None`` (can set later through ``mesh.value = value``.

        dtype : string, np.dtype, defaut=None
            Type for :attr:`value` array. Defaults to 'f8'.

        info : MeshInfo, default=None
            Mesh information (boxsize, boxcenter, nmesh, etc.),
            copied and updated with ``kwargs``.

        nthreads : int
            Number of threads to use in mesh calculations.

        attrs : dict
            Dictionary of other attributes.

        kwargs : dict
            Arguments for :class:`MeshInfo`.
        """
        if dtype is None and (value is None or np.ndim(value) == 0): dtype = 'f8' # accept single float as input
        super(RealMesh, self).__init__(value=value, info=info, nthreads=nthreads, attrs=attrs, dtype=dtype, **kwargs)
        if 'float' not in self.dtype.name:
            raise MeshError('Provide float dtype')

    def coords(self):
        """Return array of coordinates along each axis."""
        toret = []
        for idim,(n,o,d) in enumerate(zip(self.nmesh,self.offset,self.boxsize/self.nmesh)):
            toret.append(o + d*np.arange(n))
        return tuple(toret)

    def assign_cic(self, positions, weights=None, wrap=False):
        """
        Assign (paint) positions to mesh with Cloud-in-Cell scheme.

        Parameters
        ----------
        positions : array of shape (N,3)
            Cartesian positions.

        weights : array of shape (N,), default=None
            Weights; default to 1.

        wrap : boolean, default=False
            If ``True``, wrap input particle positions into the box.
        """
        size = len(positions)
        if weights is None: weights = np.ones_like(positions,shape=size,dtype=self._type_float)
        if wrap: positions = self.info.wrap(positions)
        positions = ((positions - self.boxcenter)/self.boxsize + 0.5)*self.nmesh
        positions = positions.astype(self._type_float, copy=False).ravel(order='C')
        weights = weights.astype(self._type_float, copy=False).ravel(order='C')
        if self.value is None: self.value = 0.
        type_positions = ctypeslib.ndpointer(dtype=self._type_float,shape=positions.size,flags='C')
        type_weights = ctypeslib.ndpointer(dtype=self._type_float,shape=weights.size,flags='C')
        type_nmesh = ctypeslib.ndpointer(dtype=ctypes.c_int,shape=self.ndim,flags='C')
        func = self._lib.assign_cic
        func.argtypes = (self._type_float_mesh,type_nmesh,type_positions,type_weights,ctypes.c_size_t)
        func.restype = ctypes.c_int
        self.value.shape = -1
        flag = func(self.value,self.nmesh.astype(ctypes.c_int,copy=False),positions,weights,size)
        if (flag != 0):
            raise MeshError('Issue with assign_cic')
        self.value.shape = self.shape

    def read_cic(self, positions, wrap=False):
        """
        Read mesh values interpolated at input positions with Cloud-in-Cell scheme.

        Parameters
        ----------
        positions : array of shape (N,3)
            Cartesian positions.

        wrap : boolean, default=False
            If ``True``, wrap input particle positions into the box.

        Returns
        -------
        values : array of shape (N,)
            Mesh values interpolated at input positions.
        """
        size = len(positions)
        dtype = positions.dtype
        if wrap: positions = self.info.wrap(positions)
        positions = ((positions - self.boxcenter)/self.boxsize + 0.5)*self.nmesh
        positions = positions.astype(self._type_float,copy=False).ravel(order='C')
        values = np.empty_like(positions,shape=size,order='C')
        type_positions = ctypeslib.ndpointer(dtype=self._type_float,shape=positions.size,flags='C')
        type_values = ctypeslib.ndpointer(dtype=self._type_float,shape=values.size,flags='C')
        type_nmesh = ctypeslib.ndpointer(dtype=ctypes.c_int,shape=self.ndim,flags='C')
        func = self._lib.read_cic
        func.argtypes = (self._type_float_mesh,type_nmesh,type_positions,type_values,ctypes.c_size_t)
        func.restype = ctypes.c_int
        flag = func(self.value.ravel(order='C'),self.nmesh.astype(ctypes.c_int,copy=False),positions,values,size)
        if (flag != 0):
            raise MeshError('Issue with read_cic')
        return values.astype(dtype=dtype,copy=False)

    def read_finite_difference_cic(self, positions, wrap=False):
        """
        Read derivative (finite difference scheme) of mesh values along each axis interpolated at input positions with Cloud-in-Cell scheme.

        Parameters
        ----------
        positions : array of shape (N,3)
            Cartesian positions.

        wrap : boolean, default=False
            If ``True``, wrap input particle positions into the box.

        Returns
        -------
        values : array of shape (N,)
            Derivative of mesh values interpolated at input positions.
        """
        size = len(positions)
        dtype = positions.dtype
        if wrap: positions = self.info.wrap(positions)
        positions = ((positions - self.boxcenter)/self.boxsize + 0.5)*self.nmesh
        positions = positions.astype(self._type_float,copy=False).ravel(order='C')
        values = np.empty_like(positions,order='C')
        type_positions = ctypeslib.ndpointer(dtype=self._type_float,shape=positions.size,flags='C')
        type_nmesh = ctypeslib.ndpointer(dtype=ctypes.c_int,shape=self.ndim,flags='C')
        type_boxsize = ctypeslib.ndpointer(dtype=self._type_float,shape=self.ndim,flags='C')
        func = self._lib.read_finite_difference_cic
        func.argtypes = (self._type_float_mesh,type_nmesh,type_boxsize,type_positions,type_positions,ctypes.c_size_t)
        func.restype = ctypes.c_int
        flag = func(self.value.ravel(order='C'),self.nmesh.astype(ctypes.c_int,copy=False),self.boxsize.astype(self._type_float,copy=False),positions,values,size)
        if (flag != 0):
            raise MeshError('Issue with read_finite_difference_cic')
        values.shape = (size,self.ndim)
        return values.astype(dtype=dtype,copy=False)

    def smooth_gaussian(self, radius, method='fft', nsigmas=2.5, **kwargs):
        """
        Apply Gaussian smoothing to mesh.

        Parameters
        ----------
        radius : array, float
            Smoothing scale (along each axis, or same for all axes).

        method : string, default='fft'
            Perform Gaussian smoothing in real space ('real') or using FFT ('fft').

        nsigmas : float, default=2.5
            If ``method`` is 'real', number of Gaussian sigmas where to stop convolution.

        kwargs : dict
            Optional arguments for :meth:`get_fft_engine`.
        """
        radius_ = np.empty_like(self.boxsize,order='C')
        radius_[:] = radius
        if method == 'fft':
            if kwargs or self.fft_engine is None: self.set_fft_engine(**kwargs)
            valuek = self.to_complex()
            k2 = sum(-0.5*(r*k)**2 for r,k in zip(radius_,utils.broadcast_arrays(*valuek.coords())))
            valuek *= np.exp(k2)
            self.value = valuek.to_real().value
            #func = self._lib.smooth_fft_gaussian
            #func.argtypes = (self._type_float_mesh,type_nmesh,type_boxsize)
            #func(self.value.ravel(order='C'),self.nmesh,radius/self.boxsize)
        else:
            radius = radius_/self.boxsize
            type_nmesh = ctypeslib.ndpointer(dtype=ctypes.c_int,shape=self.ndim,flags='C')
            type_boxsize = ctypeslib.ndpointer(dtype=self._type_float,shape=self.ndim,flags='C')
            func = self._lib.smooth_gaussian
            func.argtypes = (self._type_float_mesh,type_nmesh,type_boxsize,self._type_float)
            self.value.shape = -1
            func.restype = ctypes.c_int
            flag = func(self.value, self.nmesh.astype(ctypes.c_int,copy=False),radius.astype(self._type_float,copy=False),nsigmas)
            if (flag != 0):
                raise MeshError('Issue with read_finite_difference_cic')
            self.value.shape = self.shape

    def to_complex(self, *args, **kwargs):
        """
        Return :class:`ComplexMesh` computed with fast Fourier transforms.
        See :meth:`get_fft_engine` for arguments.
        """
        if kwargs or self.fft_engine is None: self.set_fft_engine(**kwargs)
        toret = ComplexMesh(self.fft_engine.forward(self.value),info=self.info,nthreads=self.nthreads,hermitian=self.fft_engine.hermitian,attrs=self.attrs)
        toret.fft_engine = self.fft_engine
        return toret

    def prod_sum(self, arrays, exp=1):
        """
        Multiply mesh by ``(arrays[0][:,None,None] + arrays[1][None,:,None] + arrays[2][None,None,:]) ** exp``

        Parameters
        ----------
        arrays : sequence of 3 float arrays
            Arrays to multiply mesh by.

        exp : int, default=1
            Exponent to raise broadcast sum of arrays to.
        """
        if len(arrays) != 3:
            raise MeshError('Provide a sequence of 3 arrays')
        arrays = list(arrays)
        #arrays = np.concatenate(arrays[::-1], axis=0, dtype=self._type_float) # ::-1 for prod_sum
        # dtype keyword for np.concatenate appears in version 1.20.0.
        arrays = np.concatenate([np.asarray(array, dtype=self._type_float) for array in arrays[::-1]], axis=0)
        if arrays.size != sum(self.shape):
            raise MeshError('Length of input arrays must match shape')
        type_nmesh = ctypeslib.ndpointer(dtype=ctypes.c_int, shape=self.ndim, flags='C')
        type_arrays = ctypeslib.ndpointer(dtype=self._type_float, shape=arrays.size, flags='C')
        func = self._lib.prod_sum
        func.argtypes = (self._type_float_mesh, type_nmesh, type_arrays, ctypes.c_int)
        func.restype = ctypes.c_int
        self.value.shape = -1
        flag = func(self.value, self.nmesh.astype(ctypes.c_int,copy=False), arrays, exp)
        self.value.shape = self.shape
        if (flag != 0):
            raise MeshError('Issue with prod_sum')

class ComplexMesh(BaseMesh):
    """
    Class holding a 3D complex mesh.

    Parameters
    ----------
    hermitian : bool
        Whether mesh has Hermitian symmetry, i.e. is real when Fourier transformed.
        In this case, :attr:`shape` is half :attr:`nmesh` on the last axis.
    """
    _attrs = BaseMesh._attrs + ['hermitian']

    def __init__(self, value=None, dtype=None, info=None, hermitian=True, nthreads=None, attrs=None, **kwargs):
        """
        Initialize :class:`ComplexMesh`.

        Parameters
        ----------
        value : array, default=None
            Numpy array holding mesh values, or ``None`` (can set later through ``mesh.value = value``.

        dtype : string, np.dtype, defaut=None
            Type for :attr:`value` array. Defaults to 'c16'.

        info : MeshInfo, default=None
            Mesh information (boxsize, boxcenter, nmesh, etc.).

        hermitian : bool
            Whether mesh has Hermitian symmetry, i.e. is real when Fourier transformed.
            In this case, :attr:`shape` is half :attr:`nmesh` on the last axis.

        nthreads : int
            Number of threads to use in mesh calculations.

        attrs : dict
            Dictionary of other attributes.

        kwargs : dict
            Arguments for :class:`MeshInfo`.
        """
        self.hermitian = hermitian
        if dtype is None and (value is None or np.ndim(value) == 0): dtype = 'c16' # accept single float as input
        super(ComplexMesh, self).__init__(value=value, info=info, nthreads=nthreads, attrs=attrs, dtype=dtype, **kwargs)
        if 'complex' not in self.dtype.name:
            raise MeshError('Provide complex dtype')

    def _copy_value(self, out=None):
        if out is None: out = np.empty_like(self.value)
        func = self._lib.copy
        type_mesh = ctypeslib.ndpointer(dtype=self._type_float, shape=2*self.size, flags='C')
        func.argtypes = (type_mesh, type_mesh, ctypes.c_size_t)
        func.restype = ctypes.c_int
        value_view = self.value.view(dtype=self._type_float)
        out_view = out.view(dtype=self._type_float)
        value_view.shape = out_view.shape = -1
        flag = func(value_view, out_view, 2*self.size)
        out = out_view.view(dtype=self.dtype)
        out.shape = self.shape
        if (flag != 0):
            raise MeshError('Issue with _copy_value')
        return out

    @property
    def shape(self):
        if self.hermitian:
            return tuple(self.nmesh[:-1]) + (self.nmesh[-1]//2 + 1,)
        return tuple(self.nmesh)

    @property
    def fundamental_freq(self):
        """Fundamental frequency of the mesh along each axis."""
        return 2.*np.pi/self.boxsize

    def coords(self):
        """Return array of frequency (wavenumbers) along each axis."""
        toret = []
        for idim,(n,d) in enumerate(zip(self.nmesh, self.boxsize/self.nmesh)):
            if (not self.hermitian) or (idim < self.ndim - 1):
                toret.append(2*np.pi*np.fft.fftfreq(n,d=d))
            else:
                toret.append(2*np.pi*np.fft.rfftfreq(n,d=d))
        return tuple(toret)

    def get_fft_engine(self, engine='numpy', **kwargs):
        """Same as :meth:`RealMesh.get_fft_engine`."""
        kwargs.setdefault('nthreads', self.nthreads)
        return get_fft_engine(engine, shape=self.nmesh, type_complex=self.dtype, hermitian=self.hermitian, **kwargs)

    def to_real(self, *args, **kwargs):
        """
        Return :class:`RealMesh` computed with fast Fourier transforms.
        See :meth:`get_fft_engine` for arguments.
        Raises a :class:`MeshError` if FFT engine has not same Hermitian symmetry.
        """
        if kwargs or self.fft_engine is None: self.set_fft_engine(**kwargs)
        if self.fft_engine.hermitian != self.hermitian:
            raise MeshError('ComplexMesh has hermitian = {} but provided FFT engine has hermitian = {}'.format(self.hermitian,self.fft_engine.hermitian))
        value = self.value
        kwargs = {}
        if isinstance(self.fft_engine, FFTWEngine):
            if self.fft_engine.hermitian: # input destroyed only when hermitian
                value = self._copy_value()
            kwargs = {'destroy_input':True}
        toret = RealMesh(self.fft_engine.backward(value, **kwargs).real, info=self.info, nthreads=self.nthreads, attrs=self.attrs)
        toret.fft_engine = self.fft_engine
        return toret

    def prod_sum(self, arrays, exp=1):
        """
        Multiply mesh by ``(arrays[0][:,None,None] + arrays[1][None,:,None] + arrays[2][None,None,:]) ** exp``

        Parameters
        ----------
        arrays : sequence of 3 float arrays
            Arrays to multiply mesh by.

        exp : int, default=1
            Exponent to raise broadcast sum of arrays to.
        """
        if len(arrays) != 3:
            raise MeshError('Provide a sequence of 3 arrays')
        arrays = list(arrays)
        arrays[-1] = np.repeat(arrays[-1], 2)
        #arrays = np.concatenate(arrays[::-1], axis=0, dtype=self._type_float) # ::-1 for prod_sum
        # dtype keyword for np.concatenate appears in version 1.20.0.
        arrays = np.concatenate([np.asarray(array, dtype=self._type_float) for array in arrays[::-1]], axis=0)
        if arrays.size != sum(self.shape) + self.shape[-1]:
            raise MeshError('Length of input arrays must match shape')
        shape = np.asarray(self.shape, dtype=ctypes.c_int)
        shape[-1] *= 2
        type_mesh = ctypeslib.ndpointer(dtype=self._type_float, shape=np.prod(shape), flags='C')
        type_nmesh = ctypeslib.ndpointer(dtype=ctypes.c_int, shape=self.ndim, flags='C')
        type_arrays = ctypeslib.ndpointer(dtype=self._type_float, shape=arrays.size, flags='C')
        func = self._lib.prod_sum
        func.argtypes = (type_mesh, type_nmesh, type_arrays, ctypes.c_int)
        func.restype = ctypes.c_int
        #value = np.array(self.value)
        #print(value.shape, value.size)
        #value.shape = (value.size,)
        self.value.shape = -1
        value_view = self.value.view(dtype=self._type_float)
        flag = func(value_view, shape, arrays, exp)
        self.value = value_view.view(dtype=self.dtype)
        if (flag != 0):
            raise MeshError('Issue with prod_sum')

class BaseFFTEngine(BaseClass):
    """
    Base engine for fast Fourier transforms.
    FFT engines should extend this class, by (at least) implementing:

    - :meth:`forward`
    - :meth:`backward`

    Attributes
    ----------
    shape : tuple
        Shape of array (in real-space, i.e. not accounting for Hermitian symmetry) to transform.

    nthreads : int
        Number of threads.

    type_real : np.dtype
        Type for real values.

    type_complex : np.dtype
        Type for complex values. Twice larger than :attr:`type_float`.

    hermitian : bool
        Whether complex array has Hermitian symmetry, i.e. is real when Fourier transformed.
    """
    def __init__(self, shape, nthreads=None, type_complex=None, type_real=None, hermitian=True, **kwargs):
        """
        Initialize FFT engine.
        Default types are 'c16' for :attr:`type_complex` and 'f8' for :attr:`type_float`.

        Parameters
        ----------
        shape : list, tuple
            Array shape.

        nthreads : int, default=None
            Number of threads.

        type_complex : string, np.dtype, default=None
            Type for complex values.
            If not provided, use ``type_real`` instead.

        type_real : string, np.dtype, default=None
            Type for real values.
            If not provided, use ``type_complex`` instead.
        """
        if nthreads is None:
            self.nthreads = int(os.environ.get('OMP_NUM_THREADS', '1'))
        else:
            self.nthreads = nthreads
        self.shape = tuple(shape)
        if type_complex is not None:
            self.type_complex = np.dtype(type_complex)
            itemsize = np.dtype(self.type_complex).itemsize
            self.type_real = np.dtype('f{:d}'.format(itemsize//2))
        else:
            if type_real is None: type_real = 'f8'
            self.type_real = np.dtype(type_real)
            itemsize = np.dtype(self.type_real).itemsize
            self.type_complex = np.dtype('c{:d}'.format(itemsize*2))
        self.hermitian = hermitian

    @property
    def ndim(self):
        """Number of dimensions."""
        return len(self.shape)

    @property
    def size(self):
        """Size of array (in real-space, i.e. not accounting for Hermitian symmetry) to transform."""
        return np.prod(self.shape)

    @property
    def hshape(self):
        """Shape in Fourier-space, accounting for Hermitian symmetry."""
        if self.hermitian:
            return self.shape[:-1] + (self.shape[-1]//2 + 1,)
        return tuple(self.shape)

    def forward(self, fun):
        """Return forward transform of ``fun``."""
        raise NotImplementedError('Implement "forward" method in your "BaseFFTEngine"-inherited FFT engine.')

    def backward(self, fun):
        """Return backward transform of ``fun``."""
        raise NotImplementedError('Implement "backward" method in your "BaseFFTEngine"-inherited FFT engine.')

class NumpyFFTEngine(BaseFFTEngine):

    """FFT engine based on :mod:`numpy.fft`."""

    def forward(self, fun):
        """Return forward transform of ``fun``."""
        if self.hermitian:
            return np.fft.rfftn(fun).astype(self.type_complex, copy=False)
        return np.fft.fftn(fun).astype(self.type_complex, copy=False)

    def backward(self, fun):
        """Return backward transform of ``fun``."""
        if self.hermitian:
            return np.fft.irfftn(fun).astype(self.type_real, copy=False)
        return np.fft.ifftn(fun).astype(self.type_complex, copy=False)

try: import pyfftw
except ImportError: pyfftw = None

class FFTWEngine(BaseFFTEngine):

    """FFT engine based on :mod:`pyfftw`."""

    def __init__(self, shape, nthreads=None, wisdom=None, save_wisdom=None, plan='measure', **kwargs):
        """
        Initialize :mod:`pyfftw` engine.

        Note
        ----
        :class:`pyfftw.FFTW` internally stores :attr:`pyfftw.FFTW._input_array` and :attr:`pyfftw.FFTW._output_array`,
        which is a waste of memory if one does not want to save them.
        e.g. performing ``engine.backward(engine.forward(array))`` would take as much as 3 times
        (2 for the forward transform, and 1 output array in the backward transform) the memory footprint of ``array``.
        As no access is provided to :attr:`pyfftw.FFTW._input_array` and :attr:`pyfftw.FFTW._output_array` attributes,
        we choose to destroy and rebuild :class:`pyfftw.FFTW` for each transform, thereby allowing Python to destroy
        undesired arrays, at a relatively modest overhead (~ 0.5 s).

        Parameters
        ----------
        shape : list, tuple
            Array shape.

        nthreads : int, default=None
            Number of threads.

        wisdom : string, tuple, default=None
            Precomputed :mod:`pyfftw` wisdom, used to accelerate FFTs.
            If a string, should be a path to previously saved FFT wisdom (with :func:`numpy.save`).
            If a tuple, directly corresponds to the wisdom.
            By default the wisdom given in ``save_wisdom`` will be loaded, if exists.

        save_wisdom : bool, string, default=None
            If not ``None``, path where to save the wisdom.
            If ``True``, the wisdom will be saved in the default path:
            'wisdom.shape-{shape[0]}-{shape[1]}-{shape[2]}.type-{type}.nthreads-{nthreads}.npy'.

        plan : string, default='measure'
            Choices are ['estimate', 'measure', 'patient', 'exhaustive'].
            The increasing amount of effort spent during the planning stage to create the fastest possible transform.
            Usually 'measure' is a good compromise.

        kwargs : dict
            Optional arguments for :class:`BaseFFTEngine`.
        """
        if pyfftw is None:
            raise NotImplementedError('Install pyfftw to use {}'.format(self.__class__.__name__))
        super(FFTWEngine, self).__init__(shape, nthreads=nthreads, **kwargs)
        plan = plan.lower()
        allowed_plans = ['estimate', 'measure', 'patient', 'exhaustive']
        if plan not in allowed_plans:
            raise MeshError('Plan {} unknown'.format(plan))
        plan = 'FFTW_{}'.format(plan.upper())

        dtype = self.type_real if self.hermitian else self.type_complex
        wisdom_fn = 'wisdom.shape-{}.type-{}.nthreads-{:d}.npy'.format('-'.join(['{:d}'.format(s) for s in self.shape]), dtype.name, self.nthreads)
        # Should we save wisdom?
        if save_wisdom and isinstance(save_wisdom, str):
            wisdom_fn = save_wisdom
        save_wisdom = bool(save_wisdom)

        if wisdom is None:
            try:
                wisdom = np.load(wisdom_fn)
                pyfftw.import_wisdom(wisdom)
            except:
                pass
            else:
                self.log_info('Loading wisdom from {}.'.format(wisdom_fn))
        elif isinstance(wisdom, str):
            self.log_info('Loading wisdom from {}.'.format(wisdom))
            wisdom = tuple(np.load(wisdom))
        else:
            pyfftw.import_wisdom(wisdom)

        fftw_f = pyfftw.empty_aligned(self.shape, dtype=dtype, order='C')
        fftw_fk = pyfftw.empty_aligned(self.hshape, dtype=self.type_complex, order='C')
        self.flags = (plan,)
        self.fftw_forward_object = pyfftw.FFTW(fftw_f, fftw_fk, axes=range(self.ndim), direction='FFTW_FORWARD', flags=self.flags, threads=self.nthreads)
        self.fftw_backward_object = pyfftw.FFTW(fftw_fk, fftw_f, axes=range(self.ndim), direction='FFTW_BACKWARD', flags=self.flags, threads=self.nthreads)
        # We delete these instances to save memory, see note above
        self.fftw_forward_object, self.fftw_backward_object = None, None
        # Allow the wisdom to be accessed from outside
        self.wisdom = pyfftw.export_wisdom()
        if save_wisdom:
            self.log_info('Saving wisdom to {}.'.format(wisdom_fn))
            np.save(wisdom_fn, self.wisdom)

    def forward(self, fun):
        """Return forward transform of ``fun``."""
        output_array = pyfftw.empty_aligned(self.hshape, dtype=self.type_complex, order='C')
        #if self.hermitian:
        #    input_array = pyfftw.empty_aligned(self.shape,dtype=self.type_real,order='C')
        #else:
        #    input_array = pyfftw.empty_aligned(self.shape,dtype=self.type_complex,order='C')
        if self.hermitian:
            fun = fun.astype(self.type_real, copy=False)
        else:
            fun = fun.astype(self.type_complex, copy=False)
        if self.fftw_forward_object is None:
            fftw_forward_object = pyfftw.FFTW(fun, output_array, axes=range(self.ndim), direction='FFTW_FORWARD', flags=self.flags, threads=self.nthreads)
            #input_array[...] = fun
            toret = fftw_forward_object(normalise_idft=True)
        else:
            toret = self.fftw_forward_object(input_array=fun, output_array=output_array, normalise_idft=True)
        return toret

    def backward(self, fun, destroy_input=True):
        """Return backward transform of ``fun``; ``destroy_input = True`` to allow destroy ``fun`` (in case dimension > 1 and hermitian)."""
        if destroy_input:
            input_array = fun
        else:
            input_array = pyfftw.empty_aligned(self.hshape, dtype=self.type_complex, order='C')
            input_array[...] = fun
        if self.hermitian:
            output_array = pyfftw.empty_aligned(self.shape, dtype=self.type_real, order='C')
        else:
            output_array = pyfftw.empty_aligned(self.shape, dtype=self.type_complex, order='C')
        if self.fftw_backward_object is None:
            fftw_backward_object = pyfftw.FFTW(input_array, output_array, axes=range(self.ndim), direction='FFTW_BACKWARD', flags=self.flags, threads=self.nthreads)
            toret = fftw_backward_object(normalise_idft=True)
        else:
            toret = self.fftw_backward_object(input_array=fun, output_array=output_array, normalise_idft=True)
        return toret

def get_fft_engine(engine, *args, **kwargs):
    """
    Return FFT engine.

    Parameters
    ----------
    engine : BaseFFTEngine, string
        FFT engine, or one of ['numpy', 'fftw'].

    args, kwargs : tuple, dict
        Arguments for FFT engine.

    Returns
    -------
    engine : BaseFFTEngine
    """
    if isinstance(engine, str):
        if engine.lower() == 'numpy':
            return NumpyFFTEngine(*args, **kwargs)
        if engine.lower() == 'fftw':
            return FFTWEngine(*args, **kwargs)
        raise ValueError('FFT engine {} is unknown'.format(engine))
    return engine


# Standard library
import os, abc
import pickle
import numpy as np
from copy import deepcopy

# Third-party
from scipy.integrate import quad
from scipy.interpolate import interp1d
from astropy.table import Table
from astropy import units as u

# Project
from .. import MIST_PATH
from ..util import check_random_state, check_units
from ..log import logger
from ..filters import *
from .imf import sample_imf, build_galaxy, imf_dict, IMFIntegrator
from .isochrones import MISTIsochrone

__all__ = ['SSP', 'MISTSSP', 'constant_sb_stars_per_pix']

class StellarPopulation(metaclass=abc.ABCMeta):
    """
    Stellar population base class.

    Parameters
    ----------
    distance : float or `~astropy.units.Quantity`, optional
        Distance to source. If float is given, the units are assumed
        to be `~astropy.units.Mpc`. Default distance is 10 `~astropy.units.pc`
        (i.e., the mags are in absolute units).
    imf : str, optional
        The initial stellar mass function. Default is `'kroupa'`.
    """

    def __init__(self, distance=10.0 * u.pc, imf='kroupa', imf_kw={}):
        self.imf = imf
        self.imf_kw = imf_kw
        self.distance = check_units(distance, 'Mpc')

    def build_pop(self, num_stars=None, **kwargs):
        """Build stellar population."""
        return NotImplementedError()

    @property
    def dist_mod(self):
        """The distance modulus."""
        return 5 * np.log10(self.distance.to('pc').value) - 5

    @property
    def num_pops(self):
        """
        The number of simple stellar populations that composes this pop.
        """
        return 1 if type(self.isochrone) != list else len(self.isochrone)

    @property
    def total_initial_live_mass(self):
        """Total initial stellar mass in solar units."""
        return self.initial_masses.sum() * u.M_sun

    @property
    def num_stars(self):
        """Number stars in population."""
        return len(self.star_masses)

    @property
    def abs_mag_table(self):
        """Absolute magnitudes in a `~astropy.table.Table` object."""
        return Table(self.abs_mags)

    @property
    def mag_table(self):
        """Apparent magnitudes in a `~astropy.table.Table` object."""
        _mags = {}
        for filt in self.filters:
            _mags[filt] = self.star_mags(filt)
        return Table(_mags)

    def to_pickle(self, file_name):
        """Pickle stellar population object."""
        pkl_file = open(file_name, 'wb')
        pickle.dump(self, pkl_file)
        pkl_file.close()

    @staticmethod
    def from_pickle(file_name):
        """Load pickle of stellar population object."""
        pkl_file = open(file_name, 'rb')
        data = pickle.load(pkl_file)
        pkl_file.close()
        return data

    def _remnants_factor(self, **kwargs):
        """
        Correction factor to account for stellar remnants in the final mass.
        The mass in luminous stars is given by M_total * factor.
        """
        m_without_remnants = self.isochrone.ssp_surviving_mass(
            self.imf, add_remnants=False, **kwargs)
        m_with_remnants = self.isochrone.ssp_surviving_mass(
            self.imf, add_remnants=True, **kwargs)
        factor = m_without_remnants / m_with_remnants
        return factor

    def copy(self):
        return deepcopy(self)

    def set_distance(self, distance):
        """
        Change the distance to the stellar population.

        Parameters
        ----------
        distance : float or `~astropy.units.Quantity`
            Distance to source. If float is given, the units are assumed
            to be `~astropy.units.Mpc`.
        """
        self.distance = check_units(distance, 'Mpc')

    def star_mags(self, bandpass, select=None):
        """
        Get the stellar apparent magnitudes.

        Parameters
        ----------
        bandpass : str
            Filter of observation. Must be a filter in the given
            photometric system(s).
        select : `~numpy.ndarray`, optional
            Boolean mask for selecting stars (True for stars to include).

        Returns
        -------
        mags : `~numpy.ndarray`
            The stellar apparent magnitudes in the given bandpass.
        """
        mags = self.abs_mags[bandpass] + self.dist_mod
        if select is not None:
            mags = mags[select]
        return mags

    def mag_integrated_component(self, bandpass):
        """
        Get the magnitude of the integrated component of the population if
        it exists.

        Parameters
        ----------
        bandpass : str
            Filter of observation. Must be a filter in the given
            photometric system(s).

        Returns
        -------
        mag : float
            The integrated magnitude if it exists. Otherwise None is returned.
        """
        if hasattr(self, 'integrated_abs_mags'):
            mag = self.integrated_abs_mags[bandpass] + self.dist_mod
        else:
            mag = None
        return mag

    def sbf_mag(self, bandpass):
        """
        Calculate the apparent SBF magnitude of the stellar population.

        Parameters
        ----------
        bandpass : str
            Filter of observation. Must be a filter in the given
            photometric system(s).

        Returns
        -------
        mbar : float
            The apparent SBF magnitude of the stellar population in the given
            bandpass.
        """
        integrated = self.mag_integrated_component(bandpass)
        if integrated is not None:
            f_int = 10**(-0.4*integrated)
            log_dddd = 4 * np.log10(self.distance.to('cm').value)
            ff_int = 10**(self._integrated_log_lumlum[bandpass] - log_dddd)
        else:
            f_int = 0.0
            ff_int = 0.0
        f_i = 10**(-0.4 * (self.star_mags(bandpass)))
        ff = np.sum(f_i**2) + ff_int
        f = np.sum(f_i) + f_int
        mbar = -2.5 * np.log10(ff / f)
        return mbar

    def integrated_color(self, blue, red, select=None):
        """
        Calculate the population's integrated color.

        Parameters
        ----------
        blue : str
            The blue bandpass. Must be a filter in the
            given photometric system(s).
        red : str
            The red bandpass. Must be a filter in the
            given photometric system(s).
        select : `~numpy.ndarray`, optional
            Boolean mask for selecting stars (True for stars to include).

        Returns
        -------
        color : float
            The integrated color.
        """
        blue_int = self.mag_integrated_component(blue)
        if select is None and blue_int is not None:
            f_blue_int = 10**(-0.4*blue_int)
            f_red_int = 10**(-0.4*self.mag_integrated_component(red))
        else:
            f_blue_int = 0.0
            f_red_int = 0.0
        blue_mag = self.star_mags(blue, select)
        F_blue = np.sum(10**(-0.4 * blue_mag)) + f_blue_int
        red_mag = self.star_mags(red, select)
        F_red = np.sum(10**(-0.4 * red_mag)) + f_red_int
        color = -2.5 * np.log10(F_blue / F_red)
        return color

    def mean_mag(self, bandpass, select=None):
        """
        Calculate the population's mean magnitude.

        Parameters
        ----------
        bandpass : str
            Filter of observation. Must be a filter in the given
            photometric system(s).
        select : `~numpy.ndarray`, optional
            Boolean mask for selecting stars (True for stars to include).

        Returns
        -------
        mag : float
            The mean magnitude in the given bandpass.
        """
        integrated = self.mag_integrated_component(bandpass)
        if select is None and integrated is not None:
            f_int = 10**(-0.4*integrated)
            n_int = self.num_stars_integrated
        else:
            f_int = 0.0
            n_int = 0.0
        mags = self.star_mags(bandpass, select)
        mean_flux = ((10**(-0.4*mags)).sum() + f_int) / (len(mags) + n_int)
        mag = -2.5 * np.log10(mean_flux)
        return mag

    def total_mag(self, bandpass, select=None):
        """
        Calculate the population's total magnitude.

        Parameters
        ----------
        bandpass : str
            Filter of observation. Must be a filter in the given
            photometric system(s).
        select : `~numpy.ndarray`, optional
            Boolean mask for selecting stars (True for stars to include).

        Returns
        -------
        mag : float
            The total magnitude in the given bandpass.
        """
        integrated = self.mag_integrated_component(bandpass)
        if select is None and integrated is not None:
            f_int = 10**(-0.4*integrated)
        else:
            f_int = 0.0
        mags = self.star_mags(bandpass, select)
        total_flux = (10**(-0.4*mags)).sum() + f_int
        mag = -2.5 * np.log10(total_flux)
        return mag

class SSP(StellarPopulation):
    """
    Generic Simple Stellar Population (SSP).

    .. note::
        You must give `total_mass` *or* `num_stars`.

    Parameters
    ----------
    isochrone : `~artpop.stars.Isochrone`
        Isochrone object.
    num_stars : int or `None`
        Number of stars in source. If `None`, then must give `total_mass`.
    total_mass : float or `~astropy.units.Quantity` or `None`
        Stellar mass of the source. If `None`, then must give `num_stars`. This
        mass accounts for stellar remnants when ``add_remnants = True``, which
        means the actual sampled mass will be less than the given value. If
        float is given, the units are assumed to be solar masses.
    distance : float or `~astropy.units.Quantity`, optional
        Distance to source. If float is given, the units are assumed
        to be `~astropy.units.Mpc`. Default distance is 10 `~astropy.units.pc`.
    mag_limit : float, optional
        Only sample individual stars that are brighter than this magnitude. All
        fainter stars will be combined into an integrated component. Otherwise,
        all stars in the population will be sampled. You must also give the
        `mag_limit_band` if you use this parameter.
    mag_limit_band : str, optional
        Bandpass of the limiting magnitude. You must give this parameter if
        you use the `mag_limit` parameter.
    imf : str, optional
        The initial stellar mass function. Default is `'kroupa'`.
    imf_kw : dict, optional
        Optional keyword arguments for sampling the stellar mass function.
    mass_tolerance : float, optional
        Tolerance in the fractional difference between the input mass and the
        final mass of the population. The parameter is only used when
        `total_mass` is given.
    add_remnants : bool, optional
        If True (default), apply scaling factor to total mass to account for
        stellar remnants in the form of white dwarfs, neutron stars,
        and black holes.
    random_state : `None`, int, list of ints, or `~numpy.random.RandomState`
        If `None`, return the `~numpy.random.RandomState` singleton used by
        ``numpy.random``. If `int`, return a new `~numpy.random.RandomState`
        instance seeded with the `int`.  If `~numpy.random.RandomState`,
        return it. Otherwise raise ``ValueError``.
    """

    def __init__(self, isochrone, num_stars=None, total_mass=None,
                 distance=10*u.pc, mag_limit=None, mag_limit_band=None,
                 imf='kroupa', imf_kw={}, mass_tolerance=0.01,
                 add_remnants=True, random_state=None):
        super(SSP, self).__init__(distance=distance, imf=imf, imf_kw=imf_kw)
        self.isochrone = isochrone
        self.filters = isochrone.filters
        self.mag_limit = mag_limit
        self.mag_limit_band = mag_limit_band
        self.rng = check_random_state(random_state)
        self.build_pop(num_stars, total_mass, mass_tolerance, add_remnants)
        self._r = {'M_star': f'{self.total_mass.value:.2e} M_sun'}

    def build_pop(self, num_stars=None, total_mass=None, mass_tolerance=0.01,
                  add_remnants=True):
        """
        Build the stellar population. You must give `total_mass`
        *or* `num_stars` as an argument.

        Parameters
        ----------
        num_stars : int or `None`
            Number of stars in source. If `None`, then must give `total_mass`.
        total_mass : float or `~astropy.units.Quantity` or `None`
            Stellar mass of the source. If `None`, then must give `num_stars`.
            This mass accounts for stellar remnants when ``add_remnants = True``
            which means the actual sampled mass will be less than the given
            value. If float is given, the units are assumed to be solar masses.
        mass_tolerance : float, optional
            Tolerance in the fractional difference between the input mass and
            the final mass of the population. The parameter is only used when
            `total_mass` is given.
        add_remnants : bool, optional
            If True (default), apply scaling factor to total mass to account
            for stellar remnants in the form of white dwarfs, neutron stars,
            and black holes.
        """

        # get isochrone object and info
        m_min, m_max = self.isochrone.m_min, self.isochrone.m_max
        imf_kw = self.imf_kw.copy()
        iso = self.isochrone
        imfint = IMFIntegrator(self.imf, m_min=m_min, m_max=m_max)
        remnants_factor = self._remnants_factor() if add_remnants else 1.0

        # calculate the fraction of stars we will sample
        m_lim, f_num_sampled, f_mass_sampled = self.sample_fraction(
            self.mag_limit,
            self.mag_limit_band
        )

        # the limiting mass is min mass if 100% of stars will be sampled
        self.has_integrated_component = m_lim != m_min

        m_min = m_lim
        self.sampled_mass_lower_limit = m_lim
        self.frac_num_sampled = f_num_sampled
        self.frac_mass_sampled = f_mass_sampled

        if num_stars is not None:

            # sample imf
            num_stars_sample = int(num_stars * f_num_sampled)
            self.initial_masses = sample_imf(
                num_stars_sample, m_min=m_min, m_max=m_max, imf=self.imf,
                random_state=self.rng, imf_kw=imf_kw)

            # star masses are interpolated from "actual" mass
            self.star_masses = iso.interpolate('mact', self.initial_masses)

            # will be < num_stars if f_num_sampled < 1.0.
            self.num_stars_integrated = int(num_stars - len(self.star_masses))
            self.sampled_mass = self.star_masses.sum()

        elif total_mass is not None:

            total_mass = check_units(total_mass, 'Msun').to('Msun').value

            # calculate fraction of mass that remains after mass loss
            mass_loss = iso.ssp_surviving_mass(
                imf=self.imf, m_min=iso.m_min, m_max=iso.m_max,
                add_remnants=False)

            # we sample less mass than total_mass to account for
            # stellar remnants and the sample fraction
            sampled_mass = total_mass * remnants_factor * f_mass_sampled

            # we increase the sampled mass to account for mass loss
            sampled_mass /= mass_loss

            # sample initial masses
            mean_mass = imfint.m_integrate(m_min, m_max)
            mean_mass /= imfint.integrate(m_min, m_max)
            num_stars_iter = int(mass_tolerance * sampled_mass / mean_mass)
            self.initial_masses = build_galaxy(
                sampled_mass, m_min=m_min, m_max=m_max, imf=self.imf,
                random_state=self.rng, num_stars_iter=num_stars_iter, **imf_kw)

            # star masses are interpolated from "actual" mass
            self.star_masses = iso.interpolate('mact', self.initial_masses)
            self.sampled_mass = self.star_masses.sum()

            # calculate approximate number of stars in integrated component
            factor = (1 - f_num_sampled) / f_num_sampled
            self.num_stars_integrated = int(self.num_stars * factor)

        else:

            raise Exception('you must give total mass *or* number of stars')

        self.abs_mags = {}
        for filt in self.filters:
            self.abs_mags[filt] = iso.interpolate(filt, self.initial_masses)

        # update masses and mags if there is an integrated component
        if self.has_integrated_component:

            # find evolved stars that are fainter than mag_limit
            sampled_mags = self.abs_mags[self.mag_limit_band] + self.dist_mod
            evolved_faint = sampled_mags > self.mag_limit

            evolved_mags = {}
            for filt in self.filters:
                evolved_mags[filt] = self.abs_mags[filt][evolved_faint]
                self.abs_mags[filt] = self.abs_mags[filt][~evolved_faint]

            # calculate evolved mass and update integrated star count
            evolved_mass = self.star_masses[evolved_faint].sum()
            self.num_stars_integrated += evolved_faint.sum()

            # update initial and sampled masses
            self.initial_masses = self.initial_masses[~evolved_faint]
            self.star_masses = self.star_masses[~evolved_faint]
            self.sampled_mass = self.star_masses.sum()

            # calculate normalized IMF weights
            w = iso.imf_weights(self.imf, m_max_norm=m_max, norm_type='number')
            _, arg = iso.nearest_mini(m_lim)

            # update total mass
            num_stars = self.num_stars_integrated + self.num_stars
            _mass = num_stars * np.sum(iso.mact[:arg] * w[:arg])
            _mass += evolved_mass
            self.total_mass = (self.sampled_mass + _mass) / remnants_factor

            # updated integrated absolute magnitudes and luminosity variances
            self.integrated_abs_mags = {}
            self._integrated_log_lumlum = {}
            for filt in iso.filters:
                mag = iso.mag_table[filt]
                flux  = num_stars * np.sum(10**(-0.4 * mag[: arg]) * w[: arg])
                flux += np.sum(10**(-0.4 * evolved_mags[filt]))
                self.integrated_abs_mags[filt] = -2.5 * np.log10(flux)
                ff = num_stars * np.sum(10**(-0.8 * mag[: arg]) * w[: arg])
                ff += np.sum(10**(-0.8 * evolved_mags[filt]))
                log_dddd = 4 * np.log10((10 * u.pc).to('cm').value)
                self._integrated_log_lumlum[filt] = np.log10(ff) + log_dddd
        else:
            # calculate total_mass with stellar remnants
            self.total_mass = self.sampled_mass / remnants_factor

        self.live_star_mass = self.total_mass * remnants_factor
        self.ssp_labels = np.ones(len(self.star_masses), dtype=int)

        self.total_mass *= u.Msun
        self.sampled_mass *= u.Msun
        self.live_star_mass *= u.Msun

        for attr in ['eep', 'log_L', 'log_Teff']:
            if hasattr(iso, attr):
                if getattr(iso, attr) is not None:
                    vals_interp = iso.interpolate(attr, self.initial_masses)
                    setattr(self, attr, vals_interp)

    def sample_fraction(self, mag_limit, mag_limit_band):
        """
        Calculate the fraction of stars by mass and number that will be
        sampled with the give limiting magnitude.

        Parameters
        ----------
        mag_limit : float, optional
            Only sample individual stars that are brighter than this magnitude.
        mag_limit_band : str, optional
            Bandpass of the limiting magnitude.

        Returns
        -------
        m_lim : float
            Initial stellar mass associated with `mag_limit`.
        f_num_sampled : float
            Fraction of stars that will be sampled by number.
        f_mass_sampled : float
            Fraction of stars that will be sampled by mass.
        """
        iso = self.isochrone
        imfint = IMFIntegrator(self.imf, iso.m_min, iso.m_max)
        m_lim = iso.m_min
        f_num_sampled = 1.0
        f_mass_sampled = 1.0
        if mag_limit is not None:
            if  mag_limit_band is None:
                raise Exception('Must give bandpass of limiting magnitude.')
            mags = iso.mag_table[mag_limit_band] + self.dist_mod
            if mag_limit < mags.max() and mag_limit > mags.min():
                m_lim = iso.mag_to_mass(
                    mag_limit - self.dist_mod, mag_limit_band).min()
                f_num_sampled  = imfint.integrate(m_lim, iso.m_max, True)
                f_mass_sampled = imfint.m_integrate(m_lim, iso.m_max, True)
            else:
                logger.warning(f'mag_lim = {mag_limit} is outside mag range.')
        return m_lim, f_num_sampled, f_mass_sampled

    def __add__(self, ssp):
        assert StellarPopulation in ssp.__class__.__mro__, 'invalid type(s)'
        assert self.filters == ssp.filters, 'must have same filters'
        assert self.distance == ssp.distance, 'SSPs must have same distance'
        new = deepcopy(self)
        if type(new.isochrone) != list:
            new.isochrone = [new.isochrone]
        if type(ssp.isochrone) != list:
            new.isochrone.append(ssp.isochrone)
        else:
            new.isochrone.extend(ssp.isochrone)

        if not hasattr(new, 'ssp_total_masses'):
            new.ssp_total_masses = [new.total_mass]
        if not hasattr(ssp, 'ssp_total_masses'):
            ssp.ssp_total_masses = [ssp.total_mass]
        new.ssp_total_masses.extend(ssp.ssp_total_masses)

        if not hasattr(new, 'ssp_total_num_stars'):
            _n = new.num_stars + new.num_stars_integrated
            new.ssp_total_num_stars = [_n]
        if not hasattr(ssp, 'ssp_total_num_stars'):
            _n = ssp.num_stars + ssp.num_stars_integrated
            ssp.ssp_total_num_stars = [_n]
        new.ssp_total_num_stars.extend(ssp.ssp_total_num_stars)

        new.total_mass = new.total_mass + ssp.total_mass
        new.sampled_mass = new.sampled_mass + ssp.sampled_mass
        new.live_star_mass = new.live_star_mass + ssp.live_star_mass
        new.frac_mass_sampled = new.sampled_mass / new.live_star_mass

        new_num_stars_total = new.num_stars + new.num_stars_integrated
        ssp_num_stars_total = ssp.num_stars + ssp.num_stars_integrated
        total_num_stars = new_num_stars_total + ssp_num_stars_total
        new.num_stars_integrated += ssp.num_stars_integrated

        new.initial_masses = np.concatenate(
            [new.initial_masses, ssp.initial_masses])
        new.star_masses = np.concatenate([new.star_masses, ssp.star_masses])
        new.frac_num_sampled = new.num_stars / total_num_stars

        new.ssp_num_fracs = []
        new.ssp_mass_fracs = []
        for n, m in zip(new.ssp_total_num_stars, new.ssp_total_masses):
            new.ssp_num_fracs.append(n / total_num_stars)
            new.ssp_mass_fracs.append(m / new.total_mass)

        # Loop over optional attributes.
        # Both SSPs must have the arrtibute to add them.
        for attr in ['eep', 'log_L', 'log_Teff']:
            if hasattr(new, attr) and hasattr(ssp, attr):
                new_attr = getattr(new, attr)
                ssp_attr = getattr(ssp, attr)
                setattr(new, attr, np.concatenate([new_attr, ssp_attr]))

        new_label = np.ones(len(ssp.star_masses), dtype=int)
        new_label *= len(new.isochrone)
        new.ssp_labels = np.concatenate([new.ssp_labels, new_label])

        if hasattr(new, 'log_age'):
            if type(new.log_age) != list:
                new.log_age = [new.log_age]
            new.log_age.append(ssp.log_age)
        if hasattr(new, 'feh'):
            if type(new.feh) != list:
                new.feh = [new.feh]
            new.feh.append(ssp.feh)

        for filt in new.filters:
            _mags = [new.abs_mags[filt], ssp.abs_mags[filt]]
            new.abs_mags[filt] = np.concatenate(_mags)
            if new.has_integrated_component and ssp.has_integrated_component:
                new_flux = 10**(-0.4 * new.integrated_abs_mags[filt])
                ssp_flux = 10**(-0.4 * ssp.integrated_abs_mags[filt])
                flux = new_flux + ssp_flux
                new.integrated_abs_mags[filt] = -2.5 * np.log10(flux)
                new_lumlum = 10**new._integrated_log_lumlum[filt]
                ssp_lumlum = 10**ssp._integrated_log_lumlum[filt]
                log_lumlum = np.log10(new_lumlum + ssp_lumlum)
                new._integrated_log_lumlum[filt] = log_lumlum
            elif ssp.has_integrated_component:
                new.integrated_abs_magw[filt] = ssp.integrated_abs_mags[filt]
                log_lumlum =  ssp._integrated_log_lumlum[filt]
                new._integrated_log_lumlum[filt] = log_lumlum

        return CompositePopulation(new)

    def __repr__(self):
        r = [f'{k} = {v}' for k, v in self._r.items()]
        t = 'Simple Stellar Population\n-------------------------\n'
        return t + '\n'.join(r)

class MISTSSP(SSP):
    """
    MIST Simple Stellar Population.

    .. note::
        You must give `total_mass` *or* `num_stars`.

    Parameters
    ----------
    log_age : float
        Log (base 10) of the simple stellar population age in years.
    feh : float
        Metallicity [Fe/H] of the simple stellar population.
    phot_system : str or list-like
        Name of the photometric system(s).
    num_stars : int or `None`
        Number of stars in source. If `None`, then must give `total_mass`.
    total_mass : float or `~astropy.units.Quantity` or `None`
        Stellar mass of the source. If `None`, then must give `num_stars`. This
        mass accounts for stellar remnants when ``add_remnants = True``, which
        means the actual sampled mass will be less than the given value. If
        float is given, the units are assumed to be solar masses.
    distance : float or `~astropy.units.Quantity`, optional
        Distance to source. If float is given, the units are assumed
        to be `~astropy.units.Mpc`. Default distance is 10 `~astropy.units.pc`.
    mag_limit : float, optional
        Only sample individual stars that are brighter than this magnitude. All
        fainter stars will be combined into an integrated component. Otherwise,
        all stars in the population will be sampled. You must also give the
        `mag_limit_band` if you use this parameter.
    mag_limit_band : str, optional
        Bandpass of the limiting magnitude. You must give this parameter if
        you use the `mag_limit` parameter.
    imf : str, optional
        The initial stellar mass function. Default is `'kroupa'`.
    mist_path : str, optional
        Path to MIST isochrone grids. Use this if you want to use a different
        path from the `MIST_PATH` environment variable.
    imf_kw : dict, optional
        Optional keyword arguments for sampling the stellar mass function.
    mass_tolerance : float, optional
        Tolerance in the fractional difference between the input mass and the
        final mass of the population. The parameter is only used when
        `total_mass` is given.
    add_remnants : bool, optional
        If True (default), apply scaling factor to total mass to account for
        stellar remnants in the form of white dwarfs, neutron stars,
        and black holes.
    random_state : `None`, int, list of ints, or `~numpy.random.RandomState`
        If `None`, return the `~numpy.random.RandomState` singleton used by
        ``numpy.random``. If `int`, return a new `~numpy.random.RandomState`
        instance seeded with the `int`.  If `~numpy.random.RandomState`,
        return it. Otherwise raise ``ValueError``.
    """

    phases = ['PMS', 'MS', 'giants', 'RGB', 'CHeB', 'AGB',
              'EAGB', 'TPAGB', 'postAGB', 'WDCS']

    def __init__(self, log_age, feh, phot_system, num_stars=None,
                 total_mass=None, distance=10*u.pc, mag_limit=None,
                 mag_limit_band=None, imf='kroupa', mist_path=MIST_PATH,
                 imf_kw={}, mass_tolerance=0.05, add_remnants=True,
                 random_state=None, **kwargs):

        self.feh = feh
        self.log_age = log_age
        self.phot_system = phot_system
        self.mist_path = mist_path
        _iso = MISTIsochrone(log_age, feh, phot_system,  mist_path, **kwargs)

        super(MISTSSP, self).__init__(
            isochrone=_iso,
            num_stars=num_stars,
            total_mass=total_mass,
            distance=distance,
            mag_limit=mag_limit,
            mag_limit_band=mag_limit_band,
            imf=imf,
            imf_kw=imf_kw,
            mass_tolerance=mass_tolerance,
            add_remnants=add_remnants,
            random_state=random_state
        )

        self._r.update({'log(age/yr)': self.log_age,
                        '[Fe/H]': self.feh,
                        'photometric system': self.phot_system})

    def select_phase(self, phase):
        """
        Generate stellar evolutionary phase mask. The mask will be `True` for
        sources that are in the give phase according to the MIST EEPs.

        Parameters
        ----------
        phase : str
            Evolutionary phase to select. Options are 'all', 'MS', 'giants',
            'RGB', 'CHeB', 'AGB', 'EAGB', 'TPAGB', 'postAGB', or 'WDCS'.

        Returns
        -------
        mask : `~numpy.ndarray`
            Mask that is `True` for stars in input phase and `False` otherwise.

        Notes
        -----
        The MIST EEP phases were taken from Table II: Primary Equivalent
        Evolutionary Points (EEPs):
        http://waps.cfa.harvard.edu/MIST/README_tables.pdf
        """
        if phase == 'all':
            mask = np.ones_like(self.eep, dtype=bool)
        elif phase == 'PMS':
            mask = self.eep < 202
        elif phase == 'MS':
            mask = (self.eep >= 202) & (self.eep < 454)
        elif phase == 'giants':
            mask = (self.eep >= 454) & (self.eep < 1409)
        elif phase == 'RGB':
            mask = (self.eep >= 454) & (self.eep <= 605)
        elif phase == 'CHeB':
            mask = (self.eep > 605) & (self.eep < 707)
        elif phase == 'AGB':
            mask = (self.eep >= 707) & (self.eep < 1409)
        elif phase == 'EAGB':
            mask = (self.eep >= 707) & (self.eep < 808)
        elif phase == 'TPAGB':
            mask = (self.eep >= 808) & (self.eep < 1409)
        elif phase == 'postAGB':
            mask = (self.eep >= 1409) & (self.eep <= 1710)
        elif phase == 'WDCS':
            mask = self.eep > 1710
        else:
            raise Exception('Uh, what phase u want?')

        return mask

    def get_star_phases(self):
        """Returns the stellar phases (as defined by the MIST EEPs)."""
        phase_list = ['PMS', 'MS', 'RGB', 'CHeB', 'EAGB',
                      'TPAGB', 'postAGB', 'WDCS']
        star_phases = np.array([''] * self.num_stars, dtype='<U8')
        for phase in phase_list:
            star_phases[self.select_phase(phase)] = phase
        return star_phases

class CompositePopulation(SSP):
    """
    Composite stellar populations.
    """

    def __init__(self, pop):

        for name, attr in pop.__dict__.items():
            setattr(self, name, attr)

    def __repr__(self):
        num_fracs = self.ssp_num_fracs
        mass_fracs = self.ssp_mass_fracs
        r = {'N_pops': self.num_pops,
             'M_star': f'{self.total_mass.value:.2e} M_sun',
             'number fractions': [f'{p * 100:.2f}%' for p in num_fracs],
             'mass fractions': [f'{p * 100:.2f}%' for p in mass_fracs]}
        if hasattr(self, 'log_age'):
            r['log(age/yr)'] = self.log_age
        if hasattr(self, 'feh'):
            r['[Fe/H]'] = self.feh
        if hasattr(self, 'phot_system'):
            r['photometric system'] = self.phot_system
        r = [f'{k} = {v}' for k, v in r.items()]
        t = 'Composite Population\n--------------------\n'
        return t + '\n'.join(r)

def constant_sb_stars_per_pix(sb, mean_mag, distance=10*u.pc, pixel_scale=0.2):
    """
    Calculate the number of stars per pixel for a uniform
    distribution (i.e., constant surface brightness) of stars.

    Parameters
    ----------
    sb : float
        Surface brightness of stellar population.
    mean_mag : float
        Mean stellar magnitude of the stellar population.
    distance : float or `~astropy.units.Quantity`
        Distance to source. If float is given, the units are assumed
        to be `~astropy.units.Mpc`.
    pixel_scale : float or `~astropy.units.Quantity`, optional
        The pixel scale of the mock image. If a float is given, the units will
        be assumed to be `~astropy.units.arcsec` per `~astropy.units.pixels`.

    Returns
    -------
    num_stars_per_pix : float
        The number of stars per pixel.
    """
    distance = check_units(distance, 'Mpc').to('pc').value
    pixel_scale = check_units(pixel_scale, u.arcsec / u.pixel).value
    dist_mod = 5 * np.log10(distance) - 5
    num_stars_per_arsec_sq = 10**(0.4 * (mean_mag + dist_mod -  sb))
    num_stars_per_pix = num_stars_per_arsec_sq * pixel_scale**2
    return num_stars_per_pix


# Onufriev-Bashford-Case Generalized Born term
# can be added to any other force field

import numpy as np

from MMTK import ParticleScalar
from MMTK.ForceFields.ForceField import ForceField

class OBCForceField(ForceField):

    """
    Onufriev-Bashford-Case Generalized Born
    """

    def __init__(self,
          prmtopFN=None,
          inv_prmtop_atom_order=None,
          desolvationGridFN=None,
          r_min = 0.14,
          r_max = 1.0,
          strength=1.0):
        """
        @param prmtopFN: an AMBER parameter and topology file
        @type strength:  C{str}
        r_min and r_max should be in units of nanometers
        """
        # Initialize the ForceField class, giving a name to this one.
        ForceField.__init__(self, 'OBC')

        # Store arguments that recreate the force field from a pickled
        # universe or from a trajectory.
        self.arguments = (prmtopFN, inv_prmtop_atom_order, \
          desolvationGridFN, r_min, r_max, strength)

        # Load the desolvation grid
        if desolvationGridFN is not None:
          import AlGDock.IO
          IO_Grid = AlGDock.IO.Grid()
          self.grid_data = IO_Grid.read(desolvationGridFN, multiplier=0.1)
          if not (self.grid_data['origin']==0.0).all():
            raise Exception('Trilinear grid origin in %s not at (0, 0, 0)!'%FN)
          self.useDesolvationGrid = True
        else:
          self.grid_data = {'spacing':np.array([0., 0., 0.]), \
                            'counts':np.array([0, 0, 0]), \
                            'vals':np.array([])}
          self.useDesolvationGrid = False

        # Store arguments as class variables
        self.prmtopFN = prmtopFN
        self.inv_prmtop_atom_order = inv_prmtop_atom_order
        self.desolvationGridFN = desolvationGridFN
        self.r_min = r_min
        self.r_max = r_max
        self.strength = strength

    def set_strength(self, strength):
      self.strength = strength

    # The following method is called by the energy evaluation engine
    # to inquire if this force field term has all the parameters it
    # requires. This is necessary for interdependent force field
    # terms. In our case, we just say "yes" immediately.
    def ready(self, global_data):
        return True

    # The following method is called by the energy evaluation engine
    # to obtain a list of the low-level evaluator objects (the C routines)
    # that handle the calculations.
    def evaluatorTerms(self, universe, subset1, subset2, global_data):
        # The energy for subsets is defined as consisting only
        # of interactions within that subset, so the contribution
        # of an external field is zero. Therefore we just return
        # an empty list of energy terms.
        if subset1 is not None or subset2 is not None:
            return []

        import numpy as np
        if (self.prmtopFN is not None) and \
           (self.inv_prmtop_atom_order is not None):
          # Get charges, radii, and scale factors from OpenMM
          import simtk.openmm
          import simtk.openmm.app as OpenMM_app

          prmtop = OpenMM_app.AmberPrmtopFile(self.prmtopFN)
          OMM_system = prmtop.createSystem(\
            nonbondedMethod=OpenMM_app.CutoffNonPeriodic, \
            nonbondedCutoff=1.5, \
            constraints=None, \
            implicitSolvent=OpenMM_app.OBC2)
          f = OMM_system.getForces()[-2]

          numParticles = f.getNumParticles()
          charges = np.zeros(numParticles)
          atomicRadii = np.zeros(numParticles)
          scaleFactors = np.zeros(numParticles)
          for n in range(numParticles):
            (charge, radius, scaleFactor) = f.getParticleParameters(n)
            charges[n] = charge/simtk.unit.elementary_charge
            atomicRadii[n] = radius/simtk.unit.nanometer
            scaleFactors[n] = scaleFactor

          charges = charges[self.inv_prmtop_atom_order]
          atomicRadii = atomicRadii[self.inv_prmtop_atom_order]
          scaleFactors = scaleFactors[self.inv_prmtop_atom_order]
        else:
          # Get charges, radii, and scale factors from the ligand database (preferred)
          charges_ps = ParticleScalar(universe)
          atomicRadii_ps = ParticleScalar(universe)
          scaleFactors_ps = ParticleScalar(universe)
          for o in universe:
            for a in o.atomList():
              charges_ps[a] = o.getAtomProperty(a, 'amber_charge')
              atomicRadii_ps[a] = o.getAtomProperty(a, 'scaling_factor_BornRadii')
              scaleFactors_ps[a] = o.getAtomProperty(a, 'scaling_factor_BornScreening')
          charges = charges_ps.array
          atomicRadii = atomicRadii_ps.array
          scaleFactors = scaleFactors_ps.array
          numParticles = charges.shape[0]

#        import time
#        import os.path
#        import MMTK_OBC
#        OBCpath = MMTK_OBC.__file__
#        print """
#        in {0}
#        last modified {1}
#            """.format(OBCpath, time.ctime(os.path.getmtime(OBCpath)))

        # Here we pass all the parameters as "simple" data types to
        # the C code that handles energy calculations.
        if self.useDesolvationGrid:
          # With desolvation grid
          from MMTK_OBC_desolv import OBCDesolvTerm
          return [OBCDesolvTerm(universe._spec, numParticles, self.strength, \
            charges, atomicRadii, scaleFactors, \
            self.grid_data['spacing'], self.grid_data['counts'], \
            self.grid_data['vals'], self.r_min, self.r_max)]
        else:
          # No desolvation grid
          from MMTK_OBC import OBCTerm
          return [OBCTerm(universe._spec, numParticles, self.strength, \
            charges, atomicRadii, scaleFactors)]


# Copyright 2016, FBPIC contributors
# Authors: Remi Lehe, Manuel Kirchen
# License: 3-Clause-BSD-LBNL
"""
This file is part of the Fourier-Bessel Particle-In-Cell code (FB-PIC)
It defines the structure and methods associated with atomic ionization.

The implemented ionization model is the ADK model. The implementation
is fully relativistic (i.e. it works in the boosted-frame as well).

Ionization is implemented by keeping ions at different ionization states in
the same Particles object, so that the number of macroparticles in this object
remains constant and arrays do not have to be reallocated (except when the
moving window creates new particles.) An array `ionization_level` keeps track
of the ionization state of each macroparticle.

On the other hand, the electrons generated by ionization do need to be added to
an existing Particles object, and this implies that the number of
macroparticles in the object does not remain constant and that the arrays are
reallocated.

In addition, at each PIC iteration, the number of new electrons need to
be counted in order to reallocate the array and copy the electrons to
the right indices. This basically involves a cumulative sum operation, which
difficult to implement on the GPU. For this reason, the corresponding array is
sent to the CPU, which performs the cumulative sum and sends the data back to
the GPU. In order, to limit the amount of data to be transfered, particles are
handled in batches of 10 particles, so that only the cumulative sum of the
number of particles in each batch need to be performed.
"""
import numpy as np
from scipy.constants import c, e, m_e, physical_constants
from scipy.special import gamma
from .read_atomic_data import get_ionization_energies
from .numba_methods import ionize_ions_numba, copy_ionized_electrons_numba
from ..cuda_numba_utils import allocate_empty, reallocate_and_copy_old, \
                                perform_cumsum_2d, generate_new_ids

# Check if CUDA is available, then import CUDA functions
from fbpic.utils.cuda import cuda_installed
from fbpic.utils.printing import catch_gpu_memory_error
if cuda_installed:
    import cupy
    from fbpic.utils.cuda import cuda_tpb_bpg_1d
    from .cuda_methods import ionize_ions_cuda, copy_ionized_electrons_cuda

class Ionizer(object):
    """
    Class that contains the data associated with ionization (on the ions side)
    and has method to calculate the ionization probability.

    The implemented ionization model is the ADK model. The implementation
    is fully relativistic (i.e. it works in the boosted-frame as well).

    Main attributes
    ---------------
    - ionization_level: 1darray of integers (one element per particle)
      which contains the ionization state of each particle
    - w_times_level: 1darray of floats (one element per particle)
      which contains the number of physical particle that correspond to each
      macroparticle, multiplied by the ionization level. (This is updated
      whenever further ionization happens, and is passed to the deposition
      kernel as the effective weight of the particles)
    """
    def __init__(self, element, ionizable_species, target_species,
                 level_start, level_max=None):
        """
        Initialize an Ionizer instance

        Parameters
        ----------
        element: string
            The atomic symbol of the considered ionizable species
            (e.g. 'He', 'N' ;  do not use 'Helium' or 'Nitrogen')

        ionizable_species: an fbpic.Particles object
            This object is not modified or registered.
            It is only used in order to pass a number of additional argument.

        target_species: a `Particles` object, or a dictionary of `Particles`
            Stores the electron macroparticles that are created in
            the ionization process.

            - If a single `Particles` object is passed, than electrons from all
            ionization levels are stored into this object.

            - If a dictionary is passed, then its keys should be integers
            (corresponding to the ionizable levels of `element`, starting
            at `level_start`), and its values should be `Particles` objects.
            In this case, the electrons from each distinct ionizable level
            will be stored into these separate objects. Note that using
            separate objects will typically require longer computing time.

            These objects are not modified when creating the class, but
            they are when ionization occurs (i.e. more particles are created)

        level_start: int
            The ionization level at which the macroparticles are initially
            (e.g. 0 for initially neutral atoms)

        level_max: int, optional
            If not None, defines the maximum ionization level that
            macroparticles can reach. Should not exceed the physical
            limit for the chosen element.
        """
        # Register a few parameters
        self.level_start = level_start
        self.level_max = level_max
        self.use_cuda = ionizable_species.use_cuda
        # Process ionized particles into batches
        self.batch_size = 10

        # Initialize ionization-relevant meta-data
        self.initialize_ADK_parameters( element, ionizable_species.dt )

        # Initialize the required arrays
        Ntot = ionizable_species.Ntot
        self.ionization_level = np.ones( Ntot, dtype=np.uint64 ) * level_start
        self.w_times_level = ionizable_species.w * self.ionization_level

        # Check if electrons from different ionization levels should
        # be stored into separate species
        if type(target_species) is dict:
            # When passing a dictionary
            # Check that the keys are the right integers
            for level in range(self.level_start, self.level_max):
                if level not in target_species.keys():
                    raise ValueError(
                    'When passing a dictionary for `target_species`, its keys '
                    'should be\nthe integers corresponding to the ionizable '
                    'levels.\n (i.e. the integers from %d to %d'
                    'for %s with level_start=%d.)' %(self.level_start,
                    self.level_max, element, self.level_start))
                # Check that the dictionary contains Particles objects
                assert isinstance(target_species[level], type(ionizable_species))
            # Convert to a list internally: the dictionary input is
            # just for less error-prone user input.
            self.target_species = [ target_species[level] \
                for level in range(self.level_start, self.level_max) ]
            self.store_electrons_per_level = True
        elif isinstance(target_species, type(ionizable_species)):
            # When passing a single Particles object
            self.target_species = [target_species]  # List of one element
            self.store_electrons_per_level = False
        else:
            raise ValueError(
                "Unexpected type for target_species: %s\n"
                "Please pass a `Particles` object, or a dictionary"
                %type(target_species))

        # Check that the target species are indeed electrons
        for species in self.target_species:
            assert species.q == -e
            assert species.m == m_e

    def initialize_ADK_parameters( self, element, dt ):
        """
        Initialize parameters needed for the calculation of ADK ionization rate

        Parameters
        ----------
        element: string
            The atomic symbol of the considered ionizable species
            (e.g. 'He', 'N' ;  do not use 'Helium' or 'Nitrogen')

        dt: float (in seconds)
            The timestep of the simulation. (The calculated ionization
            probability is a probability *per timestep*.)

        See Chen, JCP 236 (2013), equation (2) for the ionization rate formula
        """
        # Get the array of energies
        Uion = get_ionization_energies( element )
        # Check whether the element string was valid
        if Uion is None:
            raise ValueError("Unknown ionizable element %s.\n" %element + \
            "Please use atomic symbol (e.g. 'He') not full name (e.g. Helium)")
        else:
            self.element = element

        # Determine and set the maximum level of ionization
        if self.level_max is None:
            self.level_max = len(Uion)
        else:
            assert type(self.level_max) is int, "level_max must be integer"
            if self.level_max>len(Uion):
                raise ValueError("Chosen level_max for {}".format(element) + \
                                 " cannot exceed {}".format(len(Uion)))

        # Calculate the ADK prefactors (See Chen, JCP 236 (2013), equation (2))
        # - Scalars
        alpha = physical_constants['fine-structure constant'][0]
        r_e = physical_constants['classical electron radius'][0]
        wa = alpha**3 * c / r_e
        Ea = m_e*c**2/e * alpha**4/r_e
        # - Arrays (one element per ionization level)
        UH = get_ionization_energies('H')[0]
        Z = np.arange( len(Uion) ) + 1
        n_eff = Z * np.sqrt( UH/Uion )
        l_eff = n_eff[0] - 1
        C2 = 2**(2*n_eff) / (n_eff * gamma(n_eff+l_eff+1) * gamma(n_eff-l_eff))
        # For now, we assume l=0, m=0
        self.adk_power = - (2*n_eff - 1)
        self.adk_prefactor = dt * wa * C2 * ( Uion/(2*UH) ) \
            * ( 2*(Uion/UH)**(3./2)*Ea )**(2*n_eff - 1)
        self.adk_exp_prefactor = -2./3 * ( Uion/UH )**(3./2) * Ea

    @catch_gpu_memory_error
    def handle_ionization( self, ion ):
        """
        Handle ionization, either on CPU or GPU

        - For each ion macroparticle, decide whether it is going to
          be further ionized during this timestep, based on the ADK rate.
        - Add the electrons created from ionization to the `target_species`

        Parameters:
        -----------
        ion: an fbpic.Particles object
            The ionizable species, from which new electrons are created.
        """
        # Skip this function if there are no ions
        if ion.Ntot == 0:
            return

        # Process particles in batches (of typically 10, 20 particles)
        N_batch = int( ion.Ntot / self.batch_size ) + 1
        # Short-cuts
        use_cuda = self.use_cuda

        # Set the number of levels that should be distinguished
        if self.store_electrons_per_level:
            n_levels = self.level_max - self.level_start
        else:
            n_levels = 1

        # Create temporary arrays (on CPU or GPU, depending on `use_cuda`)
        ionized_from = allocate_empty( ion.Ntot, use_cuda, dtype=np.int16 )
        n_ionized = allocate_empty( (n_levels, N_batch), use_cuda,
                                    dtype=np.int64 )
        # Draw random numbers
        if self.use_cuda:
            random_draw = cupy.random.rand( ion.Ntot, dtype=cupy.float32 )
        else:
            random_draw = np.random.rand( ion.Ntot )

        # Determine the ions that are ionized, and count them in each batch
        # (one thread per batch on GPU; parallel loop over batches on CPU)
        if use_cuda:
            batch_grid_1d, batch_block_1d = cuda_tpb_bpg_1d( N_batch )
            ionize_ions_cuda[ batch_grid_1d, batch_block_1d ](
                N_batch, self.batch_size, ion.Ntot,
                self.level_start, self.level_max, n_levels,
                n_ionized, ionized_from, self.ionization_level, random_draw,
                self.adk_prefactor, self.adk_power, self.adk_exp_prefactor,
                ion.ux, ion.uy, ion.uz, ion.Ex, ion.Ey, ion.Ez,
                ion.Bx, ion.By, ion.Bz, ion.w, self.w_times_level )
        else:
            ionize_ions_numba(
                N_batch, self.batch_size, ion.Ntot,
                self.level_start, self.level_max, n_levels,
                n_ionized, ionized_from, self.ionization_level, random_draw,
                self.adk_prefactor, self.adk_power, self.adk_exp_prefactor,
                ion.ux, ion.uy, ion.uz, ion.Ex, ion.Ey, ion.Ez,
                ion.Bx, ion.By, ion.Bz, ion.w, self.w_times_level )

        # Count the total number of new electrons
        cumulative_n_ionized = perform_cumsum_2d( n_ionized, use_cuda )
        # If no new particle was created, skip the rest of this function
        if use_cuda:
            if cupy.all( cumulative_n_ionized[:,-1] == 0 ):
                return
        else:
            if np.all( cumulative_n_ionized[:,-1] == 0 ):
                return

        # Loop over the electron species associated to each level
        # (when store_electrons_per_level is False, there is a single species)
        # Reallocate electron species (on CPU or GPU depending on `use_cuda`),
        # to accomodate the electrons produced by ionization,
        # and copy the old electrons to the new arrays
        assert len(self.target_species) == n_levels
        for i_level, elec in enumerate(self.target_species):
            old_Ntot = elec.Ntot
            # Cast to int transfers the data from the GPU if needed
            new_Ntot = old_Ntot + int( cumulative_n_ionized[i_level,-1] )
            reallocate_and_copy_old( elec, use_cuda, old_Ntot, new_Ntot )
            # Create the new electrons from ionization (one thread per batch)
            if use_cuda:
                copy_ionized_electrons_cuda[ batch_grid_1d, batch_block_1d ](
                    N_batch, self.batch_size, old_Ntot, ion.Ntot,
                    cumulative_n_ionized, ionized_from,
                    i_level, self.store_electrons_per_level,
                    elec.x, elec.y, elec.z, elec.inv_gamma,
                    elec.ux, elec.uy, elec.uz, elec.w,
                    elec.Ex, elec.Ey, elec.Ez, elec.Bx, elec.By, elec.Bz,
                    ion.x, ion.y, ion.z, ion.inv_gamma,
                    ion.ux, ion.uy, ion.uz, ion.w,
                    ion.Ex, ion.Ey, ion.Ez, ion.Bx, ion.By, ion.Bz )
                # Mark the new electrons as unsorted
                elec.sorted = False
            else:
                copy_ionized_electrons_numba(
                    N_batch, self.batch_size, old_Ntot, ion.Ntot,
                    cumulative_n_ionized, ionized_from,
                    i_level, self.store_electrons_per_level,
                    elec.x, elec.y, elec.z, elec.inv_gamma,
                    elec.ux, elec.uy, elec.uz, elec.w,
                    elec.Ex, elec.Ey, elec.Ez, elec.Bx, elec.By, elec.Bz,
                    ion.x, ion.y, ion.z, ion.inv_gamma,
                    ion.ux, ion.uy, ion.uz, ion.w,
                    ion.Ex, ion.Ey, ion.Ez, ion.Bx, ion.By, ion.Bz )

            # If the electrons are tracked, generate new ids
            # (on GPU or GPU depending on `use_cuda`)
            generate_new_ids( elec, old_Ntot, new_Ntot )

    def send_to_gpu( self ):
        """
        Copy the ionization data to the GPU.
        """
        if self.use_cuda:
            # Arrays with one element per macroparticles
            self.ionization_level = cupy.asarray( self.ionization_level )
            self.w_times_level = cupy.asarray( self.w_times_level )
            # Small-size arrays with ADK parameters
            # (One element per ionization level)
            self.adk_power = cupy.asarray( self.adk_power )
            self.adk_prefactor = cupy.asarray( self.adk_prefactor )
            self.adk_exp_prefactor = cupy.asarray( self.adk_exp_prefactor )

    def receive_from_gpu( self ):
        """
        Receive the ionization data from the GPU.
        """
        if self.use_cuda:
            # Arrays with one element per macroparticles
            self.ionization_level = self.ionization_level.get()
            self.w_times_level = self.w_times_level.get()
            # Small-size arrays with ADK parameters
            # (One element per ionization level)
            self.adk_power = self.adk_power.get()
            self.adk_prefactor = self.adk_prefactor.get()
            self.adk_exp_prefactor = self.adk_exp_prefactor.get()