import pandas as pd
import pytest
import socceraction.atomic.spadl as atomicspadl
import socceraction.atomic.spadl.config as spadlconfig
import socceraction.atomic.vaep.labels as lab
from socceraction.atomic.vaep import AtomicVAEP
from socceraction.atomic.vaep import features as fs


@pytest.fixture
def test_goal_df() -> pd.DataFrame:
    return pd.DataFrame(
        [spadlconfig.actiontypes.index("shot"), spadlconfig.actiontypes.index("goal")],
        columns=["type_id"],
    )


def test_atomic_goal_from_shot_label(test_goal_df: pd.DataFrame) -> None:
    assert (lab.goal_from_shot(test_goal_df) == pd.DataFrame([[True], [False]], columns=["goal"]))[
        "goal"
    ].all()


@pytest.mark.e2e
def test_predict(sb_worldcup_data: pd.HDFStore) -> None:
    # Convert to atomic actions
    games = sb_worldcup_data["games"]
    atomic_actions = {
        game.game_id: atomicspadl.convert_to_atomic(
            sb_worldcup_data[f"actions/game_{game.game_id}"]
        )
        for game in games.itertuples()
    }
    # Test the vAEP framework on the StatsBomb World Cup data
    model = AtomicVAEP(nb_prev_actions=1)
    # comppute features and labels
    features = pd.concat(
        [
            model.compute_features(game, atomic_actions[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, atomic_actions[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)
    # rate a game
    game = games.iloc[-1]
    actions = atomic_actions[game.game_id]
    ratings = model.rate(game, actions)
    expected_rating_columns = {"offensive_value", "defensive_value", "vaep_value"}
    assert set(ratings.columns) == expected_rating_columns