from itertools import chain
import numpy as np
from numpy.polynomial.chebyshev import chebval
from scipy.spatial import Delaunay

from ..utils.smoothing import smoothspec
from .constants import lightspeed, lsun, jansky_cgs, to_cgs_at_10pc

try:
    from sklearn.neighbors import KDTree
except(ImportError):
    from scipy.spatial import cKDTree as KDTree

try:
    from sedpy.observate import getSED, vac2air, air2vac
except(ImportError):
    pass

__all__ = ["StarBasis", "BigStarBasis"]

# Useful constants
# value to go from L_sun to erg/s/cm^2 at 10pc
to_cgs = to_cgs_at_10pc

# for converting Kurucz spectral units
log4pi = np.log10(4 * np.pi)
log_rsun_cgs = np.log10(6.955) + 10
log_lsun_cgs = np.log10(lsun)
log_SB_cgs = np.log10(5.6704e-5)
log_SB_solar = log_SB_cgs + 2 * log_rsun_cgs - log_lsun_cgs

class StarBasis(object):

    _spectra = None

    def __init__(self, libname='ckc14_deimos.h5', verbose=False,
                 n_neighbors=0, log_interp=True, logify_Z=False,
                 use_params=None, rescale_libparams=False, in_memory=True,
                 norm_spec = False,
                 **kwargs):
        """An object which holds the stellar spectral library, performs
        interpolations of that library, and has methods to return attenuated,
        normalized, smoothed stellar spectra.  The interpolations are performed
        using barycenter coordinates of the enclosing simplex found from the
        Delauynay triangulation.  This is not tractable for large dimension
        (see BigStarBasis for that case).

        :param libname:
            Path to the hdf5 file to use for the spectral library.

        :param verbose:
            If True, print information about the parameters used when a point
            is outside the convex hull.

        :param n_neighbors: (default:0)
            Number of nearest neighbors to use when requested parameters are
            outside the convex hull of the library prameters.  If ``0`` then a
            ValueError is raised instead of the nearest spectrum.  If greater
            than 1 then the neighbors are combined using inverse distance
            weights.

        :param log_interp: (default:True)
            Switch to interpolate in log(flux) instead of linear flux.

        :param use_params:
            Sequence of strings. If given, only use the listed parameters
            (which must be present in the `_libparams` structure) to build the
            grid and construct spectra.  Otherwise all fields of `_libparams`
            will be used.

        :param rescale_libparams: (default: False)
            If True, rescale the parameters to the unit cube before generating
            the triangulation (and kd-tree).  Note that the `param_vector`
            method will also rescale the input parameters in this case.  This
            can help for nearest neighbor lookup and in the triangulation based
            weights when your variables have very different scales, assuming
            that the ranges give a reasonable relative distance metric.

        :param in_memory: (default: True)
            Switch to keep the spectral library in memory or access it through
            the h5py File object.  Note if the latter, then zeroed spectra are
            *not* filtered out.
        """
        # Cache initialization variables
        self.verbose = verbose
        self.logarithmic = log_interp
        self.logify_Z = logify_Z
        self._in_memory = in_memory
        self._libname = libname
        self.n_neighbors = n_neighbors
        self._rescale = rescale_libparams
        self.norm_spec = norm_spec

        # Load the library
        self.load_lib(libname)

        # Do some important bookkeeping
        if use_params is None:
            self.stellar_pars = self._libparams.dtype.names
        else:
            self.stellar_pars = tuple(use_params)
        self.ndim = len(self.stellar_pars)

        # Build the triangulation and kdtree (after rescaling)
        if self._rescale:
            ranges = [[self._libparams[d].min(), self._libparams[d].max()]
                      for d in self.stellar_pars]
            self.parameter_range = np.array(ranges).T
        self.triangulate()
        try:
            self.build_kdtree()
        except NameError:
            pass

        self.params = {}

    def load_lib(self, libname='', driver=None):
        """Read a CKC library which has been pre-convolved to be close to your
        resolution.  This library should be stored as an HDF5 file, with the
        datasets ``wavelengths``, ``parameters`` and ``spectra``.  These are
        ndarrays of shape (nwave,), (nmodels,), and (nmodels, nwave)
        respecitvely.  The ``parameters`` array is a structured array.  Spectra
        with no fluxes > 1e-32 are removed from the library if the librarty is
        kept in memory.
        """
        import h5py
        f = h5py.File(libname, "r", driver=driver)
        self._wave = np.array(f['wavelengths'])
        self._libparams = np.array(f['parameters'])

        if self._in_memory:

            if not self.norm_spec:
                self._spectra = np.array(f['spectra'])

            elif self.norm_spec:
                self._spectra = np.array(f['spectra']) / np.array(f['continuua'])

            f.close()
            # Filter library so that only existing spectra are included
            maxf = np.max(self._spectra, axis=1)
            good = maxf > 1e-32
            self._libparams = self._libparams[good]
            self._spectra = self._spectra[good, :]
        else:

            if not self.norm_spec:
                self._spectra = f['spectra']
            elif self.norm_spec:
                self._spectra = f['spectra'] / f['continuua']

        if self.logify_Z:
            from numpy.lib import recfunctions as rfn
            self._libparams['Z'] = np.log10(self._libparams['Z'])
            self._libparams = rfn.rename_fields(self._libparams, {'Z': 'logZ'})

    def update(self, **kwargs):
        """Update the `params` dictionary, turning length 1 arrays into scalars
        and pull out functions from length one arrays
        """
        for k, val in list(kwargs.items()):
            v = np.atleast_1d(val)
            try:
                if (len(v) == 1) and callable(v[0]):
                    self.params[k] = v[0]
                else:
                    self.params[k] = np.squeeze(v)
            except(KeyError):
                pass

    def get_spectrum(self, outwave=None, filters=None, peraa=False, **kwargs):
        """Return an attenuated, smoothed, distance dimmed stellar spectrum and SED.

        :returns spec:
            The spectrum on the outwave grid (assumed in air), in AB maggies.
            If peraa is True then the spectrum is erg/s/cm^2/AA.

        :returns phot:
            Observed frame photometry in units of AB maggies.  If ``lumdist``
            is not present in the parameters then these are absolute maggies,
            otherwise they are apparent.

        :returns x:
            A blob of extra quantities (e.g. mass, uncertainty)
        """
        self.update(**kwargs)

        # star spectrum (in Lsun/Hz)
        wave, spec, unc = self.get_star_spectrum(**self.params)
        spec *= self.normalize()

        # dust
        if 'dust_curve' in self.params:
            att = self.params['dust_curve'](self._wave, **self.params)
            spec *= np.exp(-att)

        # Redshifting + Wavelength solution.  We also convert to in-air.
        a = 1 + self.params.get('zred', 0)
        b = 0.0

        if 'wavecal_coeffs' in self.params:
            x = wave - wave.min()
            x = 2.0 * (x / x.max()) - 1.0
            c = np.insert(self.params['wavecal_coeffs'], 0, 0)
            # assume coeeficients give shifts in km/s
            b = chebval(x, c) / (lightspeed*1e-13)

        wa, sa = vac2air(wave) * (a + b), spec * a
        if outwave is None:
            outwave = wa

        # Broadening, interpolation onto output wavelength grid
        if 'sigma_smooth' in self.params:
            smspec = self.smoothspec(wa, sa, self.params['sigma_smooth'],
                                     outwave=outwave, **self.params)
        elif outwave is not wa:
            smspec = np.interp(outwave, wa, sa, left=0, right=0)
        else:
            smspec = sa

        # Photometry (observed frame absolute maggies)
        if filters is not None:
            mags = getSED(wa, sa * lightspeed / wa**2 * to_cgs, filters)
            phot = np.atleast_1d(10**(-0.4 * mags))
        else:
            phot = 0.0

        # Distance dimming.  Default to 10pc distance (i.e. absolute)
        dfactor = (self.params.get('lumdist', 1e-5) * 1e5)**2
        if peraa:
            # spectrum will be in erg/s/cm^2/AA
            smspec *= to_cgs / dfactor * lightspeed / outwave**2
        else:
            # Spectrum will be in maggies
            smspec *= to_cgs / dfactor / (3631*jansky_cgs)

        # Convert from absolute maggies to apparent maggies
        phot /= dfactor

        return smspec, phot, None

    def get_star_spectrum(self, **kwargs):
        """Given stellar parameters, obtain an interpolated spectrum at those
        parameters.

        :param **kwargs:
            Keyword arguments must include values for the parameters listed in
            ``stellar_pars``.

        :returns wave:
            The wavelengths at which the spectrum is defined.

        :returns spec:
            The spectrum interpolated to the requested parameters.  This has
            the same units as the supplied library spectra.

        :returns unc:
            The uncertainty spectrum, where the uncertainty is due to
            interpolation error.  Curently unimplemented (i.e. it is a None
            type object).
        """
        inds, wghts = self.weights(**kwargs)
        if self.logarithmic:
            spec = np.exp(np.dot(wghts, np.log(self._spectra[inds, :])))
        else:
            spec = np.dot(wghts, self._spectra[inds, :])
        spec_unc = None
        return self._wave, spec, spec_unc

    def smoothspec(self, wave, spec, sigma, outwave=None, **kwargs):
        outspec = smoothspec(wave, spec, sigma, outwave=outwave, **kwargs)
        return outspec

    def normalize(self):
        """Use either `logr` or `logl` to normalize the spectrum.  Both should
        be in solar units.  `logr` is checked first.  If neither is present
        then 1.0 is returned.

        :returns norm:
            Factor by which the CKC spectrum should be multiplied to get units
            of L_sun/Hz.  This assumes the native library spectrum is in units
            of erg/s/cm^2/Hz/sr.
        """
        if 'logr' in self.params:
            twologr = 2. * (self.params['logr'] + log_rsun_cgs)
        elif 'logl' in self.params:
            twologr = ((self.params['logl'] + log_lsun_cgs) -
                       4 * self.params['logt'] - log_SB_cgs - log4pi)
        else:
            return 1.0

        norm = 10**(twologr + 2 * log4pi - log_lsun_cgs)
        return norm

    def weights(self, **kwargs):
        """Delauynay weighting.  Return indices of the models forming the
        enclosing simplex, as well as the barycentric coordinates of the point
        within this simplex to use as weights.  If point is outside the convex
        hull then fallback to nearest neighbor unless ``n_neighbors`` is 0.
        """
        inparams = np.squeeze(self.param_vector(**kwargs))
        triangle_ind = self._dtri.find_simplex(inparams)
        if triangle_ind == -1:
            self.edge_flag = True
            if self.n_neighbors == 0:
                pstring = ', '.join(self.ndim * ['{}={}'])
                pstring = pstring.format(*chain(*zip(self.stellar_pars, inparams)))
                raise ValueError("Requested spectrum ({}) outside convex hull,"
                                 " and nearest neighbor interpolation turned "
                                 "off.".format(*pstring))
            ind, wght = self.weights_knn(inparams, k=self.n_neighbors)
            if self.verbose:
                print("Parameters {0} outside model convex hull. "
                      "Using model index {1} instead. ".format(inparams, ind))
            return ind, wght

        inds = self._dtri.simplices[triangle_ind, :]
        transform = self._dtri.transform[triangle_ind, :, :]
        Tinv = transform[:self.ndim, :]
        x_r = inparams - transform[self.ndim, :]
        bary = np.dot(Tinv, x_r)
        last = np.clip(1.0 - bary.sum(), 0.0, 1.0)
        wghts = np.append(bary, last)
        oo = inds.argsort()
        return inds[oo], wghts[oo]

    def rescale_params(self, points):
        """Rescale the given parameters to the unit cube, if the ``_rescale`` attribute is ``True``

        :param points:
            An array of parameter values, of shape (npoint, ndim)

        :returns x:
            An array of parameter values rescaled to the unit cube, ndarray of
            shape (npoint, ndim)
        """
        if self._rescale:
            x = np.atleast_2d(points)
            x = (x - self.parameter_range[0, :]) / np.diff(self.parameter_range, axis=0)
            return np.squeeze(x)
        else:
            return points

    def triangulate(self):
        """Build the Delauynay Triangulation of the model library.
        """
        # slow.  should use a view based method
        model_points = np.array([list(self._libparams[d])
                                 for d in self.stellar_pars]).T
        self._dtri = Delaunay(self.rescale_params(model_points))

    def build_kdtree(self):
        """Build the kdtree of the model points.
        """
        # slow.  should use a view based method
        model_points = np.array([list(self._libparams[d])
                                 for d in self.stellar_pars])
        self._kdt = KDTree(self.rescale_params(model_points.T))

    def weights_knn(self, target_points, k=1):
        """The interpolation weights are determined from the inverse distance
        to the k nearest neighbors.

        :param target_points: ndarray, shape(ntarg,npar)
            The coordinates to which you wish to interpolate.

        :param k:
            The number of nearest neighbors to use.

        :returns inds: ndarray, shape(ntarg,npar+1)
             The model indices of the interpolates.

        :returns weights: narray, shape (ntarg,npar+1)
             The weights of each model given by ind in the interpolates.
        """
        try:
            dists, inds = self._kdt.query(np.atleast_2d(target_points), k=k,
                                          return_distance=True)
        except:
            return [0], [0]
        inds = np.atleast_1d(np.squeeze(inds))
        if k == 1:
            return inds, np.ones(inds.shape)
        weights = 1 / dists
        # weights[np.isinf(weights)] = large_number
        weights = weights/weights.sum(axis=-1)
        return inds, np.atleast_1d(np.squeeze(weights))

    def param_vector(self, **kwargs):
        """Take a dictionary of parameters and return the stellar library
        parameter vector corresponding to these parameters as an ndarray.
        Raises a KeyError if the dictionary does not contain *all* of the
        required stellar parameters.
        """
        pvec = [kwargs[n] for n in self.stellar_pars]
        return self.rescale_params(np.array(pvec))

    @property
    def wavelengths(self):
        return self._wave

