from asteval import Interpreter

import multiprocessing as mp
from multiprocessing import Pool

import numpy as np

import astropy.units as u

from specutils import Spectrum1D
from specutils.fitting import fit_lines

__all__ = ['fit_model_to_spectrum']


def fit_model_to_spectrum(spectrum, component_list, expression,
                          run_fitter=False, window=None, n_cpu=None):
    """Fits a `~astropy.modeling.CompoundModel` to a
    `~specutils.Spectrum1D` instance.

    If the input spectrum represents a spectral cube, then fits
    the model to every spaxel in the cube, using
    a multiprocessor pool running in parallel (if ``n_cpu`` is
    larger than 1).

    Parameters
    ----------
    spectrum : `~specutils.Spectrum1D`
        The spectrum to be fitted.

    component_list : list
        Spectral model subcomponents stored in a list.
        Their ``'name'`` attribute must be unique. Each subcomponent
        should be an initialized object from `~astropy.modeling.Model`.

    expression : str
        The arithmetic expression that combines together
        the model subcomponents. The subcomponents are
        referred via their ``'name'`` attribute.

    run_fitter : bool
        **This is currently being ignored for 3D fits.**
        When `False` (the default), the function composes the compound
        model and returns it without fitting.

    window : `None` or `~specutils.SpectralRegion`
        See :func:`specutils.fitting.fit_lines`.

    n_cpu : `None` or int
        **This is only used for spectral cube fitting.**
        Number of cores to use for multiprocessing.
        Using all the cores at once is not recommended.
        If `None`, it will use max cores minus one.
        Set this to 1 for debugging.

    Returns
    -------
    output_model : `~astropy.modeling.CompoundModel` or list
        The model resulting from the fit. In the case of a 1D input
        spectrum, a single model instance is returned. In case of a
        3D spectral cube input, instead o model instances for every
        spaxel, a list with 2D arrays, each one storing fitted parameter
        values for all spaxels, is returned.

    output_spectrum : `~specutils.Spectrum1D`
        The realization of the fitted model as a spectrum. The spectrum
        will be 1D or 3D depending on the shape of input spectrum.
    """
    # Initial guess for the fit.
    initial_model = _build_model(component_list, expression)

    if len(spectrum.shape) > 1:
        return _fit_3D(initial_model, spectrum, window=window, n_cpu=n_cpu)
    else:
        return _fit_1D(initial_model, spectrum, run_fitter, window=window)


def _fit_1D(initial_model, spectrum, run_fitter, filter_non_finite=True, window=None):
    """
    Fits an astropy CompoundModel to a Spectrum1D instance.

    Parameters
    ----------
    initial_model : :class: `astropy.modeling.CompoundModel`
        Initial guess for the model to be fitted.
    spectrum : :class:`specutils.Spectrum1D`
        The spectrum to be fitted.
    run_fitter : bool
        When False (the default), the function composes the compound
        model and returns it without fitting.
    window : `None` or :class:`specutils.spectra.SpectralRegion`
        See :func:`specutils.fitting.fitmodels.fit_lines`.

    Returns
    -------
    output_model : :class: `astropy.modeling.CompoundModel`
        The model resulting from the fit.
    output_spectrum : :class:`specutils.Spectrum1D`
        The realization of the fitted model as a spectrum.

    """
    if run_fitter:
        if spectrum.uncertainty and not np.all(spectrum.uncertainty.array == 0):
            weights = 'unc'
        else:
            weights = None
        output_model = fit_lines(spectrum, initial_model, weights=weights,
                                 filter_non_finite=filter_non_finite, window=window)
        output_values = output_model(spectrum.spectral_axis)
    else:
        # Return without fitting.
        output_model = initial_model
        output_values = initial_model(spectrum.spectral_axis)

    # Build return spectrum
    output_spectrum = Spectrum1D(spectral_axis=spectrum.spectral_axis,
                                 flux=output_values,
                                 mask=spectrum.mask)

    return output_model, output_spectrum


