Model/scripts/data_exports.py

"""Principal-pitch data export — new DDD model edition.

Replaces sfr/principal_pitch/2_export_data.py, which read the retired
``plan_recommendations`` m2m and ``recommendation_materials`` table. In the
current model:
  * a Recommendation links to its Plan directly (``recommendation.plan_id``),
  * materials are inline on the Recommendation (``material_id`` etc.),
  * the chosen Plan per (scenario, property) is the one with ``is_default``,
  * post-works SAP/EPC + savings live on the Plan row (the new SAP calculator's
    output), so we read them directly rather than summing recommendation uplifts.

Give it a portfolio id; it resolves every *modelled* scenario for that portfolio
(scenarios that have plans) and writes ONE workbook with a ``properties`` sheet
per scenario. EPC descriptive fields (walls/roof/heating/windows/floor area/
lodgement) come live from the EPC service, because ``property_details_epc`` is
dead under the new backend.

    python scripts/data_exports.py --portfolio 814
    python scripts/data_exports.py --portfolio 814 --out "sfr/principal_pitch/Durkan.xlsx"

Reads DB_* + OPEN_EPC_API_TOKEN from backend/.env. Run from the worktree root.
"""

from __future__ import annotations

import argparse
import re
from datetime import date, datetime
from pathlib import Path
from typing import Any, Optional

import numpy as np
import pandas as pd
from sqlalchemy import text
from sqlalchemy.engine import Engine

import sys

_REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO_ROOT))

from backend.app.utils import sap_to_epc  # noqa: E402
from infrastructure.epc_client.epc_client_service import EpcClientService  # noqa: E402
from scripts.e2e_common import ENV_PATH, build_engine, load_env  # noqa: E402
from backend.app.config import get_settings  # noqa: E402

# Measure columns always present in the wide sheet (stable column set across runs).
EXPECTED_MEASURE_COLUMNS: tuple[str, ...] = (
    "suspended_floor_insulation",
    "solid_floor_insulation",
    "external_wall_insulation",
    "internal_wall_insulation",
    "cavity_wall_insulation",
    "loft_insulation",
    "flat_roof_insulation",
    "room_roof_insulation",
    "secondary_glazing",
    "double_glazing",
    "solar_pv",
    "high_heat_retention_storage_heaters",
    "air_source_heat_pump",
    "boiler_upgrade",
    "gas_boiler_upgrade",
    "roomstat_programmer_trvs",
    "time_temperature_zone_control",
    "low_energy_lighting",
    "mechanical_ventilation",
    "system_tune_up",
    "system_tune_up_zoned",
)


# --------------------------------------------------------------------------- #
# EPC descriptive fields (live from the EPC service)
# --------------------------------------------------------------------------- #
def _description_text(item: Any) -> str:
    if not isinstance(item, dict):
        return ""
    desc = item.get("description")
    if isinstance(desc, dict):
        desc = desc.get("value")
    return str(desc or "")


def _join_descriptions(value: Any) -> str:
    if isinstance(value, list):
        return "; ".join(t for t in (_description_text(d) for d in value) if t)
    return _description_text(value)


# Gov RdSAP property-type codes (the raw cert stores a code, not a word).
_PROPERTY_TYPE_CODES: dict[str, str] = {
    "0": "House", "1": "Bungalow", "2": "Flat", "3": "Maisonette", "4": "Park home",
}


def _decode_property_type(value: Any) -> Optional[str]:
    if value is None:
        return None
    s = str(value).strip()
    if s in _PROPERTY_TYPE_CODES:
        return _PROPERTY_TYPE_CODES[s]
    return s or None


def _is_expired(registration_date: Optional[str]) -> Optional[bool]:
    if not registration_date:
        return None
    try:
        lodged = datetime.fromisoformat(registration_date[:10]).date()
    except ValueError:
        return None
    return (date.today() - lodged).days > 365 * 10