class BigStarBasis(StarBasis):

    def __init__(self, libname='', verbose=False, log_interp=True,
                 n_neighbors=0, driver=None, in_memory=False,
                 use_params=None, strictness=0.0, **kwargs):
        """An object which holds the stellar spectral library, performs linear
        interpolations of that library, and has methods to return attenuated,
        normalized, smoothed stellar spoectra.

        This object is set up to work with large grids, so the models file is
        kept open for access from disk.  scikits-learn or scipy kd-trees are
        required for model access.  Ideally the grid should be regular (though
        the spacings need not be equal along a given dimension).

        :param libname:
            Path to the hdf5 file to use for the spectral library.

        :param n_neighbors: (default:0)
            Number of nearest neighbors to use when requested parameters are
            outside the convex hull of the library prameters.  If ``0`` then a
            ValueError is raised instead of the nearest spectrum.  Does not
            work, currently.

        :param verbose:
            If True, print information about the parameters used when a point
            is outside the convex hull

        :param log_interp: (default: True)
            Interpolate in log(flux) instead of flux.

        :param in_memory: (default: False)
            Switch to determine whether the grid is loaded in memory or read
            from disk each time a model is constructed (like you'd want for
            very large grids).

        :param use_params:
            Sequence of strings. If given, only use the listed parameters
            (which must be present in the `_libparams` structure) to build the
            grid and construct spectra.  Otherwise all fields of `_libparams`
            will be used.

        :param strictness: (default: 0.0)
            Float from 0.0 to 1.0 that gives the fraction of a unit hypercube
            that is required for a parameter position to be accepted.  That is,
            if the weights of the enclosing vertices sum to less than this
            number, raise an error.
        """
        self.verbose = verbose
        self.logarithmic = log_interp
        self._libname = libname
        self.n_neighbors = n_neighbors
        self._in_memory = in_memory
        self._strictness = strictness

        self.load_lib(libname, driver=driver)
        # Do some important bookkeeping
        if use_params is None:
            self.stellar_pars = self._libparams.dtype.names
        else:
            self.stellar_pars = tuple(use_params)
        self.ndim = len(self.stellar_pars)
        self.lib_as_grid()
        self.params = {}

    def load_lib(self, libname='', driver=None):
        """Read a ykc library which has been preconvolved to be close to your
        data resolution. This library should be stored as an HDF5 file, with
        the datasets ``wavelengths``, ``parameters`` and ``spectra``.  These
        are ndarrays of shape (nwave,), (nmodels,), and (nmodels, nwave)
        respecitvely.  The ``parameters`` array is a structured array.  The h5
        file object is left open so that spectra can be accessed from disk.
        """
        import h5py
        f = h5py.File(libname, "r", driver=driver)
        self._wave = np.array(f['wavelengths'])
        self._libparams = np.array(f['parameters'])
        if self._in_memory:
            self._spectra = np.array(f['spectra'])
            f.close()
        else:
            self._spectra = f['spectra']

    def get_star_spectrum(self, **kwargs):
        """Given stellar parameters, obtain an interpolated spectrum at those
        parameters.

        :param **kwargs:
            Keyword arguments must include values for the ``stellar_pars``
            parameters that are stored in ``_libparams``.

        :returns wave:
            The wavelengths at which the spectrum is defined.

        :returns spec:
            The spectrum interpolated to the requested parameters

        :returns unc:
            The uncertainty spectrum, where the uncertainty is due to
            interpolation error.  Curently unimplemented (i.e. it is a None
            type object)
        """
        inds, wghts = self.weights(**kwargs)
        if self.logarithmic:
            spec = np.exp(np.dot(wghts, np.log(self._spectra[inds, :])))
        else:
            spec = np.dot(wghts, self._spectra[inds, :])
        spec_unc = None
        return self._wave, spec, spec_unc

    def weights(self, **params):
        inds = self.knearest_inds(**params)
        wghts = self.linear_weights(inds, **params)
        if wghts.sum() <= self._strictness:
            raise ValueError("Something is wrong with the weights")
        good = wghts > 0
        # if good.sum() < 2**self.ndim:
        #     raise ValueError("Did not find all vertices of the hypercube, "
        #                      "or there is no enclosing hypercube in the library.")
        inds = inds[good]
        wghts = wghts[good]
        wghts /= wghts.sum()
        return inds, wghts

    def lib_as_grid(self):
        """Convert the library parameters to pixel indices in each dimension,
        and build and store a KDTree for the pixel coordinates.
        """
        # Get the unique gridpoints in each param
        self.gridpoints = {}
        for p in self.stellar_pars:
            self.gridpoints[p] = np.unique(self._libparams[p])
        # Digitize the library parameters
        X = np.array([np.digitize(self._libparams[p], bins=self.gridpoints[p],
                                  right=True) for p in self.stellar_pars])
        self.X = X.T
        # Build the KDTree
        self._kdt = KDTree(self.X)  # , metric='euclidean')

    def params_to_grid(self, **targ):
        """Convert a set of parameters to grid pixel coordinates.

        :param targ:
            The target parameter location, as keyword arguments.  The elements
            of ``stellar_pars`` must be present as keywords.

        :returns x:
            The target parameter location in pixel coordinates.
        """
        # bin index
        inds = np.array([np.digitize([targ[p]], bins=self.gridpoints[p], right=False) - 1
                         for p in self.stellar_pars])
        inds = np.squeeze(inds)
        # fractional index.  Could use stored denominator to be slightly faster
        try:
            find = [(targ[p] - self.gridpoints[p][i]) /
                    (self.gridpoints[p][i+1] - self.gridpoints[p][i])
                    for i, p in zip(inds, self.stellar_pars)]
        except(IndexError):
            pstring = "{0}: min={2} max={3} targ={1}\n"
            s = [pstring.format(p, targ[p], *self.gridpoints[p][[0, -1]])
                 for p in self.stellar_pars]
            raise ValueError("At least one parameter outside grid.\n{}".format(' '.join(s)))
        return inds + np.squeeze(find)

    def knearest_inds(self, **params):
        """Find all parameter ``vertices`` within a sphere of radius
        sqrt(ndim).  The parameter values are converted to pixel coordinates
        before a search of the KDTree.

        :param params:
             Keyword arguments which must include keys corresponding to
             ``stellar_pars``, the parameters of the grid.

        :returns inds:
             The sorted indices of all vertices within sqrt(ndim) of the pixel
             coordinates, corresponding to **params.
        """
        # Convert from physical space to grid index space
        xtarg = self.params_to_grid(**params)
        # Query the tree within radius sqrt(ndim)
        try:
            inds = self._kdt.query_radius(xtarg.reshape(1, -1),
                                          r=np.sqrt(self.ndim))
        except(AttributeError):
            inds = self._kdt.query_ball_point(xtarg.reshape(1, -1),
                                              np.sqrt(self.ndim))
        return np.sort(inds[0])

    def linear_weights(self, knearest, **params):
        """Use ND-linear interpolation over the knearest neighbors.

        :param knearest:
            The indices of the ``vertices`` for which to calculate weights.

        :param params:
            The target parameter location, as keyword arguments.

        :returns wght:
            The weight for each vertex, computed as the volume of the hypercube
            formed by the target parameter and each vertex.  Vertices more than
            1 away from the target in any dimension are given a weight of zero.
        """
        xtarg = self.params_to_grid(**params)
        x = self.X[knearest, :]
        dx = xtarg - x
        # Fractional pixel weights
        wght = ((1 - dx) * (dx >= 0) + (1 + dx) * (dx < 0))
        # set weights to zero if model is more than a pixel away
        wght *= (dx > -1) * (dx < 1)
        # compute hyperarea for each model and return
        return wght.prod(axis=-1)

    def triangle_weights(self, knearest, **params):
        """Triangulate the k-nearest models, then use the barycenter of the
        enclosing simplex to interpolate.
        """
        inparams = np.array([params[p] for p in self.stellar_pars])
        dtri = Delaunay(self.model_points[knearest, :])
        triangle_ind = dtri.find_simplex(inparams)
        inds = dtri.simplices[triangle_ind, :]
        transform = dtri.transform[triangle_ind, :, :]
        Tinv = transform[:self.ndim, :]
        x_r = inparams - transform[self.ndim, :]
        bary = np.dot(Tinv, x_r)
        last = 1.0 - bary.sum()
        wghts = np.append(bary, last)
        oo = inds.argsort()
        return inds[oo], wghts[oo]


