# Path: socceraction/spadl/schema.py
# class SPADLSchema(pa.SchemaModel):
#     """Definition of a SPADL dataframe."""
# 
#     game_id: Series[Any] = pa.Field()
#     original_event_id: Series[Any] = pa.Field(nullable=True)
#     action_id: Series[int] = pa.Field()
#     period_id: Series[int] = pa.Field(ge=1, le=5)
#     time_seconds: Series[float] = pa.Field(ge=0)
#     team_id: Series[Any] = pa.Field()
#     player_id: Series[Any] = pa.Field()
#     start_x: Series[float] = pa.Field(ge=0, le=spadlconfig.field_length)
#     start_y: Series[float] = pa.Field(ge=0, le=spadlconfig.field_width)
#     end_x: Series[float] = pa.Field(ge=0, le=spadlconfig.field_length)
#     end_y: Series[float] = pa.Field(ge=0, le=spadlconfig.field_width)
#     bodypart_id: Series[int] = pa.Field(isin=spadlconfig.bodyparts_df().bodypart_id)
#     bodypart_name: Optional[Series[str]] = pa.Field(isin=spadlconfig.bodyparts_df().bodypart_name)
#     type_id: Series[int] = pa.Field(isin=spadlconfig.actiontypes_df().type_id)
#     type_name: Optional[Series[str]] = pa.Field(isin=spadlconfig.actiontypes_df().type_name)
#     result_id: Series[int] = pa.Field(isin=spadlconfig.results_df().result_id)
#     result_name: Optional[Series[str]] = pa.Field(isin=spadlconfig.results_df().result_name)
# 
#     class Config:  # noqa: D106
#         strict = True
#         coerce = True
# 
# Path: socceraction/spadl/config.py
# def actiontypes_df() -> pd.DataFrame:
# def results_df() -> pd.DataFrame:
# def bodyparts_df() -> pd.DataFrame:


import json
import numpy as np
import pandas as pd
import pytest
import socceraction.spadl as spadl
import socceraction.xthreat as xt
from pathlib import Path
from pandera.typing import DataFrame, Series
from pytest_mock import MockerFixture
from sklearn.exceptions import NotFittedError
from socceraction.spadl import SPADLSchema
from socceraction.spadl.config import field_length, field_width

    Grid cells ares represented by 2D pitch coordinates. The (0,0) coordinate
    corresponds to the bottom left corner of the pitch. The 2D coordinates are
    mapped to a flat index. For a 2x2 grid, these flat indices are:
        0 1
        2 3
    """

    N = 2
    M = 2

    def test_get_cell_indexes(self) -> None:
        """It should map pitch coordinates to a 2D cell index."""
        x = Series[float]([0, field_length / 2 - 1, field_length])
        y = Series[float]([0, field_width / 2 + 1, field_width])
        xi, yi = xt._get_cell_indexes(x, y, self.N, self.M)
        pd.testing.assert_series_equal(xi, pd.Series([0, 0, 1]))
        pd.testing.assert_series_equal(yi, pd.Series([0, 1, 1]))

    def test_get_cell_indexes_out_of_bounds(self) -> None:
        """It should map out-of-bounds coordinates to the nearest cell index."""
        x = Series[float]([-10, field_length + 10])
        y = Series[float]([-10, field_width + 10])
        xi, yi = xt._get_cell_indexes(x, y, self.N, self.M)
        pd.testing.assert_series_equal(xi, pd.Series([0, 1]))
        pd.testing.assert_series_equal(yi, pd.Series([0, 1]))

    def test_get_flat_indexes(self) -> None:
        """It should map pitch coordinates to a flat index."""
        x = Series[float]([0, field_length / 2 - 1, field_length / 2 + 1, field_length])

# Path: socceraction/spadl/schema.py
# class SPADLSchema(pa.SchemaModel):
#     """Definition of a SPADL dataframe."""
# 
#     game_id: Series[Any] = pa.Field()
#     original_event_id: Series[Any] = pa.Field(nullable=True)
#     action_id: Series[int] = pa.Field()
#     period_id: Series[int] = pa.Field(ge=1, le=5)
#     time_seconds: Series[float] = pa.Field(ge=0)
#     team_id: Series[Any] = pa.Field()
#     player_id: Series[Any] = pa.Field()
#     start_x: Series[float] = pa.Field(ge=0, le=spadlconfig.field_length)
#     start_y: Series[float] = pa.Field(ge=0, le=spadlconfig.field_width)
#     end_x: Series[float] = pa.Field(ge=0, le=spadlconfig.field_length)
#     end_y: Series[float] = pa.Field(ge=0, le=spadlconfig.field_width)
#     bodypart_id: Series[int] = pa.Field(isin=spadlconfig.bodyparts_df().bodypart_id)
#     bodypart_name: Optional[Series[str]] = pa.Field(isin=spadlconfig.bodyparts_df().bodypart_name)
#     type_id: Series[int] = pa.Field(isin=spadlconfig.actiontypes_df().type_id)
#     type_name: Optional[Series[str]] = pa.Field(isin=spadlconfig.actiontypes_df().type_name)
#     result_id: Series[int] = pa.Field(isin=spadlconfig.results_df().result_id)
#     result_name: Optional[Series[str]] = pa.Field(isin=spadlconfig.results_df().result_name)
# 
#     class Config:  # noqa: D106
#         strict = True
#         coerce = True
# 
# Path: socceraction/spadl/config.py
# def actiontypes_df() -> pd.DataFrame:
# def results_df() -> pd.DataFrame:
# def bodyparts_df() -> pd.DataFrame:


# Path: tests/test_xthreat.py
import json
import numpy as np
import pandas as pd
import pytest
import socceraction.spadl as spadl
import socceraction.xthreat as xt
from pathlib import Path
from pandera.typing import DataFrame, Series
from pytest_mock import MockerFixture
from sklearn.exceptions import NotFittedError
from socceraction.spadl import SPADLSchema
from socceraction.spadl.config import field_length, field_width


    Grid cells ares represented by 2D pitch coordinates. The (0,0) coordinate
    corresponds to the bottom left corner of the pitch. The 2D coordinates are
    mapped to a flat index. For a 2x2 grid, these flat indices are:
        0 1
        2 3
    """

    N = 2
    M = 2

    def test_get_cell_indexes(self) -> None:
        """It should map pitch coordinates to a 2D cell index."""
        x = Series[float]([0, field_length / 2 - 1, field_length])
        y = Series[float]([0, field_width / 2 + 1, field_width])
        xi, yi = xt._get_cell_indexes(x, y, self.N, self.M)
        pd.testing.assert_series_equal(xi, pd.Series([0, 0, 1]))
        pd.testing.assert_series_equal(yi, pd.Series([0, 1, 1]))

    def test_get_cell_indexes_out_of_bounds(self) -> None:
        """It should map out-of-bounds coordinates to the nearest cell index."""
        x = Series[float]([-10, field_length + 10])
        y = Series[float]([-10, field_width + 10])
        xi, yi = xt._get_cell_indexes(x, y, self.N, self.M)
        pd.testing.assert_series_equal(xi, pd.Series([0, 1]))
        pd.testing.assert_series_equal(yi, pd.Series([0, 1]))

    def test_get_flat_indexes(self) -> None:
        """It should map pitch coordinates to a flat index."""
        x = Series[float]([0, field_length / 2 - 1, field_length / 2 + 1, field_length])

        y = Series[float]([0, field_width / 2 + 1, field_width / 2 - 1, field_width])