import numpy as np

from ..base import Preprocess
from ..base import Parameter


class SumMix(Preprocess):
    """Sums up multi channel source.

    Parameters
    ----------
    axis : int
    channels : int
        Number of channels will be generated by this preprocess.

    """
    axis = Parameter(-1)
    channels = Parameter(1)

    def process(self, data):
        if self.source.channels > 1:
            block = np.sum(data[0], axis=self.axis).flatten()
            data = block, *data[1:]
        return data


class MeanMix(Preprocess):
    """Averages multi channel source e.g stereo to mono mix.

    Parameters
    ----------
    axis : int
    channels : int
        Number of channels will be generated by this preprocess.

    """
    axis = Parameter(-1)
    channels = Parameter(1)

    def process(self, data):
        if self.source.channels > 1:
            block = np.mean(data[0], axis=self.axis).flatten()
            data = block, *data[1:]
        return data
