import unittest

from features.proximity import *


class ProximityFeatures(unittest.TestCase):
    df_proximity = pd.DataFrame()
    df_proximity_recoded = pd.DataFrame()
    df_proximity_features = pd.DataFrame()

    @classmethod
    def setUpClass(cls) -> None:
        cls.df_proximity = pd.read_csv("../data/example_proximity.csv")
        cls.df_proximity["participant_id"] = 99

    def test_recode_proximity(self):
        self.df_proximity_recoded = recode_proximity(self.df_proximity)
        self.assertIn("bool_prox_near", self.df_proximity_recoded)
        # Is the recoded column present?
        self.assertIn(True, self.df_proximity_recoded.bool_prox_near)
        # Are there "near" values in the data?
        self.assertIn(False, self.df_proximity_recoded.bool_prox_near)
        # Are there "far" values in the data?

    def test_count_proximity(self):
        self.df_proximity_recoded = recode_proximity(self.df_proximity)
        self.df_proximity_features = count_proximity(self.df_proximity_recoded)
        print(self.df_proximity_features.columns)
        self.assertCountEqual(
            self.df_proximity_features.columns.to_list(), FEATURES_PROXIMITY
        )