heating-mpc-ha/custom_components/heating_mpc/sensor.py

"""Diagnostics and predicted temperatures exposed as HA sensors."""
from __future__ import annotations

from homeassistant.components.sensor import SensorDeviceClass, SensorEntity, SensorStateClass
from homeassistant.const import PERCENTAGE, UnitOfTemperature
from homeassistant.helpers.update_coordinator import CoordinatorEntity

from .const import (
    DOMAIN, READY_OFF_SAMPLES, READY_ON_SAMPLES, READY_SAMPLES, SAMPLE_RESULTS,
)
from .evaluation import HORIZONS, WINDOW_DAYS


async def async_setup_entry(hass, entry, async_add_entities) -> None:
    coordinator = hass.data[DOMAIN][entry.entry_id]
    entities = [MPCStatus(coordinator, entry.entry_id), MPCWeightedTemp(coordinator, entry.entry_id), MPCOutdoorTemperature(coordinator, entry.entry_id),
                MPCLearningProgress(coordinator, entry.entry_id), MPCLastSample(coordinator, entry.entry_id)]
    entities += [MPCForecastError(coordinator, entry.entry_id, horizon) for horizon in HORIZONS]
    entities.append(MPCOperationalForecastError(coordinator, entry.entry_id))
    # Zone forecasts belong to their zone subentry, so deleting a zone removes them.
    linked = {}
    for zone in coordinator.zones.values():
        entity = MPCRoomPrediction(coordinator, entry.entry_id, zone.key)
        if zone.subentry_id:
            linked.setdefault(zone.subentry_id, []).append(entity)
        else:
            entities.append(entity)
    async_add_entities(entities)
    for subentry_id, zone_entities in linked.items():
        async_add_entities(zone_entities, config_subentry_id=subentry_id)


class MPCStatus(CoordinatorEntity, SensorEntity):
    _attr_has_entity_name = True
    _attr_name = "Heating MPC Status"
    _attr_icon = "mdi:brain"
    # Per-source report ages change every poll; recording them adds a new
    # attribute row to the database every five minutes. Current values stay visible.
    _unrecorded_attributes = frozenset({"input_quality", "weather_comparison", "heat_sources", "windows",
                                        "ventilation", "ventilation_events", "model_changes"})

    def __init__(self, coordinator, entry_id):
        super().__init__(coordinator)
        self._attr_unique_id = f"{entry_id}_status"

    @property
    def native_value(self):
        return self.coordinator.data.get("status")

    @property
    def extra_state_attributes(self):
        data = self.coordinator.data
        # Forecast details live on their own entities.
        return {key: value for key, value in data.items() if key not in (
            "status", "prediction_60m", "prediction_end", "prediction_trajectory", "forecast_accuracy",
            "zone_details", "operational_forecast")}


class MPCWeightedTemp(CoordinatorEntity, SensorEntity):
    _attr_has_entity_name = True
    _attr_name = "Heating MPC Weighted Temperature"
    _attr_device_class = SensorDeviceClass.TEMPERATURE
    _attr_native_unit_of_measurement = UnitOfTemperature.CELSIUS

    def __init__(self, coordinator, entry_id):
        super().__init__(coordinator)
        self._attr_unique_id = f"{entry_id}_weighted_temp"

    @property
    def native_value(self):
        return self.coordinator.data.get("weighted_temperature")