import numbers
from collections.abc import Iterable
import numpy as np
from numpy.random import normal
from qutip.qobjevo import QobjEvo, EvoElement
from qutip.qip.gates import (
    expand_operator, _check_qubits_oper)
from qutip.qobj import Qobj
from qutip.operators import sigmaz, destroy, identity
from qutip.tensor import tensor

__all__ = ["Noise", "DecoherenceNoise", "RelaxationNoise",
           "ControlAmpNoise", "RandomNoise", "UserNoise"]

def _dummy_qobjevo(dims, **kwargs):
    """
    Create a dummy :class":`qutip.QobjEvo` with
    a constant zero Hamiltonian. This is used since empty QobjEvo
    is not yet supported.
    """
    dummy = QobjEvo(tensor([identity(d) for d in dims]) * 0., **kwargs)
    return dummy

class Noise(object):
    """
    The base class representing noise in a processor.
    The noise object can be added to :class:`qutip.qip.Processor` and
    contributes to evolution.
    """
    def __init__(self):
        pass

    def _check_coeff_num(self, coeffs, ops_num):
        if len(coeffs) != ops_num:
            raise ValueError(
                "The length of coeffs is not {}".format(ops_num))

class DecoherenceNoise(Noise):
    """
    The decoherence noise in a processor. It generates a list of
    collapse operators.

    Parameters
    ----------
    c_ops: :class:`qutip.Qobj` or list
        The Hamiltonian representing the dynamics of the noise.
        len(ops)=len(coeffs) is required.

    targets: int or list, optional
        The indices of qubits that are acted on. Default is the first
        N qubits

    coeffs: list, optional
        A list of the coefficients for the control Hamiltonians.
        For available choice, see :class:`Qutip.QobjEvo`

    tlist: array_like, optional
        A NumPy array specifies the time of each coefficient.

    all_qubits: bool, optional
        If c_ops contains only single qubits collapse operator,
        all_qubits=True will allow it to be applied to all qubits.

    Attributes
    ----------
    c_ops: :class:`qutip.Qobj` or list
        The Hamiltonian representing the dynamics of the noise.

    targets: list
        The indices of qubits that are acted on.

    coeffs: list
        A list of the coefficients for the control Hamiltonians.
        For available choice, see :class:`Qutip.QobjEvo`

    tlist: array_like
        A NumPy array specifies the time of each coefficient.

    all_qubits: bool
        If c_ops contains only single qubits collapse operator,
        all_qubits=True will allow it to be applied to all qubits.
    """
    def __init__(self, c_ops, targets=None, coeffs=None, tlist=None,
                 all_qubits=False):
        if isinstance(c_ops, Qobj):
            self.c_ops = [c_ops]
        else:
            self.c_ops = c_ops
        self.coeffs = coeffs
        self.tlist = tlist
        self.targets = targets
        if all_qubits:
            if not all([c_op.dims == [[2], [2]] for c_op in self.c_ops]):
                raise ValueError(
                    "The operator is not a single qubit operator, "
                    "thus cannot be applied to all qubits")
        self.all_qubits = all_qubits

    def get_noise(self, N, dims=None):
        """
        Return the quantum objects representing the noise.

        Parameters
        ----------
        N: int
            The number of component systems.

        dims: list, optional
            The dimension of the components system, the default value is
            [2,2...,2] for qubits system.

        Returns
        -------
        qobjevo_list: list
            A list of :class:`qutip.Qobj` or :class:`qutip.QobjEvo`
            representing the decoherence noise.
        """
        if dims is None:
            dims = [2] * N
        qobj_list = []
        for i, c_op in enumerate(self.c_ops):
            if self.all_qubits:
                qobj_list += expand_operator(
                    oper=c_op, N=N, targets=self.targets, dims=dims,
                    cyclic_permutation=True)
            else:
                qobj_list.append(
                    expand_operator(
                        oper=c_op, N=N, targets=self.targets, dims=dims))
        # time-independent
        if self.coeffs is None:
            return qobj_list
        # time-dependent
        if self.tlist is None:
            raise ValueError("tlist is required for time-dependent noise.")
        qobjevo_list = []
        for i, temp in enumerate(qobj_list):
            self._check_coeff_num(self.coeffs, len(qobj_list))
            qobjevo_list.append(QobjEvo(
                [qobj_list[i], self.coeffs[i]],
                tlist=self.tlist))
        return qobjevo_list

class RelaxationNoise(Noise):
    """
    The decoherence on each qubit characterized by two time scales t1 and t2.

    Parameters
    ----------
    t1: float or list, optional
        Characterize the decoherence of amplitude damping for
        each qubit.

    t2: float or list, optional
        Characterize the decoherence of dephasing for
        each qubit.

    Attributes
    ----------
    t1: list
        Characterize the decoherence of amplitude damping for
        each qubit.

    t2: list
        Characterize the decoherence of dephasing for
        each qubit.
    """
    def __init__(self, t1=None, t2=None):
        self.t1 = t1
        self.t2 = t2

    def _T_to_list(self, T, N):
        """
        Check if the relaxation time is valid

        Parameters
        ----------
        T: list of float
            The relaxation time

        N: int
            The number of component systems.

        Returns
        -------
        T: list
            The relaxation time in Python list form
        """
        if (isinstance(T, numbers.Real) and T > 0) or T is None:
            return [T] * N
        elif isinstance(T, Iterable) and len(T) == N:
            if all([isinstance(t, numbers.Real) and t > 0 for t in T]):
                return T
        else:
            raise ValueError(
                "Invalid relaxation time T={},"
                "either the length is not equal to the number of qubits, "
                "or T is not a positive number.".format(T))

    def get_noise(self, N, dims=None):
        """
        Return the quantum objects representing the noise.

        Parameters
        ----------
        N: int
            The number of component systems.

        dims: list, optional
            The dimension of the components system, the default value is
            [2,2...,2] for qubits system.

        Returns
        -------
        qobjevo_list: list
            A list of :class:`qutip.Qobj` or :class:`qutip.QobjEvo`
            representing the decoherence noise.
        """
        if dims is None:
            dims = [2] * N
        self.t1 = self._T_to_list(self.t1, N)
        self.t2 = self._T_to_list(self.t2, N)
        if len(self.t1) != N or len(self.t2) != N:
            raise ValueError(
                "Length of t1 or t2 does not match N, "
                "len(t1)={}, len(t2)={}".format(
                    len(self.t1), len(self.t2)))
        qobjevo_list = []
        for qu_ind in range(N):
            t1 = self.t1[qu_ind]
            t2 = self.t2[qu_ind]
            if t1 is not None:
                qobjevo_list.append(
                    expand_operator(
                        1/np.sqrt(t1) * destroy(2), N, qu_ind, dims=dims))
            if t2 is not None:
                # Keep the total dephasing ~ exp(-t/t2)
                if t1 is not None:
                    if 2*t1 < t2:
                        raise ValueError(
                            "t1={}, t2={} does not fulfill "
                            "2*t1>t2".format(t1, t2))
                    T2_eff = 1./(1./t2-1./2./t1)
                else:
                    T2_eff = t2
                qobjevo_list.append(
                    expand_operator(
                        1/np.sqrt(2*T2_eff) * sigmaz(), N, qu_ind, dims=dims))
        return qobjevo_list

