Model/backend/app/plan/schemas.py

from pydantic import BaseModel, Field, BeforeValidator, model_validator
from typing import Annotated, List, Optional, Literal

# Example constants for validation
TYPICAL_MEASURE_TYPES = [
    "wall_insulation", "roof_insulation", "ventilation", "floor_insulation",
    "windows", "fireplace", "heating", "hot_water", "low_energy_lighting",
    "secondary_heating", "solar_pv"
]

WALL_INSULATION_MEASURES = ["internal_wall_insulation", "external_wall_insulation", "cavity_wall_insulation"]
ROOF_INSULATION_MEASURES = [
    "loft_insulation", "flat_roof_insulation", "room_roof_insulation", "sloping_ceiling_insulation"
]
WALL_INSULATION_WITH_VENTILATION_MEASURES = [
    "internal_wall_insulation+mechanical_ventilation", "external_wall_insulation+mechanical_ventilation",
    "cavity_wall_insulation+mechanical_ventilation"
]

# Both all and roof insulaiton measures are eligible for ECO4. These are the remaining fabric and heating measures
# This is based on th measures we have recommendations for
ECO4_ELIGIBILE_FABRIC_MEASURES = [
    "suspended_floor_insulation", "solid_floor_insulation", "double_glazing", "secondary_glazing"
]
ECO4_ELIGIBLE_HEATING_MEASURES = [
    "boiler_upgrade", "high_heat_retention_storage_heaters", "air_source_heat_pump", "solar_pv"
]

SPECIFIC_MEASURES = (
    WALL_INSULATION_MEASURES + ROOF_INSULATION_MEASURES + ECO4_ELIGIBILE_FABRIC_MEASURES +
    ECO4_ELIGIBLE_HEATING_MEASURES + [
        "secondary_heating", "ventilation", "low_energy_lighting", "fireplace",
        "hot_water_tank_insulation",
        "cylinder_thermostat"
    ]
)

INSULATION_MEASURES = [
    "internal_wall_insulation", "external_wall_insulation", "cavity_wall_insulation",
    "loft_insulation", "flat_roof_insulation", "room_roof_insulation", "sloping_ceiling_insulation",
    "suspended_floor_insulation", "solid_floor_insulation",
]

NON_INVASIVE_SPECIFIC_MEASURES = [
    "trickle_vents", "draught_proofing", "mixed_glazing", "cavity_extract_and_refill",
    "extension_cavity_wall_insulation"
]

# This allows us to extend high level categories for measures such as "wall_insulation" to the specific measures
# such as "external_wall_insulation", "internal_wall_insulation", "cavity_wall_insulation"
MEASURE_MAP = {
    "wall_insulation": [
        "internal_wall_insulation", "external_wall_insulation", "cavity_wall_insulation",
    ],
    "roof_insulation": [
        "loft_insulation", "flat_roof_insulation", "room_roof_insulation", "sloping_ceiling_insulation"
    ],
    "floor_insulation": ["suspended_floor_insulation", "solid_floor_insulation"],
    "heating": ["boiler_upgrade", "high_heat_retention_storage_heaters", "air_source_heat_pump"],
    "windows": ["double_glazing", "secondary_glazing"],
    "heating_controls": ["roomstat_programmer_trvs", "time_temperature_zone_control"]
}

VALID_GOALS = ["Increasing EPC", "Energy Savings", "Reducing CO2 emissions"]
VALID_HOUSING_TYPES = ["Social", "Private"]
VALID_EVENT_TYPES = ["remote_assessment", "eco_project"]


# Define the validation function for inclusions/exclusions
def check_inclusion_or_exclusion(value: str) -> str:
    if value not in TYPICAL_MEASURE_TYPES + SPECIFIC_MEASURES + NON_INVASIVE_SPECIFIC_MEASURES:
        raise ValueError(f"{value} is not an allowed inclusion")
    return value


def check_goals(value: str) -> str:
    assert value in VALID_GOALS, f"{value} is not a valid goal"
    return value


def check_housing_type(value: str) -> str:
    assert value in VALID_HOUSING_TYPES, f"{value} is not a valid housing type"
    return value


def check_event_type(value: str) -> str:
    assert value in VALID_EVENT_TYPES, f"{value} is not a valid event type"
    return value


# Use Annotated with BeforeValidator for each list item validation
InclusionOrExclusionItem = Annotated[str, BeforeValidator(check_inclusion_or_exclusion)]
Goal = Annotated[str, BeforeValidator(check_goals)]
HousingType = Annotated[str, BeforeValidator(check_housing_type)]
EventType = Annotated[str, BeforeValidator(check_event_type)]


class PlanTriggerRequest(BaseModel):
    budget: Optional[float] = None
    goal: Goal
    housing_type: HousingType
    goal_value: Optional[str] = None
    portfolio_id: int
    trigger_file_path: str
    already_installed_file_path: Optional[str] = None
    patches_file_path: Optional[str] = None
    non_invasive_recommendations_file_path: Optional[str] = None
    valuation_file_path: Optional[str] = None
    exclusions: Optional[List[InclusionOrExclusionItem]] = Field(default=None, min_length=0)
    inclusions: Optional[List[InclusionOrExclusionItem]] = Field(default=None, min_length=0)
    # This is a list of measures that we want to be included, if they are options
    # Default to empty
    required_measures: Optional[List[InclusionOrExclusionItem]] = Field(default=[], min_length=0)

    scenario_name: Optional[str] = ""
    scenario_id: Optional[str | int] = None  # Used to utilise and existing scenario for a engine run
    multi_plan: Optional[bool] = False
    default_u_values: Optional[bool] = True

    ashp_cop: Optional[float] = 2.8

    # When performing a remote assessment, if this has been set, it will allow the engine to
    # pull data from the find my epc website, to utilise as part of a remote assessment
    event_type: Optional[Literal["remote_assessment", "eco_project"]] = None

    # If true, before optimising the engine will select a slightly larger package, to account for the SAP 10 causing
    # scores to drop by a few points
    simulate_sap_10: Optional[bool] = False

    # Add in optional fields which describe the format of the asset list being used

    file_type: Optional[Literal["csv", "xlsx"]] = None
    file_format: Optional[Literal["domna_asset_list", "ara_property_list"]] = None
    sheet_name: Optional[str] = None
    sheet_count: Optional[int] = None
    # If one of index_start or index_end is set, the other must be set too
    index_start: Optional[int] = None
    index_end: Optional[int] = None

    # Task and subtask IDs
    task_id: Optional[str] = None
    subtask_id: Optional[str] = None

    # Optional flag to trigger a fabric first task
    enforce_fabric_first: Optional[bool] = False

    @model_validator(mode="after")
    def check_indexes(self):
        if (self.index_start is None) != (self.index_end is None):
            raise ValueError("Both index_start and index_end must be set or both must be None")
        return self

    @model_validator(mode="after")
    def check_goal_value_requirement(self):
        # Make sure that goal_value is set when goal is "Increasing EPC"
        if self.goal == "Increasing EPC" and not self.goal_value:
            raise ValueError("goal_value is required when goal is 'Increasing EPC'")
        return self