def _fit_3D(initial_model, spectrum, window=None, n_cpu=None):
    """
    Fits an astropy CompoundModel to every spaxel in a cube
    using a multiprocessor pool running in parallel. Computes
    realizations of the models over each spaxel.

    Parameters
    ----------
    initial_model : :class: `astropy.modeling.CompoundModel`
        Initial guess for the model to be fitted.
    spectrum : :class:`specutils.Spectrum1D`
        The spectrum that stores the cube in its 'flux' attribute.
    window : `None` or :class:`specutils.spectra.SpectralRegion`
        See :func:`specutils.fitting.fitmodels.fit_lines`.
    n_cpu : `None` or int
        Number of cores to use for multiprocessing.
        Using all the cores at once is not recommended.
        If `None`, it will use max cores minus one.
        Set this to 1 for debugging.

    Returns
    -------
    output_model : :list: a list that stores 2D arrays. Each array contains one
        parameter from `astropy.modeling.CompoundModel` instances
        fitted to every spaxel in the input cube.
    output_spectrum : :class:`specutils.Spectrum1D`
        The spectrum that stores the fitted model values in its 'flux'
        attribute.
    """
    if n_cpu is None:
        n_cpu = mp.cpu_count() - 1

    # Generate list of all spaxels to be fitted
    spaxels = _generate_spaxel_list(spectrum)

    fitted_models = []

    # Build cube with empty arrays, one per input spaxel. These
    # will store the flux values corresponding to the fitted
    # model realization over each spaxel.
    output_flux_cube = np.zeros(shape=spectrum.flux.shape)

    # Callback to collect results from workers into the cubes
    def collect_result(results):
        for i in range(len(results['x'])):
            x = results['x'][i]
            y = results['y'][i]
            model = results['fitted_model'][i]
            fitted_values = results['fitted_values'][i]

            # Store fitted model parameters
            fitted_models.append({"x": x, "y": y, "model": model})

            # Store fitted values
            output_flux_cube[x, y, :] = fitted_values

    # Run multiprocessor pool to fit each spaxel and
    # compute model values on that same spaxel.
    if n_cpu > 1:
        results = []
        pool = Pool(n_cpu)

        # The communication overhead of spawning a process for each *individual*
        # parameter set is prohibitively high (it's actually faster to run things
        # sequentially). Instead, chunk the spaxel list based on the number of
        # available processors, and have each processor do the model fitting
        # on the entire subset of spaxel tuples, then return the set of results.
        for spx in np.array_split(spaxels, n_cpu):
            # Worker for the multiprocess pool.
            worker = SpaxelWorker(spectrum.flux,
                                  spectrum.spectral_axis,
                                  initial_model,
                                  param_set=spx,
                                  window=window,
                                  mask=spectrum.mask)
            r = pool.apply_async(worker, callback=collect_result)
            results.append(r)
        for r in results:
            r.wait()

        pool.close()

    # This route is only for dev debugging because it is very slow
    # but exceptions will not get swallowed up by multiprocessing.
    else:  # pragma: no cover
        worker = SpaxelWorker(spectrum.flux,
                              spectrum.spectral_axis,
                              initial_model,
                              param_set=spaxels,
                              window=window,
                              mask=spectrum.mask)
        collect_result(worker())

    # Build output 3D spectrum
    funit = spectrum.flux.unit
    output_spectrum = Spectrum1D(wcs=spectrum.wcs,
                                 flux=output_flux_cube * funit,
                                 mask=spectrum.mask)

    return fitted_models, output_spectrum


