"""Read existing HA entities, learn thermal response, produce shadow MPC decisions."""
from __future__ import annotations

import logging
from datetime import datetime, timedelta, timezone
from functools import partial
from math import exp, isfinite
from typing import Any

from homeassistant.core import HomeAssistant
from homeassistant.config_entries import ConfigEntry
from homeassistant.helpers.storage import Store
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator

from .const import (
    CONF_ROOM_MAX_AGE, CONF_WEATHER_MAX_AGE,
    DOMAIN, CONF_CLIMATE, CONF_HEATER, CONF_ROOMS, CONF_WEIGHTS, WEATHER_SOURCES,
    CONF_HORIZON, CONF_MIN_ON, CONF_MIN_OFF, CONF_SWITCH_COST, CONF_HEAT_COST,
    CONF_BAND, DEFAULTS,
)
from .data_quality import freshness
from .physics import (
    LAG_LEARN_MAX_UNCERTAINTY, LAG_TAU_SECONDS, Learner, advance_lag,
    lag_from_time_in_state, optimize,
)
from .weather_sources import choose_weather_sources, sources_for, normalize

_LOGGER = logging.getLogger(__name__)


class HeatingMPCCoordinator(DataUpdateCoordinator[dict[str, Any]]):
    """The existing generic thermostat stays the only actuator owner."""

    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
        super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=timedelta(minutes=5))
        self.entry = entry
        self.config = {**DEFAULTS, **entry.data, **entry.options}
        self.rooms: list[str] = self.config[CONF_ROOMS]
        self.learners: dict[str, Learner] = {room: Learner() for room in self.rooms}
        self._store = Store(hass, 1, f"{DOMAIN}.{entry.entry_id}")
        self._previous: dict[str, Any] | None = None
        # Radiator-lag proxy at the last poll with a valid relay state. Tracked
        # independently of room sensors so lag survives room outages and transitions.
        self._lag: dict[str, Any] | None = None
        self._on_samples = 0
        self._off_samples = 0

    async def async_restore_model(self) -> None:
        stored = await self._store.async_load() or {}
        if not isinstance(stored, dict) or stored.get("schema_version", 1) not in (1, 2):
            _LOGGER.warning("Ignoring unsupported Heating MPC model data")
            return
        # Replacing the actual heating system invalidates heating coefficients.
        # Switching among weather providers does not discard months of learning.
        if stored.get("heater_entity", self.config[CONF_HEATER]) != self.config[CONF_HEATER]:
            _LOGGER.warning("Heating MPC heater changed; starting a fresh thermal model")
            return
        # Models carry room identity, so changing the room list is safe.
        saved = stored.get("learners", {})
        if not isinstance(saved, dict):
            return
        for room, data in saved.items():
            if room in self.learners:
                self.learners[room] = Learner.from_dict(data)
        self._update_counts()

    def _update_counts(self) -> None:
        self._on_samples = min(m.on_count for m in self.learners.values())
        self._off_samples = min(m.off_count for m in self.learners.values())

    def _sensor_quality(self, entity_id: str, now: datetime, limit: float) -> dict:
        state = self.hass.states.get(entity_id)
        if state is None:
            return {"status": "missing", "age_seconds": None}
        if state.state in ("unknown", "unavailable"):
            return {"status": state.state, "age_seconds": None}
        # last_updated is a conservative compatibility fallback for older HA.
        return freshness(getattr(state, "last_reported", getattr(state, "last_updated", None)),
                         now, limit * 60)

    def _track_lag(self, now: datetime, heater_on: bool, last_changed: datetime) -> bool:
        """Advance the lag proxy to `now`; return True if the relay held one state throughout.

        A single relay change since the last poll is integrated at its timestamp.
        Anything ambiguous (several changes, an outage, clock jumps) restarts the
        estimate from the time spent in the current state, with an explicit error bound.
        """
        tracker = self._lag
        steady = False
        lag = uncertainty = None
        if tracker is not None:
            dt = (now - tracker["time"]).total_seconds()
            if dt > 0 and heater_on == tracker["heater_on"] and last_changed <= tracker["time"]:
                lag = advance_lag(tracker["lag"], heater_on, dt)
                steady = True
            elif dt > 0 and heater_on != tracker["heater_on"] and tracker["time"] < last_changed <= now:
                before = (last_changed - tracker["time"]).total_seconds()
                lag = advance_lag(tracker["lag"], tracker["heater_on"], before)
                lag = advance_lag(lag, heater_on, (now - last_changed).total_seconds())
            if lag is not None:
                # The lag update is a contraction, so an initial error decays the same way.
                uncertainty = tracker["uncertainty"] * exp(-dt / LAG_TAU_SECONDS)
        if lag is None:
            lag, uncertainty = lag_from_time_in_state(
                heater_on, (now - last_changed).total_seconds())
        self._lag = {"time": now, "heater_on": heater_on, "lag": lag, "uncertainty": uncertainty}
        return steady

    async def _async_update_data(self) -> dict[str, Any]:
        now = datetime.now(timezone.utc)
        heater = self.hass.states.get(self.config[CONF_HEATER])
        climate = self.hass.states.get(self.config[CONF_CLIMATE])
        # Choose each weather variable independently from ordered sensor sources.
        # Sensor priorities are editable in Options without touching integration code.
        weather_entity_ids = {
            entity_id for metric in WEATHER_SOURCES
            for entity_id in sources_for(self.config, metric)
        }
        watched = set(self.rooms) | weather_entity_ids | {
            self.config[CONF_CLIMATE], self.config[CONF_HEATER]}
        snapshot = {entity: self.hass.states.get(entity) for entity in watched}
        room_quality = {}
        temps = {}
        for room in self.rooms:
            quality = self._sensor_quality(room, now, self.config[CONF_ROOM_MAX_AGE])
            value = None
            if quality["status"] == "valid":
                state = snapshot[room]
                value = normalize("temperature", state.state, state.attributes.get("unit_of_measurement"))
                if value is None or not 0 <= value <= 45:
                    quality["status"] = "invalid_value_or_unit"
                    value = None
            room_quality[room] = quality
            temps[room] = value
        weather_states = {}
        weather_quality = {}
        for entity_id in weather_entity_ids:
            quality = self._sensor_quality(entity_id, now, self.config[CONF_WEATHER_MAX_AGE])
            weather_quality[entity_id] = quality
            if quality["status"] == "valid":
                state = snapshot[entity_id]
                weather_states[entity_id] = (state.state, state.attributes.get("unit_of_measurement"))
        weather, active_sources, weather_comparisons = choose_weather_sources(
            self.config, weather_states
        )
        metric_quality = {}
        for metric in WEATHER_SOURCES:
            metric_quality[metric] = {}
            for entity_id in sources_for(self.config, metric):
                quality = dict(weather_quality[entity_id])
                if quality["status"] == "valid" and entity_id not in weather_comparisons[metric]:
                    quality["status"] = "invalid_value_or_unit"
                metric_quality[metric][entity_id] = quality
        outside = weather["temperature"]
        status: dict[str, Any] = {
            "input_quality": {"rooms": room_quality, "weather": metric_quality},
            "status": "waiting_for_sensors", "proposal": None, "prediction_60m": {},
            "temperatures": {r: t for r, t in temps.items() if t is not None},
            "on_samples": self._on_samples, "off_samples": self._off_samples,
            "samples": min(learner.count for learner in self.learners.values()),
            "mode": climate.state if climate else "unavailable",
            "weather": {k: v for k, v in weather.items() if v is not None},
            "weather_active_sources": {k: v for k, v in active_sources.items() if v},
            "weather_comparison": {k: v for k, v in weather_comparisons.items() if v},
            "weather_model_inputs": ["temperature"],
        }
        heater_valid = heater is not None and heater.state in ("on", "off")
        heater_on = heater_valid and heater.state == "on"
        steady = False
        if heater_valid:
            steady = self._track_lag(now, heater_on, heater.last_changed)
            status["heat_lag"] = round(self._lag["lag"], 3)
            status["heat_lag_uncertainty"] = round(self._lag["uncertainty"], 3)
        else:
            self._lag = None
        if (not heater_valid or climate is None or outside is None or
                any(temp is None for temp in temps.values())):
            self._previous = None
            return status

        # No casting surprises: all rooms are valid here.
        vals: dict[str, float] = {room: float(temps[room]) for room in self.rooms}
        if self._previous is None:
            status["sample_result"] = "first_sample"
        else:
            dt = (now - self._previous["time"]).total_seconds()
            # Do not fit across a relay change, a provider/fallback switch or an
            # uncertain lag estimate (restart, relay outage, several changes).
            if not 270 <= dt <= 330:
                status["sample_result"] = "interval_out_of_range"
            elif not steady or heater_on != self._previous["heater_on"]:
                status["sample_result"] = "relay_changed"
            elif self._previous["outside_source"] != active_sources["temperature"]:
                status["sample_result"] = "outside_source_changed"
            elif self._lag["uncertainty"] > LAG_LEARN_MAX_UNCERTAINTY:
                status["sample_result"] = "lag_uncertain"
            else:
                accepted = [
                    self.learners[room].learn(
                        self._previous["temps"][room], vals[room],
                        self._previous["outside"], self._lag["lag"], dt, heater_on)
                    for room in self.rooms
                ]
                status["sample_result"] = "learned" if all(accepted) else "partly_rejected"
                self._update_counts()
                # Persist at each sample so restart does not erase hours of training.
                await self._store.async_save({
                    "schema_version": 2,
                    "learners": {room: learner.as_dict() for room, learner in self.learners.items()},
                    "on_samples": self._on_samples, "off_samples": self._off_samples,
                    "heater_entity": self.config[CONF_HEATER],
                    "outside_source": active_sources["temperature"],
                })

        self._previous = {
            "time": now, "temps": vals, "outside": outside, "heater_on": heater_on,
            "outside_source": active_sources["temperature"],
        }
        weights = self.config.get(CONF_WEIGHTS, {})
        try:
            weights = {r: float(weights.get(r, 1.0)) for r in self.rooms}
            if any(not isfinite(w) or w <= 0 for w in weights.values()) or not isfinite(sum(weights.values())):
                raise ValueError
        except (AttributeError, TypeError, ValueError, OverflowError):
            status["status"] = "invalid_weights"
            return status
        weighted_temp = sum(weights[r] * vals[r] for r in self.rooms) / sum(weights.values())
        status.update({
            "weighted_temperature": round(weighted_temp, 3),
            "outside_temperature": outside,
            "reference_heat": heater_on,
            "on_samples": self._on_samples,
            "off_samples": self._off_samples,
            "samples": min(learner.count for learner in self.learners.values()),
            "models": {
                room: {"loss": round(model.theta[0], 6),
                       "heat": round(model.theta[1], 6),
                       "bias": round(model.theta[2], 6),
                       "samples": model.count,
                       "on_samples": model.on_count, "off_samples": model.off_count,
                       "mae_5m": round(model.mae, 4) if model.mae is not None else None}
                for room, model in self.learners.items()
            },
        })
        if climate.state != "heat":
            status["status"] = "paused"
            return status
        try:
            target = normalize("temperature", climate.attributes["temperature"],
                               climate.attributes.get("temperature_unit",
                                   self.hass.config.units.temperature_unit))
            if target is None or not 5 <= target <= 35:
                raise ValueError
        except (TypeError, KeyError, ValueError):
            status["status"] = "waiting_for_target"
            return status
        status["target"] = target
        # Do not act on a cold model: need ~17 hours of samples plus both conditions.
        if status["samples"] < 200 or self._on_samples < 30 or self._off_samples < 30:
            status["status"] = "learning"
            return status
        try:
            # 180-minute exhaustive search can take several seconds on RPi5;
            # never block Home Assistant's event loop with CPU work.
            job = partial(
                optimize,
                vals, outside, self._lag["lag"], target, weights, self.learners,
                heater_on, (now - heater.last_changed).total_seconds(),
                horizon_minutes=int(self.config[CONF_HORIZON]),
                minimum_on_minutes=int(self.config[CONF_MIN_ON]),
                minimum_off_minutes=int(self.config[CONF_MIN_OFF]),
                switching_cost=float(self.config[CONF_SWITCH_COST]),
                heating_cost=float(self.config[CONF_HEAT_COST]),
                comfort_band=float(self.config[CONF_BAND]),
            )
            # HA State objects are immutable snapshots. Discard a result if any
            # input changed while the worker was calculating the hypothetical plan.
            output = await self.hass.async_add_executor_job(job)
            if any(self.hass.states.get(entity) is not state for entity, state in snapshot.items()):
                status["status"] = "inputs_changed"
                return status
        except (ValueError, RuntimeError, OverflowError) as err:
            _LOGGER.warning("MPC prediction failed: %s", err)
            status["status"] = "prediction_error"
            return status
        status.update({
            "status": "ready", "proposal": output["proposal"],
            "prediction_60m": output["predicted_60m"],
            "prediction_end": output["predicted_end"],
            "plan_15min": output["plan_15min"], "objective": output["cost"],
            "model_quality_note": "5-minute in-sample MAE; not a validated 60-minute forecast error",
        })
        return status