from __future__ import division
import numpy as np
import bisect


# TODO: What commonalities can be factored out of this class and TimeSlice?
# TODO: Deprecate stuff in psychoacoustics.py in favor of these classes
class FrequencyBand(object):
    """
    Represents an interval, or band of frequencies in hertz (cycles per second)

    Args:
        start_hz (float): The lower bound of the frequency band in hertz
        stop_hz (float): The upper bound of the frequency band in hertz

    Examples::
        >>> import zounds
        >>> band = zounds.FrequencyBand(500, 1000)
        >>> band.center_frequency
        750.0
        >>> band.bandwidth
        500
    """

    def __init__(self, start_hz, stop_hz):
        super(FrequencyBand, self).__init__()
        if stop_hz <= start_hz:
            raise ValueError('stop_hz must be greater than start_hz')
        self.stop_hz = stop_hz
        self.start_hz = start_hz

    def __eq__(self, other):
        try:
            return \
                self.start_hz == other.start_hz \
                and self.stop_hz == other.stop_hz
        except AttributeError:
            return super(FrequencyBand, self).__eq__(other)

    def __hash__(self):
        return (self.__class__.__name__, self.start_hz, self.stop_hz).__hash__()

    def intersect(self, other):
        """
        Return the intersection between this frequency band and another.

        Args:
            other (FrequencyBand): the instance to intersect with

        Examples::
            >>> import zounds
            >>> b1 = zounds.FrequencyBand(500, 1000)
            >>> b2 = zounds.FrequencyBand(900, 2000)
            >>> intersection = b1.intersect(b2)
            >>> intersection.start_hz, intersection.stop_hz
            (900, 1000)
        """
        lowest_stop = min(self.stop_hz, other.stop_hz)
        highest_start = max(self.start_hz, other.start_hz)
        return FrequencyBand(highest_start, lowest_stop)

    def bandwidth_ratio(self, other):
        return other.bandwidth / self.bandwidth

    def intersection_ratio(self, other):
        intersection = self.intersect(other)
        return self.bandwidth_ratio(intersection)

    @staticmethod
    def from_start(start_hz, bandwidth_hz):
        """
        Produce a :class:`FrequencyBand` instance from a lower bound and
        bandwidth

        Args:
            start_hz (float): the lower bound of the desired FrequencyBand
            bandwidth_hz (float): the bandwidth of the desired FrequencyBand

        """
        return FrequencyBand(start_hz, start_hz + bandwidth_hz)

    @staticmethod
    def from_center(center_hz, bandwidth_hz):
        half_bandwidth = bandwidth_hz / 2
        return FrequencyBand(
            center_hz - half_bandwidth, center_hz + half_bandwidth)

    @property
    def bandwidth(self):
        """
        The span of this frequency band, in hertz
        """
        return self.stop_hz - self.start_hz

    @property
    def center_frequency(self):
        return self.start_hz + (self.bandwidth / 2)

    def __repr__(self):
        return '''FrequencyBand(
start_hz={start_hz},
stop_hz={stop_hz},
center={center},
bandwidth={bandwidth})'''.format(
            start_hz=self.start_hz,
            stop_hz=self.stop_hz,
            center=self.center_frequency,
            bandwidth=self.bandwidth)


