Model/backend/export/property_scenarios/main.py

import json
from typing import Optional, Any, Mapping, Dict, Union

import pandas as pd
from sqlalchemy.orm import Session

from backend.export.property_scenarios.input_schema import ExportRequest
from backend.export.property_scenarios.db_functions import DbMethods
from backend.app.db.connection import db_read_session
from backend.app.utils import sap_to_epc
from utils.logger import setup_logger

logger = setup_logger()


def process_export(payload: ExportRequest, session: Session) -> Dict[Union[str, int], pd.DataFrame]:
    export_files: Dict[Union[str, int], pd.DataFrame] = {}

    db_methods = DbMethods(session)

    properties_df = db_methods.get_properties(payload.portfolio_id)

    logger.info("Retrieved %s properties for export", len(properties_df))

    plans_df = db_methods.get_latest_plans(
        portfolio_id=payload.portfolio_id,
        scenario_ids=payload.scenario_ids,
        default_only=payload.default_plans_only,
    )

    logger.info("Retrieved %s plans for export", len(plans_df))

    if plans_df.empty:
        return export_files

    recommendations_df = db_methods.get_recommendations(
        plans_df["id"].tolist()
    )

    recommendations_df = db_methods.attach_materials(recommendations_df)

    if payload.default_plans_only:
        group_keys = [None]  # Single export, no scenario grouping
    else:
        group_keys = payload.scenario_ids

    for group_key in group_keys:

        if payload.default_plans_only:
            scenario_recs = recommendations_df
            export_label = "default_plans"
        else:
            scenario_recs = recommendations_df[
                recommendations_df["scenario_id"] == group_key
                ]
            export_label = group_key

        if scenario_recs.empty:
            continue

        measures_df: pd.DataFrame = scenario_recs[
            ["property_id", "measure_type", "estimated_cost"]
        ].drop_duplicates()

        pivot = measures_df.pivot(
            index="property_id",
            columns="measure_type",
            values="estimated_cost",
        ).reset_index()

        pivot["total_retrofit_cost"] = (
            pivot.drop(columns=["property_id"]).sum(axis=1)
        )

        post_sap = (
            scenario_recs.groupby("property_id")[["sap_points"]]
            .sum()
            .reset_index()
        )

        df = (
            properties_df.rename(columns={"solar_pv": "existing_solar_pv"})
            .merge(pivot, how="left", on="property_id")
            .merge(post_sap, how="left", on="property_id")
        )

        df["sap_points"] = df["sap_points"].fillna(0)
        df["predicted_post_works_sap"] = df["current_sap_points"] + df["sap_points"]
        df["predicted_post_works_epc"] = df[
            "predicted_post_works_sap"
        ].apply(sap_to_epc)

        export_files[export_label] = df

    return export_files


# ============================================================
# Lambda Handler
# ============================================================

def handler(event: dict, context: Optional[Any]) -> Mapping[str, int | str]:
    """
    Lambda event should have the following structure:
    1) task id - unique identifier for the export task (optional, can be used for tracking/logging)
    2) subtask id - unique identifier for the specific export operation (optional, can be used for tracking/logging)
    2) portfolio id - id of the portfolio to export
    3) scenario ids - list of scenario ids to export
    4) default_plans_only - flag indicating if we should only consider default plans for export (optional,
    defaults to False)

    Exxample event:
    body_dict = {
        "task_id": "test",
        "subtask_id": "test",
        "portfolio_id": 569,
        "scenario_ids": [],
        "default_plans_only": True,
    }
    :param event: Lambda event containing export request details
    :param context: Lambda context (not used in this handler but included for completeness)
    :return: HTTP response indicating success or failure of the export operation
    """
    for record in event.get("Records", []):
        try:
            body_dict = json.loads(record["body"])

            logger.debug("Validating request body")
            payload = ExportRequest.model_validate(body_dict)

            if payload._scenario_ids_ignored:
                logger.warning(
                    "Received scenario_ids in request body but they will be ignored "
                    "because default_plans_only is set to True"
                )

            logger.debug("Successfully validated request body")
            with db_read_session() as session:
                exported_files = process_export(payload, session)

            # TODO: Need to handle the exported files - e.g. upload to s3 and email a presigned url

            return {
                "statusCode": 200,
                "body": json.dumps({}),
            }

        except Exception as e:
            logger.error(f"Failed to process record: {e}")
            return {
                "statusCode": 500,
                "body": json.dumps({"message": "Failed to process export request"}),
            }

    return {
        "statusCode": 201,
        "body": json.dumps({"message": "No records to process"}),
    }