Model/backend/Funding.py

from enum import Enum
import pandas as pd
import numpy as np
from typing import List

from backend.app.plan.schemas import HousingType, WALL_INSULATION_MEASURES, ROOF_INSULATION_MEASURES


class EligibilityCaveats(Enum):
    TENANT_ON_BENEFITS_OR_LOW_INCOME = "tenant_on_benefits_or_low_income"
    INNOVATION_REQUIRED = "innovation_required"
    SOLAR_NEEDS_HEATING = "solar_needs_heating"
    NEEDS_INSULATION_TO_MINIMUM_STANDARDS = "needs_insulation_to_minimum_standards"


class Funding:
    """
    Handles eligibility and funding calculations for ECO4 & GBIS (PRS + Social Housing).
    """

    def __init__(
        self,
        tenure: str,  # 'Private' or 'Social'
        social_cavity_abs_rate: float,
        social_solid_abs_rate: float,
        private_cavity_abs_rate: float,
        private_solid_abs_rate: float,
        project_scores_matrix,
        partial_project_scores_matrix,
        whlg_eligible_postcodes
    ):
        self.tenure = tenure
        self.social_cavity_abs_rate = social_cavity_abs_rate
        self.social_solid_abs_rate = social_solid_abs_rate
        self.private_cavity_abs_rate = private_cavity_abs_rate
        self.private_solid_abs_rate = private_solid_abs_rate

        self.starting_sap_band = None
        self.ending_sap_band = None
        self.floor_area_band = None
        self.project_scores_matrix = project_scores_matrix
        self.partial_project_scores_matrix = partial_project_scores_matrix
        self.whlg_eligible_postcodes = whlg_eligible_postcodes

        self.eco4_eligible = False
        self.eco4_eligibility_caveats = []
        self.gbis_eligible = False
        self.gbis_eligibility_caveats = []

        # Funding calculation variables
        self.full_project_abs = None
        self.eco4_funding = None

        self.partial_project_abs = None

    # -----------------------
    # Utility Helpers
    # -----------------------

    @staticmethod
    def get_sap_band(sap_score_number):
        bands = [
            ("High_A", 96, float("inf")),
            ("Low_A", 92, 96),
            ("High_B", 86, 92),
            ("Low_B", 81, 86),
            ("High_C", 74.5, 81),
            ("Low_C", 69, 74.5),
            ("High_D", 61.5, 69),
            ("Low_D", 55, 61.5),
            ("High_E", 46.5, 55),
            ("Low_E", 39, 46.5),
            ("High_F", 29.5, 39),
            ("Low_F", 21, 29.5),
            ("High_G", 10.5, 21),
            ("Low_G", 1, 10.5),
        ]
        for band, lower, upper in bands:
            if lower <= sap_score_number < upper:
                return band
        return None

    @staticmethod
    def get_floor_area_band(floor_area):
        if floor_area <= 72:
            return "0-72"
        if floor_area <= 97:
            return "73-97"
        if floor_area <= 199:
            return "98-199"
        return "200"

    @staticmethod
    def _split_measures(measures: List[dict]):
        """
        Extracts measure types and flags innovation.
        measures: list of dicts like {"type": "solar_pv", "is_innovation": True}
        """
        measure_types = [m["type"] for m in measures]
        has_innovation = any(m.get("is_innovation", False) for m in measures)
        innovation_measures = [m["type"] for m in measures if m.get("is_innovation", False)]
        return measure_types, has_innovation, innovation_measures

    @staticmethod
    def _meets_upgrade_target(starting_sap: int, ending_sap: int) -> bool:
        """
        ECO4 Upgrade Requirement:
        - EPC E/D (SAP ≥ 39): must upgrade to EPC C (SAP ≥ 69)
        - EPC F/G (SAP < 39): must upgrade to EPC D (SAP ≥ 55)
        """
        if starting_sap >= 39 and ending_sap >= 69:
            return True
        if starting_sap < 39 and ending_sap >= 55:
            return True
        return False

    # -----------------------
    # Private Rented Sector
    # -----------------------

    def eco4_prs_eligibility(
        self, starting_sap: int, ending_sap: int, measure_types: List, has_solar: bool, solar_eligible: bool
    ):

        """
        ECO4 PRS eligibility:
        - EPC E–G
        - Must include SWI, FTCH, renewable heating, or DHC
        - Tenant must be on benefits (flagged)
        """
        meets_epc = starting_sap <= 54  # EPC E–G
        meets_upgrade_target = self._meets_upgrade_target(starting_sap, ending_sap)

        if not meets_epc or not meets_upgrade_target:
            self.eco4_eligible = False
            self.eco4_eligibility_caveats = []
            return

        if has_solar and not solar_eligible:
            self.eco4_eligible = False
            self.eco4_eligibility_caveats.append(EligibilityCaveats.SOLAR_NEEDS_HEATING)
            return

        has_swi = "internal_wall_insulation" in measure_types or "external_wall_insulation" in measure_types
        has_renewable = "air_source_heat_pump" in measure_types or "ground_source_heat_pump" in measure_types
        has_ftch = "first_time_central_heating" in measure_types
        has_dhc = "district_heating_connection" in measure_types

        if meets_upgrade_target and meets_epc and (
            has_swi or has_renewable or has_ftch or has_dhc or solar_eligible
        ):
            self.eco4_eligible = True
            self.eco4_eligibility_caveats.append(EligibilityCaveats.TENANT_ON_BENEFITS_OR_LOW_INCOME)
            return

        self.eco4_eligible = False
        self.eco4_eligibility_caveats = []

    def gbis_prs_eligibility(self, starting_sap: int, council_tax_band: str, measure_types: List):
        """
        GBIS PRS eligibility:
        - General route: Council Tax Band & EPC D–G
        - Low-income route: tenant on benefits (flagged)
        """
        gbis_measures = {
            "general": [
                # Cannot do CWI
                "internal_wall_insulation", "external_wall_insulation",
                "flat_roof_insulation", "suspended_floor_insulation",
                "room_roof_insulation", "solid_floor_insulation"
            ],
            "low_income": [
                "internal_wall_insulation", "external_wall_insulation",
                "flat_roof_insulation", "suspended_floor_insulation",
                "room_roof_insulation", "solid_floor_insulation",
                "cavity_wall_insulation", "loft_insulation"
            ]
        }

        meets_epc = starting_sap <= 69  # EPC D–G
        is_single_measure = len(measure_types) == 1

        if not is_single_measure or not meets_epc:
            self.gbis_eligible = False
            self.gbis_eligibility_caveats = []
            return

        # General route
        if meets_epc and council_tax_band in ["A", "B", "C", "D", "E"]:
            if any(m in gbis_measures["general"] for m in measure_types):
                self.gbis_eligible = True

        # Low-income route
        if meets_epc and any(m in gbis_measures["low_income"] for m in measure_types):
            self.gbis_eligible = True
            self.gbis_eligibility_caveats.append(EligibilityCaveats.TENANT_ON_BENEFITS_OR_LOW_INCOME)

    # -----------------------
    # Social Housing
    # -----------------------

    def eco4_sh_eligibility(
        self,
        starting_sap: int,
        ending_sap: int,
        has_innovation: bool,
        has_solar: bool,
        solar_eligible: bool
    ):
        """
        ECO4 Social Housing eligibility.
        - EPC E–G: eligible, no income check.
        - EPC D: innovation measure required.
          If solar PV is the innovation measure, must also install ASHP or HHRSH.
        """
        if has_solar and not solar_eligible:
            # The package contins solar PV but it doesn't meet the eligibility requirements
            return

        meets_epc = starting_sap <= 69
        meets_upgrade_target = self._meets_upgrade_target(starting_sap, ending_sap)

        if not meets_epc or not meets_upgrade_target:
            return

        # EPC D innovation rule
        if 55 <= starting_sap <= 68:  # EPC D

            # If we don't meet the innovation requirements, we're not eligible
            if not has_innovation:
                self.eco4_eligible = False
                self.eco4_eligibility_caveats.append(EligibilityCaveats.INNOVATION_REQUIRED)
                return

            self.eco4_eligible = True
            self.eco4_eligibility_caveats = []

        self.eco4_eligible = True
        self.eco4_eligibility_caveats = []

    def gbis_sh_eligibility(self, starting_sap: int, measure_types: List, has_innovation: bool):
        """
        GBIS Social Housing eligibility.
        - EPC E–G: single insulation measure
        - EPC D: single insulation, innovation measure
        """

        meets_epc = starting_sap <= 69  # EPC D–G
        is_single_measure = len(measure_types) == 1
        # Check if has a valid measure
        insulation_measures = [
            'internal_wall_insulation', 'external_wall_insulation', 'cavity_wall_insulation',
            'loft_insulation', 'flat_roof_insulation', 'room_roof_insulation',
            'suspended_floor_insulation', 'solid_floor_insulation',
        ]
        has_valid_measures = any(m in measure_types for m in insulation_measures)

        if not is_single_measure or not meets_epc or not has_valid_measures:
            self.gbis_eligible = False
            self.gbis_eligibility_caveats = []
            return

        if 55 <= starting_sap <= 68:  # EPC D
            # Since it's single measure if has_innovation is true, the single insulation measure
            # must be the innovation measure
            if not has_innovation:
                self.gbis_eligible = False
                self.gbis_eligibility_caveats.append(EligibilityCaveats.INNOVATION_REQUIRED)
                return

            # If we don't have an innovation measure, we're not eligible
            self.gbis_eligible = False
            self.gbis_eligibility_caveats.append(EligibilityCaveats.SOLAR_NEEDS_HEATING)
            return

        self.gbis_eligible = True

    # -----------------------
    # Score Lookup
    # -----------------------

    def calculate_full_project_abs(self):
        """Look up ABS score for full projects."""
        data = self.project_scores_matrix[
            (self.project_scores_matrix["Floor Area Segment"] == self.floor_area_band) &
            (self.project_scores_matrix["Starting Band"] == self.starting_sap_band) &
            (self.project_scores_matrix["Finishing Band"] == self.ending_sap_band)
            ]

        if data.empty:
            raise ValueError("Missing ABS rate, check the project scores matrix")

        return data["Cost Savings"].values[0]

    @staticmethod
    def get_starting_ending_uvalues(current_uvalue: float) -> tuple[str, str]:
        """
        Returns the closest starting U-value and appropriate ending U-value for solid wall insulation.

        - If current_uvalue is closest to 0.45, assume the improvement target is 0.21.
        - Otherwise, assume the target is 0.30.
        """
        possible_starting_u_values = [2.00, 1.70, 1.00, 0.60, 0.45]
        closest_starting = min(possible_starting_u_values, key=lambda x: abs(x - current_uvalue))
        ending_uvalue = 0.21 if closest_starting == 0.45 else 0.30

        return f"{closest_starting:.2f}", f"{ending_uvalue:.2f}"

    def calculate_partial_project_abs(
        self,
        measure_type: str,
        current_wall_uvalue: float = None,
        is_partial: bool = False,
        existing_li_thickness: float = None,
        # is_roof_insulated: bool = False
    ):
        """
        Calculate the partial project ABS score for a single measure.
        """
        df = self.partial_project_scores_matrix[
            (self.partial_project_scores_matrix["Total Floor Area Band"] == self.floor_area_band) &
            (self.partial_project_scores_matrix["Starting Band"] == self.starting_sap_band)
            ]

        if measure_type == "internal_wall_insulation":
            if current_wall_uvalue is None:
                raise ValueError("current_wall_uvalue is required for IWI")

            starting_str, ending_str = self.get_starting_ending_uvalues(current_wall_uvalue)
            measure_code = f"IWI_solid_{starting_str}_{ending_str}"
            pps = df[df["Measure_Type"] == measure_code]

            if pps.shape[0] != 1:
                raise ValueError(f"Invalid IWI category: {measure_code}")
            return pps.squeeze()["Cost Savings"]

        if measure_type == "external_wall_insulation":
            if current_wall_uvalue is None:
                raise ValueError("current_wall_uvalue is required for EWI")

            starting_str, ending_str = self.get_starting_ending_uvalues(current_wall_uvalue)
            measure_code = f"EWI_solid_{starting_str}_{ending_str}"
            pps = df[df["Measure_Type"] == measure_code]

            if pps.shape[0] != 1:
                raise ValueError(f"Invalid EWI category: {measure_code}")
            return pps.squeeze()["Cost Savings"]

        if measure_type == "cavity_wall_insulation":
            measure_code = "CWI_partial_fill" if is_partial else "CWI_0.033"
            pps = df[df["Measure_Type"] == measure_code]

            if pps.shape[0] != 1:
                raise ValueError(f"Invalid CWI category: {measure_code}")
            return pps.squeeze()["Cost Savings"]

        if measure_type == "loft_insulation":
            if existing_li_thickness is None:
                raise ValueError("existing_li_thickness is required for LI")

            measure_code = "LI_lessequal100" if existing_li_thickness <= 100 else "LI_greater100"
            pps = df[df["Measure_Type"] == measure_code]

            if pps.shape[0] != 1:
                raise ValueError(f"Invalid LI category: {measure_code}")
            return pps.squeeze()["Cost Savings"]

        if measure_type == "flat_roof_insulation":
            pps = df[df["Measure_Type"] == "FRI"]
            if pps.shape[0] != 1:
                raise ValueError("Invalid FRI category")
            return pps.squeeze()["Cost Savings"]

        if measure_type == "room_roof_insulation":
            # Use the more conservative score (unin is usually lower)
            # code = "RIRI_res_unin" if not is_roof_insulated else "RIRI_res_in"
            code = "RIRI_res_unin"
            pps = df[df["Measure_Type"] == code]
            if pps.shape[0] != 1:
                raise ValueError(f"Invalid RIRI category: {code}")
            return pps.squeeze()["Cost Savings"]

        if measure_type == "suspended_floor_insulation":
            pps = df[df["Measure_Type"] == "UFI"]
            if pps.shape[0] != 1:
                raise ValueError("Invalid UFI category")
            return pps.squeeze()["Cost Savings"]

        if measure_type == "solid_floor_insulation":
            pps = df[df["Measure_Type"] == "SFI"]
            if pps.shape[0] != 1:
                raise ValueError("Invalid SFI category")
            return pps.squeeze()["Cost Savings"]

        raise ValueError(f"Invalid measure type for partial project ABS calculation: {measure_type}")

    # -----------------------
    # Main Entry Point
    # -----------------------

    @staticmethod
    def check_solar_eligible_heating_system(mainheat_description, heating_control_description):
        """
        Checks if the main heating system is eligible for solar PV funding.
        :param mainheat_description: Describes the primary heating system
        :param heating_control_description: Heating controls associated to the primary heating system
        :return:
        """

        return (
            any(x in mainheat_description.lower() for x in [
                "air source heat pump", "ground source heat pump", "boiler and radiators, electric"
            ]) or (
                ("electric storage heaters" in mainheat_description)
                and (heating_control_description.lower() == "controls for high heat retention storage heaters")
            )
        )

    def check_solar_eligibility(
        self,
        measure_types,
        mainheat_description,
        heating_control_description,
        has_wall_insulation_recommendation: bool = False,
        has_roof_insulation_recommendation: bool = False,
    ):
        """
        Because of the various pre-requisites for solar, we have a self-contained function to check for
        eligibility

        Returns a tuple of booleans (has_solar, solar_eligible)
        """

        if "solar_pv" not in measure_types:
            return False, False

        # 1) We check if there is an eligible heating system in place
        has_eligibile_heating = self.check_solar_eligible_heating_system(
            mainheat_description, heating_control_description
        )

        if not has_eligibile_heating:
            # We check if there is a recommendation for an ASHP or HHRSH
            if ("air_source_heat_pump" not in measure_types) and (
                "high_heat_retention_storage_heater" not in measure_types):
                return True, False

        # 2) We check if there is a wall insulation measure for this property. If so, we make sure
        # we have a wall insulation recommendation in this package
        if has_wall_insulation_recommendation:
            # Make sure we have a wall insulation recommendation
            if not any(m in measure_types for m in WALL_INSULATION_MEASURES):
                return True, False

        # 3) We check if there is a roof insulation measure for this property. If so, we make sure
        # we have a roof insulation recommendation in this package
        if has_roof_insulation_recommendation:
            # Make sure we have a roof insulation recommendation
            if not any(m in measure_types for m in ROOF_INSULATION_MEASURES):
                return True, False

        return True, True

    def check_funding(
        self,
        measures: List[dict],
        starting_sap: int,
        ending_sap: int,
        floor_area: float,
        mainheat_description: str,
        heating_control_description: str,
        is_cavity: bool,
        current_wall_uvalue: float,
        is_partial: False,
        existing_li_thickness: float,
        council_tax_band: str = None,
        has_wall_insulation_recommendation: bool = False,
        has_roof_insulation_recommendation: bool = False,
    ):
        """
        Given a list of measures, check ECO4/GBIS eligibility.

        Because measures like solar PV are subject to the minimum insulation requirements and we can get
        exemptions on floor insulation recommendations, if has_wall_insulation_recommendation or
        has_roof_insulation_recommendation are true, we check that the measures package contain a wall or roof
        insulation measure otherwise solar PV isn't eligible
        """

        # Normalize measures
        measure_types, has_innovation, innovation_measures = self._split_measures(measures)

        # Determine if we have a solar eligible heating system
        has_solar, solar_eligible = self.check_solar_eligibility(
            measure_types,
            mainheat_description,
            heating_control_description,
            has_wall_insulation_recommendation,
            has_roof_insulation_recommendation,
        )

        # Track EPC bands and floor area
        self.starting_sap_band = self.get_sap_band(starting_sap)
        self.ending_sap_band = self.get_sap_band(ending_sap)
        self.floor_area_band = self.get_floor_area_band(floor_area)

        if self.tenure == "Private":
            # ECO4 PRS
            self.eco4_prs_eligibility(starting_sap, ending_sap, measure_types, has_solar, solar_eligible)
            # GBIS PRS
            self.gbis_prs_eligibility(starting_sap, council_tax_band or "", measure_types)

            if self.eco4_eligible:
                self.full_project_abs = self.calculate_full_project_abs()
                self.eco4_funding = self.full_project_abs * (
                    self.private_cavity_abs_rate if is_cavity else self.private_solid_abs_rate)


        elif self.tenure == "Social":
            # ECO4 Social
            self.eco4_sh_eligibility(starting_sap, ending_sap, has_innovation, has_solar, solar_eligible)

            # GBIS Social
            self.gbis_sh_eligibility(starting_sap, measure_types, has_innovation)

            if self.eco4_eligible:
                # Calculate the full project ABS for ECO4
                self.full_project_abs = self.calculate_full_project_abs()
                self.eco4_funding = self.full_project_abs * (
                    self.social_cavity_abs_rate if is_cavity else self.social_solid_abs_rate
                )

            if self.gbis_eligible:
                # Calculate the partial project score - this is dependent on the measure
                self.partial_project_abs = self.calculate_partial_project_abs(
                    measure_types[0], current_wall_uvalue, is_partial, existing_li_thickness,
                )


        else:
            raise NotImplementedError("Only 'Private' and 'Social' tenures are supported.")