from datetime import datetime
import pandas as pd
import msgpack
from utils.s3 import read_dataframe_from_s3_parquet, read_from_s3
import pytest
from backend.Property import Property
from etl.epc.Record import EPCRecord
from etl.bill_savings.KwhData import KwhData
from recommendations.HeatingRecommender import HeatingRecommender
from recommendations.tests.test_data.heating_recommendations_data import testing_examples


class TestHeatingRecommendations:

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

    @pytest.fixture
    def cleaned(self):
        df = read_from_s3(
            s3_file_name="cleaned_epc_data/cleaned.bson",
            bucket_name="retrofit-data-dev"
        )

        df = msgpack.unpackb(df, raw=False)
        return df

    @pytest.fixture
    def kwh_client(self):
        client = KwhData(bucket="retrofit-data-dev", read_consumption_data=False)
        # We fix this pricing table for these tests
        client.retail_price_comparison = pd.DataFrame(
            [
                {
                    "Date": datetime.today().strftime("%Y-%m-%d"),
                    'Average standard variable tariff (Large legacy suppliers)': 1
                }
            ]
        )
        client.retail_price_comparison["Date"] = pd.to_datetime(client.retail_price_comparison["Date"])
        return client

    @pytest.mark.parametrize(
        "test_case",
        testing_examples
    )
    def test_recommend(self, test_case, cleaning_data, cleaned, kwh_client):
        """
        With this function, we test out multiple heating descriptions and check which recomendations
        we retrieve alongside them
        :return:
        """

        epc_records = {"original_epc": test_case["epc"].copy(), "full_sap_epc": {}, "old_data": []}

        epc_record = EPCRecord(
            epc_records=epc_records,
            run_mode="newdata",
            cleaning_data=cleaning_data
        )

        p = Property(
            id=0,
            postcode=test_case["epc"]["postcode"],
            address=test_case["epc"]["address"],
            epc_record=epc_record,
            energy_assessment={
                "condition": {},
                "energy_assessment_is_newer": False
            }
        )

        # For these tests, this can be fixed
        kwh_predictions = {
            "heating_kwh_predictions": pd.DataFrame(
                [
                    {"id": p.uprn, "predictions": 12000}
                ]
            ),
            "hotwater_kwh_predictions": pd.DataFrame(
                [
                    {"id": p.uprn, "predictions": 3000}
                ]
            ),
        }

        p.set_features(cleaned=cleaned, kwh_client=kwh_client, kwh_predictions=kwh_predictions)

        recommender = HeatingRecommender(property_instance=p)
        # Check they're empty
        assert not recommender.heating_recommendations
        assert not recommender.heating_control_recommendations

        recommender.recommend(has_cavity_or_loft_recommendations=False)

        assert len(recommender.heating_recommendations) == len(test_case["heating_measure_types"])
        assert (
            len(recommender.heating_control_recommendations) ==
            len(test_case["heating_controls_measure_types"])
        )

        # Check the exact measure types
        assert (
            {x["measure_type"] for x in recommender.heating_recommendations} ==
            set(test_case["heating_measure_types"])
        )

        assert (
            {x["measure_type"] for x in recommender.heating_control_recommendations} ==
            set(test_case["heating_controls_measure_types"])
        )