import pandas as pd
from tidyfit.data import (
    train_val_test_split,
    stratified_split,
    check_nulls,
    validate_schema,
    imbalance_report,
)


def _df():
    return pd.DataFrame(
        {"x1": [1, 2, 3, 4, 5, 6], "x2": [1, 1, 2, 2, 3, 3], "y": [0, 1, 0, 1, 0, 1]}
    )


def test_split_shapes():
    (xtr, ytr), (xv, yv), (xt, yt) = train_val_test_split(
        _df(), target="y", random_state=1
    )
    assert len(xtr) + len(xv) + len(xt) == 6
    assert len(ytr) == len(xtr)


def test_stratified_split_shapes():
    (xtr, ytr), (xv, yv), (xt, yt) = stratified_split(_df(), target="y", random_state=1)
    assert len(xtr) + len(xv) + len(xt) == 6
    assert ytr.value_counts().sum() == len(ytr)


def test_nulls_schema_imbalance():
    df = _df()
    df.loc[0, "x1"] = None
    assert check_nulls(df)["x1"] == 1
    ok, errs = validate_schema(_df(), {"x1": "int64", "y": "int64"})
    assert ok and not errs
    rep = imbalance_report(_df()["y"])
    assert rep["count"].sum() == 6
