import numpy as np
from lmfit import minimize, Parameters
from .functions import *  # noqa: F403


class Fitter():
    """
    A class for multi-component fitting to spectroscopic data of ices.

    This is the heart of Omnifit, which receives spectra from the spectrum
    module, and is then capable of fitting an arbitrary number of different
    components to the target spectrum thus designated.

    Attributes
    ----------
    target_x : `numpy.ndarray`
        The x axis of the target spectrum, e.g. the wavelength.
    target_y : `numpy.ndarray`
        The y axis of the target spectrum, e.g. the optical depth.
    target_dy : `float`
        A single number expressing the average uncertainty of the y
        axis data.
    modelname : `string`
        A human-readable name for the model being fitted.
    psf : `Nonetype`, `numpy.ndarray`, or `astropy.convolution.Kernel`
        If set, this attribute can be used to give a kernel which should
        be used to convolve all the fitted data with.
    fitrange : `Nonetype` or `list`
        If set, this specifies the inclusive limits to which
        the fitting should be performed in x axis coordinates.
        For example a fitrange of [[200,250],[300,350]] sets
        two fitting windows of 200 to 250, and 300 to 350.
    color : `string`
        A string inidcating the desired plotting color of the target
        data, in a format understandable by matplotlib.
    funclist : `list`
        A list containing all the fittable functions. Each list entry
        is a dictionary containing the following keys and values:

            * 'name' : A human-readable name for the function being fitted,
                in string format.
            * 'color' : A string inidcating the desired plotting color of
                the data, in a format understandable by matplotlib.
            * 'type' : A string indicating what type of data the function
                consists of. It can be either 'analytical' or 'empirical',
                indicating an analytical function or empirical spectrum,
                respectively.
            * 'shape' : The shape of the function being fitted. In the case
                of an analytical function, this is a string indicating the
                callable name of the function. In the case of an empirical
                spectrum, this is the y-axis data from the spectrum.
            * 'params' : an lmfit `Parameters` instance containing the fitting
                parameters appropriate to the data being fitted.

    fitpars : `Parameters`
        This is where the fitting parameters are stored during and after
        minimization.
    fitres : `Minimizer`
        The fitting results are stored in this class, as documented in
        lmfit.
    """
    def __init__(
            self,
            x, y, dy=1.0,
            modelname='Unknown model',
            psf=None, fitrange=None, color='black'):
        """
        Fitter(
            x, y, dy=1.0,
            modelname='Unknown model',
            psf=None, fitrange=None, color='black')

        Constructor for the Fitter class. Initialisation happens by
        designating the target spectrum.

        Parameters
        ----------
        x : `numpy.ndarray`
            The x axis of the target spectrum, e.g. the wavelength.
        y : `numpy.ndarray`
            The y axis of the target spectrum, e.g. the optical depth.
        dy : `float`, optional
            A single number expressing the average uncertainty of the y
            axis data.
        modelname : `string`, optional
            A human-readable name for the model being fitted.
        psf : Nonetype or numpy.ndarray or astropy.convolution.Kernel, optional
            This attribute can be used to give a kernel which should be
            used to convolve all the fitted data with.
        fitrange : `Nonetype` or `list`, optional
            If set, this specifies the inclusive limits to which
            the fitting should be performed in x axis coordinates.
            For example a fitrange of [[200,250],[300,350]] sets
            two fitting windows of 200 to 250, and 300 to 350.
        color : `string`, optional
            A string inidcating the desired plotting color of the target
            data, in a format understandable by matplotlib.
        """
        if len(x) != len(y):
            raise RuntimeError('Input arrays have different sizes.')
        self.target_x = x
        self.target_y = y
        self.target_dy = dy
        self.modelname = modelname
        self.psf = psf
        self.fitrange = fitrange
        self.color = color
        self.funclist = []

    @classmethod
    def fromspectrum(cls, spectrum, **kwargs):
        """
        Fitter.fromspectrum(spectrum, **kwargs)

        An alternate way to initialise Fitter, by directly giving it
        a spectrum. Extracted data from the spectrum are the x, y,
        and (if the spectrum has been baselined) dy parameters.

        Parameters
        ----------
        spectrum : `omnifit.spectrum.BaseSpectrum` or its child class
            The input spectrum.
        **kwargs : Arguments, optional
            Additional initialisation arguments can be passed to `Fitter`
            using this. Note that x and y (and dy, if applicable) are defined
            using the data contained in the input spectrum.
        """
        if spectrum.baselined:
            return cls(
                spectrum.x.value,
                spectrum.y.value,
                spectrum.dy,
                **kwargs)
        else:
            return cls(
                spectrum.x.value,
                spectrum.y.value,
                **kwargs)

    def add_empirical(self, spectrum, params, funcname=None, color='red'):
        """
        add_empirical(spectrum, params, funcname=None, color='red')

        Add empirical data in the form of a spectrum to the fitting list.
        The spectrum must be interpolated to match the target x axis.

        Parameters
        ----------
        spectrum : `spectrum.BaseSpectrum`
            The input spectrum.
        params : `Parameters`
            The input parameters. Specifically this must contain
            the 'mul' parameter, which indicates what value the
            spectrum will be multiplied with during fitting.
        funcname : `Nonetype` or `string`, optional
            A human-readable name for the data being fitted.
            If this is left as None, the name of the spectrum will
            be used.
        color : `string`, optional
            A string inidcating the desired plotting color of the
            data, in a format understandable by matplotlib.
        """
        if not(funcname):
            funcname = spectrum.name
        if not np.all(spectrum.x.value == self.target_x):
            raise RuntimeError(
                'Input spectrum x axis does not ' +
                'match the target spectrum x axis.')
        self.funclist.append({
            'type': 'empirical',
            'shape': spectrum.y.value,
            'params': params,
            'name': funcname,
            'color': color})

    def add_analytical(
            self,
            shape, params,
            funcname='Unknown function',
            color='red'):
        """
        add_analytical(
            shape, params,
            funcname='Unknown function',
            color='red')

        Add analytical data in the form of a callable function to the
        fitting list.

        Parameters
        ----------
        shape : `string`
            The callable name of the function to be fitted.
        params : `Parameters`
            The input parameters. These should be formatted in a way that
            the function defined by shape can understand them, and that
            function should be created in such a way that it can make use
            of lmfit parameters.
        funcname : `string`, optional
            A human-readable name for the data being fitted.
        color : `string`, optional
            A string inidcating the desired plotting color of the
            data, in a format understandable by matplotlib.
        """
        self.funclist.append({
            'type': 'analytical',
            'shape': shape,
            'params': params,
            'name': funcname,
            'color': color})

    def perform_fit(self, **kwargs):
        """
        perform_fit(**kwargs)

        Uses `minimize` in lmfit to perform least-squares fitting of all the
        functions in the function list to the target data.

        Parameters
        ----------
        **kwargs : Arguments, optional
            This can be used to give additional arguments for `minimize`.
        """
        self.fitpars = self.__extract_pars()
        self.fitres = minimize(
            self.__fit_residual,
            self.fitpars,
            **kwargs)
        if not(self.fitres.success):
            raise RuntimeError('Fitting failed!')
        else:
            self.fitpars = self.fitres.params

    def __fit_residual(self, params, custrange=None):
        """
        __fit_residual(params, custrange=None)

        This is an internal function used for calculating the total
        residual of the data against the fittings function(s), given
        a set of lmfit parameters. The residual calculation can also
        be limited to a specific x axis range.

        Parameters
        ----------
        params : `Parameters`
            The parameters used for calculating the residual.
        custrange : `Nonetype` or `list`, optional
            If set, this specifies the inclusive range within which
            the residual is calculated. Otherwise the fitting range
            specified during Initialisation is used.

        Returns
        -------
        The residual function within the fitting range with the given
        lmfit parameters.
        """
        if custrange is None:
            fitrange = self.fitrange
        else:
            fitrange = custrange
        residual = 1.0*self.target_y
        for indFunc, cFunc in enumerate(self.funclist):
            oPar = Parameters()
            cParlist = cFunc['params']
            for cPar in cParlist.values():
                cParams = params[self.__func_ident(indFunc) + cPar.name]
                oPar.add(cPar.name,
                         value=cParams.value, vary=cParams.vary,
                         min=cParams.min, max=cParams.max,
                         expr=cParams.expr)
            residual -= self.__parse_function(oPar, cFunc)
        # Crop out not-numbers and fitting range exterior if necessary
        if np.any(fitrange):
            fitInd = np.isinf(residual)
            for cRange in fitrange:
                fitInd = np.logical_or(fitInd, np.logical_and(
                             np.less_equal(cRange[0], self.target_x),
                             np.greater_equal(cRange[1], self.target_x)))
        else:
            fitInd = np.isfinite(residual)
        return residual[fitInd]

    def chisq(self, checkrange=None):
        """
        chisq(checkrange=None)

        Return chi squared of fit, either in a custom range
        or in the range used by the fit.

        Parameters
        ----------
        checkrange : `Nonetype` or `list`, optional
            If set, this specifies the inclusive range within which
            the chi squared value is calculated. Otherwise the fitting
            range specified during Initialisation is used.

        Returns
        -------
        The chi squared within the desired ranged.
        """
        residual = self.__fit_residual(self.fitpars, custrange=checkrange)
        return np.sum((residual**2.0)/(self.target_dy**2.0))

    def plot_fitresults(
            self,
            ax,
            lw=[1, 2, 3], color_total='blue', legend=True,
            **kwargs):
        """
        plot_fitresults(
            ax,
            lw=[1,2,3], color_total='blue', legend=True,
            **kwargs)

        Plot the fitting results to the given matplotlib axis, with a
        number of optional parameters specifying how the different plottable
        components are presented.

        Parameters
        ----------
        axis : `matplotlib.axis`
            The axis which the plot will be generated in.
        lw : `list`, optional
            This list of 3 numbers specifies the line widths of the target
            spectrum, the fitted functions, and the total fit, respectively.
        color_total : `string`, optional
            A string inidcating the desired plotting color of the total sum
            of the fit results, in a format understandable by matplotlib.
            The colors of the target spectrum and the fitted functions are
            specified during their initialisation and addition.
        legend : `bool`, optional
            If set to True, a legend is automatically created using the
            target spectrum and fitted function names.
        **kwargs : Arguments, optional
            This can be used to pass additional arguments
            to `matplotlib.pyplot.plot`, which is used by this
            method for its plotting.
        """
        ax.plot(
            self.target_x, self.target_y,
            color=self.color, lw=lw[0],
            **kwargs)
        legList = [self.modelname]
        totRes = np.zeros(len(self.target_y))
        for indFunc, cFunc in enumerate(self.funclist):
            oPar = Parameters()
            cParList = cFunc['params']
            cCol = cFunc['color']
            for cPar in cParList.values():
                cFitPar = self.fitpars[self.__func_ident(indFunc)+cPar.name]
                oPar.add(cPar.name,
                         value=cFitPar.value, vary=cFitPar.vary,
                         min=cFitPar.min, max=cFitPar.max,
                         expr=cFitPar.expr)
            funcRes = self.__parse_function(oPar, cFunc)
            totRes += funcRes
            ax.plot(self.target_x, funcRes, lw=lw[1], color=cCol, **kwargs)
            legList.append(cFunc['name'])
        legList.append('Total fit')
        ax.plot(self.target_x, totRes, lw=lw[2], color=color_total, **kwargs)
        if legend:
            ax.legend(legList, shadow=True)

    def fitresults_tofile(self, filename, detection_threshold=5.0):
        """
        fitresults_tofile(filename)

        Export fit results to two output files which are intended to be
        easily readable and paraseable with other software.

        The first file is filename.csv, which contains x and y data of
        the fitted models, as would be visualized in a plotted fit result.
        The first column of the csv is the x value, which is shared by all
        models.
        The second column is the y value of data that was being fitted to.
        The third column is total sum of fitted models.
        The fourth to Nth columns are the individual models, in the order
        described in the second file, filename.xml.

        The second file, filename.xml is an XML file containing additional
        information about the fitted data and the fit results which are not
        easily representable in a csv-formatted file. This data is
        formatted using the following XML elements:

            * INFO : Contains all the other elements described below, and has
                the attribute "file", which is the name of the csv file pair of
                this xml file.
            * MODELNAME : Contains the name of the model.
            * HAVEPSF : A boolean value indicating whether there is a PSF
                associated with the model.
            * RMS_DATA : The uncertainty of the data.
            * NUMBER_FUNCTIONS : An integer indicating how many functions
                have been fitted to the total data.

        In addition to the above elements, each fitted function has its own
        element, designated FUNCTION, having the attribute "name" which is
        the name of the function. FUNCTION contains the following elements:

            * TYPE : If the function is an empirical one, this contains the
                string "empirical". Otherwise it contains the name of the
                called analytical function.
            * DETECTION : When generating the contents of this element,
                The method is_nondet with the detection threshold designated
                by the parameter detection_threshold. The result given by
                the method is indicated here with a "True" or "False"
                depending on whether the result is considered a detection.
            * CSV_COLUMN : Indicates which column in the CSV contains the
                fitted data for this function.
            * NUMBER_PARAMS : Inidicates how many parameters are used by
                this function i.e. the number of PARAMETER elements.

        Finally, contained within each FUNCTION element is a number of
        PARAMETER elements, which list the best-fit data for each fitted
        parameter pertaining to that function. Each PARAMETER element
        contains the attribute "name", which tells the name of the
        parameter. In addition the following elements are contained by
        each PARAMETER element:

            * VALUE : The best-fit value for this parameter.

        Parameters
        ----------
        filename : `string`
            The extensionless version of the desired filename which the
            data should be exported to. As a result the files
            "filename.csv" and "filename.xml" are created.
        detection_threshold : `float`, optional
            The threshold of detection to be used in determining whether
            the value contained by the DETECTION element is true or not.
        """
        filename_csv = filename+'.csv'
        filename_xml = filename+'.xml'
        file_xml = open(filename_xml, 'w')
        file_xml.write(
            '<!-- Automatically generated information file for csv file ' +
            filename_csv +
            '-->\n')
        file_xml.write('<INFO file="'+filename_csv+'">\n')
        file_xml.write('<MODELNAME>'+self.modelname+'</MODELNAME>\n')
        file_xml.write('<HAVEPSF>'+str(self.psf is not None)+'</HAVEPSF>\n')
        file_xml.write('<RMS_DATA>'+str(self.target_dy)+'</RMS_DATA>\n')
        file_xml.write(
            '<NUMBER_FUNCTIONS>' +
            str(len(self.funclist)) +
            '</NUMBER_FUNCTIONS>\n')
        outdata_csv = np.vstack([self.target_x, self.target_y])
        outdata_functions = np.empty([0, len(self.target_x)])
        totRes = np.zeros(len(self.target_x))
        for indFunc, cFunc in enumerate(self.funclist):
            file_xml.write('<FUNCTION name="'+cFunc['name']+'">\n')
            file_xml.write('<TYPE>')
            if cFunc['type'] == 'analytical':
                file_xml.write(cFunc['shape'])
            elif cFunc['type'] == 'empirical':
                file_xml.write('empirical')
            else:
                file_xml.write('unknown'+'\n')
            file_xml.write('</TYPE>\n')
            file_xml.write(
                '<DETECTION>' +
                str(not self.is_nondet(
                    sigma=detection_threshold
                    )[cFunc['name']]) +
                '</DETECTION>\n')
            file_xml.write('<CSV_COLUMN>'+str(indFunc+3)+'</CSV_COLUMN>\n')
            cParlist = cFunc['params']
            file_xml.write(
                '<NUMBER_PARAMS>' +
                str(len(cParlist)) +
                '</NUMBER_PARAMS>\n')
            oPar = Parameters()
            for cPar in cParlist.values():
                file_xml.write('<PARAMETER name="'+cPar.name+'">\n')
                cFitPar = self.fitpars[self.__func_ident(indFunc)+cPar.name]
                oPar.add(cPar.name,
                         value=cFitPar.value, vary=cFitPar.vary,
                         min=cFitPar.min, max=cFitPar.max,
                         expr=cFitPar.expr)
                file_xml.write('<VALUE>'+str(cFitPar.value)+'</VALUE>\n')
                file_xml.write('</PARAMETER>\n')
            funcRes = self.__parse_function(oPar, cFunc)
            outdata_functions = np.vstack([outdata_functions, funcRes])
            totRes += funcRes
            file_xml.write('</FUNCTION>\n')
        file_xml.write('</INFO>')
        file_xml.close()
        outdata_csv = np.vstack([outdata_csv, totRes, outdata_functions])
        np.savetxt(
            filename_csv,
            outdata_csv.transpose(),
            delimiter=',', header='For info, see '+filename_xml)

    def is_nondet(self, sigma=5.0):
        """
        is_nondet(sigma=5.0)

        Determines whether the fitted functions in the function list can
        be considered detections or non-detections using the given detection
        thereshold. This is done by comparing the peak of the fitted function
        within the fitting range to a multiple (set by the parameter sigma)
        of the RMS noise in the target data.
        It should be emphasized that unless the dy attribute has been set
        during the fitter class initialisation, the results returned by this
        method are meaningless.

        Parameters
        ----------
        sigma : `float`, optional
            The multiplier that should be applied to the noise when comparing
            it against the fitted function peaks.

        Returns
        -------
        A dictionary containing boolean values for each function (with
        their names as the keys) and the total fit (key 'total'), with
        True indicating that the function is considered a non-detection
        using the criteria outlined above.
        """
        minY = sigma*self.target_dy
        out = {}
        totRes = np.zeros(len(self.target_x))
        for indFunc, cFunc in enumerate(self.funclist):
            cParlist = cFunc['params']
            oPar = Parameters()
            for cPar in cParlist.values():
                cFitPar = self.fitpars[self.__func_ident(indFunc)+cPar.name]
                oPar.add(cPar.name,
                         value=cFitPar.value, vary=cFitPar.vary,
                         min=cFitPar.min, max=cFitPar.max,
                         expr=cFitPar.expr)
            funcRes = self.__parse_function(oPar, cFunc)
            if np.max(funcRes) < minY:
                out[cFunc['name']] = True
            else:
                out[cFunc['name']] = False
            totRes += funcRes
        if np.max(totRes) < minY:
            out['total'] = True
        else:
            out['total'] = False
        return out

    def fit_results(self):
        """
        fit_results()

        Return the fitting results as a dictionary.

        Parameters
        ----------
        None

        Returns
        -------
        A dictionary containing all the individual functions which were
        fitted. The key-value combinations of this dictionary consist of
        the function name, and its lmfit Parameters instance, which
        contains the best-fit results.
        """
        oResults = {}
        for indFunc, cFunc in enumerate(self.funclist):
            oKeyname_base = cFunc['name']
            oKeyind = 0
            oKeyname = oKeyname_base
            while oResults.__contains__(oKeyname):  # In case of duplicates
                oKeyind += 1
                oKeyname = oKeyname_base+'(duplicate '+str(oKeyind)+')'
            oResults[cFunc['name']] = self.__fit_result(indFunc)
        return oResults

    def __fit_result(self, index):
        """
        __fit_result(index)

        Return fitting results for a specific function in the internal
        function list.

        Parameters
        ----------
        index : `int`
            Desired index of the function to fetch from the function lsit.

        Returns
        -------
        An `Parameters` instance containing the fitting
        results for the desired function.
        """
        oParlist = self.funclist[index]['params']
        for cParname in oParlist.keys():
            coPar = self.fitpars[self.__func_ident(index)+cParname]
            coPar.name = cParname
            oParlist[cParname] = coPar
        return oParlist

    def __parse_function(self, params, function):
        """
        __parse_function(params,function)

        Parse the input function, insert parameters, return result.

        Parameters
        ----------
        params : `Parameters`
            The lmfit `Parameters` instance to use as input parameters.
        function : `dict`
            A dictionary formatted in the style that the entries inside
            funclist are formatted

        Returns
        -------
        The result of the given function with given parameters.
        """
        if function['type'] == 'empirical':
            funcres = muldata(function['shape'], params['mul'].value)  # noqa
        elif function['type'] == 'analytical':
            funcres = globals()[function['shape']](
                self.target_x,
                params, self.psf)
        else:
            raise RuntimeError('Unknown function type!')
        return funcres

    def __extract_pars(self):
        """
        __extract_pars()

        Extracts the paramers from the function list and converts them to
        a single lmfit Parameters instance, which can then be manipulated
        by the residual minimization routines.

        Parameters
        ----------
        None

        Returns
        -------
        An lmfit `Parameters` instance containing the parameters
        of *all* the fittable functions in a single place.
        """
        oPars = Parameters()
        for indFunc, cFunc in enumerate(self.funclist):
            cParlist = cFunc['params']
            for cPar in cParlist.values():
                oPars.add(self.__func_ident(indFunc)+cPar.name,
                          value=cPar.value, vary=cPar.vary,
                          min=cPar.min, max=cPar.max,
                          expr=cPar.expr)
        return oPars

    def __func_ident(self, index):
        """
        __func_ident(index)

        Generate a unique prefix string for a function, which can be
        used by `__extract_pars` to generate its master Parameters list.

        Parameters
        ----------
        index : `int`
            The index of the function.

        Returns
        -------
        A unique identifier string pertaining to that function, which
        can be used to generate unique parameter names.
        """
        return '__Func'+str(index)+'__'
