import unittest
import pandas as pd
import numpy as np
from src.aiva.representative_component import RepresentativeComponent

class TestRepresentativeComponent(unittest.TestCase):

    def test_hysteresis_downscaling(self):
        # Create a sample time series
        data = [0, 1, 2, 1.5, 0.5, 2.5, 1.8, 0.8, 3, 0]
        time_index = pd.date_range(start='2023-01-01', periods=len(data), freq='D')
        time_series = pd.Series(data, index=time_index)

        # Instantiate the component with hysteresis thresholds
        component = RepresentativeComponent(high_threshold=2.0, low_threshold=1.0)

        # Process the time series
        downscaled_series = component.process_time_series(time_series)

        # Expected output based on the hysteresis logic
        expected_output = pd.Series([
            False, False, True, True, False, True, True, False, True, False
        ], index=time_index)

        # Assert that the downscaled series matches the expected output
        pd.testing.assert_series_equal(downscaled_series, expected_output)

    def test_initial_state(self):
        # Test that the initial state is correctly set.
        data = [0.5, 1.5, 0.5]
        time_index = pd.date_range(start='2023-01-01', periods=len(data), freq='D')
        time_series = pd.Series(data, index=time_index)

        component = RepresentativeComponent(high_threshold=2.0, low_threshold=1.0, initial_state=True)
        downscaled_series = component.process_time_series(time_series)

        expected_output = pd.Series([True, True, True], index=time_index)

        pd.testing.assert_series_equal(downscaled_series, expected_output)


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