Model/infrastructure/postgres/property_baseline_performance_table.py

from __future__ import annotations

from typing import ClassVar, Optional, cast

from sqlmodel import Field, SQLModel

from datatypes.epc.domain.epc import Epc
from domain.property_baseline.bill import Bill, BillSection, BillSectionCost
from domain.property_baseline.property_baseline_performance import PropertyBaselinePerformance
from domain.property_baseline.performance import Performance
from domain.property_baseline.rebaseliner import RebaselineReason

# Each Bill section's flat-column stem (``bill_{stem}_kwh`` / ``bill_{stem}_cost_gbp``).
_SECTION_COLUMN_STEM: dict[BillSection, str] = {
    BillSection.HEATING: "heating",
    BillSection.HOT_WATER: "hot_water",
    BillSection.LIGHTING: "lighting",
    BillSection.APPLIANCES: "appliances",
    BillSection.COOKING: "cooking",
    BillSection.PUMPS_FANS: "pumps_fans",
    BillSection.COOLING: "cooling",
}


class PropertyBaselinePerformanceModel(SQLModel, table=True):
    """The ``property_baseline_performance`` row — one per Property (ADR-0004).

    Flat typed columns (not a JSONB blob) so the FE can both surface the block
    and query the lodged-vs-effective pair. The production migration is FE-owned
    (Drizzle); see docs/migrations/property-baseline-performance-table.md.
    """

    __tablename__: ClassVar[str] = "property_baseline_performance"  # pyright: ignore[reportIncompatibleVariableOverride]

    id: Optional[int] = Field(default=None, primary_key=True)
    property_id: int = Field(unique=True, index=True)

    lodged_sap_score: int
    lodged_epc_band: str
    lodged_co2_emissions_t_per_yr: float
    lodged_primary_energy_intensity_kwh_per_m2_yr: int

    effective_sap_score: int
    effective_epc_band: str
    effective_co2_emissions_t_per_yr: float
    effective_primary_energy_intensity_kwh_per_m2_yr: int

    rebaseline_reason: str

    space_heating_kwh: float
    water_heating_kwh: float

    # Bill Derivation block (ADR-0014 §6). Nullable: all None when no calculator
    # ran (stub path). The ``bill_`` prefix avoids clashing with the
    # recorded-demand ``space_heating_kwh`` / ``water_heating_kwh`` above.
    bill_heating_kwh: Optional[float] = Field(default=None)
    bill_heating_cost_gbp: Optional[float] = Field(default=None)
    bill_hot_water_kwh: Optional[float] = Field(default=None)
    bill_hot_water_cost_gbp: Optional[float] = Field(default=None)
    bill_lighting_kwh: Optional[float] = Field(default=None)
    bill_lighting_cost_gbp: Optional[float] = Field(default=None)
    bill_appliances_kwh: Optional[float] = Field(default=None)
    bill_appliances_cost_gbp: Optional[float] = Field(default=None)
    bill_cooking_kwh: Optional[float] = Field(default=None)
    bill_cooking_cost_gbp: Optional[float] = Field(default=None)
    bill_pumps_fans_kwh: Optional[float] = Field(default=None)
    bill_pumps_fans_cost_gbp: Optional[float] = Field(default=None)
    bill_cooling_kwh: Optional[float] = Field(default=None)
    bill_cooling_cost_gbp: Optional[float] = Field(default=None)
    bill_standing_charges_gbp: Optional[float] = Field(default=None)
    bill_seg_credit_gbp: Optional[float] = Field(default=None)
    bill_total_annual_bill_gbp: Optional[float] = Field(default=None)

    @classmethod
    def from_domain(
        cls, baseline: PropertyBaselinePerformance, property_id: int
    ) -> "PropertyBaselinePerformanceModel":
        model = cls(
            property_id=property_id,
            lodged_sap_score=baseline.lodged.sap_score,
            lodged_epc_band=baseline.lodged.epc_band.value,
            lodged_co2_emissions_t_per_yr=baseline.lodged.co2_emissions,
            lodged_primary_energy_intensity_kwh_per_m2_yr=baseline.lodged.primary_energy_intensity,
            effective_sap_score=baseline.effective.sap_score,
            effective_epc_band=baseline.effective.epc_band.value,
            effective_co2_emissions_t_per_yr=baseline.effective.co2_emissions,
            effective_primary_energy_intensity_kwh_per_m2_yr=baseline.effective.primary_energy_intensity,
            rebaseline_reason=baseline.rebaseline_reason,
            space_heating_kwh=baseline.space_heating_kwh,
            water_heating_kwh=baseline.water_heating_kwh,
        )
        model._write_bill(baseline.bill)
        return model

    def _write_bill(self, bill: Optional[Bill]) -> None:
        """Flatten the Bill onto the ``bill_*`` columns. When ``bill`` is None
        (no calculator ran) every bill column is left None; a section absent from
        the mapping leaves its two columns None (None != 0 — it was not billed)."""
        if bill is None:
            return
        for section, stem in _SECTION_COLUMN_STEM.items():
            cost = bill.sections.get(section)
            setattr(self, f"bill_{stem}_kwh", cost.kwh if cost is not None else None)
            setattr(
                self,
                f"bill_{stem}_cost_gbp",
                cost.cost_gbp if cost is not None else None,
            )
        self.bill_standing_charges_gbp = bill.standing_charges_gbp
        self.bill_seg_credit_gbp = bill.seg_credit_gbp
        self.bill_total_annual_bill_gbp = bill.total_gbp

    def to_domain(self) -> PropertyBaselinePerformance:
        return PropertyBaselinePerformance(
            lodged=Performance(
                sap_score=self.lodged_sap_score,
                epc_band=Epc(self.lodged_epc_band),
                co2_emissions=self.lodged_co2_emissions_t_per_yr,
                primary_energy_intensity=self.lodged_primary_energy_intensity_kwh_per_m2_yr,
            ),
            effective=Performance(
                sap_score=self.effective_sap_score,
                epc_band=Epc(self.effective_epc_band),
                co2_emissions=self.effective_co2_emissions_t_per_yr,
                primary_energy_intensity=self.effective_primary_energy_intensity_kwh_per_m2_yr,
            ),
            rebaseline_reason=cast(RebaselineReason, self.rebaseline_reason),
            space_heating_kwh=self.space_heating_kwh,
            water_heating_kwh=self.water_heating_kwh,
            bill=self._read_bill(),
        )

    def _read_bill(self) -> Optional[Bill]:
        """Reconstruct the Bill from the ``bill_*`` columns. The total is the
        not-None discriminator: a persisted bill always sets it, so its absence
        means no calculator ran and the bill was None. A section is rebuilt only
        when its kWh column is not None (paired with its cost)."""
        if self.bill_total_annual_bill_gbp is None:
            return None
        sections: dict[BillSection, BillSectionCost] = {}
        for section, stem in _SECTION_COLUMN_STEM.items():
            kwh = cast(Optional[float], getattr(self, f"bill_{stem}_kwh"))
            if kwh is None:
                continue
            cost_gbp = cast(float, getattr(self, f"bill_{stem}_cost_gbp"))
            sections[section] = BillSectionCost(kwh=kwh, cost_gbp=cost_gbp)
        return Bill(
            sections=sections,
            standing_charges_gbp=cast(float, self.bill_standing_charges_gbp),
            seg_credit_gbp=cast(float, self.bill_seg_credit_gbp),
            total_gbp=self.bill_total_annual_bill_gbp,
        )