"""Per-zone thermal model: base model plus optional extended features.

Pure Python; no Home Assistant imports. Per 5-minute step, for zone i:

    dT_i = loss x (T_out - T_i) + heat x lag + bias          (base, always learned)
         + internal x q_i                                    (electrical heat above baseline, kW)
         + sum_j couple:j x (T_j - T_i)                      (adjacent zones)

The base model is the 3-parameter model of earlier versions, kept as fallback
and as the comparison that shows whether the extra features earn their keep.
The extended model learns the same terms plus the configured extras, only from
samples where all of its inputs are known.
"""
from __future__ import annotations

from typing import Any, Mapping

from .const import READY_OFF_SAMPLES, READY_ON_SAMPLES, READY_SAMPLES
from .learning import FeatureLearner
from .physics import Learner
from .windows import ventilation_vector

BASE_FEATURES = ["loss", "heat", "bias"]


def ready(model: Learner | FeatureLearner) -> bool:
    return (model.count >= READY_SAMPLES and model.on_count >= READY_ON_SAMPLES
            and model.off_count >= READY_OFF_SAMPLES)


def extended_features(has_internal: bool, neighbors: tuple[str, ...]) -> list[str]:
    """Features of a zone's extended model; empty when it would equal the base."""
    extra = (["internal"] if has_internal else []) + [f"couple:{n}" for n in neighbors]
    return BASE_FEATURES + extra if extra else []


def feature_vector(features: list[str], temp: float, inputs: Mapping[str, Any]) -> list[float] | None:
    """Values of `features` at the start of a step; None if an input is unknown.

    `inputs`: outside, lag, internal (kW, may be None) and temps (all zones).
    """
    values = []
    for name in features:
        if name == "loss":
            values.append(inputs["outside"] - temp)
        elif name == "heat":
            values.append(inputs["lag"])
        elif name == "bias":
            values.append(1.0)
        elif name == "internal":
            if inputs.get("internal") is None:
                return None
            values.append(inputs["internal"])
        elif name.startswith("couple:"):
            neighbor = inputs.get("temps", {}).get(name.split(":", 1)[1])
            if neighbor is None:
                return None
            values.append(neighbor - temp)
        else:
            raise KeyError(name)
    return values


class ZoneModel:
    def __init__(self, base: Learner | None = None, extended: FeatureLearner | None = None,
                 ventilation: FeatureLearner | None = None) -> None:
        self.base = base or Learner()
        self.extended = extended
        self.ventilation = ventilation

    @property
    def uses_extended(self) -> bool:
        """Forecast with the extended model once it is as trained as the base needs to be."""
        return self.extended is not None and ready(self.extended)

    def main_delta(self, temp: float, inputs: Mapping[str, Any]) -> float:
        """Closed-window change. Unknown extended inputs fall back to the base model."""
        if self.uses_extended:
            x = feature_vector(self.extended.features, temp, inputs)
            if x is not None:
                return self.extended.predict(x)
        return self.base.predict_delta(temp, inputs["outside"], inputs["lag"])

    def ventilation_delta(self, temp: float, inputs: Mapping[str, Any]) -> float:
        """Extra change from open windows and recovery; 0 without windows or inputs."""
        ventilation = inputs.get("ventilation")
        if self.ventilation is None or not ventilation:
            return 0.0
        x = ventilation_vector(self.ventilation.features, temp, inputs["outside"], ventilation)
        return self.ventilation.predict(x) if x is not None else 0.0

    def delta(self, temp: float, inputs: Mapping[str, Any]) -> float:
        """Predicted 5-minute change of the zone."""
        return self.main_delta(temp, inputs) + self.ventilation_delta(temp, inputs)

    def contributions(self, temp: float, inputs: Mapping[str, Any]) -> dict[str, float]:
        """Each term's share of the predicted change, in °C per hour."""
        terms = self._main_contributions(temp, inputs)
        ventilation = inputs.get("ventilation")
        if self.ventilation is not None and ventilation:
            x = ventilation_vector(self.ventilation.features, temp, inputs["outside"], ventilation)
            if x is not None:
                for name, theta, value in zip(self.ventilation.features, self.ventilation.theta, x):
                    group = "recovery" if name == "recovery" else "ventilation"
                    terms[group] = round(terms.get(group, 0.0) + 12 * theta * value, 3)
        return terms

    def _main_contributions(self, temp: float, inputs: Mapping[str, Any]) -> dict[str, float]:
        if self.uses_extended:
            x = feature_vector(self.extended.features, temp, inputs)
            if x is not None:
                terms: dict[str, float] = {}
                for name, theta, value in zip(self.extended.features, self.extended.theta, x):
                    group = "neighbors" if name.startswith("couple:") else name
                    terms[group] = terms.get(group, 0.0) + 12 * theta * value
                return {name: round(value, 3) for name, value in terms.items()}
        a, b, c = self.base.theta
        return {"loss": round(12 * a * (inputs["outside"] - temp), 3),
                "heat": round(12 * b * inputs["lag"], 3), "bias": round(12 * c, 3)}