"""Adapter for `owrf1d` (or a local fallback).

The upstream project `owrf1d` is listed as a dependency of *zen-fronts*, but it
may be unavailable in some development environments (for example, in pure
Python CI that does not build optional native extensions).

This module provides a thin wrapper that:

* Uses `owrf1d` when it can be imported.
* Falls back to a small, deterministic online model when it cannot.

The contract required by :class:`~zen_fronts.crits.store.CriteriaStore` is:

```
model.update(y: float, t: float) -> None
model.stats() -> tuple[mu: float, trend: float, sigma2: float, t_last: float]
```

Where:

* ``mu`` is the estimated mean *at* ``t_last``.
* ``trend`` is d(mu)/dt.
* ``sigma2`` is an estimated observation variance (non-negative).
"""

from __future__ import annotations

from dataclasses import dataclass

import math


class OwrfModelProtocol:
    """Structural protocol used at runtime (no typing dependency)."""

    def update(self, y: float, t: float) -> None:  # pragma: no cover
        raise NotImplementedError

    def stats(self) -> tuple[float, float, float, float]:  # pragma: no cover
        raise NotImplementedError


@dataclass
class _Rls1D(OwrfModelProtocol):
    """A small deterministic online linear model with forgetting.

    This is a practical fallback, not a claim of being the best filter.
    It is designed to be:

    * deterministic,
    * numerically stable for typical PBT scales,
    * dependency-free.

    Model: y(t) ≈ a + b * (t - t0).
    """

    forgetting: float = 0.99
    p0: float = 1e6
    var_forgetting: float = 0.99

    # internal state
    t0: float | None = None
    t_last: float | None = None
    a: float = 0.0
    b: float = 0.0

    # 2x2 covariance (P)
    p00: float = 1e6
    p01: float = 0.0
    p10: float = 0.0
    p11: float = 1e6

    # EWMA of residual^2 as a variance proxy
    sigma2: float = 1.0
    n: int = 0

    def __post_init__(self) -> None:
        self.p00 = float(self.p0)
        self.p11 = float(self.p0)
        self.p01 = 0.0
        self.p10 = 0.0

    def update(self, y: float, t: float) -> None:
        y = float(y)
        t = float(t)
        if self.t0 is None:
            self.t0 = t
            self.t_last = t
            self.a = y
            self.b = 0.0
            self.sigma2 = 1.0
            self.n = 1
            return

        assert self.t0 is not None
        dt = t - self.t0

        # Feature x = [1, dt]
        x0 = 1.0
        x1 = dt

        # y_hat = x^T theta
        y_hat = self.a * x0 + self.b * x1
        err = y - y_hat

        # Compute P x
        px0 = self.p00 * x0 + self.p01 * x1
        px1 = self.p10 * x0 + self.p11 * x1

        denom = self.forgetting + x0 * px0 + x1 * px1
        if denom <= 0.0:
            # Extremely unlikely with sane parameters; keep state unchanged.
            self.t_last = t
            return

        k0 = px0 / denom
        k1 = px1 / denom

        # Update theta
        self.a += k0 * err
        self.b += k1 * err

        # Update covariance: P = (P - k x^T P) / forgetting
        p00_new = (self.p00 - k0 * px0) / self.forgetting
        p01_new = (self.p01 - k0 * px1) / self.forgetting
        p10_new = (self.p10 - k1 * px0) / self.forgetting
        p11_new = (self.p11 - k1 * px1) / self.forgetting

        self.p00, self.p01, self.p10, self.p11 = p00_new, p01_new, p10_new, p11_new

        # Variance proxy: EWMA of residual^2
        err2 = err * err
        if self.n == 1:
            self.sigma2 = max(err2, 1e-12)
        else:
            self.sigma2 = (
                self.var_forgetting * self.sigma2 + (1.0 - self.var_forgetting) * err2
            )
            self.sigma2 = max(self.sigma2, 1e-12)

        self.n += 1
        self.t_last = t

    def stats(self) -> tuple[float, float, float, float]:
        if self.t0 is None or self.t_last is None:
            return (0.0, 0.0, float("inf"), 0.0)
        dt_last = self.t_last - self.t0
        mu = self.a + self.b * dt_last
        return (float(mu), float(self.b), float(self.sigma2), float(self.t_last))


def make_owrf_model() -> OwrfModelProtocol:
    """Create an online model.

    If `owrf1d` is importable, use it. Otherwise, use the internal fallback.
    """

    # We do the import lazily to avoid import-time hard dependency.
    try:  # pragma: no cover
        # The exact upstream API may evolve; keep integration defensive.
        from owrf1d import Owrf1D  # type: ignore

        class _Upstream(OwrfModelProtocol):
            def __init__(self) -> None:
                self._m = Owrf1D()

            def update(self, y: float, t: float) -> None:
                # Many filters use implicit time step. If upstream expects
                # explicit time, it likely accepts it; if not, it can ignore.
                try:
                    self._m.update(y, t)
                except TypeError:
                    self._m.update(y)

            def stats(self) -> tuple[float, float, float, float]:
                s = self._m.stats()  # expected: dict-like or tuple
                if isinstance(s, dict):
                    mu = float(s.get("mu", 0.0))
                    trend = float(s.get("trend", 0.0))
                    sigma2 = float(s.get("sigma2", 0.0))
                    t_last = float(s.get("t_last", 0.0))
                    return mu, trend, max(sigma2, 0.0), t_last
                if isinstance(s, (list, tuple)) and len(s) >= 3:
                    mu = float(s[0])
                    trend = float(s[1])
                    sigma2 = float(s[2])
                    t_last = float(s[3]) if len(s) > 3 else 0.0
                    return mu, trend, max(sigma2, 0.0), t_last
                raise ValueError("Unexpected owrf1d stats format")

        return _Upstream()
    except Exception:
        return _Rls1D()


def project_mu(mu: float, trend: float, t_last: float, now: float) -> float:
    """Project mean using a linear trend."""

    dt = float(now) - float(t_last)
    if not math.isfinite(dt):
        return float(mu)
    return float(mu) + float(trend) * dt