class SpaxelWorker:
    """
    A class with callable instances that perform fitting over a
    spaxel. It provides the callable for the `Pool.apply_async`
    function, and also holds everything necessary to perform the
    fit over one spaxel.

    Additionally, the callable computes the realization of the
    model just fitted, over that same spaxel. We cannot do these
    two steps (fit and compute) separately, since we cannot
    modify parameter values in an already built CompoundModel
    instance. We need to use the current model instance while
    it still exists.
    """
    def __init__(self, flux_cube, wave_array, initial_model, param_set, window=None, mask=None):
        self.cube = flux_cube
        self.wave = wave_array
        self.model = initial_model
        self.param_set = param_set
        self.window = window
        self.mask = mask

    def __call__(self):
        results = {'x': [], 'y': [], 'fitted_model': [], 'fitted_values': []}

        for parameters in self.param_set:
            x = parameters[0]
            y = parameters[1]

            # Calling the Spectrum1D constructor for every spaxel
            # turned out to be less expensive than expected. Experiments
            # show that the cost amounts to a couple percent additional
            # running time in comparison with a version that uses a 3D
            # spectrum as input. Besides, letting an externally-created
            # spectrum reference into the callable somehow prevents it
            # to execute. This behavior was seen also with other functions
            # passed to the callable.
            flux = self.cube[x, y, :]
            if self.mask is not None:
                mask = self.mask[x, y, :]
            else:
                # If no mask is provided:
                mask = np.zeros_like(flux.value).astype(bool)

            sp = Spectrum1D(spectral_axis=self.wave, flux=flux, mask=mask)

            if sp.uncertainty and not np.all(sp.uncertainty.array == 0):
                weights = 'unc'
            else:
                weights = None
            fitted_model = fit_lines(sp, self.model, window=self.window, weights=weights)

            fitted_values = fitted_model(self.wave)

            results['x'].append(x)
            results['y'].append(y)
            results['fitted_model'].append(fitted_model)
            results['fitted_values'].append(fitted_values)

        return results


def _build_model(component_list, expression):
    """
    Builds an astropy CompoundModel from a list of components
    and an expression that links them to each other.

    Parameters
    ----------
    component_list : list
        Spectral model subcomponents stored in a list.
        Their ``'name'`` attribute must be unique. Each subcomponent
        should be an initialized object from `astropy.modeling.models'
    expression : str
        The arithmetic expression that combines together
        the model subcomponents. The subcomponents are
        referred via their ``'name'`` attribute.

    Returns
    -------
    model : :class:`astropy.modeling.CompoundModel`
        The model resulting from the fit.
    """
    # The expression parser needs the subcomponents stored in a
    # dict, with keys taken from their names. This mechanism can
    # be augmented with status feedback to the UI. This requires
    # the code in here to be refactored to provide the necessary
    # hooks for the GUI (see
    # specviz/specviz/plugins/model_editor/models.py around lines
    # 200-230).

    model_dict = {}

    for component in component_list:
        model_dict[component.name] = component

    aeval = Interpreter(usersyms=model_dict)
    model = aeval(expression)

    return model


def _handle_parameter_units(model, fitted_parameters_cube, param_units):
    """
    Extracts parameter units from a CompoundModel and parameter values
    from a list of 2D numpy arrays, and returns a dict of 2D Quantity
    arrays. The dict values are keyed by the parameter names.

    Parameters
    ----------
    model : :class:`astropy.modeling.CompoundModel`
        An instance of a model.
    fitted_parameters_cube : ndarray
        A 3D array in which the 1st dimension addresses model
        parameters, and the 2nd and 3rd dimensions address
        spaxels.
    param_units : list
        A list with each parameter's units (this is not available
        in the model instance itself).

    Returns
    -------
    fitted_parameters_dict : dict
        2D Quantity arrays keyed by parameter name
    """

    fitted_parameters_dict = {}

    for index in range(len(model.parameters)):
        key = model.param_names[index]
        _ary = fitted_parameters_cube[index, :, :]
        fitted_parameters_dict[key] = u.Quantity(_ary, param_units[index])

    return fitted_parameters_dict


def _generate_spaxel_list(spectrum):
    """
    Generates a list with tuples, each one addressing the (x,y)
    coordinates of a spaxel in a 3-D spectrum cube. If a mask is available,
    skip masked indices.

    Parameters
    ----------
    spectrum : :class:`specutils.Spectrum1D`
        The spectrum that stores the cube in its ``'flux'`` attribute.

    Returns
    -------
    spaxels : list
        List with spaxels
    """
    n_x, n_y, _ = spectrum.flux.shape

    if spectrum.mask is None:
        spx = [[(x, y) for x in range(n_x)] for y in range(n_y)]
    else:
        # return only non-masked spaxels
        spx = [[(x, y) for x in range(n_x) if np.any(~spectrum.mask[x, y])]
               for y in range(n_y) if np.any(~spectrum.mask[:, y])]

    spaxels = [item for sublist in spx for item in sublist]

    return spaxels
