"""High-dimensional horseshoe regression -- from scratch (p >> n).

Global-local shrinkage selection when there are far more predictors than samples --
the setting of genomics (thousands of genes, tens of samples), where enumeration
(2^p models) and discrete spike-and-slab both break down.  The horseshoe prior
(Carvalho, Polson & Scott 2010) puts almost all coefficients essentially at zero while
letting a few escape, via a heavy-tailed local scale:

    beta_j | lambda_j, tau, sigma ~ N(0, lambda_j^2 tau^2 sigma^2),
    lambda_j ~ C+(0,1),   tau ~ C+(0,1),   sigma^2 ~ 1/sigma^2.

Sampling uses the Makalic-Schmidt (2016) auxiliary-variable Gibbs for the half-Cauchys,
and -- crucially for p >> n -- the Bhattacharya, Chakraborty & Mallick (2016) exact
sampler for the Gaussian coefficient update, which costs O(n^2 p) rather than the naive
O(p^3): it never forms or factorises the p x p precision matrix, solving an n x n system
instead.  With p=4088, n=71 that is the difference between feasible and not.

The shrinkage weight kappa_j = 1 / (1 + lambda_j^2 tau^2) in [0,1] measures how strongly
gene j is pulled to zero (kappa ~ 1 = shrunk away, kappa ~ 0 = a real signal).
"""
import numpy as np


def _invgamma(rng, shape, scale):
    """Inverse-Gamma(shape, scale) draw (scale = rate); shape/scale may be arrays."""
    return 1.0 / rng.gamma(shape, 1.0 / scale)


def horseshoe_hd(y, X, ndraw=3000, burn=1500, seed=0):
    """Gibbs sampler for high-dimensional horseshoe regression.

    y, X are centered/standardized internally. Returns posterior-mean coefficients,
    mean shrinkage weights kappa_j, and their draws.
    """
    rng = np.random.default_rng(seed)
    y = np.asarray(y, float)
    X = np.asarray(X, float)
    n, p = X.shape
    ybar = y.mean(); yc = y - ybar
    Xm = X.mean(0); Xs = X.std(0); Xs[Xs == 0] = 1.0
    Xz = (X - Xm) / Xs                                       # standardized genes (unit variance)

    beta = np.zeros(p)
    lam2 = np.ones(p); nu = np.ones(p)
    tau2 = 1.0; xi = 1.0
    sig2 = float(np.var(yc))
    In = np.eye(n)

    Bsum = np.zeros(p); Ksum = np.zeros(p); B2sum = np.zeros(p); nkeep = 0

    for it in range(ndraw + burn):
        # ---- (1) beta | rest  via the Bhattacharya et al. (2016) fast sampler ----
        sig = np.sqrt(sig2)
        D = lam2 * tau2 * sig2                               # prior variance of each beta_j
        u = np.sqrt(D) * rng.standard_normal(p)
        delta = rng.standard_normal(n)
        v = Xz @ u / sig + delta
        M = (Xz * D) @ Xz.T / sig2 + In                     # n x n
        w = np.linalg.solve(M, yc / sig - v)
        beta = u + D * (Xz.T @ w) / sig
        b2 = beta ** 2

        # ---- (2) local scales lambda_j^2 and aux nu_j (Makalic-Schmidt) ----
        lam2 = _invgamma(rng, 1.0, 1.0 / nu + b2 / (2 * tau2 * sig2))
        nu = _invgamma(rng, 1.0, 1.0 + 1.0 / lam2)

        # ---- (3) global scale tau^2 and aux xi ----
        tau2 = _invgamma(rng, (p + 1) / 2.0, 1.0 / xi + np.sum(b2 / lam2) / (2 * sig2))
        xi = _invgamma(rng, 1.0, 1.0 + 1.0 / tau2)

        # ---- (4) error variance sigma^2 ----
        rss = np.sum((yc - Xz @ beta) ** 2)
        sig2 = _invgamma(rng, (n + p) / 2.0, 0.5 * rss + np.sum(b2 / lam2) / (2 * tau2))

        if it >= burn:
            Bsum += beta; B2sum += b2; Ksum += 1.0 / (1.0 + lam2 * tau2); nkeep += 1

    return dict(beta=Bsum / nkeep, beta2=B2sum / nkeep, kappa=Ksum / nkeep,
                Xstd=Xs, ybar=ybar, n=n, p=p)