def epc_details_from_service(svc: EpcClientService, uprn: Optional[int]) -> dict[str, Any]:
    """Flatten the UPRN's latest raw certificate into the descriptive fields the
    export needs. Returns ``{}`` when the UPRN has no EPC (blank columns)."""
    if uprn is None:
        return {}
    results = svc._search(uprn=uprn)  # pyright: ignore[reportPrivateUsage]
    if not results:
        return {}
    latest = max(results, key=lambda r: r.registration_date)
    raw = svc._fetch_certificate(latest.certificate_number)  # pyright: ignore[reportPrivateUsage]

    def _to_int(value: Any) -> Optional[int]:
        try:
            return int(value)
        except (TypeError, ValueError):
            return None

    current_sap = _to_int(raw.get("energy_rating_current"))
    return {
        "property_type": _decode_property_type(raw.get("property_type")),
        "walls": _join_descriptions(raw.get("walls")),
        "roof": _join_descriptions(raw.get("roofs")),
        "floor": _join_descriptions(raw.get("floors")),
        "windows": _join_descriptions(raw.get("window")),
        "heating": _join_descriptions(raw.get("main_heating")),
        "hot_water": _join_descriptions(raw.get("hot_water")),
        "lighting": _join_descriptions(raw.get("lighting")),
        "total_floor_area": raw.get("total_floor_area"),
        "lodgement_date": raw.get("registration_date"),
        "is_expired": _is_expired(raw.get("registration_date")),
        "current_epc_rating": raw.get("current_energy_efficiency_band"),
        "current_sap_points": current_sap,
        "original_sap_points": current_sap,
    }


# --------------------------------------------------------------------------- #
# DB reads (new model: scenario -> plan(is_default) -> recommendation)
# --------------------------------------------------------------------------- #
def modelled_scenarios(engine: Engine, portfolio_id: int) -> list[dict[str, Any]]:
    """Scenarios for the portfolio that actually have plans, newest first."""
    with engine.connect() as conn:
        rows = conn.execute(
            text(
                """
                SELECT s.id, s.name
                FROM scenario s
                WHERE s.portfolio_id = :p
                  AND EXISTS (SELECT 1 FROM plan pl WHERE pl.scenario_id = s.id)
                ORDER BY s.id
                """
            ),
            {"p": portfolio_id},
        ).mappings().all()
    return [dict(r) for r in rows]


def load_properties(engine: Engine, portfolio_id: int, svc: EpcClientService) -> pd.DataFrame:
    """Base property identity (property_type falls back to the landlord override)
    plus live EPC descriptive fields."""
    with engine.connect() as conn:
        rows = conn.execute(
            text(
                """
                SELECT p.id AS property_id, p.id AS id, p.uprn, p.address, p.postcode,
                       p.landlord_property_id, p.number_of_rooms,
                       COALESCE(p.property_type, po.override_value) AS property_type
                FROM property p
                LEFT JOIN property_overrides po
                  ON po.property_id = p.id
                 AND po.override_component = 'property_type'
                 AND po.building_part = 0
                WHERE p.portfolio_id = :p
                ORDER BY p.id
                """
            ),
            {"p": portfolio_id},
        ).mappings().all()

    records: list[dict[str, Any]] = []
    for i, r in enumerate(rows, 1):
        base: dict[str, Any] = dict(r)
        uprn = int(base["uprn"]) if base.get("uprn") is not None else None
        for key, value in epc_details_from_service(svc, uprn).items():
            if base.get(key) is None:
                base[key] = value
        records.append(base)
        if i % 50 == 0:
            print(f"  EPC fetched {i}/{len(rows)}")
    df = pd.DataFrame(records)
    df["uprn"] = df["uprn"].astype("string")
    return df


