import unittest

from datarobot.helpers import partitioning_methods as pm
from datarobot.helpers.partitioning_methods import BasePartitioningMethod


class TestPartitionMethodDispatch(unittest.TestCase):
    def test_random_cv(self):
        partition_data = dict(cv_method="random", validation_type="CV", holdout_pct=20, reps=5)
        part = BasePartitioningMethod.from_data(partition_data)
        self.assertIsInstance(part, pm.RandomCV)

    def test_random_tvh(self):
        partition_data = dict(
            cv_method="random", validation_type="TVH", holdout_pct=20, validation_pct=16
        )
        part = BasePartitioningMethod.from_data(partition_data)
        self.assertIsInstance(part, pm.RandomTVH)

    def test_stratified_cv(self):
        partition_data = dict(cv_method="stratified", validation_type="CV", holdout_pct=20, reps=5)
        part = BasePartitioningMethod.from_data(partition_data)
        self.assertIsInstance(part, pm.StratifiedCV)

    def test_stratified_tvh(self):
        partition_data = dict(
            cv_method="stratified", validation_type="TVH", holdout_pct=20, validation_pct=16
        )
        part = BasePartitioningMethod.from_data(partition_data)
        self.assertIsInstance(part, pm.StratifiedTVH)

    def test_user_cv(self):
        partition_data = dict(
            cv_method="user",
            validation_type="CV",
            user_partition_col="partitioner",
            cv_holdout_level="holdout_level",
        )
        part = BasePartitioningMethod.from_data(partition_data)
        self.assertIsInstance(part, pm.UserCV)

    def test_user_tvh(self):
        partition_data = dict(
            cv_method="user",
            validation_type="TVH",
            user_partition_col="partitioner",
            training_level="training",
            validation_level="validation",
            holdout_level="holdout",
        )
        part = BasePartitioningMethod.from_data(partition_data)
        self.assertIsInstance(part, pm.UserTVH)

    def test_group_cv(self):
        partition_data = dict(
            cv_method="group",
            validation_type="CV",
            partition_key_cols=["groupby1", "groupby2"],
            holdout_pct=20,
            reps=5,
        )
        part = BasePartitioningMethod.from_data(partition_data)
        self.assertIsInstance(part, pm.GroupCV)

    def test_group_tvh(self):
        partition_data = dict(
            cv_method="group",
            validation_type="TVH",
            partition_key_cols=["groupby1", "groupby2"],
            holdout_pct=20,
            validation_pct=16,
        )
        part = BasePartitioningMethod.from_data(partition_data)
        self.assertIsInstance(part, pm.GroupTVH)

    def test_you_screwed_up(self):
        with self.assertRaises(ValueError):
            BasePartitioningMethod.from_data({"bananas": "and cream"})

    def test_none_in_none_out(self):
        part = BasePartitioningMethod.from_data(None)
        self.assertIsNone(part)
