"""
:class:`~sidpy.proc.fitter.SidFitter` class that fits the specified dimension of a sidpy.dataset using the
user-specified fit function. An extension of scipy.optimise.curve_fit that works on sidpy.dataset

Created on Mar 9, 2022
@author: Rama Vasudevan, Mani Valleti
"""

from xml.dom import NotFoundErr
from dask.distributed import Client
import numpy as np
import dask
import inspect
from ..sid import Dimension, Dataset
from ..sid.dimension import DimensionType
from ..viz.dataset_viz import SpectralImageFitVisualizer
from ..sid.dataset import DataType

try:
    from scipy.optimize import curve_fit
except ImportError:
    curve_fit = None

try:
    from sklearn.cluster import KMeans
except ModuleNotFoundError:
    KMeans = None


class SidFitter:
    # An extension of the Process Class for Functional Fitting
    def __init__(self, sidpy_dataset, fit_fn, xvec=None, ind_dims=None, guess_fn=None, num_fit_parms=None,
                 km_guess=False, n_clus=None, return_cov=False, return_std=False, return_fit=False,
                 fit_parameter_labels=None, num_workers=2, threads=2):
        """
        Parameters
        ----------
        sidpy_dataset: (sidpy.Dataset) Sidpy dataset object to be fit

        fit_fn: (function) Function used for fitting.
        Should take xvec as the first argument and parameters as the rest of the arguments.
        Should return the function value at each of the points in the xvec

        xvec: (numpy ndarray or list of numpy ndarrays) (Optional)
        Independent variable for fitting. Should be an array
        If NOT provided, the dimension arrays are assumed to be xvecs

        ind_dims: (tuple) (Optional) Tuple with integer entries of the dimensions
            over which to parallelize. These should be the independent variable for the fitting.
            If NOT provided, it is assumed that all the non-spectral dimensions are independent dimensions.

        guess_fn: (function) (optional) This optional function should be utilized to generate priors for the full fit
        It takes (xvec,yvec) as inputs and should return the fit parameters.
        If the guess_fn is NOT provided, then the user MUST input the num_fit_parms.

        num_fit_parms: (int) Number of fitting parameters. This is needed IF the guess function is not provided to set
        the priors for the parameters for the curve_fit function.

        km_guess: (bool) (default False) When set to True: Divides the spectra into clusters using
        sklearn.optimize.kMeans, applies the fitting function on the cluster centers,
        uses the results as priors to each spectrum of the cluster.

        n_clus: (int) (default None) Used only when km_guess is set to True. Determines the number of clusters to be
        formed for sklearn.optimize.kmeans. If not provided then n_clus = self.num_computations/100

        return_std: (bool) (default False) Returns the dataset with estimated standard deviation of the parameter
        values. Square roots of the diagonal of the covariance matrix.

        return_cov: (bool) (default False) Returns the estimated covariance of fitting parameters. Confer
        scipy.optimize.curve_fit for further details

        return_fit: (bool) (default False) Returns the fitted sidpy dataset using the optimal parameters
         when set to true

        fit_parameter_labels: (list) (default None) List of parameter labels

        num_workers: (int) (default =2) Number of workers to use when setting up Dask client

        threads: (int) (default =2) Number of threads to use when setting up Dask client

        Returns:
        -------
        sidpy.dataset: if return_cov and return_fit are both set to False
        List: containing sidpy.dataset objects, if either of return_cov or return fit is set to True

        If multiple datasets are expected, the order of the returned datasets is

        [sidpy.dataset with mean parameter values,
        sidpy.dataset with estimated covariances of the fitting parameters,
        sidpy.dataset that is fit with the parameters obtained after fitting]

        """

        if guess_fn is None:
            if num_fit_parms is None:
                raise ValueError("You did not supply a guess function, you must at least provide number of fit "
                                 "parameters to set the priors for scipy.optimize.curve_fit")
        self.dataset = sidpy_dataset  # Sidpy dataset
        self.fit_fn = fit_fn  # function that takes xvec, *parameters and returns yvec at each value of xvec
        self.num_fit_parms = num_fit_parms  # int: number of fitting parameters
        self._complex_data = False  # if data is complex. Will be checked during guess/fit as needed.

        if ind_dims is not None:
            self.ind_dims = tuple(ind_dims)  # Tuple: containing indices of independent dimensions
        else:
            # All the dimensions that are not spectral will be considered as independent dimensions
            ind_dims = []
            for i, dim in self.dataset._axes.items():
                if dim.dimension_type != DimensionType.SPECTRAL:
                    ind_dims.extend([i])
            self.ind_dims = tuple(ind_dims)

        # Make sure there is at least one spectral dimension
        if len(self.ind_dims) == len(self.dataset.shape):
            raise NotImplementedError('No Spectral (dependent) dimensions found to fit')

        # Let's get the dependent dims here
        dep_dims = []  # Tuple: contains all the dependent dimensions. ind_dims+dep_dims = all_dims
        for d in np.arange(len(self.dataset.shape)):
            if d not in self.ind_dims:
                dep_dims.extend([d])
        self.dep_dims = tuple(dep_dims)

        # xvec is not provided
        if xvec is None:
            # 1D fit
            if len(self.dep_dims) == 1:
                dep_vec = np.array(self.dataset._axes[self.dep_dims[0]])
            # Multidimensional fit
            else:
                dep_vec = []
                for d in self.dep_dims:
                    dep_vec.append(np.array(self.dataset._axes[d]))

        # xvec is provided
        if xvec is not None:
            # 1D fit
            if len(self.dep_dims) == 1:
                if isinstance(xvec, np.ndarray):
                    dep_vec = xvec
                elif isinstance(xvec, list):
                    dep_vec = np.array(xvec)
                else:
                    raise TypeError('Please provide a np.ndarray or a list of independent vector values')
            # Multidimensional fit
            else:
                if isinstance(xvec, list) and len(xvec) == len(self.dep_dims):
                    dep_vec = xvec
                elif isinstance(xvec, list) and len(xvec) != len(self.dep_dims):
                    raise ValueError('The number of independent dimensions provided in the xvec do not match '
                                     'with the number of dependent dimensions of the dataset')
                else:
                    raise TypeError('Please provide a list of value-arrays corresponding to each dependent dimension')

        # Dealing with the meshgrid part of multidimensional fitting
        if len(self.dep_dims) > 1:
            self.dep_vec = [ar.ravel() for ar in np.meshgrid(*dep_vec, indexing='ij')]
        else:
            self.dep_vec = dep_vec

        self.km_guess = km_guess
        if self.km_guess:
            self.km_priors = None
            self.km_labels = None
            self.n_clus = n_clus

        self._setup_calc()
        self.guess_fn = guess_fn
        self.prior = None  # shape = [num_computations, num_fitting_parms]
        self.fit_labels = fit_parameter_labels
        self.num_workers = num_workers
        self.threads = threads
        self.guess_completed = False
        self.return_std = return_std
        self.return_cov = return_cov
        self.return_fit = return_fit
        self.fitted_dset = None


        self.mean_fit_results = []
        if self.return_cov:
            self.cov_fit_results = None
        if self.return_std:
            self.std_fit_results = None

        if 'complex' in self.dataset.dtype.name:
            self._complex_data = True
        # set up dask client
        self.client = Client(threads_per_worker=self.threads, n_workers=self.num_workers)

    def _setup_calc(self):
        self.fold_order = [[]]  # All the independent dimensions go into the first element and will be collapsed
        self.num_computations = 1

        # Here we have to come up with a way that treats the spatial dimensions as the independent dimensions
        # In other words make the argument 'ind_dims' optional
        # if self.ind_dims is not None:

        for i in np.arange(self.dataset.ndim):
            if i in self.ind_dims:
                self.fold_order[0].extend([i])
                self.num_computations *= self.dataset.shape[i]
            else:
                self.fold_order.append([i])

        self.folded_dataset = self.dataset.fold(dim_order=self.fold_order)
        self.folded_dataset_numpy = np.array(self.folded_dataset)
        self.dep_vec = np.array(self.dep_vec)

        # Here is the tricky part, dataset.unfold is designed to get back the original dataset with minimal loss of
        # information. To do this, unfold utilizes the saved information while folding the original dataset.
        # Here, we are going to tweak that information and use the unfold method on the dataset with fitted parameters.

        self._unfold_attr = {
            'dim_order_flattened': list(np.arange(len(self.fold_order[0]))) + [len(self.fold_order[0])],
            'shape_transposed': [self.dataset.shape[i] for i in self.fold_order[0]] + [-1]}
        axes, j = {}, 0
        for i, dim in self.dataset._axes.items():
            if not i in self.dep_dims:
                axes[j] = dim
                j += 1
        self._unfold_attr['_axes'] = axes

    def do_guess(self):
        """
        If a guess_fn is provided: Applies the guess_fn to get priors for the fitting parameters.
        self.prior is set as the output of guess function at each of the ind_dims

        Returns:
        None
        -------

        """
        guess_results = []
        for ind in range(self.num_computations):
            ydata = self.folded_dataset_numpy
            lazy_result = dask.delayed(self.guess_fn)(self.dep_vec, ydata[ind, :])
            guess_results.append(lazy_result)

        guess_results = dask.compute(*guess_results)
        self.prior = np.squeeze(np.array(guess_results))
        self.num_fit_parms = self.prior.shape[-1]
        self.guess_completed = True

    def do_fit(self, **kwargs):
        """
        Perform the fit.
        **kwargs: extra parameters passed to scipy.optimize.curve_fit, e.g. bounds, type of lsq algorithm, etc.
        """
        

        if self.guess_fn is not None:
            guess_function_str = inspect.getsource(self.guess_fn)
        else:
            guess_function_str = 'Not Provided'

        fit_results = []
        if not self.km_guess:
            if not self.guess_completed and self.guess_fn is not None:
                self.do_guess()

            for ind in range(self.num_computations):
                if self.prior is None:
                    p0 = np.random.normal(loc=0.5, scale=0.1, size=self.num_fit_parms)
                else:
                    p0 = self.prior[ind, :]
                ydata = self.folded_dataset_numpy[ind, :]
                if self._complex_data:
                    ydata = np.array(np.hstack([np.real(ydata), np.imag(ydata)]))

                lazy_result = dask.delayed(SidFitter.default_curve_fit)(self.fit_fn, self.dep_vec,
                                                                        ydata, self.num_fit_parms,
                                                                        return_cov=(self.return_cov or self.return_std),
                                                                        p0=p0, **kwargs)
                fit_results.append(lazy_result)

            fit_results_comp = dask.compute(*fit_results)
            self.client.close()

        else:
            self.get_km_priors(**kwargs)
            for ind in range(self.num_computations):
                ydata = self.folded_dataset_numpy[ind, :]
                if self._complex_data:
                    #ydata = ydata.flatten_complex()
                    ydata = np.array(np.hstack([np.real(ydata), np.imag(ydata)]))

                lazy_result = dask.delayed(SidFitter.default_curve_fit)(self.fit_fn, self.dep_vec,
                                                                        ydata, self.num_fit_parms,
                                                                        return_cov=(self.return_cov or self.return_std),
                                                                        p0=self.km_priors[self.km_labels[ind]],
                                                                        **kwargs)
                fit_results.append(lazy_result)

            fit_results_comp = dask.compute(*fit_results)
            self.client.close()

        if self.return_cov or self.return_std:
            # here we get back both: the parameter means and the covariance matrix!
            self.mean_fit_results = np.squeeze(
                np.array([fit_results_comp[ind][0] for ind in range(len(fit_results_comp))]))
            self.cov_fit_results = np.squeeze(
                np.array([fit_results_comp[ind][1] for ind in range(len(fit_results_comp))]))

        else:
            # in this case we can just dump it to an array because we only got the parameters back
            self.mean_fit_results = np.squeeze(np.array(fit_results_comp))

        # Here we have either the mean fit results or both mean and cov arrays. We make 2 sidpy dataset out of them
        # Make a sidpy dataset

        mean_sid_dset = Dataset.from_array(self.mean_fit_results, title='Fitting_Map')
        mean_sid_dset.metadata['fold_attr'] = self._unfold_attr.copy()
        mean_sid_dset = mean_sid_dset.unfold()

        # Set the data type
        mean_sid_dset.data_type = 'image_stack'  # We may want to pass a new type - fit map

        # We set the last dimension, i.e., the dimension with the fit parameters
        fit_dim = Dimension(np.arange(self.num_fit_parms),
                            name='fit_parms', units='a.u.',
                            quantity='fit_parameters',
                            dimension_type='temporal')
        mean_sid_dset.set_dimension(len(mean_sid_dset.shape) - 1, fit_dim)

        fit_parms_dict = {'fit_parameters_labels': self.fit_labels,
                          'fitting_function': inspect.getsource(self.fit_fn),
                          'guess_function': guess_function_str,
                          'ind_dims': self.ind_dims
                          }
        mean_sid_dset.metadata = self.dataset.metadata.copy()
        mean_sid_dset.metadata['fit_parms_dict'] = fit_parms_dict.copy()
        mean_sid_dset.original_metadata = self.dataset.original_metadata.copy()

        cov_sid_dset, std_fit_dset, fit_dset = None, None, None

        # Here we deal with the covariance dataset
        if self.return_cov:
            # Make a sidpy dataset
            cov_sid_dset = Dataset.from_array(self.cov_fit_results, title='Fitting_Map_Covariance')
            fold_attr = self._unfold_attr.copy()
            fold_attr['dim_order_flattened'] = fold_attr['dim_order_flattened'] + [
                len(fold_attr['dim_order_flattened'])]
            fold_attr['shape_transposed'] = fold_attr['shape_transposed'][:-1] + [self.num_fit_parms] + \
                                            [self.num_fit_parms]

            cov_sid_dset.metadata['fold_attr'] = fold_attr
            cov_sid_dset = cov_sid_dset.unfold()

            # Set the data type
            cov_sid_dset.data_type = 'IMAGE_4D'  # We may want to pass a new type - fit map

            cov_dims = [Dimension(np.arange(self.num_fit_parms),
                                  name='fit_cov_parms_x', units='a.u.',
                                  quantity='fit_cov_parameters',
                                  dimension_type='spectral'),
                        Dimension(np.arange(self.num_fit_parms),
                                  name='fit_cov_parms_y', units='a.u.',
                                  quantity='fit_cov_parameters',
                                  dimension_type='spectral')]

            for i, dim in enumerate(cov_dims):
                cov_sid_dset.set_dimension(i - 2 + len(cov_sid_dset.shape), dim)

            cov_sid_dset.metadata = self.dataset.metadata.copy()
            cov_sid_dset.metadata['fit_parms_dict'] = fit_parms_dict.copy()
            cov_sid_dset.original_metadata = self.dataset.original_metadata.copy()

        # Here is the std_dev dataset
        if self.return_std:
            self.std_fit_results = np.diagonal(self.cov_fit_results, axis1=-2, axis2=-1)
            std_fit_dset = Dataset.from_array(self.std_fit_results, title='Fitting_Map_std_dev')
            std_fit_dset.metadata['fold_attr'] = self._unfold_attr.copy()
            std_fit_dset = std_fit_dset.unfold()

            # Set the data type
            std_fit_dset.data_type = 'image_stack'  # We may want to pass a new type - fit map

            # We set the last dimension, i.e., the dimension with the fit parameters
            fit_dim = Dimension(np.arange(self.num_fit_parms),
                                name='std_dev', units='a.u.',
                                quantity='std_dev_fit_parms',
                                dimension_type='temporal')
            std_fit_dset.set_dimension(len(std_fit_dset.shape) - 1, fit_dim)

            std_fit_dset.metadata = self.dataset.metadata.copy()
            std_fit_dset.metadata['fit_parms_dict'] = fit_parms_dict.copy()
            std_fit_dset.original_metadata = self.dataset.original_metadata.copy()

        # Fitted dset
        if self.return_fit:
            fit_dset = self.get_fitted_dataset()
            fit_dset.metadata['fit_parms_dict'] = fit_parms_dict.copy()

        results = [mean_sid_dset, cov_sid_dset, std_fit_dset, fit_dset]
        inds = [True, self.return_cov, self.return_std, self.return_fit]
        results = [results[i] for i in range(len(inds)) if inds[i]]

        if len(results) == 0:
            return results[0]
        else:
            return results

    def get_fitted_dataset(self):
        """This method returns the fitted dataset using the parameters generated by the fit function"""
        fitted_dset = self.dataset.like_data(np.zeros_like(self.dataset.compute()),
                                             title_prefix='fitted_')

        fitted_dset_fold = fitted_dset.fold(dim_order=self.fold_order)
        output_shape = np.prod(fitted_dset_fold.shape[1:])
        user_folding = False
        ydata_fit = self.fit_fn(self.dep_vec, *self.mean_fit_results[0])

        # print(r"ydata shape is {} and squeezed is {}".format(ydata_fit.shape, ydata_fit.squeeze().shape))
        if ydata_fit.squeeze().shape[0] != output_shape:
            print('Shapes of output of fitting function is {} and original data is {} \
                  Reshaping output dataset. You are responsible for reshaping'.format(ydata_fit.shape[0],
                                                                                      output_shape,
                                                                                      ))

            fitted_dset_fold = self.dataset.like_data(np.zeros((fitted_dset_fold.shape[0], ydata_fit.shape[0])),
                                                      title_prefix='fitted_')
            user_folding = True
        # Here we make a roundtrip to numpy as earlier versions of dask did not support the assignments
        # of the form dask_array[2] = 1

        np_folded_arr = fitted_dset_fold.compute()
        for i in range(np_folded_arr.shape[0]):
            # ydata_fit = self.fit_fn(self.dep_vec, *self.mean_fit_results[i])
            # print('dep vec is {} and mean fit results are {}'.format(self.dep_vec,self.mean_fit_results[i]))
            fit_output = self.fit_fn(self.dep_vec, *self.mean_fit_results[i])
            # print('ydata output from fitting fn is {}'.format(fit_output))
            if fit_output.shape != np_folded_arr[i].shape:
                try:
                    np_folded_arr[i] = fit_output.reshape(np_folded_arr[i].shape)
                except:
                    print("Cannot reshape function output to retrieve fitted dataset")
            else:
                np_folded_arr[i] = fit_output

        if not user_folding:
            fitted_sid_dset_folded = fitted_dset_fold.like_data(np_folded_arr, title=fitted_dset_fold.title)
            fitted_sid_dset = fitted_sid_dset_folded.unfold()
            fitted_sid_dset.original_metadata = self.dataset.original_metadata.copy()
        else:
            fitted_sid_dset = fitted_dset_fold.like_data(np_folded_arr, title=fitted_dset_fold.title)
            fitted_sid_dset.original_metadata = self.dataset.original_metadata.copy()
        self.fitted_dset = fitted_sid_dset
        return fitted_sid_dset

    def get_km_priors(self, **kwargs):
        kwargs['maxfev'] = 100  # give a large number of tries for fitting the kmeans cluster centers

        shape = self.folded_dataset.shape  # We get the shape of the folded dataset
        # Our prior_dset will have the same shape except for the last dimension whose size will be equal to number of
        # fitting parameters
        dim_order = [[0], [i + 1 for i in range(len(shape) - 1)]]
        # We are using the fold function in case we have a multidimensional fit.
        # In that case we need all the spectral dimensions collapsed into a single dimension for kMeans
        # In case of a 1D fit the next line essentially does nothing.
        km_dset = self.folded_dataset.fold(dim_order)

        if self._complex_data:
            print('Warning: complex dataset detected. For Kmeans priors, we will treat real part only')
            km_dset = km_dset.real

        if KMeans is None:
            raise ModuleNotFoundError("sklearn is not installed")
        else:
            if self.n_clus is None:
                self.n_clus = int(self.num_computations / 100)
            km = KMeans(n_clusters=self.n_clus, random_state=0).fit(km_dset.compute())

        self.km_labels, self.km_centers = km.labels_, km.cluster_centers_
    
        if self._complex_data:
            km_dset = np.array(self.folded_dataset.fold(dim_order))
            self.km_centers = []
            # in the case of complex data, the centers have to be recomputed based on the labels
            for ind_l in range(self.n_clus):
                cent = km_dset[self.km_labels == ind_l, :]
                centroid = cent.real.mean(axis=0) + 1j*cent.imag.mean(axis=0)
                self.km_centers.append(centroid)
            self.km_centers = np.array(self.km_centers)
        print('---Finished KMeans, onto fiting each KM Center---')
        km_priors = []
        for i, cen in enumerate(self.km_centers):
            print('Fitting center {}'.format(i))
            num_start = 100 #number of times to restart the fit. For now this is fixed.

            if self.guess_fn is not None:
                p0 = self.guess_fn(self.dep_vec, cen)
            else:
                p0 = np.random.normal(loc=0.5, scale=0.1, size=self.num_fit_parms)
            if self._complex_data:
                cen = np.hstack([np.real(cen), np.imag(cen)])
            
            residuals = []
            for _ in range(num_start):
                
                popt = SidFitter.default_curve_fit(self.fit_fn, self.dep_vec, cen, self.num_fit_parms,
                                        return_cov=False,  p0 = p0,  **kwargs)
                temp_fit = self.fit_fn(self.dep_vec, *popt)
                #temp_fit = temp_fit[:len(temp_fit)//2] + 1j* temp_fit[len(temp_fit)//2 :]
                #temp_fit = np.hstack([np.real(cen), np.imag(cen)])
                #print(cen, temp_fit, cen.shape, temp_fit.shape)
                resid = cen - temp_fit
                resid_ss = np.sum(np.abs(resid@resid))
                residuals.append((popt, resid_ss))
                
            residuals = np.array(residuals, dtype = object)
            self.residuals = residuals
            min_idx = np.argmin(residuals[:,1])
            best_popt = residuals[min_idx,0]
            km_priors.append(best_popt)

        self.km_priors = np.array(km_priors)
        self.num_fit_parms = self.km_priors.shape[-1]

    def visualize_fit_results(self, figure=None, horizontal=True):
        '''
        Calls the interactive visualizer for comparing raw and fit datasets.

        Inputs:
            - figure: (Optional, default None) - handle to existing figure
            - horiziontal: (Optional, default True) - whether spectrum should be plotted horizontally

        '''
        dset_type = self.dataset.data_type
        supported_types = ['SPECTRAL_IMAGE']
        if self.fitted_dset == None:
            raise NotFoundErr("No fitted dataset found. Re-run with return_fit=True to use this feature")
        if dset_type == DataType.SPECTRAL_IMAGE:
            visualizer = SpectralImageFitVisualizer(self.dataset, self.fitted_dset,
                                                    figure=figure, horizontal=horizontal)
        else:
            raise NotImplementedError(
                "Data type is {} but currently we only support types {}".format(dset_type, supported_types))

        return visualizer

    @staticmethod
    def default_curve_fit(fit_fn, xvec, yvec, num_fit_parms, return_cov=True, **kwargs):

        yvec = np.array(yvec).ravel()
        if curve_fit is None:
            raise ModuleNotFoundError("scipy is not installed")
        else:
            try:
                popt, pcov = curve_fit(fit_fn, xvec, yvec, **kwargs)
            except:
                popt = np.zeros(num_fit_parms)
                pcov = np.zeros((num_fit_parms, num_fit_parms))
        if return_cov:
            return popt, pcov
        else:
            return popt