def load_recommendations(engine: Engine, scenario_id: int) -> pd.DataFrame:
    """Default, not-already-installed recommendations on each property's default
    plan for the scenario, with the material type/battery flag joined."""
    with engine.connect() as conn:
        rows = conn.execute(
            text(
                """
                SELECT pl.property_id,
                       r.measure_type,
                       r.description,
                       r.estimated_cost,
                       r.sap_points,
                       r.co2_equivalent_savings,
                       r.kwh_savings,
                       r.energy_cost_savings,
                       m.type AS material_type,
                       COALESCE(m.includes_battery, FALSE) AS includes_battery
                FROM recommendation r
                JOIN plan pl ON pl.id = r.plan_id
                LEFT JOIN material m ON m.id = r.material_id
                WHERE pl.scenario_id = :s
                  AND pl.is_default = TRUE
                  AND r.default = TRUE
                  AND r.already_installed = FALSE
                """
            ),
            {"s": scenario_id},
        ).mappings().all()
    return pd.DataFrame([dict(r) for r in rows])


def load_default_plans(engine: Engine, scenario_id: int) -> pd.DataFrame:
    """The chosen (is_default) plan per property — the new SAP calculator's
    post-works results."""
    with engine.connect() as conn:
        rows = conn.execute(
            text(
                """
                SELECT property_id, post_sap_points, post_epc_rating,
                       cost_of_works, contingency_cost, co2_savings,
                       energy_bill_savings, energy_consumption_savings,
                       valuation_increase
                FROM plan
                WHERE scenario_id = :s AND is_default = TRUE
                """
            ),
            {"s": scenario_id},
        ).mappings().all()
    return pd.DataFrame([dict(r) for r in rows])


# --------------------------------------------------------------------------- #
# Sheet building
# --------------------------------------------------------------------------- #
def _apply_battery_suffix(recs: pd.DataFrame) -> pd.DataFrame:
    """solar_pv recommendations that carry a battery material become
    solar_pv_with_battery (mirrors the old export)."""
    if recs.empty:
        return recs
    is_solar_battery = (recs["material_type"] == "solar_pv") & (recs["includes_battery"])
    recs = recs.copy()
    recs["measure_type"] = np.where(
        is_solar_battery,
        recs["measure_type"].astype(str) + "_with_battery",
        recs["measure_type"],
    )
    return recs


def build_scenario_sheet(
    properties_df: pd.DataFrame, recs: pd.DataFrame, plans: pd.DataFrame
) -> pd.DataFrame:
    recs = _apply_battery_suffix(recs)

    # Pivot: one column per measure_type holding its estimated_cost.
    if not recs.empty:
        deduped = recs.drop_duplicates(subset=["property_id", "measure_type"], keep="first")
        cost_pivot = deduped.pivot(
            index="property_id", columns="measure_type", values="estimated_cost"
        ).reset_index()
        sap_uplift = (
            recs.groupby("property_id")["sap_points"].sum().reset_index(name="sap_points")
        )
        savings = (
            recs.groupby("property_id")[
                ["co2_equivalent_savings", "kwh_savings", "energy_cost_savings"]
            ]
            .sum()
            .reset_index()
        )
    else:
        cost_pivot = pd.DataFrame({"property_id": []})
        sap_uplift = pd.DataFrame({"property_id": [], "sap_points": []})
        savings = pd.DataFrame(
            {"property_id": [], "co2_equivalent_savings": [], "kwh_savings": [], "energy_cost_savings": []}
        )

    id_cols = [
        c
        for c in [
            "landlord_property_id", "property_id", "uprn", "address", "postcode",
            "property_type", "walls", "roof", "heating", "windows",
            "current_epc_rating", "current_sap_points", "original_sap_points",
            "total_floor_area", "number_of_rooms", "lodgement_date", "is_expired", "id",
        ]
        if c in properties_df.columns
    ]

    df = (
        properties_df[id_cols]
        .merge(cost_pivot, how="left", on="property_id")
        .merge(sap_uplift, how="left", on="property_id")
        .merge(savings, how="left", on="property_id")
        .merge(plans, how="left", on="property_id")
    )

    # total retrofit cost = sum of the per-measure cost columns
    measure_cols_present = [c for c in df.columns if c in set(EXPECTED_MEASURE_COLUMNS)
                           or c.endswith("_with_battery")]
    df["total_retrofit_cost"] = df[measure_cols_present].sum(axis=1) if measure_cols_present else 0.0

    df["sap_points"] = df["sap_points"].fillna(0)
    # Post-works SAP/EPC straight from the new SAP calculator's plan row;
    # fall back to current + uplift / sap_to_epc only when the plan lacks them.
    df["predicted_post_works_sap"] = df["post_sap_points"].where(
        df["post_sap_points"].notna(), df.get("current_sap_points", 0) + df["sap_points"]
    )
    df["predicted_post_works_epc"] = df["post_epc_rating"].where(
        df["post_epc_rating"].notna(),
        df["predicted_post_works_sap"].apply(lambda x: sap_to_epc(x) if pd.notna(x) else None),
    )

    # ensure the stable measure column set exists
    for col in EXPECTED_MEASURE_COLUMNS:
        if col not in df.columns:
            df[col] = ""
    return df


