import pytest
from utils.s3 import read_dataframe_from_s3_parquet
from etl.epc.Record import EPCRecord
from etl.epc.settings import DATA_ANOMALY_MATCHES
import random


class TestEpcRecord:

    @pytest.fixture()
    def cleaning_data(self):
        cleaning_data = read_dataframe_from_s3_parquet(
            bucket_name="retrofit-data-dev", file_key="sap_change_model/cleaning_dataset.parquet",
        )

        return cleaning_data

    @pytest.fixture()
    def epc_records_1(self):
        epc_records_1 = {
            'original_epc': {
                'low-energy-fixed-light-count': '', 'address': '139 School Road, Hall Green',
                'uprn-source': 'Energy Assessor', 'floor-height': '2.6', 'heating-cost-potential': '1138',
                'unheated-corridor-length': '', 'hot-water-cost-potential': '175',
                'construction-age-band': 'England and Wales: 1900-1929', 'potential-energy-rating': 'B',
                'mainheat-energy-eff': 'Good', 'windows-env-eff': 'Average', 'lighting-energy-eff': 'Very Good',
                'environment-impact-potential': '82', 'glazed-type': 'double glazing, unknown install date',
                'heating-cost-current': '2711', 'address3': '',
                'mainheatcont-description': 'Programmer, TRVs and bypass',
                'sheating-energy-eff': 'N/A', 'property-type': 'House', 'local-authority-label': 'Birmingham',
                'fixed-lighting-outlets-count': '11', 'energy-tariff': 'Single', 'mechanical-ventilation': 'natural',
                'hot-water-cost-current': '310', 'county': '', 'postcode': 'B28 8JF', 'solar-water-heating-flag': 'N',
                'constituency': 'E14000562', 'co2-emissions-potential': '2.0', 'number-heated-rooms': '4',
                'floor-description': 'Suspended, no insulation (assumed)', 'energy-consumption-potential': '107',
                'local-authority': 'E08000025', 'built-form': 'Semi-Detached', 'number-open-fireplaces': '0',
                'windows-description': 'Fully double glazed', 'glazed-area': 'Normal', 'inspection-date': '2023-07-05',
                'mains-gas-flag': 'Y', 'co2-emiss-curr-per-floor-area': '65', 'address1': '139 School Road',
                'heat-loss-corridor': '', 'flat-storey-count': '', 'constituency-label': 'Birmingham, Hall Green',
                'roof-energy-eff': 'Average', 'total-floor-area': '103.0', 'building-reference-number': '10004697322',
                'environment-impact-current': '43', 'co2-emissions-current': '6.7',
                'roof-description': 'Pitched, 100 mm loft insulation', 'floor-energy-eff': 'N/A',
                'number-habitable-rooms': '4', 'address2': 'Hall Green', 'hot-water-env-eff': 'Good',
                'posttown': 'BIRMINGHAM', 'mainheatc-energy-eff': 'Average', 'main-fuel': 'mains gas (not community)',
                'lighting-env-eff': 'Very Good', 'windows-energy-eff': 'Average', 'floor-env-eff': 'N/A',
                'sheating-env-eff': 'N/A', 'lighting-description': 'Low energy lighting in 82% of fixed outlets',
                'roof-env-eff': 'Average', 'walls-energy-eff': 'Very Poor', 'photo-supply': '0.0',
                'lighting-cost-potential': '182', 'mainheat-env-eff': 'Good', 'multi-glaze-proportion': '100',
                'main-heating-controls': '', 'lodgement-datetime': '2023-07-13 08:23:07', 'flat-top-storey': '',
                'current-energy-rating': 'E', 'secondheat-description': 'None', 'walls-env-eff': 'Very Poor',
                'transaction-type': 'rental', 'uprn': '100070505235', 'current-energy-efficiency': '51',
                'energy-consumption-current': '366', 'mainheat-description': 'Boiler and radiators, mains gas',
                'lighting-cost-current': '182', 'lodgement-date': '2023-07-13', 'extension-count': '0',
                'mainheatc-env-eff': 'Average',
                'lmk-key': 'c1d137711da433fb3cced74b1a6848da8bbc1159d076455d26d7b4668982601e',
                'wind-turbine-count': '0',
                'tenure': 'Rented (social)', 'floor-level': '', 'potential-energy-efficiency': '84',
                'hot-water-energy-eff': 'Good', 'low-energy-lighting': '82',
                'walls-description': 'Solid brick, as built, no insulation (assumed)',
                'hotwater-description': 'From main system'}, 'full_sap_epc': {}, 'old_data': []
        }
        return epc_records_1

    def test_clean_mechanical_ventilation(self, cleaning_data, epc_records_1):
        # We have an epc with Natural ventilation - the resulting epc should also have natural ventulation

        record = EPCRecord(cleaning_data=cleaning_data)
        record.prepared_epc = {
            "mechanical-ventilation": "natural"
        }
        record._clean_ventilation()

        assert record.prepared_epc["mechanical-ventilation"] == "natural"

        record2 = EPCRecord(cleaning_data=cleaning_data)
        record2.prepared_epc = {
            "mechanical-ventilation": ""
        }

        record2._clean_ventilation()

        assert record2.prepared_epc["mechanical-ventilation"] is None

        record3 = EPCRecord(cleaning_data=cleaning_data)
        record3.prepared_epc = {
            "mechanical-ventilation": None
        }

        record3._clean_ventilation()

        assert record3.prepared_epc["mechanical-ventilation"] is None

        record4 = EPCRecord(cleaning_data=cleaning_data)
        record4.prepared_epc = {
            "mechanical-ventilation": "INVALID"
        }

        record4._clean_ventilation()

        assert record4.prepared_epc["mechanical-ventilation"] is None

    def test_clean_energy_valid_values(self, cleaning_data, epc_records_1):
        record = EPCRecord(cleaning_data=cleaning_data)
        record.prepared_epc = {
            "energy-consumption-current": "200",
            "co2-emissions-current": "5.5"
        }
        record._clean_energy()

        assert record.prepared_epc["energy-consumption-current"] == 200.0
        assert record.prepared_epc["co2-emissions-current"] == 5.5

    def test_clean_energy_empty_values(self, cleaning_data):
        # We cannot have invalid values so this should raise an exception
        record = EPCRecord(cleaning_data=cleaning_data)
        record.prepared_epc = {
            "energy-consumption-current": "",
            "co2-emissions-current": ""
        }

        with pytest.raises(ValueError):
            record._clean_energy()

    def test_clean_built_form_valid_remap(self, cleaning_data, epc_records_1):
        record = EPCRecord(cleaning_data=cleaning_data)
        # Assuming "Semi" should be remapped to "Semi-Detached"
        record.prepared_epc = {
            "built-form": "Semi-Detached",
            "property-type": "Flat"  # Assuming this affects the remapping
        }
        record._clean_built_form()

        assert record.prepared_epc["built-form"] == "Semi-Detached"

    def test_clean_built_form_anomaly(self, cleaning_data, epc_records_1):
        record = EPCRecord(cleaning_data=cleaning_data)

        record.prepared_epc = {
            "built-form": "",
            "property-type": "Flat"
        }
        record._clean_built_form()

        assert record.prepared_epc["built-form"] == "End-Terrace"

    def test_clean_floor_area_valid(self, cleaning_data):
        record = EPCRecord(cleaning_data=cleaning_data)
        record.prepared_epc = {
            "total-floor-area": "120.5"
        }
        record._clean_floor_area()

        assert record.prepared_epc["total-floor-area"] == 120.5

    def test_clean_floor_area_empty(self, cleaning_data):
        record = EPCRecord(cleaning_data=cleaning_data)
        record.prepared_epc = {
            "total-floor-area": ""
        }
        # We have no known case of missing floor area
        with pytest.raises(ValueError):
            record._clean_floor_area()

    def test_clean_heat_loss_corridor_valid(self, cleaning_data):
        record = EPCRecord(cleaning_data=cleaning_data)
        record.prepared_epc = {
            "heat-loss-corridor": "unheated corridor",
            "unheated-corridor-length": ""
        }
        record._clean_heat_loss_corridor()

        assert record.prepared_epc["heat-loss-corridor"] == "unheated corridor"

        record = EPCRecord(cleaning_data=cleaning_data)
        record.prepared_epc = {
            "heat-loss-corridor": "unheated corridor",
            "unheated-corridor-length": None
        }
        record._clean_heat_loss_corridor()

        assert record.prepared_epc["heat-loss-corridor"] == "unheated corridor"
        assert record.prepared_epc["unheated-corridor-length"] is None

    def test_clean_heat_loss_corridor_anomaly(self, cleaning_data):
        record = EPCRecord(cleaning_data=cleaning_data)
        # Assuming "InvalidCorridor" is an anomaly
        record.prepared_epc = {
            "heat-loss-corridor": "InvalidCorridor",
            "unheated-corridor-length": ""
        }
        record._clean_heat_loss_corridor()

        assert record.prepared_epc["heat-loss-corridor"] == "no corridor"

    def test_clean_mains_gas_valid(self, cleaning_data):
        record = EPCRecord(cleaning_data=cleaning_data)
        record.prepared_epc = {
            "mains-gas-flag": "Y"
        }
        record._clean_mains_gas()

        assert record.prepared_epc["mains-gas-flag"] is True

    def test_clean_mains_gas_anomaly(self, cleaning_data):
        record = EPCRecord(cleaning_data=cleaning_data)
        record.prepared_epc = {
            "mains-gas-flag": "InvalidValue"
        }
        # It should always be Y or N or an anomally value
        with pytest.raises(KeyError):
            record._clean_mains_gas()

        record = EPCRecord(cleaning_data=cleaning_data)
        record.prepared_epc = {
            "mains-gas-flag": random.choice(list(DATA_ANOMALY_MATCHES))
        }
        record._clean_mains_gas()

        assert record.prepared_epc["mains-gas-flag"] is None

    def test_clean_solar_hot_water_valid(self, cleaning_data):
        record = EPCRecord(cleaning_data=cleaning_data)
        record.prepared_epc = {
            "solar-water-heating-flag": "Y"
        }
        record._clean_solar_hot_water()

        assert record.prepared_epc["solar-water-heating-flag"] == "Y"
        assert record.solar_water_heating_flag_bool is True

    def test_clean_solar_hot_water_empty(self, cleaning_data):
        record = EPCRecord(cleaning_data=cleaning_data)
        record.prepared_epc = {
            "solar-water-heating-flag": ""
        }
        record._clean_solar_hot_water()

        assert record.prepared_epc["solar-water-heating-flag"] == "N"
        assert record.solar_water_heating_flag_bool is False

    def test_clean_number_lighting_outlets_valid(self, cleaning_data, epc_records_1):
        record = EPCRecord(cleaning_data=cleaning_data, epc_records=epc_records_1)
        record.prepared_epc = {
            "fixed-lighting-outlets-count": "5"
        }
        record._clean_number_lighting_outlets()

        assert record.prepared_epc["fixed-lighting-outlets-count"] == 5.0

    def test_clean_number_lighting_outlets_empty(self, cleaning_data, epc_records_1):
        record = EPCRecord(cleaning_data=cleaning_data)
        record.run_mode = "newdata"
        record.prepared_epc = {
            "fixed-lighting-outlets-count": "",
            "property-type": "Flat",
            "built-form": "Semi-Detached",
            "construction-age-band": "England and Wales: 1900-1929",
            "local-authority": "E08000025",
            "number-habitable-rooms": "4",
            "number-heated-rooms": "4",
        }
        record.old_data = []
        record.full_sap_epc = []
        record._clean_number_lighting_outlets()

        assert record.prepared_epc["fixed-lighting-outlets-count"] == 8.0

    def test_clean_count_variables(self, cleaning_data):
        record = EPCRecord(cleaning_data=cleaning_data)

        record.prepared_epc = {
            "number-open-fireplaces": "1",
            "extension-count": None,
            "flat-storey-count": "",
            "number-habitable-rooms": "INVALID!",
        }

        record._clean_count_variables()

        assert record.prepared_epc["number-open-fireplaces"] == 1.0
        assert record.prepared_epc["extension-count"] == 0
        assert record.prepared_epc["flat-storey-count"] is None
        assert record.prepared_epc["number-habitable-rooms"] is None

    def test_clean_floor_level(self, cleaning_data):
        record = EPCRecord(cleaning_data=cleaning_data)

        record.prepared_epc = {
            "floor-level": "1",
        }

        record._clean_floor_level()

        assert record.prepared_epc["floor-level"] == 1.0

        record = EPCRecord(cleaning_data=cleaning_data)

        record.prepared_epc = {
            "floor-level": "",
        }

        record._clean_floor_level()

        assert record.prepared_epc["floor-level"] is None

        record = EPCRecord(cleaning_data=cleaning_data)

        record.prepared_epc = {
            "floor-level": None,
        }

        record._clean_floor_level()

        assert record.prepared_epc["floor-level"] is None

    def test_clean_solar_hot_water(self, cleaning_data):
        record = EPCRecord(cleaning_data=cleaning_data)

        record.prepared_epc = {
            "solar-water-heating-flag": "Y",
        }

        record._clean_solar_hot_water()

        assert record.prepared_epc["solar-water-heating-flag"] == "Y"
        assert record.solar_water_heating_flag_bool is True

        record = EPCRecord(cleaning_data=cleaning_data)

        record.prepared_epc = {
            "solar-water-heating-flag": "N",
        }

        record._clean_solar_hot_water()

        assert record.prepared_epc["solar-water-heating-flag"] == "N"
        assert record.solar_water_heating_flag_bool is False

        record = EPCRecord(cleaning_data=cleaning_data)

        record.prepared_epc = {
            "solar-water-heating-flag": "",
        }

        record._clean_solar_hot_water()

        assert record.prepared_epc["solar-water-heating-flag"] == "N"
        assert record.solar_water_heating_flag_bool is False

        record = EPCRecord(cleaning_data=cleaning_data)

        record.prepared_epc = {
            "solar-water-heating-flag": None,
        }

        record._clean_solar_hot_water()

        assert record.prepared_epc["solar-water-heating-flag"] == "N"
        assert record.solar_water_heating_flag_bool is False
