"""Recursive least squares over named, physically bounded features.

Pure Python; no Home Assistant imports. Every coefficient is a change in room
temperature per 5-minute step per unit of its feature, so all of them share the
same target: the measured 5-minute temperature change (or, for the ventilation
layer, the part of it the closed-window model cannot explain).

New effects start at a prior of 0 ("no effect until the data shows one"). In
this unscaled RLS form a prior variance is worth 1 / (variance x typical value^2)
samples of data, so each prior below is chosen as "worth N typical samples": a
handful of samples cannot produce a confident coefficient. Confidence is the share
of a coefficient's information that comes from data rather than the prior; the
displayed standard deviation uses a running estimate of the measurement noise.
"""
from __future__ import annotations

from math import isfinite, sqrt
from typing import Any, NamedTuple

FORGETTING = 0.999  # about 3.5 days of 5-minute samples
MAX_STEP_CHANGE = 0.8  # °C per 5 minutes: larger is a glitch, not physics
# A coefficient counts as identified after this many exercising samples and once
# at least 80 % of its information comes from data.
MIN_ACTIVE_SAMPLES = 30
IDENTIFIED_CONFIDENCE = 0.8
INITIAL_NOISE_STD = 0.02  # °C per 5 minutes, about one sensor resolution step


class FeatureSpec(NamedTuple):
    prior: float
    # Unscaled RLS prior variance: worth 1 / (variance x typical^2) typical samples.
    variance: float
    lo: float
    hi: float
    # |feature value| at or above which a sample exercises this coefficient.
    active: float
    unit: str


def _worth(samples: float, typical: float) -> float:
    return 1.0 / (samples * typical * typical)


# x = feature value for one zone at the start of a 5-minute step. The first three
# match the base model's priors.
FEATURES: dict[str, FeatureSpec] = {
    # Envelope loss: x = T_out - T (°C).
    "loss": FeatureSpec(0.002, 0.0005, 0.0, 0.04, 1.0, "°C/5 min per °C outdoor difference"),
    # Boiler: x = radiator-lag proxy (0..1).
    "heat": FeatureSpec(0.045, 0.8, 0.0, 0.25, 0.1, "°C/5 min at full radiator heat"),
    # Constant gains/losses: x = 1.
    "bias": FeatureSpec(0.0, 0.08, -0.06, 0.06, 0.5, "°C/5 min"),
    # Electrical heat above baseline after its thermal response: x in kW (typ. 0.4).
    "internal": FeatureSpec(0.0, _worth(20, 0.4), 0.0, 0.5, 0.05, "°C/5 min per kW"),
    # Exchange with an adjacent zone j: x = T_j - T (°C). Shrunk hard towards 0:
    # zones heated by one boiler move together, so this is weakly identifiable.
    "couple": FeatureSpec(0.0, _worth(200, 1.0), 0.0, 0.05, 0.3, "°C/5 min per °C difference"),
    # Ventilation, fitted to what the closed-window model cannot explain. x scales
    # with the open-window factor O and T_out - T (typ. 15 °C):
    "vent_natural": FeatureSpec(0.0, _worth(20, 15), 0.0, 0.2, 0.5, "°C/5 min per °C x opening"),
    # Wind speed v in m/s (typ. 4 m/s).
    "vent_wind": FeatureSpec(0.0, _worth(20, 60), 0.0, 0.03, 0.5, "°C/5 min per °C x opening x m/s"),
    # Directional part, bounded by vent_wind so a leeward window still ventilates.
    "vent_wind_dir": FeatureSpec(0.0, _worth(20, 60), -0.03, 0.03, 0.5,
                                 "°C/5 min per °C x opening x m/s x cos(angle)"),
    "vent_gust": FeatureSpec(0.0, _worth(20, 45), 0.0, 0.03, 0.5,
                             "°C/5 min per °C x opening x m/s of gust above mean wind"),
    "vent_cross": FeatureSpec(0.0, _worth(100, 15), 0.0, 0.1, 0.5,
                              "°C/5 min per °C x opening x openings elsewhere"),
    # After closing: x = recovery weight (1 -> 0) x (T before opening - T), typ. 0.5.
    "recovery": FeatureSpec(0.0, _worth(20, 0.5), 0.0, 0.5, 0.05,
                            "°C/5 min per °C below the pre-ventilation temperature"),
}


def spec_for(name: str) -> FeatureSpec:
    """Specification of a feature; "couple:<zone>" shares the "couple" spec."""
    return FEATURES[name.split(":", 1)[0]]


