import os
import msgpack
import openpyxl
from pathlib import Path
from datetime import datetime
import pandas as pd
import numpy as np
from utils.s3 import read_from_s3, read_dataframe_from_s3_parquet
from utils.logger import setup_logger
from dotenv import load_dotenv
from tqdm import tqdm
from backend.SearchEpc import SearchEpc
from etl.eligibility.Eligibility import Eligibility
from etl.eligibility.ha_15_32.app import prepare_model_data_row
from etl.epc.DataProcessor import DataProcessor
from etl.epc.settings import COLUMNS_TO_MERGE_ON
from backend.ml_models.api import ModelApi
from etl.solar.SolarPhotoSupply import SolarPhotoSupply
from recommendations.recommendation_utils import calculate_cavity_age
from recommendation_utils import convert_thickness_to_numeric

EPC_AUTH_TOKEN = os.getenv("EPC_AUTH_TOKEN")
ENV_FILE = Path(__file__).parent / "etl" / "eligibility" / "ha_15_32" / ".env"

logger = setup_logger()
load_dotenv(ENV_FILE)


def load_data():
    workbook = openpyxl.load_workbook(f'etl/eligibility/ha_15_32/HESTIA - HA 24 ASSET LIST.xlsx')
    sheet = workbook.active
    sheet_colnames = [cell.value for cell in sheet[1]]

    rows_data = []
    rows_colors = []
    for row in sheet.iter_rows(min_row=2, values_only=False):  # Assuming the first row is headers
        row_data = [cell.value for cell in row]  # This will get you the cell values
        row_color = row[0].fill.start_color.index if row[0].fill.start_color.index != '00000000' else None
        # row_color = COLOR_INDEX[row_color]
        rows_data.append(row_data)
        rows_colors.append(row_color)

    asset_list = pd.DataFrame(rows_data, columns=sheet_colnames)
    # Remove None columns
    asset_list = asset_list.iloc[:, 0:10]
    asset_list['row_color'] = rows_colors

    asset_list["row_colour_name"] = np.where(
        asset_list["row_color"] == "FFFF0000", "red",
        np.where(asset_list["row_color"] == "FF92D050", "green", "yellow")
    )

    asset_list["row_colour_code"] = np.where(
        asset_list["row_colour_name"] == "red", "does not meet criteria",
        np.where(asset_list["row_colour_name"] == "green", "identified potential eco", "maybe in the future")
    )

    # The third column is listed as "Address" but it's actually the postcode". We have two Address columns so we
    # change just the third
    asset_list.columns.values[2] = "Postcode"

    # Split up the address on commas, which is useful for matching later
    split_addresses = asset_list['Address'].str.split(',', expand=True)
    split_addresses.columns = ['temp', 'address2', 'address3', 'address4', 'address5', 'address6']

    asset_list = pd.concat([asset_list, split_addresses], axis=1)
    # There is no commas separating house number and address 1
    split_addresses2 = asset_list['temp'].str.split(' ', expand=True)
    split_addresses2.columns = ['HouseNo', 'part1', 'part2', "part3", "part4"]
    # We could re-concatenate but we only care about HouseNo for the moment
    asset_list = pd.concat([asset_list, split_addresses2[["HouseNo"]]], axis=1)

    # Read in surveys
    survey_workbook = openpyxl.load_workbook(f'etl/eligibility/ha_15_32/HESTIA - HA 24 ECO4 SURVEY LIST.xlsx')
    survey_sheet = survey_workbook.active

    survey_rows = []
    survey_colors = []

    for row in survey_sheet.iter_rows(min_row=2, values_only=False):  # Assuming the first row is headers
        row_data = [cell.value for cell in row]  # This will get you the cell values
        row_color = row[0].fill.start_color.index if row[0].fill.start_color.index != '00000000' else None
        # row_color = COLOR_INDEX[row_color]
        survey_rows.append(row_data)
        survey_colors.append(row_color)

    survey_list = pd.DataFrame(survey_rows, columns=[cell.value for cell in survey_sheet[1]])

    survey_list["row_colour"] = survey_colors
    survey_list["survey_key"] = ["survey_" + str(i) for i in range(0, len(survey_list))]
    # Tidy up the street/block name a bit
    survey_list["Street / Block Name"] = survey_list["Street / Block Name"].str.replace("/", ", ")
    survey_list["Street / Block Name"] = survey_list["Street / Block Name"].str.lower()
    survey_list["Street / Block Name"] = survey_list["Street / Block Name"].str.strip()

    survey_list["Street / Block Name"] = survey_list["Street / Block Name"].str.replace(
        "council house, nidds lane", "nidds lane"
    )
    survey_list["Street / Block Name"] = survey_list["Street / Block Name"].str.replace(
        "wirral avenue", "wirrall avenue"
    )
    survey_list["Street / Block Name"] = survey_list["Street / Block Name"].str.replace(
        "st ives road", "st. ives crescent"
    )
    survey_list["Street / Block Name"] = survey_list["Street / Block Name"].str.replace(
        "sundringham road", "sandringham road"
    )
    survey_list["Street / Block Name"] = survey_list["Street / Block Name"].str.replace(
        "milton avenue", "milton road"
    )
    survey_list["Street / Block Name"] = survey_list["Street / Block Name"].str.replace(
        "st ives crescent", "st. ives crescent"
    )
    survey_list["Street / Block Name"] = survey_list["Street / Block Name"].str.replace(
        "council house, waterbelly lane", "waterbelly lane"
    )
    # Generally remove "councile house, " from the start of the street name
    survey_list["Street / Block Name"] = survey_list["Street / Block Name"].str.replace(
        "council house, ", ""
    )
    survey_list["Street / Block Name"] = survey_list["Street / Block Name"].str.replace(
        "st. leodegars close", "st leodegars close"
    )

    # asset_list[asset_list["Address"].str.lower().str.contains("wirral")]["Address"]

    # Drop all None rows
    survey_list = survey_list[~pd.isnull(survey_list["Street / Block Name"])]
    survey_list["survey_key"] = ["survey_" + str(i) for i in range(0, len(survey_list))]

    matched = []
    for _, row in tqdm(survey_list.iterrows(), total=len(survey_list)):
        house_number = row["NO."]
        if isinstance(house_number, str):
            house_number = house_number.lower()

        # Filter on the first line of the address
        df = asset_list[asset_list["Address"].str.lower().str.contains(row["Street / Block Name"].lower())].copy()
        # df = df[df["Postcode"].str.lower().str.contains(row["Post Code"].lower())]
        df = df[df["Address"].str.lower().str.contains(str(house_number))]
        if df.shape[0] != 1:
            df = df[df["HouseNo"] == str(house_number)]
            if df.shape[0] != 1:
                df = df[df["Postcode"].str.lower().str.contains(row["Post Code"].lower())]
                if df.shape[0] != 1:
                    print(row["Street / Block Name"])
                    print(house_number)
                    print(row["Post Code"].lower())
                    raise ValueError("Investigate")

        matched.append(
            {
                "survey_key": row["survey_key"],
                "matched_address": df["Address"].values[0],
                "survey_house_no": row["NO."],
                "survey_street_name": row["Street / Block Name"],
                "survey_postcode": row["Post Code"],
                "survey_status": row["INSTALLED OR CANCELLED"]
            }
        )

    matched = pd.DataFrame(matched)
    matched["warmfront_identified"] = True

    # Combine asset list and surveys
    data = asset_list.merge(
        matched, how="left", left_on="Address", right_on="matched_address",
    )
    data["warmfront_identified"] = data["warmfront_identified"].fillna(False)

    return data, survey_list


