from __future__ import division, print_function, absolute_import

import copy

import numpy as np
from numpy.testing import (
    assert_,
    assert_equal,
    assert_allclose,
    assert_array_equal
)
import pytest
from pytest import raises, warns

from scipy._lib.six import xrange
from scipy.signal._peak_finding import (
    argrelmax,
    argrelmin,
    peak_prominences,
    peak_widths,
    _unpack_condition_args,
    find_peaks,
    find_peaks_cwt,
    _identify_ridge_lines
)
from scipy.signal._peak_finding_utils import _local_maxima_1d, PeakPropertyWarning


def _gen_gaussians(center_locs, sigmas, total_length):
    xdata = np.arange(0, total_length).astype(float)
    out_data = np.zeros(total_length, dtype=float)
    for ind, sigma in enumerate(sigmas):
        tmp = (xdata - center_locs[ind]) / sigma
        out_data += np.exp(-(tmp**2))
    return out_data


def _gen_gaussians_even(sigmas, total_length):
    num_peaks = len(sigmas)
    delta = total_length / (num_peaks + 1)
    center_locs = np.linspace(delta, total_length - delta, num=num_peaks).astype(int)
    out_data = _gen_gaussians(center_locs, sigmas, total_length)
    return out_data, center_locs


def _gen_ridge_line(start_locs, max_locs, length, distances, gaps):
    """
    Generate coordinates for a ridge line.

    Will be a series of coordinates, starting a start_loc (length 2).
    The maximum distance between any adjacent columns will be
    `max_distance`, the max distance between adjacent rows
    will be `map_gap'.

    `max_locs` should be the size of the intended matrix. The
    ending coordinates are guaranteed to be less than `max_locs`,
    although they may not approach `max_locs` at all.
    """

    def keep_bounds(num, max_val):
        out = max(num, 0)
        out = min(out, max_val)
        return out

    gaps = copy.deepcopy(gaps)
    distances = copy.deepcopy(distances)

    locs = np.zeros([length, 2], dtype=int)
    locs[0, :] = start_locs
    total_length = max_locs[0] - start_locs[0] - sum(gaps)
    if total_length < length:
        raise ValueError('Cannot generate ridge line according to constraints')
    dist_int = length / len(distances) - 1
    gap_int = length / len(gaps) - 1
    for ind in xrange(1, length):
        nextcol = locs[ind - 1, 1]
        nextrow = locs[ind - 1, 0] + 1
        if (ind % dist_int == 0) and (len(distances) > 0):
            nextcol += ((-1)**ind)*distances.pop()
        if (ind % gap_int == 0) and (len(gaps) > 0):
            nextrow += gaps.pop()
        nextrow = keep_bounds(nextrow, max_locs[0])
        nextcol = keep_bounds(nextcol, max_locs[1])
        locs[ind, :] = [nextrow, nextcol]

    return [locs[:, 0], locs[:, 1]]


class TestLocalMaxima1d(object):

    def test_empty(self):
        """Test with empty signal."""
        x = np.array([], dtype=np.float64)
        for array in _local_maxima_1d(x):
            assert_equal(array, np.array([]))
            assert_(array.base is None)

    def test_linear(self):
        """Test with linear signal."""
        x = np.linspace(0, 100)
        for array in _local_maxima_1d(x):
            assert_equal(array, np.array([]))
            assert_(array.base is None)

    def test_simple(self):
        """Test with simple signal."""
        x = np.linspace(-10, 10, 50)
        x[2::3] += 1
        expected = np.arange(2, 50, 3)
        for array in _local_maxima_1d(x):
            # For plateaus of size 1, the edges are identical with the
            # midpoints
            assert_equal(array, expected)
            assert_(array.base is None)

    def test_flat_maxima(self):
        """Test if flat maxima are detected correctly."""
        x = np.array([-1.3, 0, 1, 0, 2, 2, 0, 3, 3, 3, 2.99, 4, 4, 4, 4, -10,
                      -5, -5, -5, -5, -5, -10])
        midpoints, left_edges, right_edges = _local_maxima_1d(x)
        assert_equal(midpoints, np.array([2, 4, 8, 12, 18]))
        assert_equal(left_edges, np.array([2, 4, 7, 11, 16]))
        assert_equal(right_edges, np.array([2, 5, 9, 14, 20]))

    @pytest.mark.parametrize('x', [
        np.array([1., 0, 2]),
        np.array([3., 3, 0, 4, 4]),
        np.array([5., 5, 5, 0, 6, 6, 6]),
    ])
    def test_signal_edges(self, x):
        """Test if behavior on signal edges is correct."""
        for array in _local_maxima_1d(x):
            assert_equal(array, np.array([]))
            assert_(array.base is None)

    def test_exceptions(self):
        """Test input validation and raised exceptions."""
        with raises(ValueError, match="wrong number of dimensions"):
            _local_maxima_1d(np.ones((1, 1)))
        with raises(ValueError, match="expected 'float64_t'"):
            _local_maxima_1d(np.ones(1, dtype=int))
        with raises(TypeError, match="list"):
            _local_maxima_1d([1., 2.])
        with raises(TypeError, match="'x' must not be None"):
            _local_maxima_1d(None)


