"""Electrical internal heat gains (UPS, PCs, servers) per zone.

Pure Python; no Home Assistant imports.

Most electrical energy ends up as heat in the room. What matters to the thermal
model is the *variation* around a persistent baseline: the constant part (for
example a UPS drawing 200 W all day) is already absorbed by the model's bias
term, so feeding it again would count it twice. Each source therefore supplies

    variable heat = fraction x (effective power - baseline)

where the baseline is a robust low percentile of the last 24 hours (or a value
set by the user), and "effective power" removes separately measured devices
that are already included in a parent measurement (a PC behind a UPS). The
variable heat reaches the room through a first-order response, because the air,
furniture and equipment take time to warm.
"""
from __future__ import annotations

from collections import deque
from dataclasses import dataclass
from math import exp, isfinite
from typing import Any, Iterable, Mapping

SUBENTRY_HEAT_SOURCE = "heat_source"
MAX_SOURCES = 16
MAX_POWER_W = 20000.0
BASELINE_WINDOW = 288  # 24 h of 5-minute samples
BASELINE_PERCENTILE = 0.2
BASELINE_MIN_SAMPLES = 12  # 1 h before the automatic baseline is trusted
RESPONSE_MINUTES_RANGE = (1.0, 180.0)
DEFAULT_RESPONSE_MINUTES = 15.0

SOURCE_KEY = "key"
SOURCE_NAME = "name"
SOURCE_SENSOR = "sensor"
SOURCE_ZONE = "zone"
SOURCE_FRACTION = "heat_fraction"
SOURCE_BASELINE_MODE = "baseline_mode"
SOURCE_BASELINE_W = "baseline_w"
SOURCE_INCLUDED_IN = "included_in"
SOURCE_RESPONSE = "response_minutes"
SOURCE_ENABLED = "enabled"
BASELINE_AUTO = "auto"
BASELINE_MANUAL = "manual"

# Active power only: VA (apparent power) and % (UPS load) are not heat.
POWER_UNITS = {"w": 1.0, "kw": 1000.0, "mw": 0.001}


def normalize_power(raw: Any, unit: Any) -> float | None:
    """Active power in W, or None when the value or unit is unusable."""
    if not isinstance(unit, str) or unit.strip().lower() not in POWER_UNITS:
        return None
    try:
        value = float(raw) * POWER_UNITS[unit.strip().lower()]
    except (TypeError, ValueError):
        return None
    return value if isfinite(value) and 0 <= value <= MAX_POWER_W else None


@dataclass(frozen=True)
class HeatSource:
    key: str
    name: str
    sensor: str
    zone: str
    fraction: float = 1.0
    baseline_mode: str = BASELINE_AUTO
    baseline_w: float = 0.0
    included_in: str | None = None
    response_minutes: float = DEFAULT_RESPONSE_MINUTES
    enabled: bool = True
    subentry_id: str | None = None


def _number(value: Any, default: float, lo: float, hi: float) -> float:
    try:
        number = float(value)
    except (TypeError, ValueError):
        return default
    return number if isfinite(number) and lo <= number <= hi else default


def source_from_data(data: Mapping[str, Any], subentry_id: str | None = None) -> HeatSource | None:
    key = data.get(SOURCE_KEY) or subentry_id
    sensor, zone = data.get(SOURCE_SENSOR), data.get(SOURCE_ZONE)
    if not (isinstance(key, str) and isinstance(sensor, str) and sensor.startswith("sensor.")
            and isinstance(zone, str)):
        return None
    name = data.get(SOURCE_NAME)
    parent = data.get(SOURCE_INCLUDED_IN)
    mode = data.get(SOURCE_BASELINE_MODE)
    return HeatSource(
        key=key, name=name if isinstance(name, str) and name.strip() else key, sensor=sensor, zone=zone,
        fraction=_number(data.get(SOURCE_FRACTION), 1.0, 0.0, 1.0),
        baseline_mode=mode if mode in (BASELINE_AUTO, BASELINE_MANUAL) else BASELINE_AUTO,
        baseline_w=_number(data.get(SOURCE_BASELINE_W), 0.0, 0.0, MAX_POWER_W),
        included_in=parent if isinstance(parent, str) and parent and parent != key else None,
        response_minutes=_number(data.get(SOURCE_RESPONSE), DEFAULT_RESPONSE_MINUTES, *RESPONSE_MINUTES_RANGE),
        enabled=data.get(SOURCE_ENABLED, True) is not False,
        subentry_id=subentry_id,
    )


