# Copyright 2026 Kemal Yaylali
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Multi-horizon forecasting: one direct head per (target, horizon).

Direct rather than recursive. A recursive one-step model iterated 288
times to reach 24 hours compounds its own bias into a beautifully smooth
lie. Direct heads cost more memory (six horizons x three targets = 18
small models, about 150 kB total) and are worth every byte.

Each head predicts a *delta from now*, then the ensemble blends three
opinions with weights that are themselves learned online:

    persistence  : it will be exactly as it is now
    climatology  : it will be whatever this hour of this day usually is
    learned RLS  : it will be now plus what the regressors imply

Persistence wins at 15 minutes. Climatology wins at 24 hours. The RLS
head wins in the middle, which is exactly the region a physical
forecaster finds hardest. The blend weights are updated by exponentiated
gradient (Hedge), so the ensemble is never worse than its best member by
more than a log factor, and it re-weights itself within a day when the
season turns.
"""

from __future__ import annotations

from typing import Dict, List, Optional, Tuple

import numpy as np

from ..features import N_FEATURES, Standardiser, supervised_pairs
from .rls import AdaptiveConformal, RecursiveLeastSquares

MEMBERS = ("persistence", "climatology", "learned", "setpoint")


class ForecastHead:
    """One target, one horizon."""

    def __init__(self, target: str, horizon_s: int, n_features: int = N_FEATURES,
                 forgetting: float = 0.9985, delta: float = 100.0,
                 alpha: float = 0.10, conformal_window: int = 400,
                 gamma: float = 0.01, hedge_eta: float = 0.35):
        self.target = target
        self.horizon_s = int(horizon_s)
        self.model = RecursiveLeastSquares(n_features, forgetting, delta)
        self.conformal = AdaptiveConformal(alpha, conformal_window, gamma)
        self.weights = np.ones(len(MEMBERS)) / len(MEMBERS)
        self.eta = float(hedge_eta)
        self.member_mae = np.zeros(len(MEMBERS))
        self.n_scored = 0

    # -------------------------------------------------------- prediction

    def predict(self, x: np.ndarray, anchor: float,
                climatology_delta: float = 0.0,
                setpoint_delta: float = 0.0) -> Dict[str, float]:
        learned_delta = self.model.predict(x)
        deltas = np.array([0.0, float(climatology_delta), float(learned_delta),
                           float(setpoint_delta)])
        blended = float(np.dot(self.weights, deltas))
        mu = float(anchor + blended)
        sigma = self.model.predict_std(x, self.model.noise_var)
        lo, hi = self.conformal.interval(mu, fallback_sigma=sigma)
        return {
            "mu": mu,
            "lo": lo,
            "hi": hi,
            "sigma": sigma,
            "delta": blended,
            "members": {m: float(anchor + d) for m, d in zip(MEMBERS, deltas)},
            "weights": {m: float(w) for m, w in zip(MEMBERS, self.weights)},
        }

    # ---------------------------------------------------------- learning

    def learn(self, x: np.ndarray, anchor: float, truth: float,
              climatology_delta: float = 0.0,
              setpoint_delta: float = 0.0) -> float:
        """One supervised step given a matured target."""
        deltas = np.array([0.0, float(climatology_delta),
                           float(self.model.predict(x)),
                           float(setpoint_delta)])
        member_pred = anchor + deltas
        losses = np.abs(member_pred - truth)

        # Hedge / exponentiated gradient on normalised losses
        scale = max(float(np.max(losses)), 1e-6)
        self.weights *= np.exp(-self.eta * losses / scale)
        self.weights = np.clip(self.weights, 1e-4, None)
        self.weights /= self.weights.sum()

        blended = float(np.dot(self.weights, member_pred))
        residual = truth - blended
        self.conformal.observe(residual)

        self.model.update(x, truth - anchor)

        self.member_mae = 0.98 * self.member_mae + 0.02 * losses
        self.n_scored += 1
        return residual

    def to_dict(self) -> Dict:
        return {"target": self.target, "horizon_s": self.horizon_s,
                "model": self.model.to_dict(), "conformal": self.conformal.to_dict(),
                "weights": self.weights.tolist(), "eta": self.eta,
                "member_mae": self.member_mae.tolist(), "n_scored": self.n_scored}

    @classmethod
    def from_dict(cls, s: Dict) -> "ForecastHead":
        h = cls(s["target"], s["horizon_s"])
        h.model = RecursiveLeastSquares.from_dict(s["model"])
        h.conformal = AdaptiveConformal.from_dict(s["conformal"])
        w = np.array(s["weights"], dtype=float)
        if w.size != len(MEMBERS):
            # A saved head from before the setpoint member existed. Reinitialise
            # uniformly rather than guessing: the Hedge weights re-converge in
            # about a day, which is far cheaper than silently mismatching a
            # member to the wrong loss and corrupting every blend until someone
            # notices.
            w = np.ones(len(MEMBERS)) / len(MEMBERS)
        h.weights = w
        h.eta = s["eta"]
        mae = np.array(s["member_mae"], dtype=float)
        # Same migration as the weights. Missing this one did not fail on load,
        # it failed later inside learn() on a shape mismatch, which is a worse
        # place to find out.
        if mae.size != len(MEMBERS):
            mae = np.zeros(len(MEMBERS))
        h.member_mae = mae
        h.n_scored = s.get("n_scored", 0)
        return h


class NowcastEnsemble:
    """The full bank of heads plus the shared feature standardiser."""

    def __init__(self, targets: Tuple[str, ...], horizons_s: Tuple[int, ...],
                 cfg_model):
        self.targets = tuple(targets)
        self.horizons = tuple(int(h) for h in horizons_s)
        self.cfg = cfg_model
        self.grid_s = int(cfg_model.grid_s)
        self.scaler = Standardiser(N_FEATURES)
        self.heads: Dict[Tuple[str, int], ForecastHead] = {
            (t, h): ForecastHead(
                t, h, N_FEATURES, cfg_model.rls_forgetting, cfg_model.rls_delta,
                cfg_model.conformal_alpha, cfg_model.conformal_window,
                cfg_model.conformal_gamma,
            )
            for t in self.targets for h in self.horizons
        }
        self.trained_rows = 0
        self.min_pairs = int(getattr(cfg_model, "min_pairs_per_head", 12))
        # Which phase of the stride this refit starts on. Rotated so that over
        # successive retrains every offset is eventually trained on, rather
        # than the model permanently seeing one sample in `steps` forever.
        self.refit_phase = 0

    # ------------------------------------------------------------ train

    def fit(self, X: np.ndarray, valid: np.ndarray, series: Dict[str, np.ndarray],
            climatology=None, grid_ts: Optional[np.ndarray] = None,
            passes: int = 1, max_pairs: int = 2500, setpoint_fn=None) -> Dict[str, int]:
        """Batch-update every head from history.

        `max_pairs` bounds the work per head to the most recent samples.
        This is not a shortcut: with a forgetting factor of 0.9985 the
        effective memory is about 11 hours, so the 4000th-most-recent
        sample carries a weight of roughly e^-6. Training on it costs
        real seconds on a Cortex-A53 and buys nothing measurable.
        """
        """Batch pass over history. Called on startup and every retrain tick."""
        if X.shape[0] < 10:
            return {"rows": 0}
        # A refit starts from the prior. Without this, every retrain tick replays
        # the same history into a live filter, and RLS with forgetting reads that
        # as new evidence each time: measured on a real station after 1.5 days,
        # 453 grid rows had produced 64,676 updates, cond(P) of 3.1e9 and a
        # weight vector of norm 1680 whose two largest entries were the annual
        # harmonics the record cannot yet resolve. The result was a six hour
        # forecast of 53 C in a 24 C room, with a plus or minus of 0.43.
        #
        # The conformal calibrators and the Hedge weights are deliberately left
        # alone: those are earned from scored forecasts, not from this regression.
        for head in self.heads.values():
            head.model.reset()
        self.scaler.partial_fit(X[valid][:: max(1, X.shape[0] // 2000)])
        Xs = self.scaler.transform(X)

        counts = {}
        for target in self.targets:
            y = series[target]
            for h in self.horizons:
                steps = max(int(round(h / self.grid_s)), 1)
                Xa, dy, anchor = supervised_pairs(Xs, valid, y, steps)
                if Xa.shape[0] < 5:
                    counts[f"{target}@{h}"] = 0
                    continue
                if Xa.shape[0] > max_pairs:
                    Xa, dy, anchor = Xa[-max_pairs:], dy[-max_pairs:], anchor[-max_pairs:]
                head = self.heads[(target, h)]
                clim = np.zeros(Xa.shape[0])
                if climatology is not None and grid_ts is not None and climatology.ready:
                    n = grid_ts.size
                    ts_a = grid_ts[:n - steps]
                    mask_len = min(ts_a.size, Xa.shape[0])
                    clim_now = climatology.predict(target, ts_a[-mask_len:])
                    clim_fut = climatology.predict(target, ts_a[-mask_len:] + h)
                    clim = np.zeros(Xa.shape[0])
                    clim[-mask_len:] = clim_fut - clim_now
                # One pair per horizon, not one per grid row. Adjacent pairs at
                # the 1 d horizon share 287 of their 288 samples, so training on
                # every row hands the filter the same outcome 288 times and RLS
                # with forgetting reads each as fresh evidence. A 400-score
                # conformal window then holds 1.4 independent outcomes while
                # believing it holds 400.
                #
                # This is not a compute shortcut that costs accuracy. Measured
                # walk-forward on four days of real station data, striding cut
                # MAE at every horizon past an hour (temperature 6h -30%,
                # humidity 6h -48%, pressure 12h -68%) with coverage unchanged,
                # and made the fit 12x faster. The redundancy was not merely
                # wasted work, it was collapsing P onto the repeated direction.
                stride = steps
                if stride > 1 and Xa.shape[0] // stride < self.min_pairs:
                    # A long horizon on a short record would otherwise train
                    # on one or two pairs, which is worse than the redundancy
                    # it avoids. The floor was chosen by sweeping it over five
                    # train splits of real data: 12 was best at every horizon,
                    # and the apparent 1 d regressions at other values were
                    # noise, since a 1 d head on four days of record is fitted
                    # and scored on well under two independent outcomes.
                    stride = max(1, Xa.shape[0] // self.min_pairs)
                idx = np.arange(self.refit_phase % stride, Xa.shape[0], stride)
                if idx.size > max_pairs:
                    idx = idx[-max_pairs:]
                for _ in range(max(int(passes), 1)):
                    for i in idx:
                        head.learn(Xa[i], anchor[i], anchor[i] + dy[i], clim[i],
                                   setpoint_fn(target, h, anchor[i]) if setpoint_fn else 0.0)
                counts[f"{target}@{h}"] = int(idx.size)
        self.trained_rows = int(X.shape[0])
        self.refit_phase += 1
        return counts

    # --------------------------------------------------------- inference

    def forecast(self, x_raw: np.ndarray, anchors: Dict[str, float], now: float,
                 climatology=None, setpoint_fn=None) -> Dict[str, Dict[int, Dict[str, float]]]:
        x = self.scaler.transform(np.atleast_2d(x_raw))[0]
        out: Dict[str, Dict[int, Dict[str, float]]] = {}
        for target in self.targets:
            anchor = float(anchors.get(target, 0.0))
            out[target] = {}
            for h in self.horizons:
                clim_delta = 0.0
                if climatology is not None and climatology.ready:
                    clim_delta = float(climatology.predict(target, np.array([now + h]))[0]
                                       - climatology.predict(target, np.array([now]))[0])
                sp = setpoint_fn(target, h, anchor) if setpoint_fn else 0.0
                out[target][h] = self.heads[(target, h)].predict(x, anchor, clim_delta, sp)
        return out

    def diagnostics(self) -> List[Dict]:
        rows = []
        for (target, h), head in sorted(self.heads.items()):
            rows.append({
                "target": target,
                "horizon_s": h,
                "n_updates": head.model.n_updates,
                "n_scored": head.n_scored,
                "weights": {m: round(float(w), 3) for m, w in zip(MEMBERS, head.weights)},
                "member_mae": {m: round(float(v), 3) for m, v in zip(MEMBERS, head.member_mae)},
                "conformal_alpha": round(head.conformal.alpha, 4),
                "conformal_halfwidth": round(float(head.conformal.quantile()), 3)
                if np.isfinite(head.conformal.quantile()) else None,
                "coverage": round(head.conformal.empirical_coverage, 3)
                if np.isfinite(head.conformal.empirical_coverage) else None,
            })
        return rows

    def to_dict(self) -> Dict:
        return {
            "targets": list(self.targets),
            "horizons": list(self.horizons),
            "grid_s": self.grid_s,
            "scaler": self.scaler.to_dict(),
            "heads": [h.to_dict() for h in self.heads.values()],
            "trained_rows": self.trained_rows,
            "refit_phase": self.refit_phase,
        }

    def load_dict(self, s: Dict) -> None:
        self.scaler = Standardiser.from_dict(s["scaler"])
        for hs in s["heads"]:
            head = ForecastHead.from_dict(hs)
            key = (head.target, head.horizon_s)
            if key not in self.heads:
                continue          # a target or horizon this build no longer has
            self.heads[key] = head
        self.trained_rows = s.get("trained_rows", 0)
        self.refit_phase = int(s.get("refit_phase", 0))
        self._apply_config_tuning()

    def _apply_config_tuning(self) -> None:
        """Tuning comes from config; only the estimate comes from the file.

        Every from_dict below this point restores its knobs alongside its data:
        lambda, delta and p_max in the RLS, alpha, gamma and the window in the
        conformal calibrator. So each of those was immutable in the field. Edit
        config.yaml, restart, and the state file quietly puts the old value
        back, which looks exactly like a change that had no effect. The same
        defect cost a Kalman retune here before it was found.
        """
        c = self.cfg
        for head in self.heads.values():
            head.model.lam = float(c.rls_forgetting)
            head.model.delta = float(c.rls_delta)
            head.conformal.retune(c.conformal_alpha, c.conformal_gamma,
                                  c.conformal_window)