def bound_covariance(p: list[list[float]], prior_variance: list[float]) -> None:
    """Keep P symmetric and no less certain than the prior, in place.

    With forgetting, unexcited directions grow by 1/forgetting each step without
    limit. Scaling row and column i together is a congruence transform, so positive
    definiteness is preserved.
    """
    n = len(p)
    for i in range(n):
        for j in range(i):
            p[i][j] = p[j][i] = (p[i][j] + p[j][i]) / 2
    for i in range(n):
        if p[i][i] > prior_variance[i]:
            scale = (prior_variance[i] / p[i][i]) ** .5
            for j in range(n):
                p[i][j] *= scale
            for j in range(n):
                p[j][i] *= scale
            p[i][i] = prior_variance[i]


def positive_definite(p: list[list[float]]) -> bool:
    """Cholesky check for a symmetric matrix."""
    n = len(p)
    lower = [[0.0] * n for _ in range(n)]
    for i in range(n):
        for j in range(i + 1):
            value = p[i][j] - sum(lower[i][k] * lower[j][k] for k in range(j))
            if i == j:
                if value <= 0 or not isfinite(value):
                    return False
                lower[i][j] = value ** .5
            else:
                lower[i][j] = value / lower[j][j]
    return True


def rls_update(theta: list[float], p: list[list[float]], x: list[float], y: float,
               forgetting: float = FORGETTING) -> tuple[list[float], float] | None:
    """One exponentially weighted RLS step, updating P in place.

    Returns the new parameters and the a priori residual, or None when the update
    would be numerically invalid (P is left untouched then).
    """
    n = len(theta)
    residual = y - sum(t * v for t, v in zip(theta, x))
    px = [sum(p[i][j] * x[j] for j in range(n)) for i in range(n)]
    denominator = forgetting + sum(x[i] * px[i] for i in range(n))
    if not isfinite(denominator) or denominator <= 0:
        return None
    gain = [value / denominator for value in px]
    for i in range(n):
        for j in range(n):
            p[i][j] = (p[i][j] - gain[i] * px[j]) / forgetting
    return [t + gain[i] * residual for i, t in enumerate(theta)], residual