def build_sources(items: Iterable[tuple[Mapping[str, Any], str | None]], zone_keys: Iterable[str]) -> list[HeatSource]:
    """Valid sources: unique keys and sensors, existing zones, one level of inclusion."""
    zone_keys = set(zone_keys)
    sources: list[HeatSource] = []
    for data, subentry_id in items:
        source = source_from_data(data, subentry_id)
        if (source is None or source.zone not in zone_keys or len(sources) == MAX_SOURCES
                or any(s.key == source.key or s.sensor == source.sensor for s in sources)):
            continue
        sources.append(source)
    keys = {s.key: s for s in sources}
    # A parent must exist and must not itself be included in another source.
    return [s if s.included_in in keys and keys[s.included_in].included_in is None
            else HeatSource(**{**s.__dict__, "included_in": None}) if s.included_in else s
            for s in sources]


def effective_power(sources: list[HeatSource], readings: Mapping[str, float | None]) -> dict[str, float | None]:
    """Power of each enabled source after removing included child sources (W).

    `readings` maps source key -> measured W (None if unavailable). A parent is
    unknown when any enabled child it contains is unknown, since the split is
    then impossible. Values below 0 (timing differences) are clamped.
    """
    result: dict[str, float | None] = {}
    for source in sources:
        if not source.enabled:
            continue
        power = readings.get(source.key)
        children = [s for s in sources if s.included_in == source.key and s.enabled]
        if power is None or any(readings.get(c.key) is None for c in children):
            result[source.key] = None
            continue
        result[source.key] = max(0.0, power - sum(readings[c.key] for c in children))
    return result


class BaselineEstimator:
    """Robust persistent load: low percentile of the last 24 hours of 5-minute samples."""

    def __init__(self, samples: Iterable[float] = ()) -> None:
        self.samples: deque[float] = deque(
            (float(v) for v in samples if isinstance(v, (int, float)) and isfinite(v) and 0 <= v <= MAX_POWER_W),
            maxlen=BASELINE_WINDOW)

    def add(self, watts: float) -> None:
        self.samples.append(watts)

    def value(self) -> float | None:
        if len(self.samples) < BASELINE_MIN_SAMPLES:
            return None
        ordered = sorted(self.samples)
        return ordered[int(BASELINE_PERCENTILE * (len(ordered) - 1))]


class ResponseFilter:
    """First-order thermal response of a heat input (kW), like the radiator lag."""

    def __init__(self, minutes: float) -> None:
        self.tau = minutes * 60.0
        self.value: float | None = None

    def advance(self, held_kw: float, seconds: float) -> float:
        """Advance over an interval during which `held_kw` was dissipated."""
        if self.value is None:
            self.value = held_kw  # assume the current load has been steady
        else:
            self.value += (1.0 - exp(-max(0.0, seconds) / self.tau)) * (held_kw - self.value)
        return self.value

    def reset(self) -> None:
        self.value = None


def forecast_internal(current_kw: float, steps: int, half_life_minutes: float | None,
                      step_minutes: float = 5.0) -> list[float]:
    """Variable heat assumed over a forecast: decays to baseline (0) with a half-life.

    A gaming PC drawing 600 W now is not assumed to keep doing so for two hours.
    `half_life_minutes=None` keeps the current value (persistence).
    """
    if half_life_minutes is None:
        return [current_kw] * steps
    if half_life_minutes <= 0:
        return [0.0] * steps
    factor = 0.5 ** (step_minutes / half_life_minutes)
    return [current_kw * factor ** (k + 1) for k in range(steps)]


def validate_source(data: Mapping[str, Any], others: Iterable[HeatSource]) -> str | None:
    """Translation key of the first problem with a heat-source form, or None."""
    others = list(others)
    if not str(data.get(SOURCE_NAME, "")).strip():
        return "source_name_required"
    if any(s.sensor == data.get(SOURCE_SENSOR) for s in others):
        return "sensor_in_other_source"
    parent = data.get(SOURCE_INCLUDED_IN)
    if parent:
        parents = {s.key: s for s in others}
        if parent not in parents or parents[parent].included_in:
            return "invalid_parent_source"
        key = data.get(SOURCE_KEY)
        if key and any(s.included_in == key for s in others):
            return "source_has_children"
    if data.get(SOURCE_BASELINE_MODE) == BASELINE_MANUAL and data.get(SOURCE_BASELINE_W) in (None, ""):
        return "baseline_required"
    if len(others) >= MAX_SOURCES:
        return "too_many_sources"
    return None


