Model/backend/export/property_scenarios/db_functions.py

from typing import List, Any, Dict, Optional
import pandas as pd
from sqlalchemy import func
from sqlalchemy.orm import Session
from collections import defaultdict

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

logger = setup_logger()


class DbMethods:

    def __init__(self, session: Session):
        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:
        """
        query = (
            self.session.query(PropertyModel, PropertyDetailsEpcModel)
            .join(
                PropertyDetailsEpcModel,
                PropertyModel.id == PropertyDetailsEpcModel.property_id,
            )
            .filter(PropertyModel.portfolio_id == portfolio_id)
            .all()
        )

        data = [
            {
                **{
                    col.name: getattr(row.PropertyModel, col.name)
                    for col in PropertyModel.__table__.columns
                },
                **{
                    col.name: getattr(row.PropertyDetailsEpcModel, col.name)
                    for col in PropertyDetailsEpcModel.__table__.columns
                },
            }
            for row in query
        ]

        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.

            plans_query = (
                self.session.query(PlanModel)
                .filter(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.

            plans_query = (
                self.session.query(PlanModel)
                .filter(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 = plans_query.all()

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

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

        if not plan_ids:
            return pd.DataFrame()

        recs_query = (
            self.session.query(
                Recommendation,
                PlanModel.scenario_id,
            )
            .join(
                PlanRecommendations,
                Recommendation.id == PlanRecommendations.recommendation_id,
            )
            .join(PlanModel, PlanModel.id == PlanRecommendations.plan_id)
            .filter(
                PlanRecommendations.plan_id.in_(plan_ids),
                Recommendation.default.is_(True),
                Recommendation.already_installed.is_(False),
            )
            .all()
        )

        data = [
            {
                **{
                    col.name: getattr(r.Recommendation, col.name)
                    for col in Recommendation.__table__.columns
                },
                "scenario_id": r.scenario_id,
            }
            for r in recs_query
        ]

        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 = recommendations_df["id"].tolist()

        materials_query = (
            self.session.query(RecommendationMaterials)
            .filter(RecommendationMaterials.recommendation_id.in_(rec_ids))
            .all()
        )

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

        for m in materials_query:
            materials_map[m.recommendation_id].append(
                {
                    "material_id": m.material_id,
                    "depth": m.depth,
                    "quantity": m.quantity,
                    "quantity_unit": m.quantity_unit,
                    "estimated_cost": m.estimated_cost,
                }
            )

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

        return recommendations_df