class ControlAmpNoise(Noise):
    """
    The noise in the amplitude of the control pulse.

    Parameters
    ----------
    coeffs: list
        A list of the coefficients for the control Hamiltonians.
        For available choices, see :class:`Qutip.QobjEvo`.

    tlist: array_like, optional
        A NumPy array specifies the time of each coefficient.

    ops: :class:`qutip.Qobj` or list
        The Hamiltonian representing the dynamics of the noise.
        len(ops)=len(coeffs) is required.

    targets: int or list, optional
        The indices of qubits that are acted on. Default is the first
        N qubits

    cyclic_permutation: boolean, optional
        If true, the Hamiltonian will be expanded for
        all cyclic permutation of the target qubits.

    Attributes
    ----------
    coeffs: list
        A list of the coefficients for the control Hamiltonians.
        For available choices, see :class:`Qutip.QobjEvo`.

    tlist: array_like
        A NumPy array specifies the time of each coefficient.

    ops: list
        The Hamiltonian representing the dynamics of the noise.

    targets: list
        The indices of qubits that are acted on.

    cyclic_permutation: boolean
        If true, the Hamiltonian will be expanded for
        all cyclic permutation of the target qubits.
    """
    def __init__(self, coeffs, tlist, ops=None, targets=None,
                 cyclic_permutation=False):
        self.coeffs = coeffs
        self.tlist = tlist
        if isinstance(ops, Qobj):
            self.ops = [ops]
        else:
            self.ops = ops
        self.targets = targets
        self.cyclic_permutation = cyclic_permutation

    def get_noise(self, N, proc_qobjevo=None, dims=None):
        """
        Return the quantum objects representing the noise.

        Parameters
        ----------
        N: int
            The number of component systems.

        proc_qobjevo: :class:`qutip.QobjEvo`, optional
            If no operator is defined in the noise object, `proc_qobjevo`
            will be used as operators, otherwise the operators in the
            object is used.

        dims: list, optional
            The dimension of the components system, the default value is
            [2,2...,2] for qubits system.

        Returns
        -------
        noise_qobjevo: :class:`qutip.QobjEvo`
            A :class:`qutip.Qobj` representing the noise.
        """
        if dims is None:
            dims = [2] * N

        # If new operators are given
        if self.ops is not None:
            if self.cyclic_permutation:
                ops = []
                for op in self.ops:
                    ops += expand_operator(
                        oper=op, N=N, targets=self.targets, dims=dims,
                        cyclic_permutation=True)
            else:
                ops = [
                    expand_operator(
                        oper=op, N=N, targets=self.targets, dims=dims)
                    for op in self.ops]
        # If no operators given, use operators in the processor
        elif proc_qobjevo is not None:
            # If there is a constant part
            if proc_qobjevo.cte.norm() > 1.e-15:
                ops = [proc_qobjevo.cte]
            else:
                ops = []
            ops += [ele.qobj for ele in proc_qobjevo.ops]
        else:
            raise ValueError(
                "No operators found.")

        if len(ops) > len(self.coeffs):
            raise ValueError("The number of coefficient has to be larger than"
                             "{}".format(len(ops)))
        return QobjEvo([[ops[i], self.coeffs[i]] for i in range(len(ops))],
                       tlist=self.tlist)

class RandomNoise(ControlAmpNoise):
    """
    Random noise in the amplitude of the control pulse. The arguments for
    the random generator need to be given as key word arguments.

    Parameters
    ----------
    rand_gen: numpy.random, optional
        A random generator in numpy.random, it has to take a ``size``
        parameter.

    dt: float, optional
        The time interval between two random amplitude. The coefficients
        of the noise are the same within this time range.

    ops: list, optional
        The Hamiltonian representing the dynamics of the noise.

    targets: list or int, optional
        The indices of qubits that are acted on.

    cyclic_permutation: boolean, optional
        If true, the Hamiltonian will be expanded for
        all cyclic permutation of the target qubits.

    kwargs:
        Key word arguments for the random number generator.

    Attributes
    ----------
    ops: list
        The Hamiltonian representing the dynamics of the noise.

    coeffs: list
        A list of the coefficients for the control Hamiltonians.
        For available choices, see :class:`Qutip.QobjEvo`.

    targets: list
        The indices of qubits that are acted on.

    cyclic_permutation: boolean
        If true, the Hamiltonian will be expanded for
        all cyclic permutation of the target qubits.

    rand_gen: numpy.random
        A random generator in numpy.random, it has to take a ``size``
        parameter.

    kwargs: dict
        Key word arguments for the random number generator.
    """
    def __init__(
            self, rand_gen=None, dt=None, ops=None, targets=None,
            cyclic_permutation=False, **kwargs):
        super(RandomNoise, self).__init__(
            coeffs=None, tlist=None, ops=ops, targets=targets,
            cyclic_permutation=cyclic_permutation)
        if rand_gen is None:
            self.rand_gen = np.random.normal
        else:
            self.rand_gen = rand_gen
        self.kwargs = kwargs
        if "size" in kwargs:
            raise ValueError("size is preditermined inside the noise object.")
        self.dt = dt

    def get_noise(self, N, proc_qobjevo=None, dims=None):
        """
        Return the quantum objects representing the noise.

        Parameters
        ----------
        N: int
            The number of component systems.

        proc_qobjevo: :class:`qutip.QobjEvo`, optional
            If no operator is defined in the noise object, `proc_qobjevo`
            wil be used as operators, otherwise the operators in the
            object is used.

        dims: list, optional
            The dimension of the components system, the default value is
            [2,2...,2] for qubits system.

        Returns
        -------
        noise_qobjevo: :class:`qutip.QobjEvo`
            A :class:`qutip.Qobj` representing the decoherence noise.
        """
        if dims is None:
            dims = [2] * N
        tlist = proc_qobjevo.tlist
        if self.ops is not None:
            if self.cyclic_permutation:
                ops_num = len(self.ops) * N
            else:
                ops_num = len(self.ops)
        elif proc_qobjevo is not None:
            # +1 for the constant part in QobjEvo,
            # if no cte part the last coeffs will be ignored
            ops_num = len(proc_qobjevo.ops) + 1
        if self.dt is not None:
            # create new tlist and random coeffs
            num_rand = int(np.floor((tlist[-1]-tlist[0])/self.dt))+1
            self.coeffs = self.rand_gen(
                **self.kwargs, size=(ops_num, num_rand))
            tlist = (np.arange(0, self.dt*num_rand, self.dt)[:num_rand] +
                     tlist[0])
            # [:num_rand] for round of error like 0.2*6=1.2000000000002
        else:
            self.coeffs = self.rand_gen(
                **self.kwargs, size=(ops_num, len(tlist)))
        self.tlist = tlist
        return super(RandomNoise, self).get_noise(
            N, proc_qobjevo=proc_qobjevo, dims=dims)

class UserNoise(Noise):
    """
    Abstract class for user defined noise. To define a noise object,
    one could overwrite the constructor and the class method `get_noise`.
    """
    def get_noise(self, N, proc_qobjevo, dims=None):
        """
        Template method. To define a noise object,
        one should over write this method and
        return the unitary evolution part as a :class: `qutip.QobjEvo`
        and a list of collapse operators in the form of either
        :class: `qutip.QobjEvo` or :class: `qutip.Qobj`.

        Parameters
        ----------
        N: int
            The number of component systems.

        proc_qobjevo: :class:`qutip.QobjEvo`
            The object representing the ideal evolution in the processor.

        dims: list, optional
            The dimension of the components system, the default value is
            [2,2...,2] for qubits system.

        Returns
        -------
        noise_qobjevo: :class:`qutip.QobjEvo`
            A :class:`qutip.Qobj` representing the decoherence noise.

        collapse_list: list
            A list of :class:`qutip.Qobj` or :class:`qutip.QobjEvo`
            representing the decoherence noise.
        """
        if dims is None:
            dims = [2] * N
        return _dummy_qobjevo(dims), []


import os.path
import glob
from collections import OrderedDict

import numpy as np
import astropy.units as u
from astropy.table import Table, Column
import scipy.io
from sunpy.io.special.genx import read_genx
from sunpy.time import parse_time
from sunpy.data import manager

__all__  = ['chianti_kev_line_common_load_light', 'chianti_kev_line_common_load',
            'chianti_kev_cont_common_load', 'read_abundance_genx', 'load_xray_abundances']

