import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import numpy as np
import pandas as pd
import warnings
from decimal import Decimal
import datetime
import logging
import itertools

from side_functions import *
from MSEval import *


class MSData():
    def __init__(self, datareader, logger=None):
        self.datareader = datareader
        self.logger = logger
        if self.logger is not None:
            self.logger.info(f'Reading data {self.datareader.filename}.')
        self.data = self.datareader()
        self.data.index = list(map(float, self.data.index))
        self.time = np.array(self.data.index)
        self.matrix = self.data.values
        self.isotope_names = np.array(self.data.columns)
        self.isotopes = {}
        self.param = None
        self.selector = None
        self.starts, self.ends = None, None
        self.laser_off, self.laser_on = None, None
        self.names = None
        self.srms = None
        self.average_peaks = None
        self.quantified = None
        self.lod = None
        self.corrected_IS = None
        self.corrected_TS = None

        self.create_isotopes()

    def __call__(self, isotopes=None, *args, **kwargs):
        plot_data(self.data, isotopes=isotopes, *args, **kwargs)

    def __repr__(self):
        res = f'{self.__class__.__name__}\n Analysis time: {self.time.max()}s\n File path: {self.datareader.filename}\n Measured isotopes: {self.isotope_names}'
        return res

    def time_to_number(self, time, integration_time=None):
        """
        takes time in seconds returns number of measured values
        depends on integration time of MS method
        """
        if not integration_time:
            integration_time = (self.time[2]-self.time[1])
        val = int(abs(time)//integration_time)
        if time < 0:
            val = -val
        return val

    def read_param(self, path):
        self.param = Param(path, self.logger)
        if self.param.peak_names:
            self.names = self.param.peak_names

    def read_srms(self, path='./SRM.xlsx'):
        self.srms = pd.ExcelFile(path).parse(index_col=0)

    def set_names(self, names):
        """
        takes list of names of peaks
        """
        self.names = names

    def create_isotopes(self):
        for key in self.isotope_names:
            self.isotopes[key] = Isotope(key, self, logger=self.logger)

    def select(self, method='treshold', s=60, sdmul=10, iolite=None):
        self.selector = Selector(
            self, s=s, sdmul=sdmul, iolite=iolite, logger=self.logger)
        self.selector.method = method
        self.starts, self.ends = self.selector()
        self.laser_on, self.laser_off = self.selector.create_on_off(
            self.starts, self.ends)

    def graph(self, ax=None, logax=False, el=None):
        """
        create matplotlib graph of intensity in time for ablation
        highlights ablation part and background signal
        """
        if ax == None:
            fig, ax = plt.subplots()

        ax.cla()
        ax.clear()
        # if element is defined, plot only one element, otherwise all
        if el:
            self.data.plot(ax=ax, y=el, kind='line', legend=False)
        else:
            self.data.plot(ax=ax, kind='line', legend=True)
            ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.05),
                      fancybox=True, shadow=True, ncol=10)

        if logax:
            ax.set_yscale('log')

        if self.starts and self.ends:
            # create lines for start and end of each ablation
            for i in range(0, len(self.starts)):
                ax.axvline(x=self.time[self.starts[i]],
                           color='blue', linewidth=2)
            for i in range(0, len(self.ends)):
                ax.axvline(x=self.time[self.ends[i]],
                           color='blue', linewidth=2)

        if self.laser_off:
            # higlights bacground
            for off in self.laser_off:
                # print(self.time[off[0]], self.time[off[1]])
                try:
                    ax.axvspan(
                        self.time[off[0]], self.time[off[1]], alpha=0.2, color='red')
                except:
                    warnings.warn('something is wrong')

        if self.laser_on:
            # higlihts ablation
            for on in self.laser_on:
                ax.axvspan(self.time[on[0]], self.time[on[1]],
                           alpha=0.2, color='green')

        plt.show()

    def average_isotopes(self, despiked=True, bcgcor_method='all', method='intensity'):
        """
        Calculate average value for all peaks for all isotopes.

        Parameters
        ----------
        bcgcor_method : str
            Method of background calculation. Possible options are [None, 'all', 'beginning', 'end']. Default is 'all'.
        despiked : bool
            State of data to be corrected. If True use despiked data, else use original. Default is True.
        method : str
            Method to get average peaks. Possible options are ['intensity', 'integral']. Default is 'intensity'.

        Returns
        -------
        data : DataFrame
            DF where index are peak names and columns are isotopes
        """
        if self.logger is not None:
            self.logger.info(
                f'Averaging of signal using despiking: {despiked}, bcgcor: {bcgcor_method}, method: {method}.')

        self.means = pd.DataFrame()
        for el, isotope in self.isotopes.items():
            isotope.average_intensity(
                despiked=despiked, bcgcor_method=bcgcor_method, method=method)
            self.means[el] = isotope.means

        if self.names:
            self.means.index = self.names
        else:
            if self.logger is not None:
                self.logger.error('Missing peak names.')
        return self.means

    def quantify_isotopes(self, srm_name='NIST610'):
        self.quantified = pd.DataFrame()
        for el, isotope in self.isotopes.items():
            isotope.quantify(srm_name=srm_name)
            self.quantified[el] = isotope.quantified

        if self.logger is not None:
            self.logger.info(
                f'Quantification of signal using SRM: {srm_name}.')
        self.quantified.index = [
            name for name in self.names if name != srm_name]
        return self.quantified

    def IS_correction(self):
        """
        Calculates correction for each element given in internal standard correction from PARAM file

        Returns
        -------
        corrected data : dict
            dict of internal standards used for correction as keys and DataFrames where index 
            are peak names and columns are isotopes with values in ppm. 
        """
        self.corrected_IS = {}

        if self.param is not None:
            if self.param.is_coef.empty:
                if self.logger is not None:
                    self.logger.error(
                        'Missing coefficients of internal standards.')
        else:
            if self.logger is not None:
                self.logger.error(
                    'Missing param file for IS correction.')

        self.correction_elements = list(self.param.is_coef.columns)

        for el in self.correction_elements:
            if el in list(map(elem_resolution, self.isotope_names)):
                self.corrected_IS[el] = (correction(
                    self.quantified, el, self.param.is_coef))
        return self.corrected_IS

    def TS_correction(self, suma=1000000, skip_isotopes=[]):
        """
        Calculates total sum correction using coefficients given in PARAM file.

        Parameters
        ----------
        suma : float
            Total sum of measured elements in ppm used for correction. 
            Normally is equal to 100% (default).
        skip_isotopes: list
            List of isotopes to be skipped, in total sum correction one element 
            can't be measured on multiple isotopes.

        Returns
        -------
        corrected data : DataFrame
            DataFrame where index are peak names and columns are isotopes with values in ppm. 
        """
        if self.param is not None:
            if not self.param.ts_coef:
                if self.logger is not None:
                    self.logger.error(
                        'Missing coeficients for total sum correction.')
            else:
                ts_coef = self.param.ts_coef
        else:
            if self.logger is not None:
                self.logger.error(
                    'Using default TS coef.')
            ts_coef = pd.ExcelFile(
                './default_sum_koef.xlsx').parse(0, index_col=0, header=None).to_dict()[1]

        self.corrected_TS = self.quantified.copy()
        self.corrected_TS.drop(skip_isotopes, axis=1, inplace=True)

        elem_lst = [element_strip(el) for el in self.corrected_TS.columns]
        if len(elem_lst) != len(set(elem_lst)):
            if self.logger is not None:
                self.logger.warning(
                    'Multiple isotopes of one element measured. Recomended to use only one isotope for total sum correction.')

        for key in ts_coef:
            elem = element_formater(key, self.corrected_TS.columns)
            if not elem:
                continue
            self.corrected_TS[elem] = self.corrected_TS[elem] / \
                ts_coef[key] * 100

        koef = suma/self.corrected_TS.sum(1)

        self.corrected_TS = self.corrected_TS.mul(list(koef), axis='rows')
        for key in ts_coef:
            elem = element_formater(key, self.corrected_TS.columns)
            if not elem:
                continue
            self.corrected_TS[elem] = self.corrected_TS[elem] * \
                ts_coef[key] / 100

        return self.corrected_TS

    def detection_limit(self, method='intensity', scale='all', ablation_time=60):
        """
        Calculate detection limit for all isotopes.

        Parameters
        ----------
        method : str
            Possible methods ['integral','intensity']. LoD must be calculates same way as average of isotopes. 
        scale : str
            Possible methods ['beginning', 'all']. LoD must be calculates same way as average of isotopes. 
        ablation_time: int
            Time in seconds of one ablation. Necessary only for method 'integral'. 
        """
        if self.logger is not None:
            self.logger.info(
                f'Calculating detection limit by method: {method}, scale: {scale}')
        self.lod = pd.Series()
        for el, isotope in self.isotopes.items():
            isotope.detection_limit(method, scale, ablation_time)
            self.lod[el] = isotope.lod
        self.lod.name = 'lod'
        return self.lod

    def report(self):
        """
        Clean all data in MSData. Replace values lower than limit of detection.
        If the value is above LoD, round to specific decimal place.
        """
        if self.corrected_TS is not None:
            self.corrected_TS = self.corrected_TS.append(self.lod)
            for column in self.corrected_TS:
                self.corrected_TS[column] = [
                    report(value, self.lod, column) for value in self.corrected_TS[column]]

        if self.corrected_IS is not None:
            self.corrected_IS = {key: df.append(self.lod)
                                 for key, df in self.corrected_IS.items()}
            for el, df in self.corrected_IS.items():
                for column in df:
                    df[column] = [report(value, self.lod, column)
                                  for value in df[column]]

        if self.quantified is not None:
            self.quantified = self.quantified.append(self.lod)
            for column in self.quantified:
                self.quantified[column] = [
                    report(value, self.lod, column) for value in self.quantified[column]]

    def export(self, path):
        """
        Export all data to excel file.

        Parameters
        ----------
        path : str
            Path to excel file where to save data. 
        """

        if self.logger is not None:
            self.logger.info(
                f'Saving data to: {path}.')

        writer = pd.ExcelWriter(path, engine='xlsxwriter')
        data = {'raw': self.data,
                'average': self.means,
                'quantified': self.quantified}

        if self.corrected_IS is not None:
            data.update(self.corrected_IS)

        if self.corrected_TS is not None:
            data['corrected_TS'] = self.corrected_TS

        for key, df in data.items():
            df.to_excel(writer, sheet_name=key)

        writer.save()

    def create_maps(self, despiked=False, bcgcor_method='all', dx=1, dy=1):
        for el, isotope in self.isotopes.items():
            isotope.elemental_distribution(despiked, bcgcor_method, dx, dy)

    def export_matrices(self, path):
        writer = pd.ExcelWriter(path, engine='xlsxwriter')
        for el, isotope in self.isotopes.items():
            df = pd.DataFrame(isotope.elmap.matrix,
                              index=isotope.elmap.y, columns=isotope.elmap.x)
            df.to_excel(writer, sheet_name=el)
        writer.save()

    def import_matrices(self, path):
        for el, isotope in self.isotopes.items():
            isotope.elmap.read_matrix(path, el)


