import pytest

from datar import stats
from datar.base import *
from datar import f
from datar.datasets import warpbreaks, state_division, state_region, airquality

from .conftest import assert_iterable_equal

def test_table():
    # https://www.rdocumentation.org/packages/base/versions/3.6.2/topics/table
    z = stats.rpois(100, 5)
    x = table(z)
    assert sum(x.values.flatten()) == 100

    #-----------------
    with data_context(warpbreaks) as _:
        tab = table(f.wool, f.tension)

    assert tab.columns.tolist() == ['H', 'L', 'M']
    assert tab.index.tolist() == ['A', 'B']
    assert_iterable_equal(tab.values.flatten(), [9] * 6)

    tab = table(warpbreaks.loc[:, ['wool', 'tension']])
    assert tab.columns.tolist() == ['H', 'L', 'M']
    assert tab.index.tolist() == ['A', 'B']
    assert_iterable_equal(tab.values.flatten(), [9] * 6)

    #-----------------
    tab = table(state_division, state_region)
    assert tab.loc['New England', 'Northeast'] == 6

    #-----------------
    with data_context(airquality) as _:
        qt = stats.quantile(f.Temp)
        ct = cut(f.Temp, qt)
        tab = table(ct, f.Month)

    assert tab.iloc[0,0] == 24

    #-----------------
    a = letters[:3]
    tab = table(a, sample(a))
    assert sum(tab.values.flatten()) == 3

    #-----------------
    tab = table(a, sample(a), dnn=['x', 'y'])
    assert tab.index.name == 'x'
    assert tab.columns.name == 'y'

    #-----------------
    a = c(NA, Inf, (1.0/(i+1) for i in range(3)))
    a = a * 10
    # tab = table(a)
    # assert_iterable_equal(tab.values.flatten(), [10] * 4)

    tab = table(a, exclude=None)
    assert_iterable_equal(tab.values.flatten(), [10] * 5)

    #------------------
    b = as_factor(rep(c("A","B","C"), 10))
    tab = table(b)
    assert tab.shape == (1, 3)
    assert_iterable_equal(tab.values.flatten(), [10] * 3)

    tab = table(b, exclude="B")
    assert tab.shape == (1, 2)
    assert_iterable_equal(tab.values.flatten(), [10] * 2)
    assert 'B' not in tab.columns

    #-------------------
    d = factor(rep(c("A","B","C"), 10), levels=c("A","B","C","D","E"))
    tab = table(d, exclude="B", dnn=['x'])
    assert_iterable_equal(tab.columns.to_list(), ["A", "C", "D", "E"])
    assert_iterable_equal(tab.values.flatten(), [10, 10, 0, 0])

    d2 = factor(rep(c("A","B","C"), 10), levels=c("A","B","C","D","E"))
    tab = table(d, d2, exclude="B")
    assert tab.shape == (4, 4)

    tab = table("abc", "cba", dnn='x')
    assert tab.shape == (3,3)
    assert sum(tab.values.flatten()) == 3

    with data_context(airquality) as _:
        tab = table(f.Ozone, f.Solar_R, exclude=None)
    assert '<NA>' in tab.columns
    assert '<NA>' in tab.index

def test_table_error():
    from datar.datasets import iris, warpbreaks
    with pytest.raises(ValueError):
        table(iris)
    with pytest.raises(ValueError):
        table(warpbreaks, iris)
    with pytest.raises(ValueError):
        table(warpbreaks.wool, iris)
    with pytest.raises(ValueError):
        table(iris.iloc[:, []])
    with pytest.raises(ValueError):
        table(iris.iloc[:, [1,2]], iris)
    with pytest.raises(ValueError):
        table(iris.iloc[:, [1]], iris, iris)
    with pytest.raises(ValueError):
        table(iris.iloc[:, [1]], iris.iloc[:, []])

[(12, 'datar.stats.rpois', 'stats.rpois', False, 'from datar import stats\n'), (35, 'datar.stats.quantile', 'stats.quantile', False, 'from datar import stats\n'), (92, 'pytest.raises', 'pytest.raises', False, 'import pytest\n'), (94, 'pytest.raises', 'pytest.raises', False, 'import pytest\n'), (96, 'pytest.raises', 'pytest.raises', False, 'import pytest\n'), (98, 'pytest.raises', 'pytest.raises', False, 'import pytest\n'), (100, 'pytest.raises', 'pytest.raises', False, 'import pytest\n'), (102, 'pytest.raises', 'pytest.raises', False, 'import pytest\n'), (104, 'pytest.raises', 'pytest.raises', False, 'import pytest\n')]