def chianti_kev_line_common_load_light():
    """
    Read only X-ray emission line info needed for the chianti_kev functions.

    Unlike chianti_kev_line_common_load which formats and returns all in the file,
    this function only returns the data required by the ChiantiKevLines class.

    Returns
    -------
    zindex: `numpy.ndarray`
        Indicies of elements as they appear in periodic table.

    line_peak_energies: `astropy.units.Quantity` of `list` of `astropy.units.Quantity`
        The energies of the line peaks.

    line_logT_bins: `numpy.ndarray` or `list` of `numpy.ndarray`
        The log10 temperature bins over which the line intensities are known.

    line_colEMs: `astropy.units.Quantity` of `list` of `astropy.units.Quantity`
        The column emission measures used to calculate the intensities for each line.

    line_element_indices: `numpy.ndarray` or `list` of `numpy.ndarray`
        The atomic number of each line minus 1.

    line_intensities: `astropy.units.Quantity`
        Intensities of each of the lines in line_properties over a temperature axis.
        The array is 2D with axes of (line, temperature axis)
    """
    # Read linefile
    contents = _read_linefile()
    out = contents["out"]

    zindex = contents["zindex"]
    line_logT_bins = _clean_array_dims(out["LOGT_ISOTHERMAL"])
    line_colEMs = 10.**_clean_array_dims(out["LOGEM_ISOTHERMAL"]) / u.cm**5
    wvl_units = _clean_units(out["WVL_UNITS"])
    int_units = _clean_units(out["INT_UNITS"])

    line_intensities = []
    line_element_indices = []
    line_peak_energies = []
    for j, lines in enumerate(out["lines"]):
        line_element_indices.append(lines["IZ"])
        line_peak_energies.append(u.Quantity(lines["WVL"], unit=wvl_units).to(u.keV, equivalencies=u.spectral()))

        # Sort lines in ascending energy.
        ordd = np.argsort(np.array(line_peak_energies[j]))
        line_element_indices[j] = line_element_indices[j][ordd]
        line_peak_energies[j] = line_peak_energies[j][ordd]

        # Extract line intensities.
        line_intensities.append(_extract_line_intensities(lines["INT"][ordd]) * int_units)

    # If there is only one element in the line properties, unpack values.
    if len(out["lines"]) == 1:
        line_element_indices = line_element_indices[0]
        line_peak_energies = line_peak_energies[0]
        line_intensities = line_intensities[0]

    return zindex, line_peak_energies, line_logT_bins, line_colEMs, line_element_indices, \
        line_intensities

def chianti_kev_line_common_load():
    """
    Read file containing X-ray emission line info needed for the chianti_kev functions.

    Returns
    -------
    zindex: `numpy.ndarray`
        Indicies of elements as they appear in periodic table.

    line_meta: `dict`
        Various metadata associated with line properties.

    line_properties: `astropy.table.Table`
        Various properties of each lines.

    line_intensities: `astropy.units.Quantity`
        Intensities of each of the lines in line_properties over a temperature axis.
        The array is 2D with axes of (line, temperature axis)
    """
    # Read linefile.
    contents = _read_linefile()
    zindex = contents["zindex"]
    out = contents["out"]

    # Repackage metadata from file.
    date = []
    for date_byte in out["DATE"]:
        date_strings = str(date_byte[3:], 'utf-8').split()
        date.append(parse_time("{0}-{1}-{2} {3}".format(date_strings[3], date_strings[0],
                                                        date_strings[1], date_strings[2])))
    if len(date) == 1:
        date = date[0]
    line_meta = {
        "IONEQ_LOGT": _clean_array_dims(out["IONEQ_LOGT"]),
        "IONEQ_NAME": _clean_string_dims(out["IONEQ_NAME"]),
        "IONEQ_REF": _combine_strings(out["IONEQ_REF"]),
        "WVL_LIMITS": _clean_array_dims(out["WVL_LIMITS"]),
        "MODEL_FILE": _clean_string_dims(out["MODEL_FILE"]),
        "MODEL_NAME": _clean_string_dims(out["MODEL_NAME"]),
        "MODEL_NE": _clean_array_dims(out["MODEL_NE"]),
        "MODEL_PE": _clean_array_dims(out["MODEL_PE"]),
        "MODEL_TE": _clean_array_dims(out["MODEL_TE"]),
        "WVL_UNITS": _clean_units(out["WVL_UNITS"]),
        "INT_UNITS": _clean_units(out["INT_UNITS"]),
        "ADD_PROTONS": _clean_array_dims(out["ADD_PROTONS"], dtype=int),
        "DATE": date,
        "VERSION": _clean_string_dims(out['VERSION']),
        "PHOTOEXCITATION": _clean_array_dims(out["PHOTOEXCITATION"], dtype=int),
        "LOGT_ISOTHERMAL": _clean_array_dims(out["LOGT_ISOTHERMAL"]),
        "LOGEM_ISOTHERMAL": _clean_array_dims(out["LOGEM_ISOTHERMAL"]),
        "chianti_doc": _clean_chianti_doc(contents["chianti_doc"])
        }

    # Repackage out["line"] into a Table with appropriate units.
    # Create a list of tables to make sure all data in file is captured.
    # Although only one iteration is expected.
    line_properties = []
    line_intensities = []
    for lines in out["lines"]:
        line_props = Table()
        line_props["IZ"] = Column(lines["IZ"], description="Atomic number of ion element.")
        line_props["ION"] = Column(lines["ION"],
            description="Integer ionization state in astronomical notation, i.e. ION-1 = negative charge of ion.")
        line_props["IDENT"] = Column(lines["IDENT"])
        line_props["IDENT_LATEX"] = Column(lines["IDENT_LATEX"])
        line_props["SNOTE"] = Column(lines["SNOTE"],
            description="Ion label in astronomical (roman numeral) notation.")
        line_props["LVL1"] = Column(lines["LVL1"])
        line_props["LVL2"] = Column(lines["LVL2"])
        line_props["TMAX"] = Column(lines["TMAX"])
        line_props["WVL"] = Column(lines["WVL"], unit=line_meta["WVL_UNITS"])
        line_props["ENERGY"] = line_props["WVL"].quantity.to(u.keV, equivalencies=u.spectral())
        line_props["FLAG"] = Column(lines["FLAG"])

        # Sort lines in ascending energy.
        ordd = np.argsort(np.array(line_props["WVL"]))[::-1]
        line_props = line_props[ordd]

        # Extract line intensities.
        line_intensities.append(_extract_line_intensities(lines["INT"][ordd]))

        # Enter outputs from this iteration into list.
        line_properties.append(line_props)
        line_intensities.append(line_ints)

    # If there is only one element in the line properties, unpack values.
    if len(out["lines"]) == 1:
        line_properties = line_properties[0]
        line_intensities = line_intensities[0]

    return zindex, line_meta, line_properties, line_intensities * line_meta["INT_UNITS"]

@manager.require('chianti_cont_1_250',
                 ['https://hesperia.gsfc.nasa.gov/ssw/packages/xray/dbase/chianti/chianti_cont_1_250_v71.sav'],
                 'aadf4355931b4c241ac2cd5669e89928615dc1b55c9fce49a155b70915a454dd')
def chianti_kev_cont_common_load(_extra=None):
    """
    Read X-ray continuum emission info needed for the chianti_kev functions.

    Returns
    -------
    zindex: `numpy.ndarray`
        Indicies of elements as they appear in periodic table.
    continuum_properties: `dict`
        Properties of continuum emission.
    """
    contfile = manager.get("chianti_cont_1_250")
    # Read file
    contents = scipy.io.readsav(contfile)
    zindex = contents["zindex"]
    edge_str = {
            "CONVERSION": _clean_array_dims(contents["edge_str"]["CONVERSION"]),
            "WVL": _clean_array_dims(contents["edge_str"]["WVL"]),
            "WVLEDGE": _clean_array_dims(contents["edge_str"]["WVLEDGE"])
                }
    continuum_properties = {
            "totcont": contents["totcont"],
            "totcont_lo": contents["totcont_lo"],
            "edge_str": edge_str,
            "ctemp": contents["ctemp"],
            "chianti_doc": _clean_chianti_doc(contents["chianti_doc"])
                            }

    return zindex, continuum_properties

@manager.require('xray_abundances',
                 ['https://hesperia.gsfc.nasa.gov/ssw/packages/xray/dbase/chianti/xray_abun_file.genx'],
                 '92c0e1f9a83da393cc38840752fda5a5b44c5b18a4946e5bf12c208771fe0fd3')
def load_xray_abundances(abundance_type=None):
    """
    Returns the abundances written in the xray_abun_file.genx

    The abundances are taken from CHIANTI and MEWE.  The source filenames are:
    cosmic sun_coronal sun_coronal_ext sun_hybrid sun_hybrid_ext sun_photospheric mewe_cosmic mewe_solar
    The first six come fron Chianti, the last two from Mewe.  They are:
    cosmic sun_coronal sun_coronal_ext sun_hybrid sun_hybrid_ext sun_photospheric mewe_cosmic mewe_solar
    These abundances are used with CHIANTI_KEV.  MEWE_KEV can only use the two mewe sourced
    abundance distributions unless using a heavily modified rel_abun structure for all of the elements.

    Parameters
    ----------
    abundance_type: `str`
        Type of abundance to be read from file.  Option are (From Chianti)
        1. cosmic
        2. sun_coronal - default abundance
        3. sun_coronal_ext
        4. sun_hybrid
        5. sun_hybrid_ext
        6. sun_photospheric
        7. mewe_cosmic
        8. mewe_solar - default for mewe_kev

    Returns
    -------
    out:
        Array of 50 abundance levels for first 50 elements.

    """
    # If kwargs not set, set defaults
    if abundance_type is None:
        abundance_type = "sun_coronal"
    xray_abundance_file = manager.get("xray_abundances")
    # Read file
    contents = read_abundance_genx(xray_abundance_file)
    # Extract relevant abundance type
    abundances = contents[abundance_type]

    return abundances

def read_abundance_genx(filename):
    # Read file.
    contents = read_genx(filename)
    # Combine data and keys from each entry in file.
    output = OrderedDict()
    for arr in contents["SAVEGEN0"]:
        output[arr["FILNAM"]] = arr["ABUND"]
    # Add header data
    output["header"] = contents["HEADER"]
    output["header"]["CHIANTI VERSION"] = float(contents["SAVEGEN1"][:3])

    return output

@manager.require('chianti_lines_1_10',
                 ['https://hesperia.gsfc.nasa.gov/ssw/packages/xray/dbase/chianti/chianti_lines_1_10_v71.sav'],
                  '2046d818efec207a83e9c5cc6ba4a5fa8574bf8c2bd8a6bb9801e4b8a2a0c677')