def get_epc_data(data, cleaned, cleaning_data, created_at, photo_supply_lookup, floor_area_decile_thresholds):
    scoring_data = []
    results = []
    nodata = []

    property_type_lookup = {
        "01 HOUSE": "House",
        "02 FLAT": "Flat",
        "03 BUNGALOW": "Bungalow",
        "05 BEDSIT": "Flat",
        "04 MAISONETTE": "Maisonette",
        "01 HOUSE MID": "House",
        "10 PBUNGALOW": "Bungalow",
        "14 SFLAT": "Flat",
        "12 SBEDSIT": "Flat",
        "11 PFLAT": "Flat",
        "13 SBUNGALOW": "Bungalow",
        " 01 HOUSE MID": "House",
        "09 PBEDSIT": "Flat"
    }

    for _, property_meta in tqdm(data.iterrows(), total=len(data)):

        searcher = SearchEpc(
            address1=property_meta["HouseNo"],
            postcode=property_meta["Postcode"],
            auth_token=EPC_AUTH_TOKEN,
            os_api_key=None,
            full_address=property_meta["Address"]
        )
        searcher.ordnance_survey_client.property_type = property_type_lookup[property_meta["Property Type"]]
        searcher.find_property(skip_os=True)

        if searcher.newest_epc is None:
            nodata.append(property_meta)
            continue

        newest_epc = searcher.newest_epc
        older_epcs = searcher.older_epcs
        full_sap_epc = searcher.full_sap_epc
        # We also want to get the penultimate epc
        penultimate_epc, _ = searcher.filter_newest_epc(older_epcs)
        if not penultimate_epc:
            penultimate_epc = newest_epc

        eligibility = Eligibility(epc=newest_epc, cleaned=cleaned)
        eligibility.check_gbis_warmfront()
        eligibility.check_eco4_warmfront()

        if (not eligibility.eco4_warmfront["eligible"]) and (not eligibility.gbis_warmfront):
            eligibility = Eligibility(epc=penultimate_epc, cleaned=cleaned)
            eligibility.check_gbis_warmfront()
            eligibility.check_eco4_warmfront()
            # If this is the case, we need to update the older epcs
            # older_epcs = [
            #     x for x in older_epcs if x["lmk-key"] not in [newest_epc["lmk-key"], penultimate_epc["lmk-key"]]
            # ]
            # If this is the case, we need to update the older epcs
            # We don't update just to make data cleaning easier
            if penultimate_epc.get("estimated") is None:
                older_epcs = [x for x in searcher.data["rows"] if x["lmk-key"] != penultimate_epc["lmk-key"]]

        # Loft MUST be suitable
        cavity_age = None
        if (
            eligibility.walls["is_cavity_wall"] and
            eligibility.walls["is_filled_cavity"] and
            eligibility.loft["suitability"] and
            eligibility.eco4_warmfront["message"] == "Failed due to full cavity - check cavity age"
        ):
            # We check the age of the cavity and if it's particularly old, we flag it
            cavity_age = calculate_cavity_age(newest_epc, older_epcs, cleaned)

        # Full checks
        eligibility.check_gbis()
        eligibility.check_eco4()

        if eligibility.eco4_warmfront["eligible"]:
            if eligibility.epc["uprn"] in ["", None]:
                eligibility.epc["uprn"] = int(property_meta["row_id"].split("_")[1])

            scoring_dictionary = prepare_model_data_row(
                property_id=property_meta["row_id"],
                modelling_epc=eligibility.epc,
                cleaned=cleaned,
                cleaning_data=cleaning_data,
                created_at=created_at,
                old_data=older_epcs,
                full_sap_epc=full_sap_epc,
                photo_supply_lookup=photo_supply_lookup,
                floor_area_decile_thresholds=floor_area_decile_thresholds
            )
            scoring_data.extend(scoring_dictionary)

        results.append(
            {
                "row_id": property_meta["row_id"],
                "uprn": eligibility.epc["uprn"],
                "Address": property_meta["Address"],
                "Postcode": property_meta["Postcode"],
                "property_type": eligibility.epc["property-type"],
                "gbis_eligible": eligibility.gbis_warmfront,
                "eco4_eligible": eligibility.eco4_warmfront["eligible"],
                "eco4_message": eligibility.eco4_warmfront["message"],
                "sap": float(eligibility.epc["current-energy-efficiency"]),
                "gbis_eligible_future": eligibility.gbis["eligible"],
                "gbis_eligible_future_message": eligibility.gbis["message"],
                "eco4_eligible_future": eligibility.eco4["eligible"],
                "eco4_eligible_future_message": eligibility.eco4["message"],
                # Property components
                "roof": eligibility.roof["clean_description"],
                "walls": eligibility.walls["clean_description"],
                "cavity_type": eligibility.cavity["type"],
                "heating": eligibility.epc["mainheat-description"],
                "tenure": eligibility.tenure,
                "date_epc": eligibility.epc["lodgement-date"],
                "cavity_age": cavity_age,
                **eligibility.walls,
                **eligibility.roof,
            }
        )

    scoring_df = pd.DataFrame(scoring_data)

    # Perform the same cleaning as in the model - first clean number of room variables though
    scoring_df = DataProcessor.apply_averages_cleaning(
        data_to_clean=scoring_df,
        cleaning_data=cleaning_data,
        cols_to_merge_on=['PROPERTY_TYPE', 'BUILT_FORM', 'CONSTRUCTION_AGE_BAND', 'LOCAL_AUTHORITY'],
        colnames=["NUMBER_HABITABLE_ROOMS", "NUMBER_HEATED_ROOMS"],
    )

    scoring_df = DataProcessor.apply_averages_cleaning(
        data_to_clean=scoring_df,
        cleaning_data=cleaning_data,
        cols_to_merge_on=COLUMNS_TO_MERGE_ON + ["LOCAL_AUTHORITY"],
    ).drop(columns=["LOCAL_AUTHORITY"])

    scoring_df = DataProcessor.clean_missings_after_description_process(
        scoring_df,
        ignore_cols=[c for c in scoring_df.columns if ("thermal_transmittance" in c) or (
            "insulation_thickness" in c) or ("ENERGY_EFF" in c)]
    )

    scoring_df = DataProcessor.clean_efficiency_variables(scoring_df)
    scoring_df["UPRN"] = scoring_df["UPRN"].astype(int)

    model_api = ModelApi(portfolio_id="ha24-eligibility", timestamp=created_at)
    all_predictions = model_api.predict_all(
        df=scoring_df,
        bucket="retrofit-data-dev",
        prediction_buckets={
            "sap_change_predictions": "retrofit-sap-predictions-dev",
            "heat_demand_predictions": "retrofit-heat-predictions-dev",
            "carbon_change_predictions": "retrofit-carbon-predictions-dev"
        }
    )

    predictions = all_predictions["sap_change_predictions"].copy()

    results_df = pd.DataFrame(results)

    predictions = predictions.rename(columns={"property_id": "row_id"}).merge(
        results_df[["row_id", "sap"]], how="left", on="row_id"
    )
    predictions["sap_uplift"] = predictions["predictions"] - predictions["sap"]
    predictions = predictions.groupby("row_id")["sap_uplift"].sum().reset_index()

    results_df = results_df.merge(
        predictions[["sap_uplift", "row_id"]],
        how="left",
        on="row_id"
    )
    results_df["post_install_sap"] = results_df["sap"] + results_df["sap_uplift"]

    eligibility_assessment = []
    for _, row in results_df[results_df["eco4_eligible"] == True].iterrows():
        # The upgrade requirements are dependent on the current SAP

        # If the property is an F or G, it only needs to upgrade to an %
        if row["sap"] <= 38:
            if row["post_install_sap"] >= 57:
                eligibility_classification = "highest confidence"
            elif row["post_install_sap"] >= 55:
                eligibility_classification = "high confidence"
            elif row["post_install_sap"] >= 53:
                eligibility_classification = "medium confidence"
            else:
                eligibility_classification = "unlikely"
        else:

            if row["post_install_sap"] >= 71:
                eligibility_classification = "highest confidence"
            elif row["post_install_sap"] >= 69:
                eligibility_classification = "high confidence"
            elif row["post_install_sap"] >= 67:
                eligibility_classification = "medium confidence"
            else:
                eligibility_classification = "unlikely"

        eligibility_assessment.append(
            {
                "row_id": row["row_id"],
                "eligibility_classification": eligibility_classification
            }
        )

    eligibility_assessment = pd.DataFrame(eligibility_assessment)

    results_df = results_df.merge(
        eligibility_assessment, how="left", on="row_id"
    )
    return results_df, scoring_data, nodata


