"""Pure-Python small grey-box thermal model and binary model predictive control.

No Home Assistant imports; this module can be tested in isolation. Step=5 minutes.
This is an experimental thermal model, not a certified boiler controller.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from itertools import product
from math import exp, isfinite
from typing import Any

STEP_MINUTES = 5
# First-order radiator lag between the relay command and the heat reaching rooms.
LAG_TAU_SECONDS = 1200.0
HEAT_LAG = 1.0 - exp(-STEP_MINUTES * 60 / LAG_TAU_SECONDS)
# Learn only when the lag proxy is known within this half-width.
LAG_LEARN_MAX_UNCERTAINTY = 0.1
# Prior covariance; also the upper bound that stops RLS windup (see _bound_covariance).
PRIOR_P = ((0.0005, 0.0, 0.0), (0.0, 0.8, 0.0), (0.0, 0.0, 0.08))


def bound(value: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, value))


def advance_lag(lag: float, command: bool, seconds: float) -> float:
    """Move the lag proxy towards the relay command over `seconds`."""
    return lag + (1.0 - exp(-seconds / LAG_TAU_SECONDS)) * (float(command) - lag)


def relay_segments(
    start_on: bool, start: float, changes: list[tuple[float, str]], end: float,
) -> list[tuple[bool, float]] | None:
    """Split [start, end] into (command, seconds) segments from ordered relay changes.

    `changes` holds (timestamp, new_state) pairs. Return None when the command is
    unknown for part of the interval: a state other than on/off, a change to the
    state already in force (so transitions were missed), or out-of-order times.
    """
    segments: list[tuple[bool, float]] = []
    command, since = start_on, start
    for when, state in changes:
        if state not in ("on", "off") or not since <= when <= end or (state == "on") == command:
            return None
        segments.append((command, when - since))
        command, since = state == "on", when
    segments.append((command, end - since))
    return segments


def lag_from_time_in_state(command: bool, seconds_in_state: float) -> tuple[float, float]:
    """Estimate the lag when only the time since the last relay change is known.

    Starting from any unknown lag in [0, 1], `seconds_in_state` of a constant command
    leaves the lag inside an interval of width exp(-t/tau). Return its midpoint and
    half-width, so the error bound is explicit rather than assuming a steady state.
    """
    width = exp(-max(0.0, seconds_in_state) / LAG_TAU_SECONDS)
    return (1.0 - width / 2) if command else width / 2, width / 2


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

    With forgetting, unexcited directions (e.g. heat input during a summer with the
    relay off) grow by 1/forgetting each step without limit. Scaling row and column i
    together is a congruence transform, so positive definiteness is preserved.
    """
    for i in range(3):
        for j in range(i):
            p[i][j] = p[j][i] = (p[i][j] + p[j][i]) / 2
    for i in range(3):
        if p[i][i] > PRIOR_P[i][i]:
            scale = (PRIOR_P[i][i] / p[i][i]) ** .5
            for j in range(3):
                p[i][j] *= scale
            for j in range(3):
                p[j][i] *= scale
            p[i][i] = PRIOR_P[i][i]


