import numpy as np
from hdmf.common.table import VectorData
from nwbwidgets.utils.dynamictable import infer_categorical_columns
from nwbwidgets.utils.testing import dicts_exact_equal
from pynwb.core import DynamicTable


def test_infer_categorical_columns():
    data1 = np.array([1, 2, 2, 3, 1, 1, 3, 2, 3])
    data2 = np.array([3, 4, 2, 4, 3, 2, 2, 4, 4])

    vd1 = VectorData('Data1', 'vector data for creating a DynamicTable', data=data1)
    vd2 = VectorData('Data2', 'vector data for creating a DynamicTable', data=data2)
    vd = [vd1, vd2]

    dynamic_table = DynamicTable(name='test table', description='This is a test table',
                                 columns=vd, colnames=['Data1', 'Data2'])

    assert dicts_exact_equal(infer_categorical_columns(dynamic_table), {'Data1': np.array([1, 2, 3]),
                                                                        'Data2': np.array([2, 3, 4])})
