Model/backend/export/property_scenarios/db_functions.py

from typing import List, Any, Dict, Optional, Tuple, Sequence
import pandas as pd
from sqlalchemy import select
from sqlalchemy.orm import Session
from sqlalchemy.engine import Row
from collections import defaultdict

from backend.app.db.models.recommendations import (
    Recommendation,
    PlanModel,
    RecommendationMaterials,
)
from backend.app.db.models.portfolio import (
    PropertyModel,
    PropertyDetailsEpcModel,
)
from backend.app.db.models.materials import Material
from utils.logger import setup_logger

logger = setup_logger()


class DbMethods:

    def __init__(self, session: Session) -> None:
        self.session = session

    def get_properties(self, portfolio_id: int) -> pd.DataFrame:
        """
        Function to fetch the property data, for property scenario exports
        :param portfolio_id:
        :return:
        """
        stmt = (
            select(PropertyModel, PropertyDetailsEpcModel)
            .join(
                PropertyDetailsEpcModel,
                PropertyModel.id == PropertyDetailsEpcModel.property_id,
            )
            .where(PropertyModel.portfolio_id == portfolio_id)
        )

        rows: Sequence[Row[Tuple[PropertyModel, PropertyDetailsEpcModel]]] = (
            self.session.execute(stmt).all()
        )

        data: List[Dict[str, Any]] = [
            {
                **{
                    col.name: getattr(property_model, col.name)
                    for col in PropertyModel.__table__.columns.values()
                },
                **{
                    col.name: getattr(epc_model, col.name)
                    for col in PropertyDetailsEpcModel.__table__.columns.values()
                },
            }
            for property_model, epc_model in rows
        ]

        return pd.DataFrame(data)

    def get_latest_plans(
        self,
        portfolio_id: int,
        scenario_ids: Optional[List[int]] = None,
        default_only: bool = False,
    ) -> pd.DataFrame:
        """
        Fetch latest plans.

        Modes:
        1) Scenario mode: latest per (scenario_id, property_id)
        2) Default mode: latest default plan per property (ignores scenario_ids)

        """

        # -----------------------------
        # Sanity checks
        # -----------------------------
        if default_only and scenario_ids:
            # Override scenario_ids to make it explicit that they will be ignored in the query
            scenario_ids = None

        if not default_only and not scenario_ids:
            raise ValueError(
                "Either scenario_ids must be provided "
                "or default_only must be True."
            )

        # -----------------------------
        # Filter on just the default plans - we ignore the scenario ids. NOTE - this is specific to postgres
        # and relies on DISTINCT ON behaviour.
        # -----------------------------
        if default_only:
            # Latest default plan per property (ignore scenarios entirely)
            # DISTINCT ON (property_id) keeps the first row per property,
            # ordered by created_at DESC so we get the newest one.

            stmt = (
                select(PlanModel)
                .where(
                    PlanModel.portfolio_id == portfolio_id,
                    PlanModel.is_default.is_(True),
                )
                .distinct(PlanModel.property_id)
                .order_by(
                    PlanModel.property_id,
                    PlanModel.created_at.desc(),
                )
            )

        else:
            # Latest plan per (scenario_id, property_id)
            # DISTINCT ON (scenario_id, property_id) keeps the newest
            # plan per scenario/property combination.

            assert scenario_ids is not None

            stmt = (
                select(PlanModel)
                .where(
                    PlanModel.portfolio_id == portfolio_id,
                    PlanModel.scenario_id.in_(scenario_ids),
                )
                .distinct(
                    PlanModel.scenario_id,
                    PlanModel.property_id,
                )
                .order_by(
                    PlanModel.scenario_id,
                    PlanModel.property_id,
                    PlanModel.created_at.desc(),
                )
            )

        logger.info("Fetching plans")

        plans: Sequence[PlanModel] = self.session.scalars(stmt).all()

        return pd.DataFrame(
            [
                {
                    col.name: getattr(plan, col.name)
                    for col in PlanModel.__table__.columns.values()
                }
                for plan in plans
            ]
        )

    def get_recommendations(self, plan_ids: List[int]) -> pd.DataFrame:

        if not plan_ids:
            logger.info("No plan ids provided")
            return pd.DataFrame()

        stmt = (
            select(Recommendation, PlanModel.scenario_id, PlanModel.name)
            .join(PlanModel, PlanModel.id == Recommendation.plan_id)
            .where(
                Recommendation.plan_id.in_(plan_ids),
                Recommendation.default.is_(True),
                Recommendation.already_installed.is_(False),
            )
        )

        rows: Sequence[Tuple[Recommendation, Optional[int], Optional[str]]] = (
            self.session.execute(stmt).tuples().all()
        )

        data: List[Dict[str, Any]] = [
            {
                **{
                    col.name: getattr(rec_model, col.name)
                    for col in Recommendation.__table__.columns.values()
                },
                "scenario_id": scenario_id,
                "plan_name": plan_name,
            }
            for rec_model, scenario_id, plan_name in rows
        ]

        return pd.DataFrame(data)

    def attach_materials(self, recommendations_df: pd.DataFrame) -> pd.DataFrame:

        if recommendations_df.empty:
            recommendations_df["materials"] = []
            return recommendations_df

        rec_ids: List[int] = recommendations_df["id"].astype(int).tolist()

        stmt = (
            select(RecommendationMaterials, Material)
            .join(Material, RecommendationMaterials.material_id == Material.id)
            .where(RecommendationMaterials.recommendation_id.in_(rec_ids))
        )

        rows: Sequence[Tuple[RecommendationMaterials, Material]] = (
            self.session.execute(stmt).tuples().all()
        )

        materials_map: Dict[int, List[Dict[str, Any]]] = defaultdict(list)

        for rec_mat, material in rows:
            materials_map[rec_mat.recommendation_id].append(
                {
                    "material_id": rec_mat.material_id,
                    "depth": rec_mat.depth,
                    "quantity": rec_mat.quantity,
                    "quantity_unit": rec_mat.quantity_unit,
                    "estimated_cost": rec_mat.estimated_cost,
                    "type": material.type.value if material.type else None,
                    "includes_battery": material.includes_battery,
                }
            )

        recommendations_df["materials"] = recommendations_df["id"].astype(int).apply(
            lambda x: materials_map.get(x, [])
        )

        return recommendations_df