import numpy as np
from sklearn.cluster import KMeans

def median_filter(a, w):
    ap = np.pad(a, (w//2, w//2), mode='edge')
    o = np.zeros(a.shape[0])
    for i in np.arange(a.shape[0]):
        sl = ap[i:i+w]
        o[i] = np.median(sl)
    return o

class Segmenter:
    seg_win_size_sec = 4.0  #: window size for stat. measures for segmentation, in sec
    seg_win_step_sec = 1.0  #: step for segmentation, in sec
    n_clusters = 8  #: clusters for KMeans algorithm
    seg_filt_win_sec = 20.0 #: median filter width for smoothing segments

    def __init__(self): pass

    def get_segment_boundaries(self, fs, guitar):
        """split the spectral power signal 'guitar' into stochastically similar segments."""
        segment_ids = self.get_segments(fs, guitar)
        stxs = np.diff(segment_ids) != 0
        i_stxs = np.where(stxs)[0]
        return i_stxs

    def get_segments(self, fs, guitar):
        """split the spectral power signal 'guitar' into stochastically similar segments."""
        seg_filt_win = int(self.seg_filt_win_sec / self.seg_win_step_sec)
        seg_guitar_data = self._sig_stochastics(fs, guitar)
        X = np.vstack((
            seg_guitar_data[:,0]*1.4,
            np.sqrt(np.sum(seg_guitar_data[:,1:]**2, axis=1))
        )).T
        # cluster by stochastic characteristics
        kmeans = KMeans(n_clusters=self.n_clusters, random_state=0, n_init="auto").fit(X)
        segment_ids = np.floor(median_filter(kmeans.labels_, seg_filt_win)).astype(int)
        # up-sample segment id assignment
        iidx = np.linspace(0, segment_ids.shape[0], guitar.shape[0], endpoint=False).astype(int)
        return segment_ids[iidx]

    def _sig_stochastics(self, fs, y):
        """compute the stochastic moments of the signal. normalized."""
        seg_win_size = int(self.seg_win_size_sec * fs)
        seg_win_step = int(self.seg_win_step_sec * fs)
        #
        seg_y_data = np.zeros((y.shape[0] // seg_win_step, 4))
        y_pad = np.pad(y, (seg_win_size // 2, seg_win_size // 2))
        y_0_max, y_0_mean = np.max(y), np.mean(y)
        y_1_max = np.max(np.mean((y - y_0_mean)**2))
        y_2_max = np.max(np.mean(np.abs((y - y_0_mean)**3)))
        y_3_max = np.max(np.mean((y - y_0_mean)**4))
        wo = int(self.seg_win_size_sec/self.seg_win_step_sec)
        for i in np.arange(wo//2, y.shape[0] // seg_win_step - wo//2):
            i_c = int((i+0.5)*seg_win_step)
            y_slice = y_pad[i_c-seg_win_size//2:i_c+seg_win_size//2]
            mean = np.mean(y_slice)
            seg_y_data[i,0] = mean / y_0_max
            seg_y_data[i,1] = np.mean((y_slice - mean)**2) / y_1_max
            seg_y_data[i,2] = np.mean(np.abs((y_slice - mean)**3)) / y_2_max / 2
            seg_y_data[i,3] = np.mean((y_slice - mean)**4) / y_3_max / 4
        #
        return seg_y_data