import numpy as np

import numpy.testing as npt

import pytest

from stumpy.chains import allc, atsc

test_data = [

(

np.array([47, 32, 1, 22, 2, 58, 3, 36, 4, -5, 5, 40], dtype=np.int64),

np.array([11, 7, 4, 7, 6, 11, 8, 11, 10, 10, 11, -1], dtype=np.int64),

np.array([-1, 0, 1, 1, 2, 0, 4, 1, 6, 2, 8, 7], dtype=np.int64),

)

]

@pytest.mark.parametrize("Value, IR, IL", test_data)

def test_atsc(Value, IR, IL):

j = 2

ref = np.array([2, 4, 6, 8, 10], np.int64)

comp = atsc(IL, IR, j)

npt.assert_equal(ref, comp)

@pytest.mark.parametrize("Value, IR, IL", test_data)

def test_allc(Value, IR, IL):

S_ref = [

np.array([1, 7, 11], dtype=np.int64),

np.array([0], dtype=np.int64),

np.array([3], dtype=np.int64),

np.array([9], dtype=np.int64),

np.array([2, 4, 6, 8, 10], dtype=np.int64),

np.array([5], dtype=np.int64),

]

C_ref = np.array([2, 4, 6, 8, 10], dtype=np.int64)

S_comp, C_comp = allc(IL, IR)

S_ref = sorted(S_ref, key=lambda x: (len(x), list(x)))

S_comp = sorted(S_comp, key=lambda x: (len(x), list(x)))

npt.assert_equal(S_ref, S_comp)

npt.assert_equal(C_ref, C_comp)