class FeatureLearner:
    """RLS over a configurable list of named features (see FEATURES)."""

    def __init__(self, features: list[str]) -> None:
        if not features or len(set(features)) != len(features):
            raise ValueError("Features must be unique and non-empty")
        self.features = list(features)
        self._specs = [spec_for(name) for name in self.features]
        self.theta = [spec.prior for spec in self._specs]
        self.p = [[spec.variance if i == j else 0.0 for j in range(len(self._specs))]
                  for i, spec in enumerate(self._specs)]
        self.count = 0
        self.error_sum = 0.0
        self.rejected_count = 0
        self.on_count = 0
        self.off_count = 0
        self.active = [0] * len(self.features)
        # Bias-corrected running mean of squared a priori residuals.
        self._residual_sq = 0.0
        self._residual_weight = 0.0

    def predict(self, x: list[float]) -> float:
        return sum(t * v for t, v in zip(self.theta, x))

    def learn(self, x: list[float], y: float, command: bool | None = None) -> bool:
        """Fit one sample of features `x` to the 5-minute change `y`."""
        if (len(x) != len(self.features) or not all(isfinite(v) for v in x)
                or not isfinite(y) or abs(y) > MAX_STEP_CHANGE):
            self.rejected_count += 1
            return False
        result = rls_update(self.theta, self.p, x, y)
        if result is None:
            self.rejected_count += 1
            return False
        theta, residual = result
        bound_covariance(self.p, [spec.variance for spec in self._specs])
        self.theta = self._bounded(theta)
        self.count += 1
        self.error_sum += abs(residual)
        self._residual_sq = FORGETTING * self._residual_sq + (1 - FORGETTING) * residual * residual
        self._residual_weight = FORGETTING * self._residual_weight + (1 - FORGETTING)
        if command is not None:
            if command:
                self.on_count += 1
            else:
                self.off_count += 1
        for i, (value, spec) in enumerate(zip(x, self._specs)):
            if abs(value) >= spec.active:
                self.active[i] += 1
        return True

    def _bounded(self, theta: list[float]) -> list[float]:
        bounded = [min(spec.hi, max(spec.lo, value)) for value, spec in zip(theta, self._specs)]
        if "vent_wind_dir" in self.features:
            # The directional term may at most cancel the isotropic wind term.
            limit = bounded[self.features.index("vent_wind")] if "vent_wind" in self.features else 0.0
            i = self.features.index("vent_wind_dir")
            bounded[i] = min(limit, max(-limit, bounded[i]))
        return bounded

    @property
    def noise_var(self) -> float:
        """Estimated variance of the unexplained 5-minute change (°C²)."""
        if self._residual_weight <= 0:
            return INITIAL_NOISE_STD ** 2
        return self._residual_sq / self._residual_weight

    @property
    def mae(self) -> float | None:
        return self.error_sum / self.count if self.count else None

    def coefficient(self, name: str) -> float:
        return self.theta[self.features.index(name)] if name in self.features else 0.0

    def coefficients(self) -> dict[str, dict[str, Any]]:
        """Value, uncertainty and support of each coefficient, for diagnostics.

        `confidence` is the share of the coefficient's information that comes from
        data rather than the prior (0 = prior only); `std` is its standard error.
        """
        result = {}
        for i, (name, spec) in enumerate(zip(self.features, self._specs)):
            variance = min(spec.variance, max(0.0, self.p[i][i]))
            confidence = 1 - variance / spec.variance
            result[name] = {
                "value": round(self.theta[i], 6),
                "std": round(sqrt(self.noise_var * variance), 6),
                "confidence": round(confidence, 3),
                "samples": self.active[i],
                "identified": (self.active[i] >= MIN_ACTIVE_SAMPLES
                               and confidence >= IDENTIFIED_CONFIDENCE),
            }
        return result

    def as_dict(self) -> dict[str, Any]:
        return {"features": self.features, "theta": self.theta, "p": self.p,
                "count": self.count, "error_sum": self.error_sum,
                "rejected_count": self.rejected_count, "on_count": self.on_count,
                "off_count": self.off_count, "active": self.active,
                "residual_sq": self._residual_sq, "residual_weight": self._residual_weight}

    @classmethod
    def from_dict(cls, source: Any, features: list[str]) -> "FeatureLearner":
        """Restore what is valid for `features`; other features start from the prior.

        Coefficients are matched by name, so adding or removing a heat source,
        window or neighbour retrains only the affected coefficients. Anything
        malformed yields a fresh learner.
        """
        fresh = cls(features)
        try:
            if not isinstance(source, dict):
                return fresh
            names = source["features"]
            theta = [float(v) for v in source["theta"]]
            p = [[float(v) for v in row] for row in source["p"]]
            n = len(names)
            if (not all(isinstance(name, str) for name in names) or len(set(names)) != n
                    or len(theta) != n or len(p) != n or any(len(row) != n for row in p)):
                return fresh
            if not all(isfinite(v) for v in theta + [v for row in p for v in row]):
                return fresh
            for i in range(n):
                for j in range(i):
                    if abs(p[i][j] - p[j][i]) > 1e-6 * max(1.0, abs(p[i][j]), abs(p[j][i])):
                        return fresh
            if not positive_definite(p):
                return fresh
            counts = [source.get(key, 0) for key in
                      ("count", "rejected_count", "on_count", "off_count")]
            active = source.get("active", [0] * n)
            if (any(type(v) is not int or v < 0 for v in counts + list(active))
                    or len(active) != n or counts[2] + counts[3] > counts[0]):
                return fresh
            error_sum = float(source.get("error_sum", 0))
            residual_sq = float(source.get("residual_sq", 0.0))
            residual_weight = float(source.get("residual_weight", 0.0))
            if (not all(isfinite(v) and v >= 0 for v in (error_sum, residual_sq, residual_weight))
                    or residual_weight > 1):
                return fresh
            index = {name: i for i, name in enumerate(names)}
            kept = [(j, index[name]) for j, name in enumerate(features) if name in index]
            obj = cls(features)
            for j, i in kept:
                obj.theta[j] = theta[i]
                obj.active[j] = active[i]
                for k, m in kept:
                    obj.p[j][k] = p[i][m]
            obj.theta = obj._bounded(obj.theta)
            bound_covariance(obj.p, [spec.variance for spec in obj._specs])
            if not positive_definite(obj.p):
                return fresh
            obj.count, obj.rejected_count, obj.on_count, obj.off_count = counts
            obj.error_sum = error_sum
            obj._residual_sq, obj._residual_weight = residual_sq, residual_weight
            return obj
        except (KeyError, TypeError, ValueError, OverflowError):
            return fresh