def _read_linefile():
    linefile = manager.get('chianti_lines_1_10')
    # Read file
    contents = scipy.io.readsav(linefile)
    zindex = contents["zindex"]
    out = contents["out"]

    return contents

def _extract_line_intensities(lines_int_sorted):
    line_ints = np.empty((lines_int_sorted.shape[0], lines_int_sorted[0].shape[0]), dtype=float)
    for i in range(line_ints.shape[0]):
        line_ints[i, :] = lines_int_sorted[i]
    return line_ints

def _clean_array_dims(arr, dtype=None):
    # Initialize a single array to hold contents of input arr.
    result = np.empty(list(arr.shape) + list(arr[0].shape))
    # Combine arrays in arr into single array.
    for i in range(arr.shape[0]):
        result[i] = arr[i]
    # Remove redundant dimensions
    result = np.squeeze(result)
    # If result is now unsized, convert to scalar.
    if result.shape == ():
        result = result.item()
        if dtype is not None:
            dtype(result)
    return result

def _clean_string_dims(arr):
    result = [str(s, 'utf-8') for s in arr]
    if len(result) == 1:
        result = result[0]
    return result

def _combine_strings(arr):
    result = [".".join([str(ss, 'utf-8') for ss in s]) for s in arr]
    if len(result) == 1:
        result = result[0]
    return result

def _clean_units(arr):
    result = []
    for a in arr:
        unit = str(a, 'utf-8')
        unit_components = unit.split()
        for i, component in enumerate(unit_components):
            # Remove plurals
            if component in ["photons", "Angstroms"]:
                component = component[:-1]
            # Insert ** for indices.
            component_minus_split = component.split("-")
            if len(component_minus_split) > 1:
                "**-".join(component_minus_split)
            component_plus_split = component.split("+")
            if len(component_plus_split) > 1:
                "**-".join(component_plus_split)
            unit_components[i] = component
        result.append("*".join(unit_components))
    if len(result) == 1:
        result = result[0]

    return u.Unit(result)

def _clean_chianti_doc(arr):
    chianti_doc = {}
    chianti_doc["ion_file"] = str(arr[0][0], 'utf-8')
    chianti_doc["ion_ref"] = "{0}.{1}.{2}".format(str(arr["ion_ref"][0][0], 'utf-8'),
                                                  str(arr["ion_ref"][0][1], 'utf-8'),
                                                  str(arr["ion_ref"][0][2], 'utf-8'))
    chianti_doc["version"] = str(arr[0][2], 'utf-8')
    return chianti_doc

import numpy,import scipy,import astropy,from astropy


#Author - Evan Leister
import eclipse_cells as ec
import numpy as np
import matplotlib.pyplot as plt

class Injector(object):
    def __init__(self, index, x, y, ratio, layer_id, end_days, mass_rate):
        self.index = index
        self.x = x
        self.y = y
        self.ratio = ratio
        self.layer_id = layer_id
        self.end_days = end_days # list of intervals in integer days
        self.mass_rate = mass_rate # list of mass rate is given in Mt/yr
        self.radius = 1.0
        if len(mass_rate) != len(end_days):
            print "Mass inflow and interval ends must match"
            return 1
    def write_injector(self, f):
        print "mass rate is"
        print self.mass_rate
        f.write(', '.join([str(self.index), str(self.x), str(self.y),\
                str(self.ratio), str(self.layer_id),\
                str(self.radius), str(len(self.mass_rate))]))
        f.write('\n')
        for i in range(len(self.mass_rate)):
            print len(self.mass_rate), len(self.end_days)
            f.write(', '.join([str(self.end_days[i]), \
                    str(self.mass_rate[i])]))
            f.write('\n')
        return f

def write_injwells(injectors):
    f = open("InjWells.txt","w")
    print "writing Injwells.txt"
    for inj in injectors:
        inj.write_injector(f)
    f.close()
    return 0

def write_system(timestep_days, output_days, simtime_years, output_control, layers):
    f = open("System.txt", "w")
    print "writing System.txt"
    f.write(''.join([str(timestep_days),'\n']))
    f.write(''.join([str(output_days),'\n']))
    years = simtime_years
    f.write(''.join([str(simtime_years),'\n']))
    f.write(''.join([output_control,'\n']))
    # mass balance output 'string' massbalance nomassbalance
    f.write('massbalance \n')
    # number of layers in the model [int]
    f.write(''.join([str(len(layers)),'\n']))
    # file names that contain the grid data
    for s in layers:
        f.write(s)
    f.close()
    return 0

class Layer(object):
    def __init__(self, layer_name, l_type, l_id, l_co2_rho, l_bri_rho, l_co2_mu,\
            l_bri_mu, sc_res, sb_res, c_co2, c_bri, c_roc, cap_rp_id, \
            nx, ny, nz = 1, gradient = 10.,\
            homogeneous = False, permval = 2000., poroval = False):
        """
        self.layer_name : Name of Text File
        l_type : Layer Type
        l_id : Layer ID
        l_co2_rho : CO2 density
        self.l_bri_rho : Brine Density
        self.l_co2_mu : CO2 viscosity
        self.l_bri_mu : brine viscosity
        self.sc_res : CO2 residual saturation
        self.sb_res : brine residual saturation
        self.c_co2 : CO2 compressibility
        self.c_bri : Brine Compressibility
        self.c_roc : Rock Compressibility
        self.cap_rp_id : Capillary - Rel/perm ID
        self.gradient : Pressure gradient
        """
        self.layer_name = layer_name
        self.l_type = l_type
        self.l_id = l_id
        self.l_co2_rho = l_co2_rho
        self.l_bri_rho = l_bri_rho
        self.l_co2_mu = l_co2_mu
        self.l_bri_mu = l_bri_mu
        self.sc_res = sc_res
        self.sb_res = sb_res
        self.c_co2 = c_co2
        self.c_bri = c_bri
        self.c_roc = c_roc
        self.cap_rp_id = cap_rp_id
        # list of GridCell objects
        self.grid_cells = []
        self.gradient = gradient #[MPa/km]
        self.nx = nx
        self.ny = ny
        self.nz = nz
        self.homogeneous = homogeneous
        self.permval = permval
        self.poroval = poroval

    def fill_uniform_grid(self, dx, dy, dz, center_depth, phi, k):
        for j in range(self.ny):
            for i in range(self.nx):
                x = (dx/2. + dx * i)
                y = (dy/2. + dy * j)
                top_b = -center_depth + dz/2.
                bottom_b = -center_depth - dz/2.

                if i == (self.nx - 1):
                    east_bc = 3
                else:
                    east_bc = 1

                if j == (self.ny - 1):
                    north_bc = 3
                else:
                    north_bc = 1

                if i == 0:
                    west_bc = 3
                else:
                    west_bc = 1

                if j == 0:
                    south_bc = 3
                else:
                    south_bc = 1

                pressure = -self.gradient * bottom_b * 1000
                gc = GridCell(top_b, bottom_b, x, y, dx, dy, phi, k,\
                        west_bc, east_bc, south_bc, north_bc, pressure)
                self.grid_cells.append(gc)

        return 0

    def plot_perm_data(self, e_cells):
        cell_ind = np.zeros(len(e_cells))
        anis = np.zeros(len(e_cells))
        depth = np.zeros(len(e_cells))
        perm = np.zeros(len(e_cells))
        poro = np.zeros(len(e_cells))
        for i in range(len(e_cells)):
            cell_ind[i] = i
            perm[i] = e_cells[i].getXPermeability()
            anis[i] = e_cells[i].getZPermeability() / \
                        e_cells[i].getXPermeability()
            depth[i] = e_cells[i].getTopZ()
            poro[i] = e_cells[i].getPorosity()
        print "plotting anisotropy ratio"
        fig1 = plt.figure()
        ax1 = fig1.add_subplot(111)
        a = ax1.plot(cell_ind, anis)
        ax1.set_xlabel('cell_index []')
        ax1.set_ylabel('anisotropy kz/kx')
        plt.savefig('ec_anis_cells.png')
        plt.close()
        print "plotting permeabilityvsdepth"
        fig2 = plt.figure()
        ax2 = fig2.add_subplot(111)
        b = ax2.scatter(depth, perm)
        ax2.set_xlabel('depth [m]')
        ax2.set_ylabel('permeability [md]')
        plt.savefig('ec_perm_depth.png')
        plt.close()
        print "plotting porosity vsdepth"
        fig3 = plt.figure()
        ax3 = fig3.add_subplot(111)
        b = ax3.scatter(depth, poro)
        ax3.set_xlabel('depth [m]')
        ax3.set_ylabel('porosity []')
        plt.savefig('ec_poro_depth.png')
        plt.close()

        return 0

    def fill_nonuniform_grid(self, e_cells):
        print "Filling nonuniform grid"
        count = 0
        k = 0
        columnz = []
        columnk = []
        columnphi = []
        check_col = self.nz
        check_plane = self.nx*self.ny
        for j in range(self.ny-1,-1,-1):
            for i in range(0,self.nx):
                for k in range(0,self.nz):
                    ind = (i + self.nx *j) + check_plane*k
                    if i == 32 and j == 77:
                        print k, e_cells[ind].getZPermeability()
                    if e_cells[ind].getXPermeability() > 1:
                        columnz.append(e_cells[ind].getTopZ())
                        columnk.append(e_cells[ind].getXPermeability())
                        columnphi.append(e_cells[ind].getPorosity())
                # spits out the averages after the column index is filled.
                kmean, kvar = stats(columnk)
                if self.homogeneous == True:
                    k_write = self.permval
                else:
                    k_write = kmean
                if self.poroval == False:
                    phimean , phivar = stats(columnphi)
                else:
                    phimean = self.poroval

                top_b = -columnz[0]
                bottom_b = -columnz[-1]

                x = e_cells[ i + self.nx * j].getCenterX()
                y = e_cells[ i + self.nx * j].getCenterY()

                # get correct dx and dy
                if i != (self.nx-1):
                    x_1 = e_cells[i+1 + self.nx*j].getCenterX()
                    x_0 = e_cells[ i + self.nx * j].getCenterX()
                    dx = (x_1 - x_0)
                else:
                    x_1 = e_cells[i + self.nx*j].getCenterX()
                    x_0 = e_cells[ i-1 + self.nx * j].getCenterX()
                    dx = (x_1 - x_0)
                if j != 0:
                    y_1 = e_cells[i + self.nx *(j-1)].getCenterY()
                    y_0 = e_cells[i + self.nx * j].getCenterY()
                    dy = y_1 - y_0
                else:
                    y_1 = e_cells[i + self.nx *j].getCenterY()
                    y_0 = e_cells[i + self.nx *(j+1)].getCenterY()
                    dy = y_1 - y_0

                #boundary condition key [integers]
                # 1 = internal
                # 3 = constant pressure
                # 4 = no flow
                if i == (self.nx - 1):
                    east_bc = 3
                else:
                    east_bc = 1

                if j == 0:
                    north_bc = 3
                else:
                    north_bc = 1

                if i == 0:
                    west_bc = 3
                else:
                    west_bc = 1

                if j == (self.ny - 1):
                    south_bc = 3
                else:
                    south_bc = 1

                pressure = -self.gradient * bottom_b * 1000
                gc = GridCell(top_b, bottom_b, x, y, dx, dy, phimean, k_write,\
                        west_bc, east_bc, south_bc, north_bc, pressure)
                self.grid_cells.append(gc)

                # increment loops
                count += 1
                columnz = []
                columnk = []
                columnphi = []
                k = 0
        return 0

    def write_layer(self):
        print "writing layer " + self.layer_name
        f = open("".join([self.layer_name,'.txt']),"w")
        g = open("thickness.txt","w")
        # layer type
        f.write(''.join([str(self.l_type) + '\n']))
        # layer id
        f.write(''.join([str(self.l_id) + '\n']))
        # fluid parameters
        f.write(''.join([str(self.l_co2_rho),'\n']))
        f.write(''.join([str(self.l_bri_rho),'\n']))
        f.write(''.join([str(self.l_co2_mu),'\n']))
        f.write(''.join([str(self.l_bri_mu),'\n']))
        f.write(''.join([str(self.sc_res),'\n']))
        f.write(''.join([str(self.sb_res),'\n']))
        f.write(''.join([str(self.c_co2),'\n']))
        f.write(''.join([str(self.c_bri),'\n']))
        f.write(''.join([str(self.c_roc),'\n']))
        if self.cap_rp_id == 0:
            f.write(''.join([str(self.cap_rp_id),'\n']))
        elif self.cap_rp_id == 1:
            lamb = 3.
            p_entry = 3000000.
            f.write(''.join([str(self.cap_rp_id), ', ',\
                    str(lamb), ', ', \
                    str(p_entry), '\n']))
            self.plot_cap_rp_bc(lamb, p_entry)
        # number of cells
        f.write(''.join([str(self.nx * self.ny), '\n']))
        for cel in self.grid_cells:
            g.write(''.join([str((cel.top_b - cel.bottom_b)),', ']))
            cel.write_cell(f)
        f.close()
        return 0
    def bc_cap(self, pentry, lamb):
        return pcap
    def plot_cap_rp_bc(self, lamb, p_entry):
        sb = np.linspace(self.sb_res,1.)
        pc = np.zeros(len(sb))
        krb = np.zeros(len(sb))
        krc = np.zeros(len(sb))
        for i in range(len(sb)):
            seff = (sb[i] - self.sb_res) / (1 - self.sb_res)
            pc[i] = p_entry * pow(seff, -1/lamb)
            krb[i] = pow(seff, (2 + 3 * lamb)/lamb)
            krc[i] = (1 - seff)**2 * (1 - pow(seff, (2 + lamb) / lamb))
        fig1 = plt.figure()
        ax1 = fig1.add_subplot(111)
        ax1.plot(sb, pc)
        ax1.set_xlabel('sb []')
        ax1.set_ylabel('pcap [Pa]')
        plt.savefig('pcap.png')
        plt.clf()
        fig2 = plt.figure()
        ax2 = fig2.add_subplot(111)
        ax2.plot(sb, krb, label = 'krb')
        ax2.plot(sb, krc, label = 'krc')
        ax2.legend(loc=1)
        ax2.set_xlabel('sb')
        plt.savefig('relperm.png')
        return 0