def _safe_sheet_name(name: str, used: set[str]) -> str:
    clean = re.sub(r"[:\\/?*\[\]]", "", name or "scenario").strip() or "scenario"
    clean = clean[:31]
    base, i = clean, 1
    while clean in used:
        suffix = f" ({i})"
        clean = base[: 31 - len(suffix)] + suffix
        i += 1
    used.add(clean)
    return clean


def export_portfolio(portfolio_id: int, out_path: Path) -> None:
    load_env(ENV_PATH)
    settings = get_settings()
    engine = build_engine()
    svc = EpcClientService(auth_token=settings.OPEN_EPC_API_TOKEN)

    with engine.connect() as conn:
        pname = conn.execute(
            text("SELECT name FROM portfolio WHERE id = :p"), {"p": portfolio_id}
        ).scalar()

    scenarios = modelled_scenarios(engine, portfolio_id)
    if not scenarios:
        raise SystemExit(f"No modelled scenarios (with plans) for portfolio {portfolio_id}.")
    print(f"Portfolio {portfolio_id} ({pname}) — {len(scenarios)} modelled scenario(s): "
          f"{[s['name'] for s in scenarios]}")

    print("Loading properties + EPC descriptive fields…")
    properties_df = load_properties(engine, portfolio_id, svc)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    used_names: set[str] = set()
    with pd.ExcelWriter(out_path) as writer:
        for s in scenarios:
            recs = load_recommendations(engine, s["id"])
            plans = load_default_plans(engine, s["id"])
            sheet_df = build_scenario_sheet(properties_df, recs, plans)
            sheet = _safe_sheet_name(s["name"] or f"scenario_{s['id']}", used_names)
            sheet_df.to_excel(writer, sheet_name=sheet, index=False)
            print(f"  sheet {sheet!r}: {len(sheet_df)} properties, "
                  f"{0 if recs.empty else len(recs)} recommendations")
    print(f"Wrote {out_path}")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--portfolio", type=int, required=True)
    ap.add_argument("--out", type=Path, default=None,
                    help="output xlsx path (default: sfr/principal_pitch/<portfolio name>.xlsx)")
    args = ap.parse_args()

    out = args.out
    if out is None:
        load_env(ENV_PATH)
        with build_engine().connect() as conn:
            nm = conn.execute(text("SELECT name FROM portfolio WHERE id=:p"), {"p": args.portfolio}).scalar()
        safe = re.sub(r"[\\/:*?\"<>|]", "_", str(nm or f"portfolio_{args.portfolio}"))
        out = _REPO_ROOT / "sfr" / "principal_pitch" / f"{safe}.xlsx"
    export_portfolio(args.portfolio, out)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())