import numpy as np

import pypsr


def test_ami():
    a = np.arange(100)
    assert np.isclose(np.exp(pypsr.ami(a, a)), 10)
    assert np.isclose(pypsr.ami(a, np.ones_like(a)), 0)


def test_lagged_ami():
    lags, ami = pypsr.lagged_ami(np.arange(10), min_lag=0, max_lag=5)
    assert lags.dtype == 'int64'
    assert ami.dtype == 'float64'
    assert np.all(lags == np.arange(5))
    assert np.all(np.isclose(np.exp(ami), np.arange(10, 5, -1)))


def test_ami_unsqueezed_vector():
    a = np.arange(10)[:, np.newaxis]
    assert np.isclose(np.exp(pypsr.ami(a, a)), 10)
    assert np.isclose(pypsr.ami(a, np.ones_like(a)), 0)


def test_lagged_ami_unsqueezed_vector():
    a = np.arange(10)[:, np.newaxis]
    lags, ami = pypsr.lagged_ami(a)
    assert np.all(lags == np.arange(5))
    assert np.all(np.isclose(np.exp(ami), np.arange(10, 5, -1)))

@@ -11,16 +11,30 @@
 py as np
+%0Aimport pytest
 %0A%0Aimport
@@ -182,32 +182,162 @@
 _like(a)), 0)%0A%0A%0A
+def test_ami_different_length_signals():%0A    a = np.arange(10)%0A    with pytest.raises(ValueError):%0A        pypsr.ami(a, a%5B:-1%5D)%0A%0A%0A
 def test_lagged_

test error when calling ami with different length vectors