class Isotope():
    """
    Extract specific Isotope data from MSData object.

    Parameters
    ----------
    isotope_name : str
        Isotope to extract from MSData (i.e. Li7)
    ms_data : MSData
        MS data to extract from.

    Returns
    -------
    data : array
        `data` interpreted as an array.

    """

    def __init__(self, isotope_name, ms_data, logger=None, *args, **kwargs):
        self.isotope_name = elem_resolution(isotope_name)
        self.ms_data = ms_data
        self.isotope_number = int(
            ''.join([c for c in self.isotope_name if c.isnumeric()]))
        self.data = ms_data.data[isotope_name].values
        self.despiked = None
        self.bcg = None
        self.bcg_corrected = None
        self.bcgcor_method = None
        self.means = None
        self.ratio = None
        self.quantified = None
        self.elmap = None
        self.logger = logger

        # TODO use self.despiked as a state, in next functions if despiked is True, use that
        # values if they exists, if not do despike if false use original

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        if isinstance(idx, (int, slice)):
            return self.data[idx]
        return [self.data[i] for i in idx]

    def __call__(self):
        plt.plot(self.data)

    def __repr__(self):
        res = f'{self.__class__.__name__} {self.isotope_name} ({len(self)} values)\n{self.data[:10]}'
        if len(self) > 10:
            res = res[:-1] + '...]'
        return res

    def __iter__(self): return iter(self.data)
    def __setitem__(self, i, o): self.data[i] = o
    def __delitem__(self, i): del(self.data[i])

    def element(self):
        return ''.join([c for c in self.isotope_name if c.isalpha()])

    def despike(self, win=3, treshold=10):
        sig = pd.Series(self.data)
        skip = int((win-1)/2)
        outliers = np.zeros(len(sig), dtype=bool)
        outliers[skip:-skip] = True
        rm = np.array(sig.rolling(win).mean().dropna())
        rs = np.array(sig.rolling(win).std().dropna())**0.5
        outliers[skip:-skip] = sig[skip:-skip] > rm+rs*treshold
        sig[outliers] = rm[outliers[skip:-skip]]
        self.despiked = sig

    def bcg_correction(self, method='all', despiked=True):
        """
        Calculate bacground correction for selected Isotope.

        Parameters
        ----------
        method : str
            Method of background calculation. Possible options are ['all', 'beginning', 'end']. Default is 'all'.
        despiked : bool
            State of data to be corrected. If True use despiked data, else use original. Default is True.
        """

        if not self.ms_data.laser_on:
            if self.logger is not None:
                self.logger.error(
                    'Bacground correction requires selection of peaks first.')
            return

        if self.bcgcor_method != method:
            self.bcg = Background(
                self, self.ms_data.laser_on, self.ms_data.laser_off)

        if not despiked:
            data = np.copy(self.data)
        elif despiked:
            if self.despiked is None:
                self.despike()
            data = np.copy(self.despiked)

        if method == 'beginning':
            data_out = data-self.bcg.bcg_means[0]
        elif method == 'end':
            data_out = data-self.bcg.bcg_means[-1]
        elif method == 'all':
            data_out = data-self.bcg.bcg_interp
        else:
            if self.logger is not None:
                self.logger.warning(
                    'Select correct method  for background correction. Possible options: [\'beginning\', \'end\', \'all\']')
            return

        data_out[data_out < 0] = 0
        self.bcg_corrected = data_out
        self.bcgcor_method = method

    def average_intensity(self, despiked=True, bcgcor_method='all', method='intensity'):
        """
        Calculate average intensity for peaks in selected Isotope.

        Parameters
        ----------
        bcgcor_method : str
            Method of background calculation. Possible options are [None, 'all', 'beginning', 'end']. Default is 'all'.
        despiked : bool
            State of data to be corrected. If True use despiked data, else use original. Default is True.
        method : str
            Method to get average peaks. Possible options are ['intensity', 'integral']. Default is 'intensity'.
        """
        laser_on = self.ms_data.laser_on

        if bcgcor_method is not None:
            if self.bcg_corrected is None:
                self.bcg_correction(method=bcgcor_method, despiked=despiked)
            data = np.copy(self.bcg_corrected)
        else:
            if not despiked:
                data = np.copy(self.data)
            elif despiked:
                if self.despiked is None:
                    self.despike()
                data = np.copy(self.despiked)

        means = []
        for on in laser_on:
            if method == 'intensity':
                means.append(np.mean(data[on[0]:on[1]]))
            elif method == 'integral':
                sample_y = data[on[0]:on[1]]
                sample_x = data.index[on[0]:on[1]]
                means.append(np.trapz(sample_y, sample_x))
            else:
                if self.logger is not None:
                    self.logger.error(
                        'Select correct method for peak averaging. Possible options: [\'intensity\', \'integral\']')
                return
        self.means = means

    def quantify(self, srm_name='NIST610'):
        """
        Calculate quantification of average peaks in Isotope.
        Parameters
        ----------
        srm_name : str
            Name of the standard reference material to be used for quantification. Must be in SRM file.
        """
        if self.ms_data.srms is not None:
            if srm_name in self.ms_data.srms.index:
                srm = self.ms_data.srms.loc[srm_name, :]
            else:
                if self.logger is not None:
                    self.logger.error(
                        f'Standard reference material {srm_name} not in file.')
                return
        else:
            if self.logger is not None:
                self.logger.error('Missing standard reference materials file.')
            return

        if self.means is None:
            if self.logger is not None:
                self.logger.error('Missing average values for peaks.')
            return

        if self.ms_data.names is None:
            if self.logger is not None:
                self.logger.error('Missing names of peaks.')
            return

        if element_strip(self.isotope_name) not in self.ms_data.srms.columns:
            if self.logger is not None:
                self.logger.error(f'Unknown isotope: {self.isotope_name}.')
            return

        spots = np.array(
            [m for m, s in zip(self.means, self.ms_data.names) if s != srm_name])
        stdsig = np.mean(
            [[m for m, s in zip(self.means, self.ms_data.names) if s == srm_name]])
        self.ratio = srm[element_strip(self.isotope_name)]/stdsig
        self.quantified = spots * self.ratio

    def detection_limit(self, method='intensity', scale='all', ablation_time=60):
        """
        Calculate limit of detection for isotope.
        Parameters
        ----------
        method : str
            Possible methods ['integral','intensity']. LoD must be calculates same way as average of isotopes. 
        scale : str
            Possible methods ['beginning', 'all']. LoD must be calculates same way as average of isotopes. 
        ablation_time: int
            Time in seconds of one ablation. Necessary only for method 'integral'. 
        """
        if scale == 'all':
            bcg = np.array(
                list(itertools.chain.from_iterable(self.bcg.bcg_all)))
        elif scale == 'beginning':
            bcg = np.array(self.bcg.bcg_all[0])
        if method == 'integral':
            self.lod = (bcg.std()*ablation_time*self.ratio)
        elif method == 'intensity':
            self.lod = (bcg.std()*3*self.ratio)

    def elemental_distribution(self, despiked=False, bcgcor_method='all', dx=1, dy=1):

        laser_on = self.ms_data.laser_on

        if bcgcor_method is not None:
            if self.bcgcor_method != bcgcor_method:
                self.bcg_correction(method=bcgcor_method, despiked=despiked)
            data = np.copy(self.bcg_corrected)
        else:
            if not despiked:
                data = np.copy(self.data)
            elif despiked:
                if self.despiked is None:
                    self.despike()
                data = np.copy(self.despiked)

        l = []
        for on in laser_on:
            arr = data[on[0]:on[1]]
            arr[arr < 0] = 0
            l.append(arr)

        lens = [len(r) for r in l]
        maxlen = max(lens)
        mat = np.empty([len(l), maxlen])
        mat[:] = np.nan
        mask = np.arange(maxlen) < np.array(lens)[:, None]
        mat[mask] = np.concatenate(l)

        self.elmap = ElementalMap(mat, dx, dy)


