"""Empirical Bayes estimators via MacKay's evidence maximization."""

from __future__ import annotations

import math
from typing import Any, Optional, Union, cast

import numpy as np
from numpy.typing import NDArray
from scipy.linalg import cho_factor, cho_solve
from scipy.sparse import csc_array
from sklearn.utils.validation import check_X_y
from typing_extensions import Self

from ._blas_helpers import compute_eta_dense, update_precision_dense
from ._empirical_bayes import (
    accumulate_sufficient_stats,
    mackay_update_normal_online,
)
from ._estimators import NormalRegressor, _invalidate_cached_properties
from ._sparse_bayesian_linear_regression import (
    DenseFactor,
    PrecisionFactor,
    create_sparse_factor,
)


class EmpiricalBayesNormalRegressor(NormalRegressor):
    """Bayesian linear regression with empirical Bayes hyperparameter tuning.

    Extends :class:`NormalRegressor` with automatic optimization of the
    prior precision (``alpha``) and noise precision (``beta``) via
    MacKay's evidence maximization framework [1]_. During ``fit``, the
    hyperparameters are iteratively updated to maximize the log marginal
    likelihood. During ``partial_fit``, a single online MacKay step is
    performed using accumulated sufficient statistics.

    When ``learning_rate < 1``, exponential forgetting is applied to the
    precision matrix. To prevent the prior contribution from collapsing
    to zero under repeated decay, *stabilized forgetting* [2]_ re-injects
    a fixed prior floor after each decay step.

    Parameters
    ----------
    alpha : float, default=1.0
        Initial prior precision of the weights. The prior is
        :math:`w \\sim \\mathcal{N}(0, \\alpha^{-1} I)`. Updated
        automatically during fitting.
    beta : float, default=1.0
        Initial noise precision. The likelihood is
        :math:`y \\mid x, w \\sim \\mathcal{N}(x^T w, \\beta^{-1})`.
        Updated automatically during fitting.
    n_eb_iter : int, default=10
        Maximum number of empirical Bayes iterations during ``fit``.
        Each iteration re-fits the posterior and runs one MacKay update.
        Set to 0 to disable EB tuning during ``fit``.
    eb_tol : float, default=1e-4
        Convergence tolerance on the change in log marginal likelihood
        between successive EB iterations.
    learning_rate : float, default=1.0
        Decay rate for sequential updates. Values less than 1
        geometrically shrink the precision matrix on each call to
        ``decay`` or ``partial_fit``, enabling adaptation to
        non-stationary environments.
    sparse : bool, default=False
        If True, use sparse matrix operations for the precision matrix.
        Input ``X`` must be a ``scipy.sparse.csc_array``. When CHOLMOD
        is available (via ``scikit-sparse``), it is used for efficient
        Cholesky factorization; otherwise falls back to SuperLU.
    random_state : int, np.random.Generator, or None, default=None
        Controls the random number generator for ``sample``. Pass an
        int for reproducible results across calls.
    trace_method : {'auto', 'diagonal'}, default='auto'
        Method for computing :math:`\\operatorname{tr}(\\Lambda^{-1})`
        in the MacKay update for ``alpha``:

        - ``'auto'``: Uses Takahashi recursion (exact, exploits
          sparsity) for sparse matrices, Cholesky for dense.
        - ``'diagonal'``: Fast :math:`O(p)` approximation
          :math:`\\operatorname{tr}(\\Lambda^{-1}) \\approx
          \\sum_i 1/\\Lambda_{ii}`.

    Attributes
    ----------
    log_evidence_ : float
        Log marginal likelihood at convergence (after ``fit``), or
        ``-inf`` if ``n_eb_iter=0``.
    n_eb_iterations_ : int
        Number of EB iterations performed during the last ``fit``.
    eb_converged_ : bool
        Whether the EB loop converged within ``eb_tol`` during
        the last ``fit``.

    See Also
    --------
    NormalRegressor : Base estimator without empirical Bayes tuning.
    BayesianGLM : Bayesian GLM for non-Gaussian likelihoods (logistic,
        Poisson).

    Notes
    -----
    **Evidence maximization (MacKay's update rules)**

    Given the posterior precision
    :math:`\\Lambda = \\alpha I + \\beta X^T X` and MAP estimate
    :math:`\\hat{w}`, the effective number of well-determined
    parameters is:

    .. math::

        \\gamma = p - \\alpha \\operatorname{tr}(\\Lambda^{-1})

    The hyperparameters are updated as [1]_:

    .. math::

        \\alpha_{\\text{new}} = \\frac{\\gamma}{\\hat{w}^T \\hat{w}},
        \\qquad
        \\beta_{\\text{new}} = \\frac{N - \\gamma}{
        \\| y - X \\hat{w} \\|^2}

    **Stabilized forgetting**

    Under exponential decay with rate :math:`\\gamma < 1`, the prior
    contribution to the precision diagonal decays as
    :math:`\\gamma^t \\alpha \\to 0`. Stabilized forgetting [2]_
    re-injects :math:`(1 - \\gamma^n) \\alpha` after each decay step,
    so the prior contribution converges to :math:`\\alpha` instead of
    zero:

    .. math::

        s_{t+n} = \\gamma^n s_t + (1 - \\gamma^n) \\alpha

    where :math:`s_t` tracks the cumulative prior scalar on the
    diagonal.

    References
    ----------
    .. [1] MacKay, D. J. C. (1992). "Bayesian Interpolation."
       *Neural Computation*, 4(3), 415-447.

    .. [2] Kulhavy, R. & Zarrop, M. B. (1993). "On a general concept of
       forgetting." *International Journal of Control*, 58(4), 905-924.

    Examples
    --------
    Batch fitting with automatic hyperparameter tuning:

    >>> import numpy as np
    >>> from bayesianbandits import EmpiricalBayesNormalRegressor
    >>> rng = np.random.default_rng(42)
    >>> X = rng.standard_normal((100, 3))
    >>> w_true = np.array([1.0, -0.5, 0.0])
    >>> y = X @ w_true + 0.1 * rng.standard_normal(100)
    >>> model = EmpiricalBayesNormalRegressor(random_state=42)
    >>> _ = model.fit(X, y)
    >>> model.eb_converged_
    True
    >>> np.round(model.coef_, 2)  # Close to [1.0, -0.5, 0.0]
    array([ 1.  , -0.49, -0.02])

    Online learning with forgetting for non-stationary environments:

    >>> model = EmpiricalBayesNormalRegressor(
    ...     learning_rate=0.99, random_state=42
    ... )
    >>> for i in range(10):  # doctest: +SKIP
    ...     X_batch = rng.standard_normal((5, 3))
    ...     y_batch = X_batch @ w_true + 0.1 * rng.standard_normal(5)
    ...     model.partial_fit(X_batch, y_batch)
    """

    def __init__(
        self,
        alpha: float = 1.0,
        beta: float = 1.0,
        *,
        n_eb_iter: int = 10,
        eb_tol: float = 1e-4,
        learning_rate: float = 1.0,
        sparse: bool = False,
        random_state: Union[int, np.random.Generator, None] = None,
        trace_method: str = "auto",
    ) -> None:
        super().__init__(
            alpha=alpha,
            beta=beta,
            learning_rate=learning_rate,
            sparse=sparse,
            random_state=random_state,
        )
        self.n_eb_iter = n_eb_iter
        self.eb_tol = eb_tol
        self.trace_method = trace_method

    def _eb_mackay_step_online(self) -> None:
        """MacKay step using accumulated sufficient statistics for beta."""
        alpha_new, beta_new, _ = mackay_update_normal_online(
            self.coef_,
            self.cov_inv_,
            self.alpha,
            self.beta,
            self._prior_scalar,
            self._effective_n,
            self._eff_yTy,
            self._eff_XTy,
            factor=self._precision_factor,
            trace_method=self.trace_method,
        )
        self.alpha = alpha_new
        self.beta = beta_new

    def _accumulate_stats(
        self,
        X: Union[NDArray[Any], csc_array],
        y: NDArray[Any],
        prior_decay: float,
    ) -> None:
        """Update decayed sufficient statistics for the beta update."""
        self._effective_n, self._eff_yTy, self._eff_XTy = accumulate_sufficient_stats(
            self._effective_n,
            self._eff_yTy,
            self._eff_XTy,
            X,
            y,
            prior_decay,
        )

    def fit(
        self,
        X_fit: Union[NDArray[Any], csc_array],
        y: NDArray[Any],
        sample_weight: Optional[NDArray[Any]] = None,
    ) -> Self:
        """
        Fit the model with empirical Bayes hyperparameter tuning.

        Iteratively optimizes ``alpha`` (prior precision) and ``beta``
        (noise precision) by maximizing the log marginal likelihood
        using MacKay's evidence framework, then performs a final
        posterior update with the converged hyperparameters.

        Parameters
        ----------
        X_fit : array-like of shape (n_samples, n_features)
            Training data.
        y : array-like of shape (n_samples,)
            Target values.
        sample_weight : array-like of shape (n_samples,), optional
            Individual weights for each sample. If None, all samples
            are given weight 1.0.

        Returns
        -------
        self : EmpiricalBayesNormalRegressor
            Fitted estimator with tuned ``alpha`` and ``beta``.

        See Also
        --------
        partial_fit : Incremental update with one online MacKay step.
        """
        X_fit, y = check_X_y(
            X_fit,  # type: ignore
            y,
            copy=True,
            ensure_2d=True,
            dtype=np.float64,
            accept_sparse="csc" if self.sparse else False,
        )

        if self.n_eb_iter > 0:
            # Sufficient stats are constant during batch fitting (no decay)
            effective_n = float(y.shape[0])
            eff_yTy = float(y @ y)
            if isinstance(X_fit, csc_array):
                eff_XTy = np.asarray(X_fit.T @ y, dtype=np.float64).ravel()
            else:
                eff_XTy = np.asarray(X_fit.T @ y, dtype=np.float64)

            prev_evidence = -math.inf
            converged = False
            iterations = 0
            log_ev = -math.inf

            for i in range(self.n_eb_iter):
                self._initialize_prior(X_fit)
                self._fit_helper(X_fit, y, sample_weight)
                # After fresh fit: Λ = α·I + β·XᵀX
                self._prior_scalar = self.alpha

                alpha_new, beta_new, log_ev = mackay_update_normal_online(
                    self.coef_,
                    self.cov_inv_,
                    self.alpha,
                    self.beta,
                    self._prior_scalar,
                    effective_n,
                    eff_yTy,
                    eff_XTy,
                    factor=self._precision_factor,
                    trace_method=self.trace_method,
                )
                self.alpha = alpha_new
                self.beta = beta_new
                iterations = i + 1

                if abs(log_ev - prev_evidence) < self.eb_tol:
                    converged = True
                    break
                prev_evidence = log_ev

            # Final refit with converged hyperparams
            self._initialize_prior(X_fit)
            self._fit_helper(X_fit, y, sample_weight)

            self.log_evidence_ = log_ev
            self.n_eb_iterations_ = iterations
            self.eb_converged_ = converged
        else:
            self._initialize_prior(X_fit)
            self._fit_helper(X_fit, y, sample_weight)
            self.log_evidence_ = -math.inf
            self.n_eb_iterations_ = 0
            self.eb_converged_ = False

        # After fit, the precision matrix is consistent with self.alpha.
        self._prior_scalar = self.alpha

        # Initialize sufficient statistics from the fit data.
        self._effective_n = float(y.shape[0])
        self._eff_yTy = float(y @ y)
        if isinstance(X_fit, csc_array):
            self._eff_XTy: NDArray[np.float64] = np.asarray(
                X_fit.T @ y, dtype=np.float64
            ).ravel()
        else:
            self._eff_XTy = np.asarray(X_fit.T @ y, dtype=np.float64)

        return self

    @_invalidate_cached_properties
    def _correct_precision(
        self,
        alpha_old: float,
        beta_old: float,
    ) -> None:
        """Correct the precision matrix after a MacKay hyperparameter update.

        The precision matrix decomposes as prior + data:
            Λ = _prior_scalar·I + data_component

        After MacKay changes α and β, rescale both components so the
        matrix is consistent with the new hyperparameters.

        After correction, eagerly refactorizes so the factor is ready
        for ``sample()`` without an extra factorization.
        """
        alpha_new = self.alpha
        beta_new = self.beta

        if alpha_new == alpha_old and beta_new == beta_old:
            return

        prior_scalar = self._prior_scalar
        beta_ratio = beta_new / beta_old
        new_prior_scalar = prior_scalar * (alpha_new / alpha_old)
        diag_correction = new_prior_scalar - beta_ratio * prior_scalar

        if self.sparse:
            cov_inv = cast(csc_array, self.cov_inv_)
            cov_inv *= beta_ratio
            if diag_correction != 0.0:
                cov_inv.setdiag(cov_inv.diagonal() + diag_correction)
            self.cov_inv_ = cov_inv
        else:
            self.cov_inv_ *= beta_ratio
            if diag_correction != 0.0:
                diag_idx = np.diag_indices_from(self.cov_inv_)
                self.cov_inv_[diag_idx] += diag_correction

        self._prior_scalar = new_prior_scalar
        if "_precision_factor" in self.__dict__:
            del self._precision_factor

    def _reinject_prior(self, prior_reinjection: float) -> None:
        """Add stabilized prior re-injection to the precision diagonal.

        After exponential decay the prior contribution shrinks toward zero.
        This adds back ``prior_reinjection`` to every diagonal entry so that
        the prior converges to ``alpha`` instead (Kulhavy & Zarrop, 1993).
        """
        if prior_reinjection == 0.0:
            return
        if self.sparse:
            cov_inv = cast(csc_array, self.cov_inv_)
            cov_inv.setdiag(cov_inv.diagonal() + prior_reinjection)
            self.cov_inv_ = cov_inv
        else:
            diag_idx = np.diag_indices_from(self.cov_inv_)
            self.cov_inv_[diag_idx] += prior_reinjection
        if "_precision_factor" in self.__dict__:
            del self._precision_factor

    @_invalidate_cached_properties
    def _fit_helper(
        self,
        X: Union[NDArray[Any], csc_array],
        y: NDArray[Any],
        sample_weight: Optional[NDArray[Any]] = None,
    ) -> None:  # type: ignore[override]
        """Override to fold prior reinjection into precision construction.

        When ``_pending_reinjection > 0``, adds ``(1 - γ) · α · I`` to
        the precision during the decay + data update.  This means the
        factorization for the linear solve already reflects the
        reinjected prior — no separate ``_reinject_prior`` call or
        refactorization needed.
        """
        reinjection = getattr(self, "_pending_reinjection", 0.0)
        if reinjection == 0.0:
            # No reinjection needed — delegate to base class.
            super()._fit_helper(X, y, sample_weight)
            return

        # -- Below: base class logic with reinjection folded in --

        if self.sparse:
            X = csc_array(X)

        assert X.shape is not None

        from ._estimators import compute_effective_weights

        if sample_weight is None:
            sample_weight = np.ones(X.shape[0], dtype=np.float64)
        else:
            sample_weight = np.asarray(sample_weight, dtype=np.float64)

        effective_weights = compute_effective_weights(
            X.shape[0], sample_weight, self.learning_rate
        )

        prior_decay = self.learning_rate ** X.shape[0]

        if self.sparse:
            # Element-wise scaling; absorb beta into weights
            assert isinstance(X, csc_array)
            w_sqrt = np.sqrt(self.beta * effective_weights)
            X_weighted = X.multiply(w_sqrt.reshape(-1, 1)).tocsc()
            cov_inv = cast(
                csc_array,
                prior_decay * self.cov_inv_ + X_weighted.T @ X_weighted,
            )
            # Fold in the prior reinjection
            cov_inv.setdiag(cov_inv.diagonal() + reinjection)
        else:
            # Fused X^T W X + prior via dsyrk (upper triangle only)
            w_sqrt = np.sqrt(effective_weights)
            X_weighted = X * w_sqrt[:, np.newaxis]
            prior_scaled = np.asfortranarray(prior_decay * self.cov_inv_)
            # Add reinjection to diagonal before dsyrk
            diag_idx = np.diag_indices_from(prior_scaled)
            prior_scaled[diag_idx] += reinjection
            cov_inv = update_precision_dense(self.beta, X_weighted, prior_scaled)

        # Apply weights to y for the linear term
        y_weighted = y * effective_weights

        if self.sparse:
            # Scale vectors instead of sparse matrices
            eta = self.cov_inv_ @ (prior_decay * self.coef_) + X.T @ (
                self.beta * y_weighted
            )
            eta = cast(NDArray[np.float64], eta)
            assert isinstance(cov_inv, csc_array)
            factor: PrecisionFactor = create_sparse_factor(cov_inv)
            coef = factor.solve(eta)
            self._precision_factor = factor
        else:
            # eta = prior_decay * cov_inv_ @ coef_ + beta * X^T @ y_weighted
            eta = compute_eta_dense(
                prior_decay, self.cov_inv_, self.coef_, self.beta, X, y_weighted
            )
            # Cache the Cholesky factor for reuse in cov_/sample
            cho = cho_factor(cov_inv, lower=False, check_finite=False)
            self._precision_factor = DenseFactor(_U=cho[0], _n_features=cho[0].shape[0])
            coef = cho_solve(cho, eta, check_finite=False)

        self.cov_inv_ = cov_inv
        self.coef_ = coef

    def partial_fit(
        self,
        X: Union[NDArray[Any], csc_array],
        y: NDArray[Any],
        sample_weight: Optional[NDArray[Any]] = None,
    ) -> Self:
        """
        Incrementally update the posterior and retune hyperparameters.

        Performs a recursive Bayesian update (delegating to
        ``NormalRegressor.partial_fit``), then runs one online MacKay
        step to adjust ``alpha`` and ``beta`` using accumulated
        sufficient statistics. The precision matrix is corrected to
        reflect the updated hyperparameters.

        When ``learning_rate < 1``, stabilized forgetting (Kulhavy &
        Zarrop 1993) re-injects ``(1 - γ^n)·α`` into the precision
        diagonal so that the prior contribution converges to ``alpha``
        instead of decaying to zero.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.
        y : array-like of shape (n_samples,)
            Target values.
        sample_weight : array-like of shape (n_samples,), optional
            Individual weights for each sample. If None, all samples
            are given weight 1.0.

        Returns
        -------
        self : EmpiricalBayesNormalRegressor
            Updated estimator with retuned hyperparameters.

        See Also
        --------
        fit : Fit from scratch with full EB iteration loop.
        decay : Increase uncertainty without observing new data.
        """
        had_prior_scalar = hasattr(self, "_prior_scalar")

        # Clear cached factor so _fit_helper creates a fresh one.
        # The sparsity pattern may change with new X, and a stale
        # permutation can cause catastrophic fill-in (O(n^2) instead
        # of O(nnz)).  Reuse is only safe inside the EB loop where X
        # is fixed.
        if "_precision_factor" in self.__dict__:
            del self._precision_factor

        n_samples = X.shape[0] if hasattr(X, "shape") else len(X)  # type: ignore[arg-type]
        prior_decay = self.learning_rate**n_samples

        if had_prior_scalar:
            # Stabilized forgetting: decay the tracked prior contribution but
            # re-inject (1 - prior_decay) * alpha so it converges to alpha
            # instead of decaying to zero.  This ensures the EB-tuned prior
            # always regularizes new or dormant coefficients.
            self._prior_scalar = (
                prior_decay * self._prior_scalar + (1 - prior_decay) * self.alpha
            )

        # Store reinjection amount for _fit_helper to fold in.
        # When > 0, _fit_helper adds this to the precision diagonal during
        # construction, avoiding a separate _reinject_prior + refactorization.
        self._pending_reinjection = (
            (1 - prior_decay) * self.alpha if had_prior_scalar else 0.0
        )

        alpha_old = self.alpha
        beta_old = self.beta

        result = super().partial_fit(X, y, sample_weight)

        self._pending_reinjection = 0.0

        if not had_prior_scalar:
            if hasattr(self, "_prior_scalar"):
                # fit() was called by super().partial_fit() and already
                # set _prior_scalar, sufficient stats, and ran the EB loop.
                return result

            # sample() previously called _initialize_prior (setting coef_),
            # so super().partial_fit() did an incremental update instead of
            # calling fit(). Initialize EB state for the first time.
            #
            # After the incremental update the precision matrix is:
            #   cov_inv_ = prior_decay * alpha_old * I + beta_old * X^T X
            # so the prior contribution to the diagonal is prior_decay * alpha_old.
            self._prior_scalar = prior_decay * alpha_old

        X_fit, y = check_X_y(
            X,  # type: ignore
            y,
            copy=True,
            ensure_2d=True,
            dtype=np.float64,
            accept_sparse="csc" if self.sparse else False,
        )

        if not had_prior_scalar:
            # First observation — initialize sufficient statistics.
            self._effective_n = float(y.shape[0])
            self._eff_yTy = float(y @ y)
            if isinstance(X_fit, csc_array):
                self._eff_XTy = np.asarray(X_fit.T @ y).ravel()
            else:
                self._eff_XTy = X_fit.T @ y
        else:
            # Accumulate sufficient statistics for the beta update.
            self._accumulate_stats(X_fit, y, prior_decay)

        self._eb_mackay_step_online()
        self._correct_precision(alpha_old, beta_old)

        return result

    def decay(
        self,
        X: Union[NDArray[Any], csc_array],
        *,
        decay_rate: Optional[float] = None,
    ) -> None:
        """
        Decay the precision matrix with stabilized prior re-injection.

        Applies exponential forgetting ``Λ_new = γ^n · Λ_old`` and
        then re-injects ``(1 - γ^n)·α`` onto the diagonal, following
        stabilized forgetting (Kulhavy & Zarrop 1993). This ensures
        the prior contribution to the precision converges to ``alpha``
        rather than decaying to zero under repeated decay steps.

        Sufficient statistics (``_effective_n``, ``_eff_yTy``,
        ``_eff_XTy``) used for the online MacKay beta update are
        also decayed.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Used only to determine the number of time steps ``n`` for
            the decay exponent ``γ^n``. The actual feature values are
            not used.
        decay_rate : float, default=None
            Decay factor :math:`\\gamma` in (0, 1]. If None, uses
            ``self.learning_rate``.

        See Also
        --------
        partial_fit : Update the model with new observations.
        """
        if not hasattr(self, "coef_"):
            return

        if decay_rate is None:
            decay_rate = self.learning_rate

        assert X.shape is not None
        prior_decay = decay_rate ** X.shape[0]

        prior_reinjection = 0.0
        if hasattr(self, "_prior_scalar"):
            self._prior_scalar = (
                prior_decay * self._prior_scalar + (1 - prior_decay) * self.alpha
            )
            prior_reinjection = (1 - prior_decay) * self.alpha
        if hasattr(self, "_effective_n"):
            self._effective_n *= prior_decay
            self._eff_yTy *= prior_decay
            self._eff_XTy *= prior_decay

        # Base class applies uniform decay: cov_inv_ *= prior_decay
        super().decay(X, decay_rate=decay_rate)

        self._reinject_prior(prior_reinjection)
