import numpy as np
import pandas as pd
from unittest import TestCase
from sklearn.model_selection import GroupKFold

import category_encoders as encoders
import tests.helpers as th
from category_encoders.wrapper import PolynomialWrapper, NestedCVWrapper


class TestMultiClassWrapper(TestCase):
    def test_invariance_to_data_types(self):
        x = np.array([
            ['a', 'b', 'c'],
            ['a', 'b', 'c'],
            ['b', 'b', 'c'],
            ['b', 'b', 'b'],
            ['b', 'b', 'b'],
            ['a', 'b', 'a'],
        ])
        y = [1, 2, 3, 3, 3, 3]
        wrapper = PolynomialWrapper(encoders.TargetEncoder())
        result = wrapper.fit_transform(x, y)
        th.verify_numeric(result)

        x = pd.DataFrame([
            ['a', 'b', 'c'],
            ['a', 'b', 'c'],
            ['b', 'b', 'c'],
            ['b', 'b', 'b'],
            ['b', 'b', 'b'],
            ['a', 'b', 'a'],
        ], columns=['f1', 'f2', 'f3'])
        y = ['bee', 'cat', 'dog', 'dog', 'dog', 'dog']
        wrapper = PolynomialWrapper(encoders.TargetEncoder())
        result2 = wrapper.fit_transform(x, y)

        self.assertTrue((result.values == result2.values).all(), 'The content should be the same regardless whether we pass Numpy or Pandas data type.')

    def test_transform_only_selected(self):
        x = pd.DataFrame([
            ['a', 'b', 'c'],
            ['a', 'a', 'c'],
            ['b', 'a', 'c'],
            ['b', 'c', 'b'],
            ['b', 'b', 'b'],
            ['a', 'b', 'a'],
        ], columns=['f1', 'f2', 'f3'])
        y = ['bee', 'cat', 'dog', 'dog', 'dog', 'dog']
        wrapper = PolynomialWrapper(encoders.LeaveOneOutEncoder(cols=['f2']))

        # combination fit() + transform()
        wrapper.fit(x, y)
        result = wrapper.transform(x, y)
        self.assertEqual(len(result.columns), 4, 'We expect 2 untouched features + f2 target encoded into 2 features')

        # directly fit_transform()
        wrapper = PolynomialWrapper(encoders.LeaveOneOutEncoder(cols=['f2']))
        result2 = wrapper.fit_transform(x, y)
        self.assertEqual(len(result2.columns), 4, 'We expect 2 untouched features + f2 target encoded into 2 features')

        pd.testing.assert_frame_equal(result, result2)

    def test_refit_stateless(self):
        # test that when the encoder is fitted multiple times no old state is carried
        x = pd.DataFrame([
            ['a', 'b', 'c'],
            ['a', 'b', 'c'],
            ['b', 'b', 'c'],
            ['b', 'b', 'b'],
            ['b', 'b', 'b'],
            ['a', 'b', 'a'],
        ], columns=['f1', 'f2', 'f3'])
        y1 = ['bee', 'cat', 'dog', 'dog', 'dog', 'dog']
        y2 = ['bee', 'cat', 'duck', 'duck', 'duck', 'duck']
        wrapper = PolynomialWrapper(encoders.TargetEncoder())
        result_first_fit = wrapper.fit_transform(x, y1)
        expected_categories_1 = {"cat", "dog"}  # 'bee' is dropped since first label is always dropped
        expected_categories_2 = {"cat", "duck"}
        self.assertEqual(set(wrapper.label_encoder.category_mapping[0]["mapping"].index), {"bee", "cat", "dog"})
        self.assertEqual(set(wrapper.feature_encoders.keys()), expected_categories_1)
        result_second_fit = wrapper.fit_transform(x, y2)
        self.assertEqual(set(wrapper.label_encoder.category_mapping[0]["mapping"].index), {"bee", "cat", "duck"})
        self.assertEqual(set(wrapper.feature_encoders.keys()), expected_categories_2)


class TestNestedCVWrapper(TestCase):
    def test_train_not_equal_to_valid(self):
        x = np.array([
            ['a', 'b', 'c'],
            ['a', 'b', 'c'],
            ['a', 'b', 'c'],
            ['a', 'b', 'c'],
            ['b', 'b', 'c'],
            ['b', 'b', 'c'],
            ['b', 'b', 'b'],
            ['b', 'b', 'b'],
            ['b', 'b', 'b'],
            ['b', 'b', 'b'],
            ['a', 'b', 'a'],
            ['a', 'b', 'a'],
        ])
        y = [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]
        wrapper = NestedCVWrapper(encoders.TargetEncoder(), cv=3)
        result_train, result_valid = wrapper.fit_transform(x, y, X_test=x)

        # We would expect result_train != result_valid since result_train has been generated using nested
        # folds and result_valid is generated by fitting the encoder on all of the x & y daya
        self.assertFalse(np.allclose(result_train, result_valid))


    def test_custom_cv(self):
        x = np.array([
            ['a', 'b', 'c'],
            ['a', 'b', 'c'],
            ['a', 'b', 'c'],
            ['a', 'b', 'c'],
            ['b', 'b', 'c'],
            ['b', 'b', 'c'],
            ['b', 'b', 'b'],
            ['b', 'b', 'b'],
            ['b', 'b', 'b'],
            ['b', 'b', 'b'],
            ['a', 'b', 'a'],
            ['a', 'b', 'a'],
        ])
        groups = [1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3]
        y = [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]
        gkfold = GroupKFold(n_splits=3)
        wrapper = NestedCVWrapper(encoders.TargetEncoder(), cv=gkfold)
        result_train, result_valid = wrapper.fit_transform(x, y, X_test=x, groups=groups)

        # We would expect result_train != result_valid since result_train has been generated using nested
        # folds and result_valid is generated by fitting the encoder on all of the x & y daya
        self.assertFalse(np.allclose(result_train, result_valid))
