import unittest
import pandas as pd
from etl.solar.SolarPhotoSupply import SolarPhotoSupply


class TestSolarPhotoSupply(unittest.TestCase):

    def setUp(self):
        # Mock data for photo_supply_lookup and floor_area_decile_thresholds
        self.photo_supply_lookup = pd.DataFrame({
            "tenure": ["leasehold", "freehold"],
            "built_form": ["detached", "semi-detached"],
            "property_type": ["house", "flat"],
            "construction_age_band": ["pre-1900", "1900-1929"],
            "is_flat": [False, True],
            "is_pitched": [True, False],
            "is_roof_room": [False, True],
            "floor_area_decile": [0, 1],
            "photo_supply": [100, 200]
        })

        self.floor_area_decile_thresholds = pd.DataFrame({
            "floor_area_decile_thresholds": [50, 100]
        })

        self.solar_photo_supply = SolarPhotoSupply([], {})

    def test_correct_filtering(self):
        result = self.solar_photo_supply.filter_photo_supply_lookup(
            self.photo_supply_lookup,
            self.floor_area_decile_thresholds,
            "leasehold",
            "detached",
            "house",
            "pre-1900",
            False,
            True,
            False,
            45
        )
        self.assertEqual(len(result), 1)
        self.assertEqual(result.iloc[0]["photo_supply"], 100)

    def test_no_matches(self):
        with self.assertRaises(ValueError):
            self.solar_photo_supply.filter_photo_supply_lookup(
                self.photo_supply_lookup,
                self.floor_area_decile_thresholds,
                "leasehold",
                "unknown",
                "house",
                "pre-1900",
                False,
                True,
                False,
                45
            )

    def test_floor_area_decile_matching(self):
        result = self.solar_photo_supply.filter_photo_supply_lookup(
            self.photo_supply_lookup,
            self.floor_area_decile_thresholds,
            "freehold",
            "semi-detached",
            "flat",
            "1900-1929",
            True,
            False,
            True,
            60
        )
        self.assertEqual(len(result), 1)
        self.assertEqual(result.iloc[0]["photo_supply"], 200)

    def test_invalid_parameters(self):
        with self.assertRaises(AttributeError):
            self.solar_photo_supply.filter_photo_supply_lookup(
                self.photo_supply_lookup,
                self.floor_area_decile_thresholds,
                123,  # Invalid type for tenure
                "detached",
                "house",
                "pre-1900",
                False,
                True,
                False,
                45
            )

    def test_classify_floor_area(self):
        # Setup
        thresholds = [10, 20, 30, 40, 50]
        solar_photo_supply = SolarPhotoSupply([], {})

        # Test Case 1: Valid floor area
        floor_area = 25
        expected_decile = 2
        result = solar_photo_supply.classify_floor_area(floor_area, thresholds)
        self.assertEqual(result, expected_decile, "Decile classification did not match expected result")

        # Test Case 2: Out of range floor area
        floor_area = 60
        expected_decile = len(thresholds)
        result = solar_photo_supply.classify_floor_area(floor_area, thresholds)
        self.assertEqual(result, expected_decile, "Decile classification for out of range value is incorrect")


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