"""
bishop.linear_models – PRML Chapters 3 & 4
===========================================
Uses sklearn for standard models; adds a pure-numpy RVM (Relevance Vector
Machine) which is not available in any major library.

    from ml_ai_library.bishop.linear_models import (
        BayesianLinearRegression,   # sklearn.linear_model.BayesianRidge
        LogisticRegressionIRLS,     # sklearn.linear_model.LogisticRegression (lbfgs)
        RelevanceVectorMachine,     # sparse Bayesian – novel gap, numpy impl
    )
"""

from __future__ import annotations
import numpy as np
from sklearn.linear_model import BayesianRidge, LogisticRegression
from sklearn.base import BaseEstimator, RegressorMixin


# Re-export sklearn models under PRML names ---------------------------------

class BayesianLinearRegression(BayesianRidge):
    """
    Bayesian linear regression (PRML §3.3).

    Wraps sklearn.linear_model.BayesianRidge with the same API.
    Adds ``predict_with_std`` convenience method.

    Example
    -------
    >>> model = BayesianLinearRegression(alpha_init=1.0, lambda_init=1.0)
    >>> model.fit(X_train, y_train)
    >>> mean, std = model.predict_with_std(X_test)
    """

    def predict_with_std(self, X: np.ndarray):
        """Return (mean, std) of predictive distribution."""
        return self.predict(X, return_std=True)


class LogisticRegressionIRLS(LogisticRegression):
    """
    Logistic regression (PRML §4.3) – IRLS / Newton-CG solver.

    Wraps sklearn.linear_model.LogisticRegression with solver='newton-cg'
    which implements the IRLS algorithm described in the book.

    Example
    -------
    >>> clf = LogisticRegressionIRLS(C=1.0, max_iter=200)
    >>> clf.fit(X_train, y_train)
    >>> probs = clf.predict_proba(X_test)
    """

    def __init__(self, C: float = 1.0, max_iter: int = 200, **kwargs):
        super().__init__(solver="newton-cg", C=C, max_iter=max_iter, **kwargs)


# Relevance Vector Machine (PRML §7.2) – not in sklearn -------------------

class RelevanceVectorMachine(BaseEstimator, RegressorMixin):
    """
    Relevance Vector Machine for regression (PRML §7.2).

    Sparse Bayesian learning: each weight w_i has its own precision
    hyperparameter alpha_i.  The EM update drives most alpha_i → ∞,
    leaving only a sparse set of 'relevance vectors'.

    Parameters
    ----------
    kernel : callable or None
        k(X, X') → (N, M) gram matrix.  None → linear kernel X @ X'.T
    beta_init : float
        Initial noise precision.
    max_iter : int
        Maximum EM steps.
    tol : float
        Convergence tolerance on beta change.

    Example
    -------
    >>> rvm = RelevanceVectorMachine(max_iter=500)
    >>> rvm.fit(X_train, y_train)
    >>> y_pred = rvm.predict(X_test)
    >>> print("Relevance vectors:", len(rvm.relevance_vectors_))
    """

    def __init__(self, kernel=None, beta_init: float = 1.0,
                 max_iter: int = 300, tol: float = 1e-5):
        self.kernel = kernel
        self.beta_init = beta_init
        self.max_iter = max_iter
        self.tol = tol

    @staticmethod
    def _linear_kernel(A, B):
        return A @ B.T

    def fit(self, X: np.ndarray, t: np.ndarray):
        kern = self.kernel or self._linear_kernel
        self._X_train = X.copy()
        Phi = kern(X, X)          # (N, N) design matrix
        N = Phi.shape[0]
        alpha = np.ones(N)
        beta = float(self.beta_init)
        active = np.ones(N, dtype=bool)

        for _ in range(self.max_iter):
            Phi_a = Phi[:, active]
            A = np.diag(alpha[active]) + beta * Phi_a.T @ Phi_a
            Sigma = np.linalg.inv(A)
            mu = beta * Sigma @ Phi_a.T @ t

            gamma = 1.0 - alpha[active] * np.diag(Sigma)
            alpha[active] = np.clip(gamma / (mu ** 2 + 1e-12), 0, 1e9)
            beta_new = float(
                (N - gamma.sum()) / (np.sum((t - Phi_a @ mu) ** 2) + 1e-12)
            )

            # Prune: mark active entries whose alpha exceeded threshold
            prev_active = active.copy()
            active[prev_active] = alpha[prev_active] < 1e8
            if not active.any():
                break
            if abs(beta_new - beta) / (abs(beta) + 1e-10) < self.tol:
                break
            beta = beta_new

        # Recompute mu for the final (possibly pruned) active set so that
        # len(mu) == active.sum() and predict() dimensions are consistent.
        if active.any():
            Phi_a = Phi[:, active]
            A     = np.diag(alpha[active]) + beta * Phi_a.T @ Phi_a
            mu    = beta * np.linalg.inv(A) @ Phi_a.T @ t

        self._mu = mu
        self._active = active
        self._beta = beta
        self.relevance_vectors_ = X[active]
        return self

    def predict(self, X: np.ndarray) -> np.ndarray:
        kern = self.kernel or self._linear_kernel
        Phi = kern(X, self._X_train[self._active])
        return Phi @ self._mu