class FrequencyScale(object):
    """
    Represents a set of frequency bands with monotonically increasing start
    frequencies

    Args:
        frequency_band (FrequencyBand): A band representing the entire span of
            this scale.  E.g., one might want to generate a scale spanning the
            entire range of human hearing by starting with
            :code:`FrequencyBand(20, 20000)`
        n_bands (int): The number of bands in this scale
        always_even (bool): when converting frequency slices to integer indices
            that numpy can understand, should the slice size always be even?

    See Also:
        :class:`~zounds.spectral.LinearScale`
        :class:`~zounds.spectral.GeometricScale`
    """

    def __init__(self, frequency_band, n_bands, always_even=False):
        super(FrequencyScale, self).__init__()
        self.always_even = always_even
        self.n_bands = n_bands
        self.frequency_band = frequency_band
        self._bands = None
        self._starts = None
        self._stops = None

    @property
    def bands(self):
        """
        An iterable of all bands in this scale
        """
        if self._bands is None:
            self._bands = self._compute_bands()
        return self._bands

    @property
    def band_starts(self):
        if self._starts is None:
            self._starts = [b.start_hz for b in self.bands]
        return self._starts

    @property
    def band_stops(self):
        if self._stops is None:
            self._stops = [b.stop_hz for b in self.bands]
        return self._stops

    def _compute_bands(self):
        raise NotImplementedError()

    def __len__(self):
        return self.n_bands

    @property
    def center_frequencies(self):
        """
        An iterable of the center frequencies of each band in this scale
        """
        return (band.center_frequency for band in self)

    @property
    def bandwidths(self):
        """
        An iterable of the bandwidths of each band in this scale
        """
        return (band.bandwidth for band in self)

    def ensure_overlap_ratio(self, required_ratio=0.5):
        """
        Ensure that every adjacent pair of frequency bands meets the overlap
        ratio criteria.  This can be helpful in scenarios where a scale is
        being used in an invertible transform, and something like the `constant
        overlap add constraint
        <https://ccrma.stanford.edu/~jos/sasp/Constant_Overlap_Add_COLA_Cases.html>`_
        must be met in order to not introduce artifacts in the reconstruction.

        Args:
            required_ratio (float): The required overlap ratio between all
                adjacent frequency band pairs

        Raises:
            AssertionError: when the overlap ratio for one or more adjacent
                frequency band pairs is not met
        """

        msg = \
            'band {i}: ratio must be at least {required_ratio} but was {ratio}'

        for i in xrange(0, len(self) - 1):
            b1 = self[i]
            b2 = self[i + 1]

            try:
                ratio = b1.intersection_ratio(b2)
            except ValueError:
                ratio = 0

            if ratio < required_ratio:
                raise AssertionError(msg.format(**locals()))

    @property
    def Q(self):
        """
        The quality factor of the scale, or, the ratio of center frequencies
        to bandwidths
        """
        return np.array(list(self.center_frequencies)) \
               / np.array(list(self.bandwidths))

    @property
    def start_hz(self):
        """
        The lower bound of this frequency scale
        """
        return self.frequency_band.start_hz

    @property
    def stop_hz(self):
        """
        The upper bound of this frequency scale
        """
        return self.frequency_band.stop_hz

    def get_slice(self, frequency_band):
        """
        Given a frequency band, and a frequency dimension comprised of
        n_samples, return a slice using integer indices that may be used to
        extract only the frequency samples that intersect with the frequency
        band
        """
        start_index = bisect.bisect_left(
            self.band_stops, frequency_band.start_hz)
        stop_index = bisect.bisect_left(
            self.band_starts, frequency_band.stop_hz)

        if self.always_even and (stop_index - start_index) % 2:
            # KLUDGE: This is simple, but it may make sense to choose move the
            # upper *or* lower bound, based on which one introduces a lower
            # error
            stop_index += 1
        return slice(start_index, stop_index)

    def __eq__(self, other):
        return \
            self.__class__ == other.__class__ \
            and self.frequency_band == other.frequency_band \
            and self.n_bands == other.n_bands

    def __iter__(self):
        return iter(self.bands)

    def _construct_scale_from_slice(self, bands):
        freq_band = FrequencyBand(bands[0].start_hz, bands[-1].stop_hz)
        return self.__class__(freq_band, len(bands))

    def __getitem__(self, index):
        try:
            # index is an integer or slice
            bands = self.bands[index]
        except TypeError:
            # index is a frequency band
            bands = self.bands[self.get_slice(index)]

        if isinstance(bands, FrequencyBand):
            return bands

        return self._construct_scale_from_slice(bands)

    def __str__(self):
        cls = self.__class__.__name__
        return '{cls}(band={self.frequency_band}, n_bands={self.n_bands})' \
            .format(**locals())

    def __repr__(self):
        return self.__str__()


class LinearScale(FrequencyScale):
    """
    A linear frequency scale with constant bandwidth.  Appropriate for use
    with transforms whose coefficients also lie on a linear frequency scale,
    e.g. the FFT or DCT transforms.

    Args:
        frequency_band (FrequencyBand): A band representing the entire span of
            this scale.  E.g., one might want to generate a scale spanning the
            entire range of human hearing by starting with
            :code:`FrequencyBand(20, 20000)`
        n_bands (int): The number of bands in this scale
        always_even (bool): when converting frequency slices to integer indices
            that numpy can understand, should the slice size always be even?

    Examples:
        >>> from zounds import FrequencyBand, LinearScale
        >>> scale = LinearScale(FrequencyBand(20, 20000), 10)
        >>> scale
        LinearScale(band=FrequencyBand(
        start_hz=20,
        stop_hz=20000,
        center=10010.0,
        bandwidth=19980), n_bands=10)
        >>> scale.Q
        array([ 0.51001001,  1.51001001,  2.51001001,  3.51001001,  4.51001001,
                5.51001001,  6.51001001,  7.51001001,  8.51001001,  9.51001001])
    """

    def __init__(self, frequency_band, n_bands, always_even=False):
        super(LinearScale, self).__init__(frequency_band, n_bands, always_even)

    @staticmethod
    def from_sample_rate(sample_rate, n_bands, always_even=False):
        """
        Return a :class:`~zounds.spectral.LinearScale` instance whose upper
        frequency bound is informed by the nyquist frequency of the sample rate.

        Args:
            sample_rate (SamplingRate): the sample rate whose nyquist frequency
                will serve as the upper frequency bound of this scale
            n_bands (int): the number of evenly-spaced frequency bands
        """
        fb = FrequencyBand(0, sample_rate.nyquist)
        return LinearScale(fb, n_bands, always_even=always_even)

    def _compute_bands(self):
        freqs = np.linspace(
            self.start_hz, self.stop_hz, self.n_bands, endpoint=False)
        # constant, non-overlapping bandwidth
        bandwidth = freqs[1] - freqs[0]
        return tuple(FrequencyBand(f, f + bandwidth) for f in freqs)