def forecast_response(response_kw: float, variable_kw: float, steps: int, tau_minutes: float,
                      half_life_minutes: float | None, step_minutes: float = 5.0) -> list[float]:
    """Thermal response (kW) at the end of each forecast step.

    The variable power returns towards baseline (see forecast_internal) and the
    response keeps following it with its own time constant.
    """
    alpha = 1.0 - exp(-step_minutes / tau_minutes)
    held = [variable_kw] + forecast_internal(variable_kw, steps - 1, half_life_minutes, step_minutes)
    result, value = [], response_kw
    for power in held:
        value += alpha * (power - value)
        result.append(value)
    return result


class HeatSourceTracker:
    """Per-poll processing of all heat sources: baselines, responses, zone inputs."""

    def __init__(self, sources: list[HeatSource], saved: Any = None,
                 fingerprints: Mapping[str, str] | None = None) -> None:
        self.sources = sources
        self.fingerprints = dict(fingerprints or {})
        saved = saved if isinstance(saved, Mapping) else {}
        stored = saved.get("baselines") if isinstance(saved.get("baselines"), Mapping) else {}
        old_prints = saved.get("fingerprints") if isinstance(saved.get("fingerprints"), Mapping) else None
        # Samples of a replaced sensor (or a changed inclusion) describe a different load.
        self.reset_baselines = [s.key for s in sources if old_prints is not None and s.key in old_prints
                                and old_prints[s.key] != self.fingerprints.get(s.key)]
        self._baselines = {s.key: BaselineEstimator(
            stored.get(s.key) if isinstance(stored.get(s.key), list) and s.key not in self.reset_baselines else ())
            for s in sources}
        self._filters = {s.key: ResponseFilter(s.response_minutes) for s in sources}
        self._held: dict[str, tuple[float, float | None]] = {}  # key -> (timestamp, variable kW)
        self._settled_since: dict[str, float] = {}
        self.state: dict[str, dict[str, Any]] = {}

    def update(self, now: float, readings: Mapping[str, float | None]) -> None:
        """Process one poll. `now` is a timestamp; `readings` maps key -> measured W or None."""
        effective = effective_power(self.sources, readings)
        self.state = {}
        for source in self.sources:
            info: dict[str, Any] = {"power_w": readings.get(source.key), "enabled": source.enabled}
            self.state[source.key] = info
            if not source.enabled:
                continue
            watts = effective.get(source.key)
            estimator = self._baselines[source.key]
            if watts is not None:
                estimator.add(watts)
            baseline = source.baseline_w if source.baseline_mode == BASELINE_MANUAL else estimator.value()
            variable = None if watts is None or baseline is None else source.fraction * (watts - baseline) / 1000
            # The heat during the interval just ended is what was drawn at its start.
            held = self._held.get(source.key)
            filt = self._filters[source.key]
            if held and held[1] is not None and 0 < now - held[0] <= 900:
                response = filt.advance(held[1], now - held[0])
            else:
                filt.reset()
                self._settled_since.pop(source.key, None)
                response = None if variable is None else filt.advance(variable, 0)
            if response is not None:
                self._settled_since.setdefault(source.key, now)
            self._held[source.key] = (now, variable)
            info.update(effective_w=watts, baseline_w=baseline, variable_kw=variable, response_kw=response,
                        settled=response is not None
                        and now - self._settled_since[source.key] >= 2 * filt.tau)

    def zone_input(self, zone: str) -> float | None:
        """Summed thermal response of a zone's sources (kW); None if unknown or unsettled."""
        total = 0.0
        for source in self.sources:
            if source.zone != zone or not source.enabled:
                continue
            info = self.state.get(source.key, {})
            if not info.get("settled"):
                return None
            total += info["response_kw"]
        return total

    def zone_forecast(self, zone: str, steps: int, half_life_minutes: float | None) -> list[float]:
        """Assumed thermal response of a zone's sources over a forecast (kW per step)."""
        total = [0.0] * steps
        for source in self.sources:
            info = self.state.get(source.key, {})
            if source.zone != zone or not source.enabled or info.get("response_kw") is None:
                continue
            variable = info["variable_kw"] if info.get("variable_kw") is not None else 0.0
            for k, value in enumerate(forecast_response(info["response_kw"], variable, steps,
                                                        source.response_minutes, half_life_minutes)):
                total[k] += value
        return total

    def as_dict(self) -> dict[str, Any]:
        return {"baselines": {key: [round(v, 1) for v in est.samples] for key, est in self._baselines.items()},
                "fingerprints": self.fingerprints}