from __future__ import annotations from bisect import bisect_left from collections.abc import Mapping from csv import reader from dataclasses import dataclass from decimal import Decimal from enum import StrEnum from functools import partial import gzip import logging import os from typing import Any, TextIO, cast from homeassistant.components import light from homeassistant.components.light import ( ATTR_BRIGHTNESS, ATTR_COLOR_MODE, ATTR_COLOR_TEMP_KELVIN, ATTR_EFFECT, ATTR_HS_COLOR, COLOR_MODES_COLOR, ColorMode, ) from homeassistant.core import HomeAssistant, State from homeassistant.util.color import color_temperature_kelvin_to_mired, color_temperature_to_hs from custom_components.powercalc.common import SourceEntity from custom_components.powercalc.errors import ( LutFileNotFoundError, StrategyConfigurationError, ) from custom_components.powercalc.power_profile.power_profile import PowerProfile from .strategy_interface import PowerCalculationStrategyInterface _LOGGER = logging.getLogger(__name__) BrightnessLutValue = float ColorTempLutValue = dict[int, float] SatLutValue = dict[int, float] HsLutValue = dict[int, SatLutValue] LookupDictValue = BrightnessLutValue | ColorTempLutValue | HsLutValue LookupDictType = dict[int, LookupDictValue] EffectTableType = dict[str, dict[int, float]] class LookupMode(StrEnum): EFFECT = "effect" BRIGHTNESS = "brightness" COLOR_TEMP = "color_temp" HS = "hs" @staticmethod def from_color_mode(color_mode: ColorMode) -> LookupMode: return LookupMode(color_mode.value) @dataclass class _LutEntry: """Holds a lookup dictionary together with its pre-sorted brightness key list.""" table: LookupDictType sorted_keys: list[int] @dataclass class _EffectEntry: """Holds an effect lookup table (str → {brightness: power}).""" table: EffectTableType # manufacturer, model, lookup mode, sub profile _CacheKey = tuple[str, str, LookupMode, str | None] class LutRegistry: def __init__(self, hass: HomeAssistant) -> None: self._hass = hass self._lut_entries: dict[_CacheKey, _LutEntry] = {} self._effect_entries: dict[_CacheKey, _EffectEntry] = {} self._supported_modes: dict[tuple[str, str, str], set[LookupMode]] = {} async def get_lookup_entry( self, power_profile: PowerProfile, lookup_mode: LookupMode, ) -> _LutEntry: """Return a cached _LutEntry for the given profile and mode.""" cache_key = self._cache_key(power_profile, lookup_mode) entry = self._lut_entries.get(cache_key) if entry is None: entry = await self._hass.async_add_executor_job(partial(self._load_lut_entry, power_profile, lookup_mode)) self._lut_entries[cache_key] = entry return entry async def get_effect_entry( self, power_profile: PowerProfile, ) -> _EffectEntry: """Return a cached _EffectEntry for the given profile.""" cache_key = self._cache_key(power_profile, LookupMode.EFFECT) entry = self._effect_entries.get(cache_key) if entry is None: entry = await self._hass.async_add_executor_job(partial(self._load_effect_entry, power_profile)) self._effect_entries[cache_key] = entry return entry async def get_supported_modes(self, power_profile: PowerProfile) -> set[LookupMode]: """Return the LUT modes supported by the profile.""" cache_key = (power_profile.manufacturer, power_profile.model, "supported_modes") supported_modes = self._supported_modes.get(cache_key) if supported_modes is None: supported_modes = set() filenames = cast( list[str], await self._hass.async_add_executor_job(os.listdir, power_profile.get_model_directory()), ) for filename in filenames: if filename.endswith((".csv.gz", ".csv")): base_name = filename.split(".", 1)[0] supported_modes.add(LookupMode(base_name)) self._supported_modes[cache_key] = supported_modes return supported_modes @staticmethod def _cache_key(power_profile: PowerProfile, lookup_mode: LookupMode) -> _CacheKey: return power_profile.manufacturer, power_profile.model, lookup_mode, power_profile.sub_profile @classmethod def _load_lut_entry(cls, power_profile: PowerProfile, lookup_mode: LookupMode) -> _LutEntry: """Load a non-effect CSV into a typed _LutEntry.""" raw: dict[int, Any] = {} csv_file = cls.get_lut_file(power_profile, lookup_mode) line_count = 0 with csv_file: csv_reader = reader(csv_file) next(csv_reader) # skip header row for row in csv_reader: if lookup_mode == LookupMode.HS: bri_key = int(row[0]) hue_key = int(row[1]) sat_key = int(row[2]) raw.setdefault(bri_key, {}).setdefault(hue_key, {})[sat_key] = float(row[3]) elif lookup_mode == LookupMode.COLOR_TEMP: bri_key = int(row[0]) ct_key = int(row[1]) raw.setdefault(bri_key, {})[ct_key] = float(row[2]) else: raw[int(row[0])] = float(row[1]) line_count += 1 _LOGGER.debug("LUT file loaded: %d lines", line_count) table = cast(LookupDictType, raw) return _LutEntry(table=table, sorted_keys=sorted(table.keys())) @classmethod def _load_effect_entry(cls, power_profile: PowerProfile) -> _EffectEntry: """Load an effect CSV into a typed _EffectEntry.""" raw: dict[str, dict[int, float]] = {} csv_file = cls.get_lut_file(power_profile, LookupMode.EFFECT) line_count = 0 with csv_file: csv_reader = reader(csv_file) next(csv_reader) # skip header row for row in csv_reader: effect_name: str = row[0] bri_key = int(row[1]) raw.setdefault(effect_name, {})[bri_key] = float(row[2]) line_count += 1 _LOGGER.debug("Effect LUT file loaded: %d lines", line_count) return _EffectEntry(table=raw) @staticmethod def get_lut_file(power_profile: PowerProfile, lookup_mode: LookupMode) -> TextIO: """ Open the LUT file for the given power profile and color mode. When the file is gzipped it will be decompressed transparently. """ path = os.path.join(power_profile.get_model_directory(), f"{lookup_mode}.csv") gzip_path = f"{path}.gz" if os.path.exists(gzip_path): _LOGGER.debug("Loading LUT data file: %s", gzip_path) return gzip.open(gzip_path, "rt") if os.path.exists(path): _LOGGER.debug("Loading LUT data file: %s", path) return open(path) raise LutFileNotFoundError(f"Data file not found: {path}") class LutStrategy(PowerCalculationStrategyInterface): def __init__( self, source_entity: SourceEntity, lut_registry: LutRegistry, profile: PowerProfile, ) -> None: self._source_entity = source_entity self._lut_registry = lut_registry self._profile = profile self._supported_modes: set[LookupMode] = set() self._effect_entry: _EffectEntry | None = None async def initialize(self) -> None: self._supported_modes = await self._lut_registry.get_supported_modes(self._profile) async def calculate(self, entity_state: State) -> Decimal | None: """Calculate the power consumption based on brightness, mired, hsl or effect.""" attrs = entity_state.attributes brightness = attrs.get(ATTR_BRIGHTNESS) if brightness is None: _LOGGER.warning( "%s: Could not calculate power. no brightness set", entity_state.entity_id, ) return None if brightness > 255: brightness = 255 color_mode = await self.get_selected_color_mode(attrs) if color_mode == ColorMode.UNKNOWN: _LOGGER.warning( "%s: Could not calculate power. color mode unknown", entity_state.entity_id, ) return None effect = attrs.get(ATTR_EFFECT) if effect and str(effect).lower() not in ("off", "none", "white"): return await self._calculate_effect_power(entity_state, str(effect), brightness) lut_mode = LookupMode.from_color_mode(color_mode) try: lut_entry = await self._lut_registry.get_lookup_entry(self._profile, lut_mode) except LutFileNotFoundError: _LOGGER.error( "%s: Lookup table not found for color mode (model: %s, color_mode: %s)", entity_state.entity_id, self._profile.model, color_mode, ) return None light_setting = self.create_light_setting(entity_state, color_mode, brightness) if light_setting is None: return None _LOGGER.debug( "%s: Looking up power usage with settings: %s", entity_state.entity_id, {attr: getattr(light_setting, attr) for attr in vars(light_setting)}, ) power = self.lookup_power(lut_entry, light_setting) _LOGGER.debug("%s: Calculated power:%s", entity_state.entity_id, power) return Decimal(power) async def _calculate_effect_power( self, entity_state: State, effect: str, brightness: int, ) -> Decimal | None: """Look up power for an active light effect.""" if LookupMode.EFFECT not in self._supported_modes: _LOGGER.warning("%s: Effects not supported for this power profile", entity_state.entity_id) return None effect_entry = await self._lut_registry.get_effect_entry(self._profile) effect_table = effect_entry.table.get(effect) if effect_table is None: _LOGGER.warning('%s: Effect "%s" not found in LUT', entity_state.entity_id, effect) return None sorted_keys = sorted(effect_table.keys()) return Decimal(self._interpolate(effect_table, sorted_keys, brightness)) def create_light_setting( self, entity_state: State, color_mode: ColorMode, brightness: int, ) -> LightSetting | None: """Create a LightSetting object based on the entity state.""" light_setting = LightSetting(color_mode=color_mode, brightness=brightness) attrs = entity_state.attributes if color_mode == ColorMode.COLOR_TEMP: color_temp = attrs.get(ATTR_COLOR_TEMP_KELVIN) if color_temp is None: _LOGGER.error( "%s: Could not calculate power. no color temp set. " "Please check the attributes of your light in the developer tools.", entity_state.entity_id, ) return None light_setting.color_temp = color_temperature_kelvin_to_mired(color_temp) return light_setting if color_mode == ColorMode.HS: try: original_color_mode = attrs.get(ATTR_COLOR_MODE) hs = ( color_temperature_to_hs(attrs[ATTR_COLOR_TEMP_KELVIN]) if original_color_mode == ColorMode.COLOR_TEMP else attrs[ATTR_HS_COLOR] ) light_setting.hue = int(hs[0] / 360 * 65535) light_setting.saturation = int(hs[1] / 100 * 255) except KeyError, TypeError, ValueError: _LOGGER.error( "%s: Could not calculate power. no hue/sat set. " "Please check the attributes of your light in the developer tools.", entity_state.entity_id, ) return None return light_setting async def get_selected_color_mode(self, attrs: Mapping[str, Any]) -> ColorMode: """Get the selected color mode for the entity.""" try: color_mode = ColorMode(str(attrs.get(ATTR_COLOR_MODE, ColorMode.UNKNOWN))) except ValueError: color_mode = ColorMode.UNKNOWN if color_mode == ColorMode.WHITE: return ColorMode.BRIGHTNESS if color_mode == ColorMode.UNKNOWN: return color_mode if color_mode in COLOR_MODES_COLOR: color_mode = ColorMode.HS lookup_mode = LookupMode.from_color_mode(color_mode) if lookup_mode not in self._supported_modes and color_mode == ColorMode.COLOR_TEMP: _LOGGER.debug("Color mode not natively supported, falling back to HS") color_mode = ColorMode.HS return color_mode @staticmethod def _nearest_key(sorted_keys: list[int], x: int) -> int: """Return the key in sorted_keys nearest to x.""" i = bisect_left(sorted_keys, x) if i == 0: return sorted_keys[0] if i >= len(sorted_keys): return sorted_keys[-1] before = sorted_keys[i - 1] after = sorted_keys[i] return before if (x - before) <= (after - x) else after @staticmethod def _interpolate(table: dict[int, float], sorted_keys: list[int], brightness: int) -> float: """Linear interpolation over a flat {brightness: power} table.""" if brightness in table: return table[brightness] i = bisect_left(sorted_keys, brightness) if i == 0: return table[sorted_keys[0]] if i >= len(sorted_keys): return table[sorted_keys[-1]] b0 = sorted_keys[i - 1] b1 = sorted_keys[i] p0 = table[b0] p1 = table[b1] return p0 + (p1 - p0) * ((brightness - b0) / (b1 - b0)) def lookup_power( self, lut_entry: _LutEntry, light_setting: LightSetting, ) -> float: lookup_table = lut_entry.table sorted_keys = lut_entry.sorted_keys brightness = light_setting.brightness # Exact brightness match — skip interpolation entirely. if brightness in lookup_table: return self.lookup_power_for_brightness(lookup_table[brightness], light_setting) i = bisect_left(lut_entry.sorted_keys, brightness) if i == 0: return self.lookup_power_for_brightness(lookup_table[sorted_keys[0]], light_setting) if i >= len(sorted_keys): return self.lookup_power_for_brightness(lookup_table[sorted_keys[-1]], light_setting) b0 = sorted_keys[i - 1] b1 = sorted_keys[i] p0 = self.lookup_power_for_brightness(lookup_table[b0], light_setting) p1 = self.lookup_power_for_brightness(lookup_table[b1], light_setting) return p0 + (p1 - p0) * ((brightness - b0) / (b1 - b0)) def lookup_power_for_brightness( self, lut_value: LookupDictValue, light_setting: LightSetting, ) -> float: if isinstance(lut_value, float): return lut_value if light_setting.color_mode == ColorMode.COLOR_TEMP: return self.get_nearest(lut_value, light_setting.color_temp or 0) # type: ignore # HS path — outer dict is hue → {saturation → power} hs_table = cast(HsLutValue, lut_value) sat_values = self.get_nearest(hs_table, light_setting.hue or 0) return self.get_nearest(sat_values, light_setting.saturation or 0) def get_nearest[NearestT: (float, dict[int, float])]( self, lookup_dict: dict[int, NearestT], search_key: int, ) -> NearestT: """Return the value mapped at search_key or the nearest neighbour key.""" value = lookup_dict.get(search_key) if value is not None: return value nearest = self._nearest_key(sorted(lookup_dict.keys()), search_key) return lookup_dict[nearest] async def validate_config(self) -> None: if self._source_entity.domain != light.DOMAIN: raise StrategyConfigurationError( "Only light entities can use the LUT mode", "lut_unsupported_color_mode", ) @dataclass class LightSetting: color_mode: ColorMode brightness: int hue: int | None = None saturation: int | None = None color_temp: int | None = None effect: str | None = None