import pandas as pd
import pytest

from delta_barth.analysis import parse
from delta_barth.errors import FeaturesMissingError


def test_check_needed_features():
    target_features = set(("feat1", "feat2", "feat3"))
    data = pd.DataFrame(
        data=[[1, 2, 3, 4, 5]], columns=["feat1", "feat2", "feat3", "feat4", "feat5"]
    )
    parse.check_needed_features(data, target_features)
    data = pd.DataFrame(
        data=[[1, 2, 3, 4, 5]], columns=["featX", "feat2", "feat3", "feat4", "feat5"]
    )
    with pytest.raises(FeaturesMissingError):
        parse.check_needed_features(data, target_features)


def test_map_features_to_targets():
    feature_map = dict(feat1="feat10", feat2="feat20", feat5="feat50")
    data = pd.DataFrame(
        data=[[1, 2, 3, 4, 5]], columns=["feat1", "feat2", "feat3", "feat4", "feat5"]
    )
    data = parse.map_features_to_targets(data, feature_map)
    assert "feat10" in data.columns
    assert "feat20" in data.columns
    assert "feat50" in data.columns
    assert "feat3" in data.columns
    assert "feat4" in data.columns
    assert "feat1" not in data.columns
    assert "feat2" not in data.columns
    assert "feat5" not in data.columns