import unittest
import pandas as pd 
import numpy as np


from data import preprocess_data


class TestData(unittest.TestCase):

    def test_filter_events(self):


        df = pd.DataFrame(
            {
                "hillas_intensity": np.array([10.0, 100.0, 40, 1000.0]),
                "hillas_width": np.array([0, 0.5, 0.5, 1]),
                "morphology_n_pixels": np.array([1, np.nan, 5, 10]),
                "leakage_intensity_width_2": np.array([0.1, 0.1, 0.1, 0.1]),
                "true_energy": np.random.rand(4),
                "RandomForestRegressor_tel_energy": np.random.rand(4),
                "type": ["LST", "MST", "SST", "LST"],
            }
        )

        training_features = ["hillas_intensity", "hillas_width", "morphology_n_pixels", "leakage_intensity_width_2"]

        filtered_df = preprocess_data(df, training_features)


        # les 3 premières lignes sont supprimées car elles ne respectent pas les conditions de filtrage:
        assert len(filtered_df) == 1
        assert filtered_df['hillas_intensity'].values[0] == 3.0
        filtered_df = preprocess_data(df, training_features)

        # les 3 premières lignes sont supprimées car elles ne respectent pas les conditions de filtrage:
        assert len(filtered_df) == 1
        assert filtered_df['hillas_intensity'].values[0] == 3.0

    

if __name__ == '__main__':
    unittest.main()