class TestRidgeLines(object):

    def test_empty(self):
        test_matr = np.zeros([20, 100])
        lines = _identify_ridge_lines(test_matr, 2*np.ones(20), 1)
        assert_(len(lines) == 0)

    def test_minimal(self):
        test_matr = np.zeros([20, 100])
        test_matr[0, 10] = 1
        lines = _identify_ridge_lines(test_matr, 2*np.ones(20), 1)
        assert_(len(lines) == 1)

        test_matr = np.zeros([20, 100])
        test_matr[0:2, 10] = 1
        lines = _identify_ridge_lines(test_matr, 2*np.ones(20), 1)
        assert_(len(lines) == 1)

    def test_single_pass(self):
        distances = [0, 1, 2, 5]
        gaps = [0, 1, 2, 0, 1]
        test_matr = np.zeros([20, 50]) + 1e-12
        length = 12
        line = _gen_ridge_line([0, 25], test_matr.shape, length, distances, gaps)
        test_matr[line[0], line[1]] = 1
        max_distances = max(distances)*np.ones(20)
        identified_lines = _identify_ridge_lines(test_matr, max_distances, max(gaps) + 1)
        assert_array_equal(identified_lines, [line])

    def test_single_bigdist(self):
        distances = [0, 1, 2, 5]
        gaps = [0, 1, 2, 4]
        test_matr = np.zeros([20, 50])
        length = 12
        line = _gen_ridge_line([0, 25], test_matr.shape, length, distances, gaps)
        test_matr[line[0], line[1]] = 1
        max_dist = 3
        max_distances = max_dist*np.ones(20)
        #This should get 2 lines, since the distance is too large
        identified_lines = _identify_ridge_lines(test_matr, max_distances, max(gaps) + 1)
        assert_(len(identified_lines) == 2)

        for iline in identified_lines:
            adists = np.diff(iline[1])
            np.testing.assert_array_less(np.abs(adists), max_dist)

            agaps = np.diff(iline[0])
            np.testing.assert_array_less(np.abs(agaps), max(gaps) + 0.1)

    def test_single_biggap(self):
        distances = [0, 1, 2, 5]
        max_gap = 3
        gaps = [0, 4, 2, 1]
        test_matr = np.zeros([20, 50])
        length = 12
        line = _gen_ridge_line([0, 25], test_matr.shape, length, distances, gaps)
        test_matr[line[0], line[1]] = 1
        max_dist = 6
        max_distances = max_dist*np.ones(20)
        #This should get 2 lines, since the gap is too large
        identified_lines = _identify_ridge_lines(test_matr, max_distances, max_gap)
        assert_(len(identified_lines) == 2)

        for iline in identified_lines:
            adists = np.diff(iline[1])
            np.testing.assert_array_less(np.abs(adists), max_dist)

            agaps = np.diff(iline[0])
            np.testing.assert_array_less(np.abs(agaps), max(gaps) + 0.1)

    def test_single_biggaps(self):
        distances = [0]
        max_gap = 1
        gaps = [3, 6]
        test_matr = np.zeros([50, 50])
        length = 30
        line = _gen_ridge_line([0, 25], test_matr.shape, length, distances, gaps)
        test_matr[line[0], line[1]] = 1
        max_dist = 1
        max_distances = max_dist*np.ones(50)
        #This should get 3 lines, since the gaps are too large
        identified_lines = _identify_ridge_lines(test_matr, max_distances, max_gap)
        assert_(len(identified_lines) == 3)

        for iline in identified_lines:
            adists = np.diff(iline[1])
            np.testing.assert_array_less(np.abs(adists), max_dist)

            agaps = np.diff(iline[0])
            np.testing.assert_array_less(np.abs(agaps), max(gaps) + 0.1)