def analyse_results(results_df, data, survey_list):
    analysis_data = data[["row_id", "survey_key", "warmfront_identified"]].merge(
        results_df, how="left", on="row_id"
    ).merge(
        survey_list[["survey_key", survey_list.columns[0]]].rename(columns={survey_list.columns[0]: "funding_scheme"}),
        how="left", on="survey_key"
    )

    # NEW

    analysis_data["roof_insulation_thickness"] = np.where(
        pd.isnull(analysis_data["roof_insulation_thickness"]), None, analysis_data["roof_insulation_thickness"]
    )
    analysis_data["roof_insulation_thickness_numeric"] = analysis_data["roof_insulation_thickness"].apply(
        lambda x: convert_thickness_to_numeric(x, is_flat=False, is_pitched=True)
    )

    warmfront_sold_eco4 = analysis_data[
        (analysis_data["warmfront_identified"] == True) & (
            analysis_data["funding_scheme"].isin(["ECO4 A/W", "AFFORDABLE WARMTH"]))
        ]

    warmfront_sold_gbis = analysis_data[
        (analysis_data["warmfront_identified"] == True) & (
            analysis_data["funding_scheme"].isin(["ECO4 GBIS (ECO+)"]))
        ]
    # 1407

    additional_eco4_warmfront_not_sold = analysis_data[
        (analysis_data["eco4_eligible"] == True) & (analysis_data["warmfront_identified"] == False) & (
            analysis_data["roof_insulation_thickness_numeric"] <= 100)
        ]

    additional_gbis_warmfront_not_sold = analysis_data[
        (analysis_data["gbis_eligible"] == True) & (analysis_data["warmfront_identified"] == False) & (
            ~analysis_data["row_id"].isin(additional_eco4_warmfront_not_sold["row_id"].values)
        )
        ]

    additional_gbis_warmfront_not_sold["walls"].value_counts()
    analysis_data["walls"].value_counts()

    # END NEW

    all_identified_eco = analysis_data[
        (analysis_data["warmfront_identified"] & analysis_data["funding_scheme"].isin(
            ["ECO4 A/W"])) |
        (analysis_data["eco4_eligible"])
        ]

    all_identified_gbis = analysis_data[
        (analysis_data["warmfront_identified"] & analysis_data["funding_scheme"].isin(
            ["ECO4 GBIS (ECO+)"])) |
        (analysis_data["gbis_eligible"] & analysis_data["eco4_eligible"].isin([False, None]))
        ]

    warmfront_identified = analysis_data[analysis_data["warmfront_identified"]]

    # Of the ECO jobs, what proportion to we get right
    warmfront_identified_eco = warmfront_identified[
        warmfront_identified["funding_scheme"].isin(["ECO4 A/W", "AFFORDABLE WARMTH"])
    ]

    eco_success_rate = warmfront_identified_eco["eco4_eligible"].sum() / warmfront_identified_eco.shape[0]

    warmfront_identified_gbis = warmfront_identified[
        warmfront_identified["funding_scheme"].isin(["ECO4 GBIS (ECO+)"])
    ]

    # No gbis for this
    # gbis_success_rate = warmfront_identified_gbis["gbis_eligible"].sum() / warmfront_identified_gbis.shape[0]

    # Additional identified
    additional_identified_eco = analysis_data[
        (analysis_data["eco4_eligible"] == True) & (analysis_data["warmfront_identified"] == False)
        ]

    additional_identified_eco["eligibility_classification"].value_counts()

    additional_identified_gbis = analysis_data[
        (analysis_data["gbis_eligible"] == True) & (analysis_data["eco4_eligible"] == False) & (
            analysis_data["warmfront_identified"] == False
        )
        ].shape[0]
    # Future
    additional_identified_eco_future = analysis_data[
        (analysis_data["eco4_eligible_future"] == True) & (analysis_data["warmfront_identified"] == False)
        ].shape[0]
    additional_identified_gbis_future = analysis_data[
        (analysis_data["gbis_eligible_future"] == True) & (analysis_data["eco4_eligible_future"] == False) & (
            analysis_data["warmfront_identified"] == False
        )
        ].shape[0]