class Peak():
    def __init__(self, isotope, bounds):
        if isotope.data is None:
            raise ValueError(
                'Missing data for isotope: {}'.format(isotope.isotope_name))
        self.data = isotope.data[bounds[0]:bounds[1]]

    def __len__(self):
        return len(self.data)

    def __call__(self, *args, **kwargs):
        plt.plot(self.data, *args, **kwargs)

    def __repr__(self):
        pass


class ElementalMap():
    def __init__(self, matrix=None, dx=None, dy=None):
        """
        Class of Elemental map holding data from MSData representing elemental distribution of one isotope.
        Parameters
        ----------
        matrix : numpy matrix
            Matrix of z values of image.
        dx: float
            Step between z values in x dirextion.
        dy: float
            Step between z values in y dirextion.
        """
        self.matrix = matrix
        self.dx = dx
        self.dy = dy
        self.create_xy()

    def __call__(self, fig=None, ax=None, units='', axis=True, clb=True, *args, **kwargs):
        """
        Show image of elemental map.
        Parameters
        ----------
        ax : matplotlib axes
            Axes where plot will be drawn, if None new axes will be created.
        units : str
            Label for colorbar
        axis : boolean
            If true then show axis with labels in the image.
        clb : boolean
            Wheather to show colorbal in image.
        """
        if fig is None or ax is None:
            fig, ax = plt.subplots()
        ax.cla()

        im = plt.imshow(self.matrix, extent=[
                        0, self.x[-1], self.y[-1], 0], *args, **kwargs)

        if not axis:
            ax.axis('off')
        if clb:
            divider = make_axes_locatable(ax)
            cax = divider.append_axes("right", size="5%", pad=0.05)
            clb = fig.colorbar(im, cax=cax)
            clb.ax.set_title(units)

    def __repr__(self):
        return f'{self.__class__.__name__}: {self.matrix.shape}'

    def create_xy(self):
        self.x = [i*self.dx for i in range(self.matrix.shape[1])]
        self.y = [i*self.dy for i in range(self.matrix.shape[0])]

    def read_matrix(self, path, el):
        """
        Import created map from excel, where each isotope is on one list.
        Parameters
        ----------
        path : str
            Path to excel file.
        el: str
            Name of isotope to import.
        """
        file = pd.ExcelFile(path)
        df = file.parse(el, index_col=0)
        self.matrix = np.matrix(df)
        self.dx = df.columns[1]-df.columns[0]
        self.dy = df.index[1]-df.index[0]
        self.create_xy()

    def write_matrix(self, path, el):
        """
        Save created map to excel, where each isotope is on one list.

        Parameters
        ----------
        path : str
            Path to excel file.
        el: str
            Name of isotope to export. Will be used as sheet name.
        """
        writer = pd.ExcelWriter(path, engine='xlsxwriter')
        df = pd.DataFrame(self.matrix, index=self.y, columns=self.x)
        df.to_excel(writer, sheet_name=el)
        writer.save()

    def rotate(self):
        """Rotate map by 90 degrees"""
        self.matrix = np.rot90(self.matrix)
        self.x, self.y = self.y, self.x

    def flip(self, axis=None):
        """
        Flip matrix of image.

        Parameters
        ----------
        axis : str
            Axis along which to flip over. The default, axis=None, 
        will flip over all of the axes, axis=0 flip matrix vertically, 
        axis=1 flip matrixhorizontally.
        """
        self.matrix = np.flip(self.matrix, axis)