class TestArgrel(object):

    def test_empty(self):
        # Regression test for gh-2832.
        # When there are no relative extrema, make sure that
        # the number of empty arrays returned matches the
        # dimension of the input.

        empty_array = np.array([], dtype=int)

        z1 = np.zeros(5)

        i = argrelmin(z1)
        assert_equal(len(i), 1)
        assert_array_equal(i[0], empty_array)

        z2 = np.zeros((3,5))

        row, col = argrelmin(z2, axis=0)
        assert_array_equal(row, empty_array)
        assert_array_equal(col, empty_array)

        row, col = argrelmin(z2, axis=1)
        assert_array_equal(row, empty_array)
        assert_array_equal(col, empty_array)

    def test_basic(self):
        # Note: the docstrings for the argrel{min,max,extrema} functions
        # do not give a guarantee of the order of the indices, so we'll
        # sort them before testing.

        x = np.array([[1, 2, 2, 3, 2],
                      [2, 1, 2, 2, 3],
                      [3, 2, 1, 2, 2],
                      [2, 3, 2, 1, 2],
                      [1, 2, 3, 2, 1]])

        row, col = argrelmax(x, axis=0)
        order = np.argsort(row)
        assert_equal(row[order], [1, 2, 3])
        assert_equal(col[order], [4, 0, 1])

        row, col = argrelmax(x, axis=1)
        order = np.argsort(row)
        assert_equal(row[order], [0, 3, 4])
        assert_equal(col[order], [3, 1, 2])

        row, col = argrelmin(x, axis=0)
        order = np.argsort(row)
        assert_equal(row[order], [1, 2, 3])
        assert_equal(col[order], [1, 2, 3])

        row, col = argrelmin(x, axis=1)
        order = np.argsort(row)
        assert_equal(row[order], [1, 2, 3])
        assert_equal(col[order], [1, 2, 3])

    def test_highorder(self):
        order = 2
        sigmas = [1.0, 2.0, 10.0, 5.0, 15.0]
        test_data, act_locs = _gen_gaussians_even(sigmas, 500)
        test_data[act_locs + order] = test_data[act_locs]*0.99999
        test_data[act_locs - order] = test_data[act_locs]*0.99999
        rel_max_locs = argrelmax(test_data, order=order, mode='clip')[0]

        assert_(len(rel_max_locs) == len(act_locs))
        assert_((rel_max_locs == act_locs).all())

    def test_2d_gaussians(self):
        sigmas = [1.0, 2.0, 10.0]
        test_data, act_locs = _gen_gaussians_even(sigmas, 100)
        rot_factor = 20
        rot_range = np.arange(0, len(test_data)) - rot_factor
        test_data_2 = np.vstack([test_data, test_data[rot_range]])
        rel_max_rows, rel_max_cols = argrelmax(test_data_2, axis=1, order=1)

        for rw in xrange(0, test_data_2.shape[0]):
            inds = (rel_max_rows == rw)

            assert_(len(rel_max_cols[inds]) == len(act_locs))
            assert_((act_locs == (rel_max_cols[inds] - rot_factor*rw)).all())


class TestPeakProminences(object):

    def test_empty(self):
        """
        Test if an empty array is returned if no peaks are provided.
        """
        out = peak_prominences([1, 2, 3], [])
        for arr, dtype in zip(out, [np.float64, np.intp, np.intp]):
            assert_(arr.size == 0)
            assert_(arr.dtype == dtype)

        out = peak_prominences([], [])
        for arr, dtype in zip(out, [np.float64, np.intp, np.intp]):
            assert_(arr.size == 0)
            assert_(arr.dtype == dtype)

    def test_basic(self):
        """
        Test if height of prominences is correctly calculated in signal with
        rising baseline (peak widths are 1 sample).
        """
        # Prepare basic signal
        x = np.array([-1, 1.2, 1.2, 1, 3.2, 1.3, 2.88, 2.1])