from collections import defaultdict from typing import Dict, List, Optional from backend.app.db.functions.recommendations_functions import ( bulk_update_plans, get_default_plans, get_most_recent_plans_by_portfolio_id, get_most_recent_plans_by_scenario_ids, get_scenarios_by_portfolio_id, ) from backend.app.db.models.recommendations import PlanModel, ScenarioModel from backend.app.domain.classes.plan import Plan from backend.app.domain.classes.scenario import Scenario from utils.logger import setup_logger logger = setup_logger() def process_portfolio( portfolio_id: int, scenarios_to_consider: Optional[List[int]] = None, scenario_priority_order: Optional[List[int]] = None, ) -> None: # TODO: make this a class logger.info(f"Processing portfolio {portfolio_id}") all_scenarios: List[Scenario] = _load_scenarios_for_portfolio(portfolio_id) plans_by_id: Dict[int, Plan] = {} # TODO: make this an in-memory repository class if scenarios_to_consider: if len(scenarios_to_consider) < 2: raise ValueError( "Cannot run auto categorisation for fewer than 2 scenarios" ) # first get all plans that we're interested in plans_for_consideration: List[Plan] = _load_plans_for_portfolio( portfolio_id, all_scenarios, scenarios_to_consider ) for plan in plans_for_consideration: if plan.id is not None: # just in case plans_by_id[plan.id] = plan # then unset existing defaults on domain objects regardless of whether they're under consideration or not default_plans: List[Plan] = _get_default_plans(portfolio_id, all_scenarios) for plan in default_plans: plan.set_default(False) if plan.id is not None: # just in case plans_by_id[plan.id] = plan logger.info(f"Successfully unset {len(default_plans)} default plan(s)") # then set new defaults on domain objects under consideration plans_for_consideration_by_property: Dict[int, List[Plan]] = ( _group_plans_by_property(plans_for_consideration) ) for property_id, property_plans in plans_for_consideration_by_property.items(): if not property_plans: raise ValueError(f"No plans for property {property_id}") try: cheapest_plan = choose_cheapest_relevant_plan( property_plans, scenario_priority_order ) except Exception: logger.error(f"Failed to find cheapest plan for property {property_id}") raise property_plans = _update_plan_objects(property_plans, cheapest_plan) for plan in property_plans: if plan.id is not None: # just in case plans_by_id[plan.id] = plan logger.info("Successfully set defaults on Plan objects in memory") # then pass all domain objects to database to update (regardless of whether they've changed) _update_plans_in_db(list(plans_by_id.values())) logger.info(f"Successfully updated {len(plans_by_id)} Plans in database") def choose_cheapest_relevant_plan( plans: List[Plan], scenario_priority_order: Optional[List[int]] = None ) -> Plan: scenario_priority_order = scenario_priority_order or [] eligible_plans: List[Plan] = [plan for plan in plans if plan.is_compliant] or plans if not eligible_plans: raise ValueError("No plans available to choose from.") for plan in eligible_plans: if plan.id is None: # This should never actually happen, but plan.id is optional to cater # for new plans. We are only working with already persisted plans here raise ValueError( f"All plans must have an ID, but found a plan with no ID: {plan}" ) min_cost: float = min(plan.cost for plan in eligible_plans) cheapest_plans: List[Plan] = [ plan for plan in eligible_plans if plan.cost == min_cost ] for priority_scenario_id in scenario_priority_order: for plan in cheapest_plans: if plan.scenario.id == priority_scenario_id: return plan return cheapest_plans[0] def _get_default_plans(portfolio_id: int, scenarios: List[Scenario]) -> List[Plan]: default_plan_models = get_default_plans(portfolio_id) scenario_map = {s.id: s for s in scenarios} return [ Plan.from_sqlalchemy(p, scenario_map[p.scenario_id]) for p in default_plan_models if p.scenario_id in scenario_map ] def _load_scenarios_for_portfolio(portfolio_id: int) -> List[Scenario]: scenario_models: List[ScenarioModel] = get_scenarios_by_portfolio_id(portfolio_id) return [Scenario.from_sqlalchemy(s) for s in scenario_models] def _load_plans_for_portfolio( portfolio_id: int, all_scenarios: List[Scenario], scenarios_to_consider: Optional[List[int]] = None, ) -> List[Plan]: if scenarios_to_consider: logger.info(f"Getting plans for {len(scenarios_to_consider)} scenarios") plan_models: List[PlanModel] = get_most_recent_plans_by_scenario_ids( scenarios_to_consider ) logger.info(f"Got {len(plan_models)} plan models from database") else: logger.info( f"No list of Plans to consider provided. Getting all Plans for portfolio {portfolio_id}" ) plan_models: List[PlanModel] = get_most_recent_plans_by_portfolio_id( portfolio_id ) plans: List[Plan] = [] if not all_scenarios: raise Exception(f"No scenarios found for Portfolio {portfolio_id}") for model in plan_models: scenario = next((s for s in all_scenarios if s.id == model.scenario_id)) if not scenario: logger.info(f"No Scenario associated with Plan of ID {model.id}") continue plans.append(Plan.from_sqlalchemy(model, scenario)) logger.info(f"Got {len(plans)} Plans") return plans def _group_plans_by_property(plans: List[Plan]) -> Dict[int, List[Plan]]: grouped: dict[int, List[Plan]] = defaultdict(list) for plan in plans: grouped[plan.record.property_id].append(plan) return grouped def _update_plan_objects(plans: List[Plan], cheapest_plan: Plan) -> List[Plan]: for plan in plans: should_be_default: bool = plan.id == cheapest_plan.id plan.set_default(should_be_default) if should_be_default: logger.debug( f"Setting Plan {plan.id} (Scenario Name: {plan.scenario.record.name}) to default" ) return plans def _update_plans_in_db(plans: List[Plan]) -> None: plan_models: List[PlanModel] = [] scenario_models: List[ScenarioModel] = [] for plan in plans: plan_model, scenario_model = plan.to_sqlalchemy() plan_models.append(plan_model) scenario_models.append(scenario_model) bulk_update_plans(plan_models, scenario_models)