import numpy as np

from sasdata.data import SasData
from sasdata.data_backing import Group
from sasdata.dataset_types import one_dim, two_dim
from sasdata.quantities.quantity import Quantity
from sasdata.quantities.units import per_angstrom, per_centimeter


def test_1d():
    q = [1, 2, 3, 4, 5]
    i = [5, 4, 3, 2, 1]

    q_quantity = Quantity(np.array(q), per_angstrom)
    i_quantity = Quantity(np.array(i), per_centimeter)

    data_contents = {
        'Q': q_quantity,
        'I': i_quantity
    }

    data = SasData('TestData', data_contents, one_dim, Group('root', {}), True)

    assert all(data.abscissae.value == np.array(q))
    assert all(data.ordinate.value == np.array(i))


def test_2d():
    # This could be autogenerated but I am hard coding to reduce the logic in
    # the test.
    qx = [1, 1, 1, 2, 2, 2, 3, 3, 3]
    qy = [1, 2, 3, 1, 2, 3, 1, 2, 3]
    i = [1, 2, 3]

    qx_quantity = Quantity(np.array(qx), per_angstrom)
    qy_quantity = Quantity(np.array(qy), per_angstrom)
    i_quantity = Quantity(np.array(i), per_centimeter)

    data_contents = {
        'Qx': qx_quantity,
        'Qy': qy_quantity,
        'I': i_quantity
    }

    data = SasData('TestData', data_contents, two_dim, Group('root', {}), True)

    assert all(data.ordinate.value == np.array(i))
    assert (data.abscissae.value == np.array([[1, 1], [1, 2], [1, 3], [2, 1], [2, 2], [2, 3], [3, 1], [3, 2], [3, 3]])).all().all()
