"""tklsq.py

汎用線形回帰・誤差評価・モデル選択ユーティリティ。

設計方針
--------
- 基本は線形最小二乗: y = X beta + eps
- パラメータ誤差は cov(beta) = sigma^2 (X^T W X)^(-1)
- 推定値誤差は Var[y_mean] = diag(X_new cov(beta) X_new^T)
- 予測誤差は Var[y_pred] = Var[y_mean] + sigma^2
- モデル選択は AIC/BIC/AICc と、必要に応じてベイズ evidence を使う

このファイルは tklib / matplotlib / openpyxl に依存しない「計算コア」です。
入出力、グラフ、Excel保存は呼び出し側スクリプトに残す想定です。
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Union
import math
import numpy as np

ArrayLike = Union[Sequence[float], np.ndarray]
CandidateType = Union[Sequence[np.ndarray], Mapping[str, np.ndarray]]


# =====================================================================
# Dataclasses
# =====================================================================

@dataclass
class LeastSquaresResult:
    """線形最小二乗の結果。"""

    beta: np.ndarray
    beta_std: Optional[np.ndarray]
    cov_beta: Optional[np.ndarray]
    sigma2_resid: Optional[float]
    residuals: np.ndarray
    y_fit: np.ndarray
    N: int
    p: int
    dof: int
    rank: int
    RSS: float
    WRSS: float
    singular_values: np.ndarray
    condition_number: float
    error_estimation_enabled: bool
    warning: str = ""
    known_sigma: bool = False

    @property
    def sigma_resid(self) -> Optional[float]:
        if self.sigma2_resid is None:
            return None
        return float(math.sqrt(max(self.sigma2_resid, 0.0)))


@dataclass
class PredictionResult:
    """線形モデルの推定値・誤差バンド。"""

    y_mean: np.ndarray
    sigma_param: Optional[np.ndarray]
    sigma_pred: Optional[np.ndarray]
    var_param: Optional[np.ndarray]
    var_pred: Optional[np.ndarray]


@dataclass
class ModelSelectionResult:
    """複数モデル比較の結果。"""

    labels: List[str]
    scores: np.ndarray
    weights: np.ndarray
    best_index: int
    criterion: str
    results: Optional[List[LeastSquaresResult]] = None
    log_evidences: Optional[np.ndarray] = None

    @property
    def best_label(self) -> str:
        return self.labels[self.best_index]


@dataclass
class BayesianLinearResult:
    """ベイズ線形回帰の事後分布。"""

    mean: np.ndarray
    cov: np.ndarray
    alpha: float
    beta: float
    sigma_noise: float
    log_evidence: Optional[float] = None
    n_iter: int = 0
    converged: bool = False


# =====================================================================
# Design matrices
# =====================================================================

def polynomial_design_matrix(
    x: ArrayLike,
    order: Optional[int] = None,
    basis_indices: Optional[Sequence[int]] = None,
) -> np.ndarray:
    """多項式基底の設計行列を作る。

    Parameters
    ----------
    x : array-like
        1次元の説明変数。
    order : int, optional
        0, 1, ..., order の全基底を使う。
    basis_indices : sequence of int, optional
        任意のべき指数を指定する。例: [0, 1, 3]。

    Returns
    -------
    X : ndarray, shape (N, p)
    """
    x = np.asarray(x, dtype=float).reshape(-1)
    if basis_indices is None:
        if order is None:
            raise ValueError("Specify either order or basis_indices")
        basis_indices = list(range(order + 1))
    return np.vstack([x ** int(i) for i in basis_indices]).T


def design_matrix_from_functions(
    x: ArrayLike,
    basis_functions: Sequence,
) -> np.ndarray:
    """任意の基底関数リストから設計行列を作る。

    basis_functions の各要素は f(x_array) を返す関数、または
    np.vectorize 可能な f(x_scalar) を想定する。
    """
    x = np.asarray(x, dtype=float).reshape(-1)
    cols = []
    for f in basis_functions:
        v = np.asarray(f(x), dtype=float)
        if v.shape == () or v.shape != x.shape:
            v = np.asarray([f(xi) for xi in x], dtype=float)
        cols.append(v.reshape(-1))
    return np.vstack(cols).T


def add_intercept(X: np.ndarray) -> np.ndarray:
    """設計行列の先頭列に定数項を追加する。"""
    X = np.asarray(X, dtype=float)
    return np.column_stack([np.ones(X.shape[0]), X])


# =====================================================================
# Least squares and uncertainty
# =====================================================================

def linear_lsq(
    X: np.ndarray,
    y: ArrayLike,
    *,
    weights: Optional[ArrayLike] = None,
    known_sigma: bool = False,
    rcond: Optional[float] = None,
    cond_warning: float = 1.0e12,
) -> LeastSquaresResult:
    """線形最小二乗を行い、可能なら共分散行列と誤差を計算する。

    Parameters
    ----------
    X : ndarray, shape (N, p)
        設計行列。
    y : array-like, shape (N,)
        観測値。
    weights : array-like, optional
        重み。通常は 1/sigma_y^2 を指定する。
        None の場合は等重み。
    known_sigma : bool
        True の場合、weights が絶対的な 1/sigma_y^2 であるとみなし、
        cov(beta) = (X^T W X)^(-1) とする。
        False の場合、残差から sigma^2 を推定して掛ける。
    rcond : float or None
        np.linalg.lstsq の cutoff。
    cond_warning : float
        条件数がこの値を超えると warning に追記する。

    Notes
    -----
    rank deficient または N-p<=0 の場合でも beta は返す。
    ただし unbiased な残差分散とパラメータ誤差は返さず、
    beta_std, cov_beta, sigma2_resid は None になる。
    """
    X = np.asarray(X, dtype=float)
    y = np.asarray(y, dtype=float).reshape(-1)
    if X.ndim != 2:
        raise ValueError("X must be a 2D array")
    if X.shape[0] != y.size:
        raise ValueError(f"X rows ({X.shape[0]}) and y size ({y.size}) mismatch")

    N, p = X.shape
    dof = N - p

    if weights is None:
        w = np.ones(N, dtype=float)
    else:
        w = np.asarray(weights, dtype=float).reshape(-1)
        if w.size != N:
            raise ValueError(f"weights size ({w.size}) and y size ({N}) mismatch")
        if np.any(w < 0):
            raise ValueError("weights must be non-negative")

    sw = np.sqrt(w)
    Xw = X * sw[:, None]
    yw = y * sw

    beta, _residuals_lstsq, rank, svals = np.linalg.lstsq(Xw, yw, rcond=rcond)
    y_fit = X @ beta
    residuals = y - y_fit
    RSS = float(residuals.T @ residuals)
    WRSS = float((sw * residuals).T @ (sw * residuals))

    if len(svals) == 0 or np.min(svals) <= 0:
        condition_number = float("inf")
    else:
        condition_number = float(np.max(svals) / np.min(svals))

    warning_parts: List[str] = []
    error_enabled = True

    if rank < p:
        error_enabled = False
        warning_parts.append(
            f"WARNING: Design matrix is rank deficient (rank={rank}, p={p}). "
            "Least-squares coefficients were calculated, but parameter "
            "errors/covariance and uncertainty bands are disabled."
        )

    if dof <= 0:
        error_enabled = False
        warning_parts.append(
            f"WARNING: Residual degrees of freedom N-p = {N}-{p} = {dof}. "
            "The unbiased residual variance is undefined. Parameter errors, "
            "prediction errors, confidence intervals, and uncertainty bands are disabled."
        )

    if condition_number > cond_warning:
        warning_parts.append(
            f"WARNING: Ill-conditioned design matrix: cond={condition_number:.3e}. "
            "Fitted coefficients may be numerically unstable."
        )

    if not error_enabled:
        return LeastSquaresResult(
            beta=beta,
            beta_std=None,
            cov_beta=None,
            sigma2_resid=None,
            residuals=residuals,
            y_fit=y_fit,
            N=N,
            p=p,
            dof=dof,
            rank=rank,
            RSS=RSS,
            WRSS=WRSS,
            singular_values=svals,
            condition_number=condition_number,
            error_estimation_enabled=False,
            warning="\n".join(warning_parts),
            known_sigma=known_sigma,
        )

    XtWX = Xw.T @ Xw
    try:
        XtWX_inv = np.linalg.inv(XtWX)
    except np.linalg.LinAlgError:
        # full rank のはずでも極端に悪条件の場合の保険。
        XtWX_inv = np.linalg.pinv(XtWX, rcond=rcond if rcond is not None else 1e-15)
        warning_parts.append(
            "WARNING: np.linalg.inv failed; covariance was calculated with pseudo-inverse."
        )

    sigma2_resid = WRSS / dof
    cov_scale = 1.0 if known_sigma else sigma2_resid
    cov_beta = cov_scale * XtWX_inv
    beta_var = np.diag(cov_beta)
    beta_std = np.sqrt(np.maximum(beta_var, 0.0))

    return LeastSquaresResult(
        beta=beta,
        beta_std=beta_std,
        cov_beta=cov_beta,
        sigma2_resid=float(sigma2_resid),
        residuals=residuals,
        y_fit=y_fit,
        N=N,
        p=p,
        dof=dof,
        rank=rank,
        RSS=RSS,
        WRSS=WRSS,
        singular_values=svals,
        condition_number=condition_number,
        error_estimation_enabled=True,
        warning="\n".join(warning_parts),
        known_sigma=known_sigma,
    )


def polynomial_lsq(
    x: ArrayLike,
    y: ArrayLike,
    order: Optional[int] = None,
    *,
    basis_indices: Optional[Sequence[int]] = None,
    weights: Optional[ArrayLike] = None,
    known_sigma: bool = False,
    rcond: Optional[float] = None,
) -> LeastSquaresResult:
    """多項式線形回帰の便利関数。"""
    X = polynomial_design_matrix(x, order=order, basis_indices=basis_indices)
    return linear_lsq(X, y, weights=weights, known_sigma=known_sigma, rcond=rcond)


def parameter_prediction_variance(X_new: np.ndarray, cov_beta: np.ndarray) -> np.ndarray:
    """Var[y_mean] = diag(X_new cov_beta X_new^T) を計算する。"""
    X_new = np.asarray(X_new, dtype=float)
    cov_beta = np.asarray(cov_beta, dtype=float)
    return np.sum((X_new @ cov_beta) * X_new, axis=1)


def predict_linear(
    X_new: np.ndarray,
    result: LeastSquaresResult,
    *,
    sigma2_noise: Optional[float] = None,
    include_noise: bool = True,
) -> PredictionResult:
    """線形回帰結果から、推定値と誤差バンドを計算する。"""
    X_new = np.asarray(X_new, dtype=float)
    y_mean = X_new @ result.beta

    if result.cov_beta is None:
        return PredictionResult(
            y_mean=y_mean,
            sigma_param=None,
            sigma_pred=None,
            var_param=None,
            var_pred=None,
        )

    var_param = parameter_prediction_variance(X_new, result.cov_beta)
    sigma_param = np.sqrt(np.maximum(var_param, 0.0))

    if include_noise:
        if sigma2_noise is None:
            sigma2_noise = 0.0 if result.sigma2_resid is None else result.sigma2_resid
        var_pred = var_param + float(sigma2_noise)
        sigma_pred = np.sqrt(np.maximum(var_pred, 0.0))
    else:
        var_pred = None
        sigma_pred = None

    return PredictionResult(
        y_mean=y_mean,
        sigma_param=sigma_param,
        sigma_pred=sigma_pred,
        var_param=var_param,
        var_pred=var_pred,
    )


def predict_polynomial(
    x_new: ArrayLike,
    result: LeastSquaresResult,
    order: Optional[int] = None,
    *,
    basis_indices: Optional[Sequence[int]] = None,
    sigma2_noise: Optional[float] = None,
    include_noise: bool = True,
) -> PredictionResult:
    """多項式回帰結果から、推定値と誤差バンドを計算する。"""
    if basis_indices is None and order is None:
        order = result.p - 1
    X_new = polynomial_design_matrix(x_new, order=order, basis_indices=basis_indices)
    return predict_linear(X_new, result, sigma2_noise=sigma2_noise, include_noise=include_noise)


def measurement_std(y: ArrayLike, ddof: int = 1) -> float:
    """観測値全体の標準偏差を測定ばらつきの目安として返す。"""
    y = np.asarray(y, dtype=float).reshape(-1)
    if y.size <= ddof:
        return 0.0
    return float(np.std(y, ddof=ddof))


def delta_method_variance(jacobian: np.ndarray, cov_beta: np.ndarray) -> np.ndarray:
    """デルタ法: Var[g(beta)] = J cov(beta) J^T の対角を返す。

    jacobian は shape (N, p) または (p,) を想定する。
    派生量、例えば活性化エネルギーや移動度の log 変換後の誤差伝播に使える。
    """
    J = np.asarray(jacobian, dtype=float)
    if J.ndim == 1:
        J = J.reshape(1, -1)
    return np.sum((J @ cov_beta) * J, axis=1)


# =====================================================================
# Scores and information criteria
# =====================================================================

def regression_scores(y: ArrayLike, y_fit: ArrayLike, p: int = 0) -> Dict[str, float]:
    """基本的な回帰スコアを返す。"""
    y = np.asarray(y, dtype=float).reshape(-1)
    y_fit = np.asarray(y_fit, dtype=float).reshape(-1)
    residuals = y - y_fit
    N = y.size
    RSS = float(residuals.T @ residuals)
    TSS = float(((y - np.mean(y)).T @ (y - np.mean(y))))
    rmse = math.sqrt(RSS / N) if N > 0 else float("nan")
    r2 = 1.0 - RSS / TSS if TSS > 0 else float("nan")
    adj_r2 = 1.0 - (1.0 - r2) * (N - 1) / (N - p) if N > p and np.isfinite(r2) else float("nan")
    return {"N": float(N), "p": float(p), "RSS": RSS, "RMSE": rmse, "R2": r2, "adj_R2": adj_r2}


def gaussian_log_likelihood_from_rss(RSS: float, N: int, sigma2: Optional[float] = None) -> float:
    """Gaussian residual を仮定した log likelihood。

    sigma2=None の場合、MLE sigma2=RSS/N を使う。
    """
    eps = np.finfo(float).tiny
    if sigma2 is None:
        sigma2 = max(float(RSS) / max(int(N), 1), eps)
    else:
        sigma2 = max(float(sigma2), eps)
    return float(-0.5 * N * (math.log(2.0 * math.pi * sigma2) + 1.0 if RSS > 0 else math.log(2.0 * math.pi * sigma2)))


def information_criteria(result: LeastSquaresResult, *, k: Optional[int] = None) -> Dict[str, float]:
    """AIC, AICc, BIC を計算する。"""
    N = result.N
    if k is None:
        k = result.p
    logL = gaussian_log_likelihood_from_rss(result.RSS, N)
    aic = 2.0 * k - 2.0 * logL
    if N - k - 1 > 0:
        aicc = aic + (2.0 * k * (k + 1.0)) / (N - k - 1.0)
    else:
        aicc = float("inf")
    bic = k * math.log(N) - 2.0 * logL if N > 0 else float("inf")
    return {"logL": logL, "AIC": aic, "AICc": aicc, "BIC": bic}


def normalized_weights_from_scores(scores: ArrayLike, *, lower_is_better: bool = True) -> np.ndarray:
    """AIC/BIC や -log evidence から相対重みを計算する。"""
    s = np.asarray(scores, dtype=float)
    if lower_is_better:
        z = -0.5 * (s - np.nanmin(s))
    else:
        z = s - np.nanmax(s)
    z = np.where(np.isfinite(z), z, -np.inf)
    w = np.exp(z)
    total = np.sum(w)
    if total <= 0:
        return np.full_like(w, 1.0 / len(w), dtype=float)
    return w / total


def _normalize_candidates(
    candidates: CandidateType,
    labels: Optional[Sequence[str]] = None,
) -> Tuple[List[str], List[np.ndarray]]:
    if isinstance(candidates, Mapping):
        out_labels = [str(k) for k in candidates.keys()]
        matrices = [np.asarray(v, dtype=float) for v in candidates.values()]
    else:
        matrices = [np.asarray(v, dtype=float) for v in candidates]
        if labels is None:
            out_labels = [f"model_{i}" for i in range(len(matrices))]
        else:
            if len(labels) != len(matrices):
                raise ValueError("labels length and candidates length mismatch")
            out_labels = [str(v) for v in labels]
    return out_labels, matrices


def select_models_by_ic(
    candidates: CandidateType,
    y: ArrayLike,
    *,
    labels: Optional[Sequence[str]] = None,
    criterion: str = "BIC",
    weights: Optional[ArrayLike] = None,
    rcond: Optional[float] = None,
) -> ModelSelectionResult:
    """候補設計行列を AIC/AICc/BIC で比較する。"""
    criterion_key = criterion.upper()
    if criterion_key not in {"AIC", "AICC", "BIC"}:
        raise ValueError("criterion must be 'AIC', 'AICc', or 'BIC'")
    if criterion_key == "AICC":
        criterion_key = "AICc"

    out_labels, matrices = _normalize_candidates(candidates, labels)
    results = [linear_lsq(X, y, weights=weights, rcond=rcond) for X in matrices]
    scores = np.array([information_criteria(r)[criterion_key] for r in results], dtype=float)
    model_weights = normalized_weights_from_scores(scores, lower_is_better=True)
    best_index = int(np.nanargmin(scores))
    return ModelSelectionResult(
        labels=out_labels,
        scores=scores,
        weights=model_weights,
        best_index=best_index,
        criterion=criterion_key,
        results=results,
    )


def polynomial_candidates(x: ArrayLike, max_order: int, *, min_order: int = 0) -> Dict[str, np.ndarray]:
    """0..max_order の多項式候補を dict として作る。"""
    return {f"poly{d}": polynomial_design_matrix(x, order=d) for d in range(min_order, max_order + 1)}


def basis_index_candidates(x: ArrayLike, models: Sequence[Sequence[int]]) -> Dict[str, np.ndarray]:
    """任意の多項式べき指数リストから候補設計行列を作る。"""
    return {str(list(m)): polynomial_design_matrix(x, basis_indices=m) for m in models}


# =====================================================================
# Bayesian linear regression / evidence
# =====================================================================

def bayesian_posterior(
    X: np.ndarray,
    y: ArrayLike,
    *,
    alpha: float = 1.0,
    beta: float = 1.0,
    mu0: Optional[ArrayLike] = None,
    prior_cov: Optional[np.ndarray] = None,
) -> BayesianLinearResult:
    """線形ガウスモデルの事後平均・共分散を計算する。

    prior: beta_vec ~ N(mu0, prior_cov / alpha)
    noise: y|beta_vec ~ N(X beta_vec, I / beta)
    """
    X = np.asarray(X, dtype=float)
    y = np.asarray(y, dtype=float).reshape(-1)
    N, p = X.shape
    if mu0 is None:
        mu0_arr = np.zeros(p)
    else:
        mu0_arr = np.asarray(mu0, dtype=float).reshape(-1)
    if mu0_arr.size != p:
        raise ValueError("mu0 size mismatch")

    if prior_cov is None:
        prior_cov_arr = np.eye(p)
    else:
        prior_cov_arr = np.asarray(prior_cov, dtype=float)
        if prior_cov_arr.shape != (p, p):
            raise ValueError("prior_cov shape mismatch")

    prior_precision = alpha * np.linalg.inv(prior_cov_arr)
    S_inv = prior_precision + beta * (X.T @ X)
    S = np.linalg.inv(S_inv)
    m = S @ (beta * X.T @ y + prior_precision @ mu0_arr)
    log_ev = bayesian_log_evidence(X, y, alpha=alpha, beta=beta, mu0=mu0_arr, prior_cov=prior_cov_arr)
    return BayesianLinearResult(
        mean=m,
        cov=S,
        alpha=float(alpha),
        beta=float(beta),
        sigma_noise=float(math.sqrt(1.0 / beta)),
        log_evidence=log_ev,
    )


def bayesian_log_evidence(
    X: np.ndarray,
    y: ArrayLike,
    *,
    alpha: float = 1.0,
    beta: float = 1.0,
    mu0: Optional[ArrayLike] = None,
    prior_cov: Optional[np.ndarray] = None,
) -> float:
    """線形ガウスモデルの log evidence を計算する。

    prior: beta_vec ~ N(mu0, prior_cov / alpha)
    noise: y|beta_vec ~ N(X beta_vec, I / beta)

    N が大きい場合は O(N^3) なので、候補数・データ数が大きい用途では
    高速版の追加を検討する。
    """
    X = np.asarray(X, dtype=float)
    y = np.asarray(y, dtype=float).reshape(-1)
    N, p = X.shape
    if mu0 is None:
        mu0_arr = np.zeros(p)
    else:
        mu0_arr = np.asarray(mu0, dtype=float).reshape(-1)
    if prior_cov is None:
        prior_cov_arr = np.eye(p)
    else:
        prior_cov_arr = np.asarray(prior_cov, dtype=float)

    actual_prior_cov = prior_cov_arr / float(alpha)
    C = (1.0 / float(beta)) * np.eye(N) + X @ actual_prior_cov @ X.T
    r = y - X @ mu0_arr
    sign, logdet = np.linalg.slogdet(C)
    if sign <= 0:
        return -float("inf")
    q = float(r.T @ np.linalg.solve(C, r))
    return float(-0.5 * (N * math.log(2.0 * math.pi) + logdet + q))


def empirical_bayes_linear(
    X: np.ndarray,
    y: ArrayLike,
    *,
    alpha0: float = 1.0,
    beta0: Optional[float] = None,
    sigma_noise0: float = 1.0,
    mu0: Optional[ArrayLike] = None,
    max_iter: int = 100,
    tol: float = 1.0e-4,
    eps: float = 1.0e-300,
) -> BayesianLinearResult:
    """MacKay型の evidence maximization で alpha, beta を推定する。

    現在は等方 prior beta_vec ~ N(mu0, I/alpha) 用。
    """
    X = np.asarray(X, dtype=float)
    y = np.asarray(y, dtype=float).reshape(-1)
    N, p = X.shape
    if mu0 is None:
        mu0_arr = np.zeros(p)
    else:
        mu0_arr = np.asarray(mu0, dtype=float).reshape(-1)
    if mu0_arr.size != p:
        raise ValueError("mu0 size mismatch")

    alpha = float(alpha0)
    beta = float(beta0 if beta0 is not None else 1.0 / (sigma_noise0 ** 2))
    eigvals = np.linalg.eigvalsh(X.T @ X)
    converged = False

    for it in range(1, max_iter + 1):
        S_inv = alpha * np.eye(p) + beta * (X.T @ X)
        S = np.linalg.inv(S_inv)
        m = S @ (beta * X.T @ y + alpha * mu0_arr)
        gamma = float(np.sum(beta * eigvals / (alpha + beta * eigvals)))
        dm2 = float((m - mu0_arr).T @ (m - mu0_arr))
        res = y - X @ m
        rss = float(res.T @ res)

        alpha_new = gamma / max(dm2, eps)
        beta_new = max(N - gamma, eps) / max(rss, eps)

        da = abs(alpha_new - alpha) / max(abs(alpha), eps)
        db = abs(beta_new - beta) / max(abs(beta), eps)
        alpha, beta = float(alpha_new), float(beta_new)
        if da < tol and db < tol:
            converged = True
            break

    posterior = bayesian_posterior(X, y, alpha=alpha, beta=beta, mu0=mu0_arr)
    posterior.n_iter = it
    posterior.converged = converged
    return posterior


def predict_bayesian(
    X_new: np.ndarray,
    result: BayesianLinearResult,
    *,
    include_noise: bool = True,
) -> PredictionResult:
    """ベイズ線形回帰結果から予測平均・分散を計算する。"""
    X_new = np.asarray(X_new, dtype=float)
    y_mean = X_new @ result.mean
    var_param = parameter_prediction_variance(X_new, result.cov)
    sigma_param = np.sqrt(np.maximum(var_param, 0.0))
    if include_noise:
        var_pred = var_param + 1.0 / result.beta
        sigma_pred = np.sqrt(np.maximum(var_pred, 0.0))
    else:
        var_pred = None
        sigma_pred = None
    return PredictionResult(y_mean, sigma_param, sigma_pred, var_param, var_pred)


def posterior_model_probabilities(log_evidences: ArrayLike, log_priors: Optional[ArrayLike] = None) -> np.ndarray:
    """log evidence と log prior から事後モデル確率を計算する。"""
    le = np.asarray(log_evidences, dtype=float)
    if log_priors is None:
        lp = np.zeros_like(le)
    else:
        lp = np.asarray(log_priors, dtype=float)
        if lp.shape != le.shape:
            raise ValueError("log_priors shape mismatch")
    z = le + lp
    z = z - np.nanmax(z)
    w = np.exp(z)
    total = np.sum(w)
    if total <= 0:
        return np.full_like(w, 1.0 / len(w), dtype=float)
    return w / total


def select_models_by_evidence(
    candidates: CandidateType,
    y: ArrayLike,
    *,
    labels: Optional[Sequence[str]] = None,
    alpha: float = 1.0,
    beta: float = 1.0,
    log_priors: Optional[ArrayLike] = None,
) -> ModelSelectionResult:
    """候補設計行列を Bayesian log evidence で比較する。"""
    out_labels, matrices = _normalize_candidates(candidates, labels)
    log_evs = np.array([bayesian_log_evidence(X, y, alpha=alpha, beta=beta) for X in matrices])
    probs = posterior_model_probabilities(log_evs, log_priors=log_priors)
    best_index = int(np.nanargmax(probs))
    return ModelSelectionResult(
        labels=out_labels,
        scores=log_evs,
        weights=probs,
        best_index=best_index,
        criterion="log_evidence",
        results=None,
        log_evidences=log_evs,
    )


# =====================================================================
# Optional sparse model selection
# =====================================================================

def ard_select(
    X: np.ndarray,
    y: ArrayLike,
    *,
    threshold: float = 1.0e-3,
    fit_intercept: bool = False,
    **kwargs,
) -> Dict[str, object]:
    """sklearn の ARDRegression を使った疎な基底選択。

    sklearn がない環境では ImportError を出す。
    ライブラリ本体の必須依存に sklearn を入れないため、ここで局所 import する。
    """
    try:
        from sklearn.linear_model import ARDRegression
    except ImportError as exc:
        raise ImportError("ard_select requires scikit-learn") from exc

    X = np.asarray(X, dtype=float)
    y = np.asarray(y, dtype=float).reshape(-1)
    ard = ARDRegression(fit_intercept=fit_intercept, compute_score=True, **kwargs)
    ard.fit(X, y)
    coef = np.asarray(ard.coef_, dtype=float)
    selected = [int(i) for i, c in enumerate(coef) if abs(c) > threshold]
    return {
        "model": ard,
        "coef": coef,
        "intercept": float(getattr(ard, "intercept_", 0.0)),
        "selected": selected,
        "scores": getattr(ard, "scores_", None),
        "alpha": getattr(ard, "alpha_", None),
        "lambda": getattr(ard, "lambda_", None),
    }