class LogScale(FrequencyScale):
    def __init__(self, frequency_band, n_bands, always_even=False):
        super(LogScale, self).__init__(
            frequency_band, n_bands, always_even=always_even)

    def _compute_bands(self):
        center_freqs = np.logspace(
            np.log10(self.start_hz),
            np.log10(self.stop_hz),
            self.n_bands + 1)
        # variable bandwidth
        bandwidths = np.diff(center_freqs)
        return tuple(FrequencyBand.from_center(cf, bw)
                     for (cf, bw) in zip(center_freqs[:-1], bandwidths))


class GeometricScale(FrequencyScale):
    """
    A constant-Q scale whose center frequencies progress geometrically rather
    than linearly

    Args:
        start_center_hz (int): the center frequency of the first band in the
            scale
        stop_center_hz (int): the center frequency of the last band in the scale
        bandwidth_ratio (float): the center frequency to bandwidth ratio
        n_bands (int): the total number of bands

    Examples:
        >>> from zounds import GeometricScale
        >>> scale = GeometricScale(20, 20000, 0.05, 10)
        >>> scale
        GeometricScale(band=FrequencyBand(
        start_hz=19.5,
        stop_hz=20500.0,
        center=10259.75,
        bandwidth=20480.5), n_bands=10)
        >>> scale.Q
        array([ 20.,  20.,  20.,  20.,  20.,  20.,  20.,  20.,  20.,  20.])
        >>> list(scale.center_frequencies)
        [20.000000000000004, 43.088693800637671, 92.831776672255558,
            200.00000000000003, 430.88693800637651, 928.31776672255558,
            2000.0000000000005, 4308.8693800637648, 9283.1776672255564,
            20000.000000000004]
    """

    def __init__(
            self,
            start_center_hz,
            stop_center_hz,
            bandwidth_ratio,
            n_bands,
            always_even=False):
        self.__bands = [
            FrequencyBand.from_center(cf, cf * bandwidth_ratio)
            for cf in np.geomspace(start_center_hz, stop_center_hz, num=n_bands)
            ]
        band = FrequencyBand(self.__bands[0].start_hz, self.__bands[-1].stop_hz)
        super(GeometricScale, self).__init__(
            band, n_bands, always_even=always_even)
        self.start_center_hz = start_center_hz
        self.stop_center_hz = stop_center_hz
        self.bandwidth_ratio = bandwidth_ratio

    def _construct_scale_from_slice(self, bands):
        return ExplicitScale(bands)

    def __eq__(self, other):
        return \
            super(GeometricScale, self).__eq__(other) \
            and self.start_center_hz == other.start_center_hz \
            and self.stop_center_hz == other.stop_center_hz \
            and self.bandwidth_ratio == other.bandwidth_ratio

    def _compute_bands(self):
        return self.__bands


class ExplicitScale(FrequencyScale):
    """
    A scale where the frequency bands are provided explicitly, rather than
    computed

    Args:
        bands (list of FrequencyBand): The explicit bands used by this scale

    See Also:
        :class:`~zounds.spectral.FrequencyAdaptive`
    """
    def __init__(self, bands):
        bands = list(bands)
        frequency_band = FrequencyBand(bands[0].start_hz, bands[-1].stop_hz)
        super(ExplicitScale, self).__init__(
            frequency_band, len(bands), always_even=False)
        self._bands = bands

    def _construct_scale_from_slice(self, bands):
        return ExplicitScale(bands)

    def _compute_bands(self):
        return self._bands

    def __eq__(self, other):
        return all([a == b for (a, b) in zip(self, other)])


class BarkScale(FrequencyScale):
    def __init__(self, frequency_band, n_bands):
        super(BarkScale, self).__init__(frequency_band, n_bands)


class MelScale(FrequencyScale):
    def __init__(self, frequency_band, n_bands):
        super(MelScale, self).__init__(frequency_band, n_bands)


class ChromaScale(FrequencyScale):
    def __init__(self, frequency_band, n_bands):
        super(ChromaScale, self).__init__(frequency_band, n_bands)
