import os

from pandas import DataFrame
from pandas.testing import assert_frame_equal

from kloppy import to_pandas, load_metrica_tracking_data, load_tracab_tracking_data, transform
from kloppy.domain import (
    Period, DatasetFlag, Point, AttackingDirection,
    TrackingDataset, PitchDimensions, Dimension,
    Orientation, Frame, EventDataset, PassEvent
)


class TestHelpers:
    def test_load_metrica_tracking_data(self):
        base_dir = os.path.dirname(__file__)
        dataset = load_metrica_tracking_data(
            f'{base_dir}/files/metrica_home.csv',
            f'{base_dir}/files/metrica_away.csv'
        )
        assert len(dataset.records) == 6
        assert len(dataset.periods) == 2

    def test_load_tracab_tracking_data(self):
        base_dir = os.path.dirname(__file__)
        dataset = load_tracab_tracking_data(
            f'{base_dir}/files/tracab_meta.xml',
            f'{base_dir}/files/tracab_raw.dat'
        )
        assert len(dataset.records) == 5  # only alive=True
        assert len(dataset.periods) == 2

    def _get_tracking_dataset(self):
        periods = [
            Period(id=1, start_timestamp=0.0, end_timestamp=10.0, attacking_direction=AttackingDirection.HOME_AWAY),
            Period(id=2, start_timestamp=15.0, end_timestamp=25.0, attacking_direction=AttackingDirection.AWAY_HOME)
        ]
        tracking_data = TrackingDataset(
            flags=~(DatasetFlag.BALL_OWNING_TEAM | DatasetFlag.BALL_STATE),
            pitch_dimensions=PitchDimensions(
                x_dim=Dimension(0, 100),
                y_dim=Dimension(-50, 50)
            ),
            orientation=Orientation.HOME_TEAM,
            frame_rate=25,
            records=[
                Frame(
                    frame_id=1,
                    timestamp=0.1,
                    ball_owning_team=None,
                    ball_state=None,
                    period=periods[0],

                    away_team_player_positions={},
                    home_team_player_positions={},
                    ball_position=Point(x=100, y=-50)
                ),
                Frame(
                    frame_id=2,
                    timestamp=0.2,
                    ball_owning_team=None,
                    ball_state=None,
                    period=periods[0],

                    away_team_player_positions={'1': Point(x=10, y=20)},
                    home_team_player_positions={'1': Point(x=15, y=35)},
                    ball_position=Point(x=0, y=50)
                )
            ],
            periods=periods
        )
        return tracking_data

    def test_transform(self):
        tracking_data = self._get_tracking_dataset()

        # orientation change AND dimension scale
        transformed_dataset = transform(
            tracking_data,
            to_orientation="AWAY_TEAM",
            to_pitch_dimensions=[[0, 1], [0, 1]]
        )

        assert transformed_dataset.frames[0].ball_position == Point(x=0, y=1)
        assert transformed_dataset.frames[1].ball_position == Point(x=1, y=0)

    def test_to_pandas(self):
        tracking_data = self._get_tracking_dataset()

        data_frame = to_pandas(tracking_data)

        expected_data_frame = DataFrame.from_dict({
            'period_id': {0: 1, 1: 1},
            'timestamp': {0: 0.1, 1: 0.2},
            'ball_state': {0: None, 1: None},
            'ball_owning_team': {0: None, 1: None},
            'ball_x': {0: 100, 1: 0},
            'ball_y': {0: -50, 1: 50},
            'player_home_1_x': {0: None, 1: 15.0},
            'player_home_1_y': {0: None, 1: 35.0},
            'player_away_1_x': {0: None, 1: 10.0},
            'player_away_1_y': {0: None, 1: 20.0}
        })

        assert_frame_equal(data_frame, expected_data_frame)