class MPCRoomPrediction(CoordinatorEntity, SensorEntity):
    _attr_has_entity_name = True
    _attr_device_class = SensorDeviceClass.TEMPERATURE
    _attr_native_unit_of_measurement = UnitOfTemperature.CELSIUS
    _attr_icon = "mdi:chart-line"

    def __init__(self, coordinator, entry_id, room):
        super().__init__(coordinator)
        self.room = room
        self._attr_unique_id = f"{entry_id}_{room}_predicted_60min"
        self._attr_name = f"MPC {coordinator.zones[room].name} 60min Prediction"

    # Change every poll; the 60-minute value itself is recorded.
    _unrecorded_attributes = frozenset({"forecast_15min", "plan_15min", "forecast", "contributions_c_per_h",
                                        "coefficients", "current_temperature", "model_in_use"})

    @property
    def native_value(self):
        return self.coordinator.data.get("prediction_60m", {}).get(self.room)

    @property
    def extra_state_attributes(self):
        data = self.coordinator.data
        trajectory = data.get("prediction_trajectory", {}).get(self.room) or []
        details = data.get("zone_details", {}).get(self.room, {})
        return {
            "current_temperature": details.get("temperature"),
            "forecast": {f"{minutes}_min": trajectory[minutes // 15 - 1]
                         for minutes in (15, 30, 60, 120) if len(trajectory) >= minutes // 15},
            "forecast_15min": trajectory or None,
            "plan_15min": data.get("plan_15min"),
            "basis": "proposed MPC plan; counterfactual while Generic Thermostat controls the relay",
            # Windows with an unknown state: the forecast assumes their state (see Status).
            "uncertain_windows": details.get("uncertain_windows") or [],
            "model_in_use": details.get("model_in_use"),
            "contributions_c_per_h": details.get("contributions_c_per_h"),
            "coefficients": details.get("coefficients"),
        }

class MPCOutdoorTemperature(CoordinatorEntity, SensorEntity):
    """Display chosen outdoor temperature and all currently valid alternatives."""
    _attr_has_entity_name = True
    _attr_name = "MPC Outdoor Temperature"
    _attr_device_class = SensorDeviceClass.TEMPERATURE
    _attr_native_unit_of_measurement = UnitOfTemperature.CELSIUS
    _attr_icon = "mdi:thermometer"

    def __init__(self, coordinator, entry_id):
        super().__init__(coordinator)
        self._attr_unique_id = f"{entry_id}_outside_temperature"

    @property
    def native_value(self):
        return self.coordinator.data.get("weather", {}).get("temperature")

    @property
    def extra_state_attributes(self):
        data = self.coordinator.data
        return {
            "active_source": data.get("weather_active_sources", {}).get("temperature"),
            "alternatives": data.get("weather_comparison", {}).get("temperature", {}),
        }


class MPCForecastError(CoordinatorEntity, SensorEntity):
    """Model accuracy at one horizon: the model replayed with the inputs that actually happened.

    Not a forecast accuracy (a real forecast does not know the future relay, power or
    windows); see MPCOperationalForecastError. The unique ID keeps its original name.
    """
    _attr_has_entity_name = True
    _attr_device_class = SensorDeviceClass.TEMPERATURE_DELTA
    _attr_native_unit_of_measurement = UnitOfTemperature.CELSIUS
    _attr_state_class = SensorStateClass.MEASUREMENT
    _attr_suggested_display_precision = 2
    _attr_icon = "mdi:target"
    _unrecorded_attributes = frozenset({"zones"})

    def __init__(self, coordinator, entry_id, horizon):
        super().__init__(coordinator)
        self.horizon = horizon
        self._attr_unique_id = f"{entry_id}_forecast_error_{horizon}min"
        self._attr_name = f"Heating MPC Model Error {horizon} min"

    @property
    def _accuracy(self):
        return self.coordinator.data.get("forecast_accuracy", {}).get(self.horizon, {})

    @property
    def native_value(self):
        accuracy = self._accuracy
        # No number from a handful of overlapping forecasts.
        return accuracy.get("mae") if accuracy.get("quality", "insufficient") != "insufficient" else None

    @property
    def extra_state_attributes(self):
        accuracy = self._accuracy
        return {
            "kind": "model accuracy: replayed with the actual relay, electrical heat and window history",
            "samples": accuracy.get("samples", 0),
            "effective_samples": accuracy.get("effective_samples", 0),
            "quality": accuracy.get("quality", "insufficient"),
            "mae": accuracy.get("mae"),
            "rmse": accuracy.get("rmse"),
            "bias": accuracy.get("bias"),
            "persistence_mae": accuracy.get("persistence_mae"),
            "skill": accuracy.get("skill"),
            "base_model_mae": accuracy.get("base_model_mae"),
            "improvement_vs_base": accuracy.get("improvement_vs_base"),
            # Per zone, by name.
            "zones": {self.coordinator.zones[key].name if key in self.coordinator.zones else key: stats
                      for key, stats in accuracy.get("rooms", {}).items()},
            "window_days": WINDOW_DAYS,
            "unit_of_attributes": "°C",
            "basis": ("all zones replayed together with the actual relay, electrical heat and "
                      "window history; outdoor temperature held at the start value"),
        }


class MPCOperationalForecastError(CoordinatorEntity, SensorEntity):
    """Accuracy of forecasts as published, using only what was known when they were made.

    A forecast assumes the MPC's plan; it is scored only when the relay actually
    followed that plan, otherwise counted as counterfactual. State: the 60-minute
    error; all horizons in the attributes.
    """
    _attr_has_entity_name = True
    _attr_name = "Heating MPC Forecast Error"
    _attr_device_class = SensorDeviceClass.TEMPERATURE_DELTA
    _attr_native_unit_of_measurement = UnitOfTemperature.CELSIUS
    _attr_state_class = SensorStateClass.MEASUREMENT
    _attr_suggested_display_precision = 2
    _attr_icon = "mdi:crystal-ball"
    _unrecorded_attributes = frozenset({"horizons"})

    def __init__(self, coordinator, entry_id):
        super().__init__(coordinator)
        self._attr_unique_id = f"{entry_id}_operational_forecast_error"

    @property
    def native_value(self):
        hour = self.coordinator.data.get("operational_forecast", {}).get(60, {})
        return hour.get("mae") if hour.get("quality", "insufficient") != "insufficient" else None

    @property
    def extra_state_attributes(self):
        summary = self.coordinator.data.get("operational_forecast", {})
        names = {key: zone.name for key, zone in self.coordinator.zones.items()}
        horizons = {}
        for horizon, stats in summary.items():
            horizons[f"{horizon}_min"] = {**{k: v for k, v in stats.items() if k != "zones"},
                                          "zones": {names.get(z, z): v for z, v in stats.get("zones", {}).items()}}
        return {
            "kind": "forecast accuracy: published forecasts, scored only where the relay followed the plan",
            "horizons": horizons,
            "window_days": WINDOW_DAYS,
            "unit_of_attributes": "°C",
        }


class MPCLearningProgress(CoordinatorEntity, SensorEntity):
    """How close the least trained room is to being used for forecasts."""
    _attr_has_entity_name = True
    _attr_name = "Heating MPC Learning Progress"
    _attr_native_unit_of_measurement = PERCENTAGE
    _attr_state_class = SensorStateClass.MEASUREMENT
    _attr_suggested_display_precision = 0
    _attr_icon = "mdi:school-outline"

    def __init__(self, coordinator, entry_id):
        super().__init__(coordinator)
        self._attr_unique_id = f"{entry_id}_learning_progress"

    @property
    def native_value(self):
        return self.coordinator.data.get("learning_progress")

    @property
    def extra_state_attributes(self):
        data = self.coordinator.data
        return {
            "samples": data.get("samples"),
            "on_samples": data.get("on_samples"),
            "off_samples": data.get("off_samples"),
            "required": {"samples": READY_SAMPLES, "on_samples": READY_ON_SAMPLES,
                         "off_samples": READY_OFF_SAMPLES},
        }


class MPCLastSample(CoordinatorEntity, SensorEntity):
    """What the last poll did for learning; its history shows when learning happens."""
    _attr_has_entity_name = True
    _attr_name = "Heating MPC Last Sample"
    _attr_translation_key = "last_sample"  # translated state names
    _attr_device_class = SensorDeviceClass.ENUM
    _attr_options = list(SAMPLE_RESULTS)
    _attr_icon = "mdi:clipboard-pulse-outline"
    _unrecorded_attributes = frozenset({"counts_since_start"})

    def __init__(self, coordinator, entry_id):
        super().__init__(coordinator)
        self._attr_unique_id = f"{entry_id}_last_sample"

    @property
    def native_value(self):
        return self.coordinator.data.get("sample_result")

    @property
    def extra_state_attributes(self):
        return {"counts_since_start": self.coordinator.data.get("sample_results_since_start", {})}