class GridCell(object):
    def __init__(self, top_b, bottom_b, x, y, dx, dy, phi, k,\
            west_bc, east_bc, south_bc, north_bc, pressure):
        self.top_b = top_b
        self.bottom_b = bottom_b
        self.x = x
        self.y = y
        self.dx = dx
        self.dy = dy
        self.phi = phi
        self.k = k
        self.west_bc = west_bc
        self.east_bc = east_bc
        self.south_bc = south_bc
        self.north_bc = north_bc
        self.pressure = pressure

    def write_cell(self, f):
        f.write(''.join([str(self.x),', ']))
        f.write(''.join([str(self.y),', ']))
        f.write(''.join([str(self.dx),', ']))
        f.write(''.join([str(self.dy),', ']))
        # porosity
        f.write(''.join(['%.4f' % self.phi ,', ']))
        # NOTE: Enters permeability for all dimensions since the formation is isotropic
        # in the planar directions and the z permeability is not used in VESA.
        xperm = self.k
        yperm = self.k
        zperm = self.k
        f.write(''.join(['%.0f' % xperm, ', ' ,'%.0f' % yperm, ', ',\
                '%.0f' % zperm, ', ']))
        f.write(''.join(['%.4f'  % self.bottom_b, ', ']))
        f.write(''.join(['%.4f' % self.top_b, ', ']))
        # initial CO2 saturation
        f.write('0.0, ')
        # past CO2 saturation
        f.write('0.0, ')
        #initial pressure at bottom of formation [Pa]
        f.write(''.join(['%.0f' % self.pressure, ', ']))
        f.write(''.join([str(self.east_bc), ', ']))
        f.write(''.join([str(self.north_bc), ', ']))
        f.write(''.join([str(self.west_bc), ', ']))
        f.write(''.join([str(self.south_bc), ', ']))
        f.write('\n')
        return f

def stats(data):
    sum_s = 0.0
    for value in data:
        sum_s += value
    mean = sum_s/len(data)
    sum_s = 0.0
    for value in data:
        sum_s += (value - mean)**2
    variance = sum_s/(len(data)-1)
    return(mean,variance)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
:mod:`milboost`
==================

.. module:: milboost
    :platform: Unix, Windows
    :synopsis:

.. moduleauthor:: hbldh <henrik.blidh@nedomkull.com>

Created on 2015-11-06, 08:48

