"""Thermal zones: rooms or areas, each with one or more temperature sensors.

Pure Python; no Home Assistant imports.

Two kinds of weight are kept apart on purpose:

* sensor weights combine several sensors in one zone into one temperature;
* the comfort weight says how much the zone matters to the MPC.

Neither is a physical coefficient. A cold bedroom can get a higher comfort
weight without its learned heat loss changing.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from math import isfinite
from typing import Any, Iterable, Mapping

SUBENTRY_ZONE = "zone"
MAX_ZONES = 12
MAX_SENSORS_PER_ZONE = 6
COMFORT_WEIGHT_RANGE = (0.1, 10.0)
SENSOR_WEIGHT_RANGE = (0.1, 10.0)
TARGET_OFFSET_RANGE = (-3.0, 3.0)

# Subentry data keys.
ZONE_KEY = "key"
ZONE_NAME = "name"
ZONE_SENSORS = "sensors"
ZONE_SENSOR_WEIGHTS = "sensor_weights"
ZONE_COMFORT_WEIGHT = "comfort_weight"
ZONE_TARGET_OFFSET = "target_offset"
ZONE_NEIGHBORS = "neighbors"


@dataclass(frozen=True)
class Zone:
    # Stable identity for learned models and entity unique IDs. Zones migrated from
    # 0.3 keep their sensor's entity ID; new zones get a generated key.
    key: str
    name: str
    sensors: tuple[str, ...]
    sensor_weights: Mapping[str, float] = field(default_factory=dict)
    comfort_weight: float = 1.0
    target_offset: float = 0.0
    neighbors: tuple[str, ...] = ()
    subentry_id: str | None = None

    def sensor_weight(self, sensor: str) -> float:
        return float(self.sensor_weights.get(sensor, 1.0))


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 zone_from_data(data: Mapping[str, Any], subentry_id: str | None = None) -> Zone | None:
    """Parse one zone's stored data; None if it cannot be used."""
    key = data.get(ZONE_KEY) or subentry_id
    sensors = data.get(ZONE_SENSORS)
    if not isinstance(key, str) or not isinstance(sensors, (list, tuple)):
        return None
    sensors = tuple(dict.fromkeys(s for s in sensors if isinstance(s, str) and s.startswith("sensor.")))
    if not sensors:
        return None
    raw_weights = data.get(ZONE_SENSOR_WEIGHTS)
    raw_weights = raw_weights if isinstance(raw_weights, Mapping) else {}
    weights = {s: _number(raw_weights.get(s), 1.0, *SENSOR_WEIGHT_RANGE) for s in sensors}
    neighbors = data.get(ZONE_NEIGHBORS)
    neighbors = tuple(n for n in neighbors if isinstance(n, str) and n != key) \
        if isinstance(neighbors, (list, tuple)) else ()
    name = data.get(ZONE_NAME)
    return Zone(
        key=key,
        name=name if isinstance(name, str) and name.strip() else key,
        sensors=sensors,
        sensor_weights=weights,
        comfort_weight=_number(data.get(ZONE_COMFORT_WEIGHT), 1.0, *COMFORT_WEIGHT_RANGE),
        target_offset=_number(data.get(ZONE_TARGET_OFFSET), 0.0, *TARGET_OFFSET_RANGE),
        neighbors=neighbors,
        subentry_id=subentry_id,
    )


def legacy_zone_data(rooms: Iterable[str], weights: Mapping[str, Any] | None,
                     names: Mapping[str, str] | None = None) -> list[dict[str, Any]]:
    """Zone data for 0.3 configurations: one zone per room sensor.

    The sensor's entity ID becomes the zone key, so learned models, forecast
    entities and accuracy statistics carry over unchanged.
    """
    weights = weights if isinstance(weights, Mapping) else {}
    names = names or {}
    result = []
    for room in dict.fromkeys(r for r in rooms if isinstance(r, str)):
        result.append({
            ZONE_KEY: room,
            ZONE_NAME: names.get(room) or room,
            ZONE_SENSORS: [room],
            ZONE_SENSOR_WEIGHTS: {room: 1.0},
            ZONE_COMFORT_WEIGHT: _number(weights.get(room), 1.0, *COMFORT_WEIGHT_RANGE),
            ZONE_TARGET_OFFSET: 0.0,
            ZONE_NEIGHBORS: [],
        })
    return result


def build_zones(items: Iterable[tuple[Mapping[str, Any], str | None]]) -> list[Zone]:
    """Valid zones from (data, subentry_id) pairs: unique keys, each sensor in one zone."""
    zones: list[Zone] = []
    keys: set[str] = set()
    taken: set[str] = set()
    for data, subentry_id in items:
        zone = zone_from_data(data, subentry_id)
        if zone is None or zone.key in keys or taken.intersection(zone.sensors):
            continue
        zones.append(zone)
        keys.add(zone.key)
        taken.update(zone.sensors)
        if len(zones) == MAX_ZONES:
            break
    # Neighbours must exist; exchange is symmetric, so record both directions.
    linked = {zone.key: set(n for n in zone.neighbors if n in keys) for zone in zones}
    for key, others in list(linked.items()):
        for other in others:
            linked[other].add(key)
    return [Zone(**{**zone.__dict__, "neighbors": tuple(sorted(linked[zone.key]))}) for zone in zones]


def zone_temperature(zone: Zone, readings: Mapping[str, float | None]) -> tuple[float | None, frozenset[str]]:
    """Sensor-weighted mean of the zone's usable readings, and which sensors it used.

    The set of sensors matters for learning: when a sensor drops out or rejoins,
    the mean jumps by the offset between sensors, which is not a real change.
    """
    used = [s for s in zone.sensors if readings.get(s) is not None]
    if not used:
        return None, frozenset()
    total = sum(zone.sensor_weight(s) for s in used)
    value = sum(zone.sensor_weight(s) * float(readings[s]) for s in used) / total
    return value, frozenset(used)


def validate_zone(data: Mapping[str, Any], others: Iterable[Zone], weather_sensors: Iterable[str]) -> str | None:
    """Translation key of the first problem with a zone form, or None."""
    sensors = data.get(ZONE_SENSORS) or []
    if not str(data.get(ZONE_NAME, "")).strip():
        return "zone_name_required"
    if not sensors or len(sensors) > MAX_SENSORS_PER_ZONE or len(set(sensors)) != len(sensors):
        return "zone_sensors"
    others = list(others)
    if any(set(sensors).intersection(zone.sensors) for zone in others):
        return "sensor_in_other_zone"
    if set(sensors).intersection(weather_sensors):
        return "sensor_is_weather_source"
    if len(others) >= MAX_ZONES:
        return "too_many_zones"
    return None