from datetime import datetime
import pandas as pd
import pickle
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):
        with open("recommendations/tests/test_data/cleaning_data.pkl", "rb") as f:
            data = pickle.load(f)

        return data

    @pytest.fixture()
    def cleaned(self):
        with open("recommendations/tests/test_data/cleaned.pkl", "rb") as f:
            df = pickle.load(f)

        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:
        """

        # We patch an old version of cleaned which is missing some attributes for 'mainheat-description'
        for x in cleaned['mainheat-description']:
            x["has_hot-water-only"] = False
            x["has_mineral_and_wood"] = False
            x["has_dual_fuel_appliance"] = False

        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

        recommender.recommend(has_cavity_or_loft_recommendations=False)

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

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