@dataclass
class Learner:
    """Recursive least squares for dT=a*(Tout-T)+b*effective_heat+c."""

    theta: list[float] = field(default_factory=lambda: [0.002, 0.045, 0.0])
    p: list[list[float]] = field(default_factory=lambda: [list(row) for row in PRIOR_P])
    count: int = 0
    error_sum: float = 0.0
    on_count: int = 0
    off_count: int = 0
    # Samples offered but not fitted (outliers, implausible inputs).
    rejected_count: int = 0

    def predict_delta(self, temp: float, outside: float, heat: float) -> float:
        a, b, c = self.theta
        return a * (outside - temp) + b * heat + c

    def learn(self, temp_before: float, temp_after: float, outside: float, effective_heat: float,
              duration_seconds: float = 300, command: bool | None = None) -> bool:
        if (not all(isfinite(x) for x in (temp_before, temp_after, outside, effective_heat))
                or not 0 <= effective_heat <= 1
                or not isfinite(duration_seconds) or not 270 <= duration_seconds <= 330):
            return self._reject()
        y = (temp_after - temp_before) * 300 / duration_seconds
        if abs(y) > 0.8:  # sensor glitch/door opened/restarted device
            return self._reject()
        x = [outside - temp_before, effective_heat, 1.0]
        predicted = sum(t * v for t, v in zip(self.theta, x))
        residual = y - predicted
        px = [sum(self.p[i][j] * x[j] for j in range(3)) for i in range(3)]
        forgetting = 0.999
        denominator = forgetting + sum(x[i] * px[i] for i in range(3))
        if not isfinite(denominator) or denominator <= 0:
            return self._reject()
        k = [value / denominator for value in px]
        self.theta = [t + k[i] * residual for i, t in enumerate(self.theta)]
        for i in range(3):
            for j in range(3):
                self.p[i][j] = (self.p[i][j] - k[i] * px[j]) / forgetting
        _bound_covariance(self.p)
        # Physical constraints; avoids unstable temperature predictions.
        self.theta = [
            bound(self.theta[0], 0.0, 0.04),
            bound(self.theta[1], 0.0, 0.25),
            bound(self.theta[2], -0.06, 0.06),
        ]
        self.count += 1
        self.error_sum += abs(residual)
        if command is not None:
            if command:
                self.on_count += 1
            else:
                self.off_count += 1
        return True

    def _reject(self) -> bool:
        self.rejected_count += 1
        return False

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

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

    @classmethod
    def from_dict(cls, source: dict[str, Any]) -> "Learner":
        obj = cls()
        try:
            if not isinstance(source, dict):
                return obj
            theta = [float(v) for v in source["theta"]]
            p = [[float(v) for v in r] for r in source["p"]]
            if len(theta) != 3 or len(p) != 3 or any(len(r) != 3 for r in p):
                return obj
            if not all(isfinite(v) for v in theta + [v for row in p for v in row]):
                return obj
            if not (0 <= theta[0] <= .04 and 0 <= theta[1] <= .25 and -.06 <= theta[2] <= .06):
                return obj
            # Tolerate float round-off relative to magnitude (older versions could
            # accumulate it during windup), then apply the current covariance bound.
            for i in range(3):
                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 obj
            _bound_covariance(p)
            # A covariance must be symmetric positive definite (Cholesky check).
            lower = [[0.0] * 3 for _ in range(3)]
            for i in range(3):
                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 obj
                        lower[i][j] = value ** .5
                    else:
                        lower[i][j] = value / lower[j][j]
            counts = [source.get(key, 0)
                      for key in ("count", "on_count", "off_count", "rejected_count")]
            if any(type(v) is not int or v < 0 for v in counts):
                return obj
            count, on_count, off_count, rejected_count = counts
            error_sum = float(source.get("error_sum", 0))
            if not isfinite(error_sum) or error_sum < 0 or on_count + off_count > count:
                return obj
            return cls(theta, p, count, error_sum, on_count, off_count, rejected_count)
        except (KeyError, TypeError, ValueError, OverflowError):
            return obj


def simulate_step(temps: dict[str, float], outside: float, lag: float,
                  command: bool, learners: dict[str, Learner]) -> tuple[dict[str, float], float]:
    next_lag = lag + HEAT_LAG * (float(command) - lag)
    new = {name: value + learners[name].predict_delta(value, outside, next_lag)
           for name, value in temps.items()}
    return new, next_lag


