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]
    # 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")}


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",
            "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):
    """Mean absolute error of the thermal model at one horizon, replayed against reality."""
    _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 Forecast Error {horizon} min"

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

    @property
    def native_value(self):
        return self._accuracy.get("mae")

    @property
    def extra_state_attributes(self):
        accuracy = self._accuracy
        return {
            "samples": accuracy.get("samples", 0),
            "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 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", {})}