import re
from functools import partial

import numpy as np
import pandas as pd
import pytest
from numpy import dtypes
from numpy.testing import assert_array_equal

from src.backend.utils import pandas_helpers as ph


class TestAddColIfNotExists:
    _sample_df = pd.DataFrame({"a": [1, 2, 3]})
    _test_col = "col"

    def test_add_col_if_not_exists_single_as_str(self):
        df = self._sample_df.copy()
        ph.add_col_if_not_exists(df, self._test_col)
        assert_array_equal(df[self._test_col].values, [np.nan] * 3)

    def test_add_col_if_not_exists_multiple_and_not_overrides(self):
        col_names = ["a", "col1", "col2"]
        df = self._sample_df.copy()
        ph.add_col_if_not_exists(df, col_names, default_value=True)
        assert set(df.columns.tolist()) == set(col_names)  # Ignore order
        for col_name in col_names:
            expected = [1, 2, 3] if col_name == "a" else [True] * 3
            assert_array_equal(df[col_name].values, expected)

    @pytest.mark.parametrize(
        "dtype,expected_dtype",
        [
            (None, dtypes.Float64DType),
            # We set dtype to be bool but all values are NaNs so it's inferred as float64
            (bool, dtypes.Float64DType),
        ],
    )
    def test_add_col_if_not_exists_dtype_all_nan(self, dtype, expected_dtype):
        df = self._sample_df.copy()
        ph.add_col_if_not_exists(df, self._test_col, dtype=dtype, default_value=np.nan)
        assert isinstance(df[self._test_col].dtype, expected_dtype)

    def test_add_col_if_not_exists_dtype_infer(self):
        df = self._sample_df.copy()
        ph.add_col_if_not_exists(df, self._test_col, dtype=bool, default_value=np.nan)
        df.at[0, self._test_col] = True
        assert df.at[0, self._test_col] is True


class TestExtractRows:
    MOCK_COL = "col"

    def test_extract_rows(self):
        """Ensure that rows fulfilling condition are extracted from the input
        DataFrame into a new one.
        """
        df = pd.DataFrame([{self.MOCK_COL: v} for v in range(5)])
        result = ph.extract_rows(df, df[self.MOCK_COL].ne(2))
        assert result[self.MOCK_COL].tolist() == [0, 1, 3, 4]
        assert df[self.MOCK_COL].tolist() == [2]

    def test_extract_rows_empty_input_preserves_schema(self):
        """Ensure that the schema of the input DataFrame is preserved when it's empty
        (i.e. the subset must have the same columns as the input).
        """
        df = pd.DataFrame(columns=[self.MOCK_COL])
        result = ph.extract_rows(df, df[self.MOCK_COL].ne(2))
        assert result.empty
        assert df.columns == result.columns


def test_group_into_set():
    mock_df = pd.DataFrame(
        [
            {"a": "x", "b": "y", "c": "ABC"},
            {"a": "x", "b": "y", "c": "DEF"},
            {"a": "x", "b": "y", "c": "GHI"},
            {"a": "q", "b": "q", "c": "GHI"},
        ]
    )

    result = ph.group_into_set(mock_df, "c")
    result = result.to_dict(orient="records")
    expected_df_dict = [
        {"a": "q", "b": "q", "c": {"GHI"}},
        {"a": "x", "b": "y", "c": {"ABC", "DEF", "GHI"}},
    ]
    assert result == expected_df_dict


class TestToSet:
    func = partial(ph.to_set, separator_regex=re.compile(r"\s*,\s*"))

    @pytest.mark.parametrize(
        "value,expected",
        [("ES, FR, IT", {"ES", "FR", "IT"}), ("ES,FR, IT", {"ES", "FR", "IT"})],
    )
    def test_to_set(self, value, expected):
        assert self.func(value) == expected

    def test_to_set_nan(self):
        assert self.func(np.nan) is np.nan

    def test_to_set_other(self):
        with pytest.raises(TypeError):
            self.func(1)