def optimize(
    temps: dict[str, float], outside: float, effective_heat: float,
    target: float, weights: dict[str, float], learners: dict[str, Learner],
    is_on: bool, seconds_since_change: float, horizon_minutes: int = 120,
    minimum_on_minutes: int = 15, minimum_off_minutes: int = 10,
    switching_cost: float = .15, heating_cost: float = .025,
    comfort_band: float = .15,
) -> dict[str, Any]:
    """Enumerate 15min ON/OFF plans; honor minimum ON/OFF durations.

    120min = 256 plans (fast on Raspberry Pi 5). All rooms share one burner demand.
    """
    if not temps or set(temps) != set(learners):
        raise ValueError("Missing room temperatures or learner")
    if not 30 <= horizon_minutes <= 180 or horizon_minutes % 15:
        raise ValueError("Horizon must be 30..180 minutes in 15-minute steps")
    if not isfinite(target) or not isfinite(outside):
        raise ValueError("Invalid temperature")
    if any(not isfinite(t) for t in temps.values()):
        raise ValueError("Invalid room temperature")
    if any(not isfinite(float(weights.get(room, 0))) or float(weights.get(room, 0)) <= 0 for room in temps):
        raise ValueError("All rooms need a positive weight")
    if not isfinite(effective_heat) or not 0 <= effective_heat <= 1:
        raise ValueError("Invalid heating lag")
    if any(not isfinite(v) or v < 0 for v in (
        seconds_since_change, minimum_on_minutes, minimum_off_minutes,
        switching_cost, heating_cost, comfort_band,
    )):
        raise ValueError("Invalid optimizer parameter")
    total_weight = sum(float(weights[room]) for room in temps)
    if not isfinite(total_weight):
        raise ValueError("Invalid total weight")
    norm = {room: float(weights[room]) / total_weight for room in temps}
    blocks = horizon_minutes // 15
    nsteps_per_block = 3
    best = float("inf")
    result = None
    min_on_s = minimum_on_minutes * 60
    min_off_s = minimum_off_minutes * 60
    elapsed_initial = max(0.0, seconds_since_change)
    # More sequences than 4096 are deliberately excluded by horizon validation.
    for bits in product((False, True), repeat=blocks):
        current = bool(is_on)
        elapsed_s = elapsed_initial
        state = dict(temps)
        lag = effective_heat
        cost = 0.0
        at_60 = None
        feasible = True
        for idx, bit in enumerate(bits):
            if bit != current:
                minimum = min_on_s if current else min_off_s
                if elapsed_s < minimum:
                    feasible = False
                    break
                cost += switching_cost
                elapsed_s = 0.0
                current = bit
            for tick in range(nsteps_per_block):
                state, lag = simulate_step(state, outside, lag, bit, learners)
                mismatch = sum(
                    norm[room] * max(0.0, abs(value - target) - comfort_band) ** 2
                    for room, value in state.items()
                )
                # A lower bound on weighted average comfort also matters.
                avg = sum(norm[r] * value for r, value in state.items())
                mismatch += max(0.0, abs(avg - target) - comfort_band) ** 2
                cost += 8.0 * mismatch + (heating_cost if bit else 0.0)
                if idx * 15 + (tick + 1) * STEP_MINUTES == 60:
                    at_60 = dict(state)
            elapsed_s += 15 * 60
        if feasible and cost < best:
            best = cost
            result = (bits, at_60 or {}, dict(state))
    if result is None:
        raise RuntimeError("No feasible schedule")
    plan, temps_60m, temps_end = result
    # Re-simulate the chosen plan for its trajectory, keeping the search loop lean.
    trajectory: dict[str, list[float]] = {room: [] for room in temps}
    state, lag = dict(temps), effective_heat
    for bit in plan:
        for _ in range(nsteps_per_block):
            state, lag = simulate_step(state, outside, lag, bit, learners)
        for room, value in state.items():
            trajectory[room].append(round(value, 3))
    return {
        "proposal": bool(plan[0]),
        "plan_15min": [int(v) for v in plan],
        "predicted_60m": {r: round(t, 3) for r, t in temps_60m.items()},
        "predicted_end": {r: round(t, 3) for r, t in temps_end.items()},
        "trajectory_15min": trajectory,
        "cost": round(best, 4),
    }