def app():
    data, survey_list = load_data()

    data["row_id"] = ["ha24_" + str(i) for i in range(0, len(data))]

    cleaned = read_from_s3(
        s3_file_name="cleaned_epc_data/cleaned.bson",
        bucket_name="retrofit-data-dev"
    )
    cleaned = msgpack.unpackb(cleaned, raw=False)

    cleaning_data = read_dataframe_from_s3_parquet(
        bucket_name="retrofit-data-dev", file_key="sap_change_model/cleaning_dataset.parquet",
    )

    created_at = datetime.now().isoformat()

    photo_supply_lookup, floor_area_decile_thresholds = SolarPhotoSupply.load(bucket="retrofit-data-dev")

    results_df, scoring_data, nodata = get_epc_data(
        data, cleaned, cleaning_data, created_at, photo_supply_lookup, floor_area_decile_thresholds
    )

    # Pickle results just in case
    # import pickle
    # with open("ha24_10_jan.pickle", "wb") as f:
    #     pickle.dump(
    #         {
    #             "scoring_data": scoring_data,
    #             "results": results_df,
    #             "nodata": nodata
    #         }, f
    #     )

    # Read in pickle
    # import pickle
    # with open("ha24_10_jan.pickle", "rb") as f:
    #     saved = pickle.load(f)
    # scoring_data = saved["scoring_data"]
    # results_df = saved["results"]
    # nodata = saved["nodata"]