"""

from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import absolute_import

import warnings

import numpy as np
from sklearn.ensemble.weight_boosting import ClassifierMixin, BaseWeightBoosting
from sklearn.tree import DecisionTreeClassifier
from sklearn.utils.validation import check_is_fitted
from scipy.optimize import fminbound

from skboost.milboost.softmax import SoftmaxFunction

__all__ = ['MILBoostClassifier', ]

class MILBoostClassifier(ClassifierMixin, BaseWeightBoosting):

    def __init__(self,
                 base_estimator=DecisionTreeClassifier(max_depth=10),
                 softmax=None,
                 n_estimators=50,
                 learning_rate=1.0,
                 random_state=None,
                 verbose=False):

        super(MILBoostClassifier, self).__init__(
            base_estimator=base_estimator,
            n_estimators=n_estimators,
            learning_rate=learning_rate,
            random_state=random_state)

        if not isinstance(softmax, SoftmaxFunction):
            raise TypeError("Softmax input must be an object of class `SoftmaxFunction`")
        self.softmax_fcn = softmax
        self._verbose = verbose

        self._bag_labels = None
        self._inferred_y = None
        self._bag_partitioning = None

    def __str__(self):
        return "{0}, with {1} {2} classifiers".format(
            self.__class__.__name__, len(self.estimators_), self.estimators_[0])

    def fit(self, X, y, sample_weight=None):
        """Build a boosted classifier from the training set (X, y).

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape = [n_samples, n_features]
            The training input samples. Sparse matrix can be CSC, CSR, COO,
            DOK, or LIL. DOK and LIL are converted to CSR.

        y : array-like of shape = [n_samples]
            The target values (class labels).

        sample_weight : array-like of shape = [n_samples], optional
            Sample weights. If None, the sample weights are initialized to
            ``1 / n_samples``.

        Returns
        -------
        self : object
            Returns self.
        """

        # Pre-compute bag labels and inferred instance labels from y.
        unique_bag_ids = np.unique(y)
        self._bag_labels = np.zeros((max(np.abs(unique_bag_ids)) + 1, ), 'int')
        self._bag_labels[np.abs(unique_bag_ids)] = np.sign(unique_bag_ids)
        self._bag_labels = self._bag_labels[1:]
        self._inferred_y = np.sign(y)
        self._bag_partitioning = np.cumsum(np.bincount(np.abs(y))[1:])

        # Fit
        out = super(MILBoostClassifier, self).fit(X, y, sample_weight)

        # Clean away stored labels.
        self._bag_labels = None
        self._inferred_y = None
        self._bag_partitioning = None

        return out

    def _boost(self, iboost, X, y, sample_weight, random_state):

        if iboost > 0:
            dv_pre = self.decision_function(X)
            instance_probabilites = self._estimate_instance_probabilities(dv_pre)
            bag_probabilites = self._estimate_bag_probabilites(instance_probabilites)
            sample_weight = self._calculate_new_weights(instance_probabilites, bag_probabilites)
        else:
            dv_pre = np.zeros(((X.shape[0]),), 'float')

        estimator = self._make_estimator()
        try:
            estimator.set_params(random_state=self.random_state)
        except ValueError:
            pass

        _weights = np.abs(sample_weight)
        estimator.fit(X, self._inferred_y, sample_weight=_weights)
        y_predict = estimator.predict(X)

        # Instances incorrectly classified
        incorrect = y_predict != self._inferred_y

        # Error fraction
        estimator_error = np.mean(
            np.average(incorrect, weights=_weights, axis=0))

        if iboost == 0:
            self.classes_ = getattr(estimator, 'classes_', None)
            self.n_classes_ = len(self.classes_)

        # Estimate alpha, the estimator weight.
        estimator_weight, nll = self._find_estimator_weight(y, dv_pre, y_predict)
        if self._verbose:
            print("[{0}] - err={1:.4f}, w={2:.4f}, -L={3:.4f}".format(iboost, estimator_error, estimator_weight, nll))

        if estimator_weight < 1e-5:
            _weights = None

        return _weights, estimator_weight, estimator_error

    def _negative_log_likelihood(self, bag_probabilities):
        positive_bags_log_prob = np.log(bag_probabilities[self._bag_labels > 0])
        positive_bags_log_prob[np.isinf(positive_bags_log_prob)] = 0.0
        negative_bags_log_prob = np.log(1 - bag_probabilities[self._bag_labels < 0])
        negative_bags_log_prob[np.isinf(negative_bags_log_prob)] = 0.0

        return -(np.sum(positive_bags_log_prob) + np.sum(negative_bags_log_prob))

    def _find_estimator_weight(self, y, dv_pre, y_pred):
        """Make line search to determine estimator weights."""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")

            def optimization_function(alpha):
                p_ij = self._estimate_instance_probabilities(dv_pre + alpha * y_pred)
                p_i = self._estimate_bag_probabilites(p_ij)
                return self._negative_log_likelihood(p_i)

            # TODO: Add option to choose optimization method.

            alpha, fval, err, n_func = fminbound(optimization_function, 0.0, 5.0, full_output=True, disp=1)
            if self.learning_rate < 1.0:
                alpha *= self.learning_rate
        return alpha, fval

    def _estimate_instance_probabilities(self, dv):
        return 1.0 / (1 + np.exp(-(2 * dv)))

    def _estimate_bag_probabilites(self, instance_probabilites):
        bags = self._bag_split(instance_probabilites)
        bag_probabilities = np.array([self.softmax_fcn.f(x) for x in bags])
        return bag_probabilities

    def _calculate_new_weights(self, instance_probabilites, bag_probabilities):
        weights = []
        for p_ij, p_i, Y_i in zip(self._bag_split(instance_probabilites),
                                  bag_probabilities,
                                  self._bag_labels):
            if Y_i > 0:
                if p_i == 0.0:
                    p_i = np.finfo(float).resolution
                term_1 = (2 * p_ij * (1 - p_ij)) / p_i
            else:
                if p_i == 1.0:
                    p_i = 1 - np.finfo(float).resolution
                term_1 = -((2 * p_ij * (1 - p_ij)) / (1 - p_i))
            weights += (term_1 * self.softmax_fcn.d_dt(p_ij)).tolist()

        return np.array(weights) / np.sum(np.abs(weights))

    def _bag_split(self, x):
        return np.split(x, self._bag_partitioning)[:-1]

    def decision_function(self, X):
        """Compute the decision function of ``X``.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape = [n_samples, n_features]
            The training input samples. Sparse matrix can be CSC, CSR, COO,
            DOK, or LIL. DOK and LIL are converted to CSR.

        Returns
        -------
        score : array, shape = [n_samples, k]
            The decision function of the input samples. The order of
            outputs is the same of that of the `classes_` attribute.
            Binary classification is a special cases with ``k == 1``,
            otherwise ``k==n_classes``. For binary classification,
            values closer to -1 or 1 mean more like the first or second
            class in ``classes_``, respectively.
        """
        check_is_fitted(self, "n_classes_")
        X = self._validate_X_predict(X)

        classes = self.classes_[:, np.newaxis]
        pred = sum((estimator.predict(X) == classes).T * w
                   for estimator, w in zip(self.estimators_,
                                           self.estimator_weights_))
        pred[:, 0] *= -1
        return pred.sum(axis=1)

    def staged_decision_function(self, X):
        """Compute decision function of ``X`` for each boosting iteration.

        This method allows monitoring (i.e. determine error on testing set)
        after each boosting iteration.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape = [n_samples, n_features]
            The training input samples. Sparse matrix can be CSC, CSR, COO,
            DOK, or LIL. DOK and LIL are converted to CSR.

        Returns
        -------
        score : generator of array, shape = [n_samples, k]
            The decision function of the input samples. The order of
            outputs is the same of that of the `classes_` attribute.
            Binary classification is a special cases with ``k == 1``,
            otherwise ``k==n_classes``. For binary classification,
            values closer to -1 or 1 mean more like the first or second
            class in ``classes_``, respectively.
        """
        check_is_fitted(self, "n_classes_")
        X = self._validate_X_predict(X)

        classes = self.classes_[:, np.newaxis]
        pred = None

        for weight, estimator in zip(self.estimator_weights_,
                                     self.estimators_):

            current_pred = estimator.predict(X)
            current_pred = (current_pred == classes).T * weight

            if pred is None:
                pred = current_pred
            else:
                pred += current_pred

            tmp_pred = np.copy(pred)
            tmp_pred[:, 0] *= -1
            yield (tmp_pred).sum(axis=1)

    def predict(self, X):
        """Predict classes for X.

        The predicted class of an input sample is computed as the weighted mean
        prediction of the classifiers in the ensemble.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape = [n_samples, n_features]
            The training input samples. Sparse matrix can be CSC, CSR, COO,
            DOK, or LIL. DOK and LIL are converted to CSR.

        Returns
        -------
        y : array of shape = [n_samples]
            The predicted classes.
        """
        pred = self.decision_function(X)

        return self.classes_.take(pred > 0, axis=0)

    def staged_predict(self, X):
        """Return staged predictions for X.

        The predicted class of an input sample is computed as the weighted mean
        prediction of the classifiers in the ensemble.

        This generator method yields the ensemble prediction after each
        iteration of boosting and therefore allows monitoring, such as to
        determine the prediction on a test set after each boost.

        Parameters
        ----------
        X : array-like of shape = [n_samples, n_features]
            The input samples.

        Returns
        -------
        y : generator of array, shape = [n_samples]
            The predicted classes.
        """
        classes = self.classes_

        for pred in self.staged_decision_function(X):
            yield np.array(classes.take(pred > 0, axis=0))

    def predict_proba(self, X):
        """Predict class probabilities for X.

        The predicted class probabilities of an input sample is computed as
        the weighted mean predicted class probabilities of the classifiers
        in the ensemble.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape = [n_samples, n_features]
            The training input samples. Sparse matrix can be CSC, CSR, COO,
            DOK, or LIL. DOK and LIL are converted to CSR.

        Returns
        -------
        p : array of shape = [n_samples]
            The class probabilities of the input samples. The order of
            outputs is the same of that of the `classes_` attribute.
        """
        check_is_fitted(self, "n_classes_")

        n_classes = self.n_classes_
        X = self._validate_X_predict(X)

        proba = sum(estimator.predict_proba(X) * w
                    for estimator, w in zip(self.estimators_,
                                            self.estimator_weights_))

        proba /= self.estimator_weights_.sum()
        proba = np.exp((1. / (n_classes - 1)) * proba)
        normalizer = proba.sum(axis=1)[:, np.newaxis]
        normalizer[normalizer == 0.0] = 1.0
        proba /= normalizer

        return proba

    def staged_predict_proba(self, X):
        """Predict class probabilities for X.

        The predicted class probabilities of an input sample is computed as
        the weighted mean predicted class probabilities of the classifiers
        in the ensemble.

        This generator method yields the ensemble predicted class probabilities
        after each iteration of boosting and therefore allows monitoring, such
        as to determine the predicted class probabilities on a test set after
        each boost.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape = [n_samples, n_features]
            The training input samples. Sparse matrix can be CSC, CSR, COO,
            DOK, or LIL. DOK and LIL are converted to CSR.

        Returns
        -------
        p : generator of array, shape = [n_samples]
            The class probabilities of the input samples. The order of
            outputs is the same of that of the `classes_` attribute.
        """
        X = self._validate_X_predict(X)

        n_classes = self.n_classes_
        proba = None
        norm = 0.

        for weight, estimator in zip(self.estimator_weights_,
                                     self.estimators_):
            norm += weight

            current_proba = estimator.predict_proba(X) * weight

            if proba is None:
                proba = current_proba
            else:
                proba += current_proba

            real_proba = np.exp((1. / (n_classes - 1)) * (proba / norm))
            normalizer = real_proba.sum(axis=1)[:, np.newaxis]
            normalizer[normalizer == 0.0] = 1.0
            real_proba /= normalizer

            yield real_proba

    def predict_log_proba(self, X):
        """Predict class log-probabilities for X.

        The predicted class log-probabilities of an input sample is computed as
        the weighted mean predicted class log-probabilities of the classifiers
        in the ensemble.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape = [n_samples, n_features]
            The training input samples. Sparse matrix can be CSC, CSR, COO,
            DOK, or LIL. DOK and LIL are converted to CSR.

        Returns
        -------
        p : array of shape = [n_samples]
            The class probabilities of the input samples. The order of
            outputs is the same of that of the `classes_` attribute.
        """
        return np.log(self.predict_proba(X))