import pandas as pd
import pytest
from socceraction.vaep import VAEP
from socceraction.vaep import features as fs


@pytest.fixture(scope="session")
def vaep_model(sb_worldcup_data: pd.HDFStore) -> VAEP:
    # Test the vAEP framework on the StatsBomb World Cup data
    model = VAEP(nb_prev_actions=1)
    # comppute features and labels
    games = sb_worldcup_data["games"]
    features = pd.concat(
        [
            model.compute_features(game, sb_worldcup_data[f"actions/game_{game.game_id}"])
            for game in games.iloc[:-1].itertuples()
        ]
    )
    expected_features = set(fs.feature_column_names(model.xfns, model.nb_prev_actions))
    assert set(features.columns) == expected_features
    labels = pd.concat(
        [
            model.compute_labels(game, sb_worldcup_data[f"actions/game_{game.game_id}"])
            for game in games.iloc[:-1].itertuples()
        ]
    )
    expected_labels = {"scores", "concedes"}
    assert set(labels.columns) == expected_labels
    assert len(features) == len(labels)
    # fit the model
    model.fit(features, labels)
    return model


@pytest.mark.e2e
def test_predict(sb_worldcup_data: pd.HDFStore, vaep_model: VAEP) -> None:
    games = sb_worldcup_data["games"]
    game = games.iloc[-1]
    actions = sb_worldcup_data[f"actions/game_{game.game_id}"]
    ratings = vaep_model.rate(game, actions)
    expected_rating_columns = {"offensive_value", "defensive_value", "vaep_value"}
    assert set(ratings.columns) == expected_rating_columns


@pytest.mark.e2e
def test_predict_with_missing_features(sb_worldcup_data: pd.HDFStore, vaep_model: VAEP) -> None:
    games = sb_worldcup_data["games"]
    game = games.iloc[-1]
    actions = sb_worldcup_data[f"actions/game_{game.game_id}"]
    X = vaep_model.compute_features(game, actions)
    del X["period_id_a0"]
    with pytest.raises(ValueError):
        vaep_model.rate(game, actions, X)