from sklearn.base import BaseEstimator
from scipy.stats import median_abs_deviation
import numpy as np

class DeSpike(BaseEstimator):
    '''
    A transformer to remove cosmic spikes from spectra as described by Whitaker & Hayes.

    Parameters
    ----------
    window: int,
        ..
    threshold: float,
        A point is considered a spike if Z-score >= threshold.
    ignore: np.ndarray,
        Array of wavelengths to ignore. Uses index if ignore_ref is not specified.
    ignore_ref: np.ndarray,
        Array of reference wavelengths used by ignore.


    D. A. Whitaker and K. Hayes, "A simple algorithm for despiking Ramann spectra",
    Chemometrics and Intelligent Laboratory Systems, vol. 179, pp. 82-84, Aug. 2018,
    doi: 10.1016/j.chemolab.2018.06.009.
    '''

    def __init__(self, *,
                 window: int = 5,
                 threshold: float = 6,
                 ignore: np.ndarray = np.array(["None"]),
                 ignore_ref: np.ndarray = np.array(["None"])):
        # super().__init__()
        self.window = window
        self.threshold = threshold
        self.ignore = ignore
        self.ignore_ref = ignore_ref

    def fit(self, X: np.ndarray, y=None) -> "DeSpike":
        '''
        Compute the modified Z-scores for all spectra.

        Parameters
        ----------
        X: np.ndarray of shape (n_spectra, n_wavelengths)
            The input array from which to remove the spikes.

        y: None
            Ignored to align with API.

        Returns
        -------
        self: Despike
            The fitted transformer, i.e. calculated the modified z-scores.
        '''

        nrow, ncol = X.shape
        self.Z = np.zeros(shape=(nrow, ncol-1))

        for spectrum in range(nrow):
            x = np.diff(X[spectrum])
            m = np.median(x)
            M = median_abs_deviation(x, scale="normal")
            self.Z[spectrum] = (x-m)/M
            
        # Insert 1 col lost through np.diff
        self.Z = np.insert(self.Z, 0, 0, axis=1)
        
        self.is_fitted_ = True
        return self
    
    def transform(self, X: np.ndarray, y=None):
        '''
        Smooth the peaks flaged as a spike.

        Parameters
        ----------
        X: np.ndarray of shape (n_spectra, n_wavelengths)
            The input array from which to remove the spikes.

        y: None
            Ignored to align with API.
        '''
        nrow, ncol = X.shape
        X_cleaned = np.copy(X)

        if self.ignore_ref.any() == str:
            self.ignore_ref = range(ncol)

        if self.ignore.any() == str:
            mask = np.repeat(True, ncol)
        else:
            mask = np.isin(self.ignore_ref, self.ignore)

        
        for spectrum in range(nrow):
            zb = abs(self.Z[spectrum]) >= self.threshold
            zb[mask] = False
            zb[0] = zb[-1] = True
        
            for i, spike in enumerate(zb):
                if spike:
                    # Get window around spike & selecte spike free for average.
                    w = np.arange(
                            np.max((0, i - self.window)),
                            np.min((ncol, i + self.window)))
                    w = w[~zb[w]]
                    # Remove empty windows
                    if len(w) != 0:
                        X_cleaned[spectrum, i] = np.mean(X[spectrum, w])
        return X_cleaned
