core/homeassistant/components/amberelectric/sensor.py

"""Amber Electric Sensor definitions."""

# There are three types of sensor: Current, Forecast and Grid
# Current and forecast will create general, controlled load and feed in as required
# At the moment renewables in the only grid sensor.


from __future__ import annotations

from collections.abc import Mapping
from typing import Any

from amberelectric.model.channel import ChannelType
from amberelectric.model.current_interval import CurrentInterval
from amberelectric.model.forecast_interval import ForecastInterval

from homeassistant.components.sensor import (
    STATE_CLASS_MEASUREMENT,
    SensorEntity,
    SensorEntityDescription,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_ATTRIBUTION, CURRENCY_DOLLAR, ENERGY_KILO_WATT_HOUR
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity

from .const import DOMAIN
from .coordinator import AmberUpdateCoordinator

ATTRIBUTION = "Data provided by Amber Electric"

ICONS = {
    "general": "mdi:transmission-tower",
    "controlled_load": "mdi:clock-outline",
    "feed_in": "mdi:solar-power",
}

UNIT = f"{CURRENCY_DOLLAR}/{ENERGY_KILO_WATT_HOUR}"


def friendly_channel_type(channel_type: str) -> str:
    """Return a human readable version of the channel type."""
    if channel_type == "controlled_load":
        return "Controlled Load"
    if channel_type == "feed_in":
        return "Feed In"
    return "General"


class AmberSensor(CoordinatorEntity, SensorEntity):
    """Amber Base Sensor."""

    def __init__(
        self,
        coordinator: AmberUpdateCoordinator,
        description: SensorEntityDescription,
        channel_type: ChannelType,
    ) -> None:
        """Initialize the Sensor."""
        super().__init__(coordinator)
        self.site_id = coordinator.site_id
        self.entity_description = description
        self.channel_type = channel_type

    @property
    def unique_id(self) -> None:
        """Return a unique id for each sensors."""
        self._attr_unique_id = (
            f"{self.site_id}-{self.entity_description.key}-{self.channel_type}"
        )


class AmberPriceSensor(AmberSensor):
    """Amber Price Sensor."""

    @property
    def native_value(self) -> str | None:
        """Return the current price in $/kWh."""
        interval = self.coordinator.data[self.entity_description.key][self.channel_type]

        if interval.channel_type == ChannelType.FEED_IN:
            return round(interval.per_kwh, 0) / 100 * -1
        return round(interval.per_kwh, 0) / 100

    @property
    def device_state_attributes(self) -> Mapping[str, Any] | None:
        """Return additional pieces of information about the price."""
        interval = self.coordinator.data[self.entity_description.key][self.channel_type]

        data: dict[str, Any] = {ATTR_ATTRIBUTION: ATTRIBUTION}
        if interval is None:
            return data

        data["duration"] = interval.duration
        data["date"] = interval.date.isoformat()
        data["per_kwh"] = round(interval.per_kwh)
        if interval.channel_type == ChannelType.FEED_IN:
            data["per_kwh"] = data["per_kwh"] * -1
        data["nem_date"] = interval.nem_time.isoformat()
        data["spot_per_kwh"] = round(interval.spot_per_kwh)
        data["start_time"] = interval.start_time.isoformat()
        data["end_time"] = interval.end_time.isoformat()
        data["renewables"] = round(interval.renewables)
        data["estimate"] = interval.estimate
        data["spike_status"] = interval.spike_status.value
        data["channel_type"] = interval.channel_type.value

        if interval.range is not None:
            data["range_min"] = interval.range.min
            data["range_max"] = interval.range.max

        return data


class AmberForecastSensor(AmberSensor):
    """Amber Forecast Sensor."""

    @property
    def native_value(self) -> str | None:
        """Return the first forecast price in $/kWh."""
        intervals = self.coordinator.data[self.entity_description.key][
            self.channel_type
        ]
        interval = intervals[0]

        if interval.channel_type == ChannelType.FEED_IN:
            return round(interval.per_kwh, 0) / 100 * -1
        return round(interval.per_kwh, 0) / 100

    @property
    def device_state_attributes(self) -> Mapping[str, Any] | None:
        """Return additional pieces of information about the price."""
        intervals = self.coordinator.data[self.entity_description.key][
            self.channel_type
        ]

        data = {
            "forecasts": [],
            "channel_type": intervals[0].channel_type.value,
            ATTR_ATTRIBUTION: ATTRIBUTION,
        }

        for interval in intervals:
            datum = {}
            datum["duration"] = interval.duration
            datum["date"] = interval.date.isoformat()
            datum["nem_date"] = interval.nem_time.isoformat()
            datum["per_kwh"] = round(interval.per_kwh)
            if interval.channel_type == ChannelType.FEED_IN:
                datum["per_kwh"] = datum["per_kwh"] * -1
            datum["spot_per_kwh"] = round(interval.spot_per_kwh)
            datum["start_time"] = interval.start_time.isoformat()
            datum["end_time"] = interval.end_time.isoformat()
            datum["renewables"] = round(interval.renewables)
            datum["spike_status"] = interval.spike_status.value

            if interval.range is not None:
                datum["range_min"] = interval.range.min
                datum["range_max"] = interval.range.max

            data["forecasts"].append(datum)

        return data


class AmberGridSensor(CoordinatorEntity, SensorEntity):
    """Sensor to show single grid specific values."""

    def __init__(
        self,
        coordinator: AmberUpdateCoordinator,
        description: SensorEntityDescription,
    ) -> None:
        """Initialize the Sensor."""
        super().__init__(coordinator)
        self.site_id = coordinator.site_id
        self.entity_description = description
        self._attr_device_state_attributes = {ATTR_ATTRIBUTION: ATTRIBUTION}
        self._attr_unique_id = f"{coordinator.site_id}-{description.key}"

    @property
    def unique_id(self) -> None:
        """Return a unique id for each sensors."""
        self._attr_unique_id = f"{self.site_id}-{self.entity_description.key}"

    @property
    def native_value(self) -> str | None:
        """Return the value of the sensor."""
        return self.coordinator.data["grid"][self.entity_description.key]


async def async_setup_entry(
    hass: HomeAssistant,
    entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up a config entry."""
    coordinator: AmberUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]

    current: dict[str, CurrentInterval] = coordinator.data["current"]
    forecasts: dict[str, list[ForecastInterval]] = coordinator.data["forecasts"]

    entities: list = []
    for channel_type in current:
        description = SensorEntityDescription(
            key="current",
            name=f"{entry.title} - {friendly_channel_type(channel_type)} Price",
            native_unit_of_measurement=UNIT,
            state_class=STATE_CLASS_MEASUREMENT,
            icon=ICONS[channel_type],
        )
        entities.append(AmberPriceSensor(coordinator, description, channel_type))

    for channel_type in forecasts:
        description = SensorEntityDescription(
            key="forecasts",
            name=f"{entry.title} - {friendly_channel_type(channel_type)} Forecast",
            native_unit_of_measurement=UNIT,
            state_class=STATE_CLASS_MEASUREMENT,
            icon=ICONS[channel_type],
        )
        entities.append(AmberForecastSensor(coordinator, description, channel_type))

    renewables_description = SensorEntityDescription(
        key="renewables",
        name=f"{entry.title} - Renewables",
        native_unit_of_measurement="%",
        state_class=STATE_CLASS_MEASUREMENT,
        icon="mdi:solar-power",
    )
    entities.append(AmberGridSensor(coordinator, renewables_description))

    async_add_entities(entities)