Model/backend/categorisation/processor.py

from collections import defaultdict
from typing import Dict, List, Optional

from backend.app.db.functions.recommendations_functions import (
    bulk_update_plans,
    get_default_plans_and_scenarios,
    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}")

    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, 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)
    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) -> List[Plan]:
    default_plan_models, default_scenario_models = get_default_plans_and_scenarios(
        portfolio_id
    )

    return [
        Plan.from_sqlalchemy(
            p,
            next(
                Scenario.from_sqlalchemy(s)
                for s in default_scenario_models
                if s.id == p.scenario_id
            ),
        )
        for p in default_plan_models
    ]


def _load_plans_for_portfolio(
    portfolio_id: int, 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] = []

    scenarios: List[ScenarioModel] = get_scenarios_by_portfolio_id(portfolio_id)

    if not scenarios:
        raise Exception(f"No scenarios found for Portfolio {portfolio_id}")

    for model in plan_models:

        scenario_model = next((s for s in scenarios if s.id == model.scenario_id))
        if not scenario_model:
            logger.info(f"No Scenario associated with Plan of ID {model.id}")
            continue

        plans.append(
            Plan.from_sqlalchemy(model, Scenario.from_sqlalchemy(scenario_model))
        )

    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)