"""From-scratch Bayesian VAR with Minnesota prior and a Gibbs sampler.

Follows Meseguer (2010): mortality is modelled in **first differences of log
rates** (Delta log m), separately by sex, over the 22 abridged age groups. Two
models (both estimated by a hand-written Gibbs sampler, general lag order p):

  M1  "Minnesota"    : reduced-form VAR with intercept, Minnesota prior on the AR
                       coefficients (own first lag centred on 0 for differenced
                       data), independent Normal-inverse-Wishart -> 2-block Gibbs.
  M2  "steady state" : Villani mean-adjusted VAR y_t = psi + sum Pi_l (y_{t-l}-psi)
                       + e, with the same Minnesota prior on Pi and an informative
                       Normal prior on the steady state psi -> 3-block Gibbs.

Minnesota prior standard deviations (Meseguer eq. 16), for coefficient on variable
j at lag l in equation i:
    j = i : lam1 / l**lam3
    j != i: lam1 * lam2(i,j) / l**lam3 * (s_i / s_j)
with s_i the residual std of an AR(p) fit of series i, and lam2(i,j) set to
0.8 * corr(Delta log y_i, Delta log y_j).
"""

import numpy as np
from scipy.stats import invwishart


# --------------------------------------------------------------------------- #
# data / prior construction
# --------------------------------------------------------------------------- #

def build_regressors(Y, p, intercept):
    """Stack a VAR(p) design. Y is (T, m). Returns (Yt, X) with X lag-blocks
    ordered [lag1 (m cols), lag2, ...], optionally prefixed by an intercept."""
    T, m = Y.shape
    rows = T - p
    lags = [Y[p - l:T - l] for l in range(1, p + 1)]     # each (rows, m)
    X = np.hstack(lags)
    if intercept:
        X = np.hstack([np.ones((rows, 1)), X])
    return Y[p:], X


def _series_scale(Y, p):
    """s_i = residual std of an AR(p) regression of each series on its own lags."""
    T, m = Y.shape
    s = np.empty(m)
    for i in range(m):
        yi = Y[p:, i]
        Xi = np.column_stack([np.ones(T - p)] + [Y[p - l:T - l, i] for l in range(1, p + 1)])
        beta, *_ = np.linalg.lstsq(Xi, yi, rcond=None)
        s[i] = (yi - Xi @ beta).std(ddof=len(beta))
    return s


def minnesota_moments(Y, p, lam1=0.2, lam2=0.5, lam3=1.0, own_mean=0.0,
                      intercept=True, intercept_sd=1e3):
    """Prior mean B0 and prior variance V0, each shaped (k, m) with k regressors
    per equation and columns = equations. `lam2` may be a scalar or an (m, m)
    matrix lam2(i,j) (e.g. 0.8 * corr)."""
    T, m = Y.shape
    s = _series_scale(Y, p)
    k = (1 if intercept else 0) + m * p
    B0 = np.zeros((k, m))
    V0 = np.zeros((k, m))
    off = 1 if intercept else 0
    if intercept:
        V0[0, :] = intercept_sd ** 2                     # diffuse intercept
    lam2_mat = lam2 if np.ndim(lam2) == 2 else np.full((m, m), lam2)
    for i in range(m):                                   # equation i
        for l in range(1, p + 1):
            for j in range(m):                           # regressor variable j
                r = off + (l - 1) * m + j
                if j == i and l == 1:
                    B0[r, i] = own_mean
                if j == i:
                    V0[r, i] = (lam1 / l ** lam3) ** 2
                else:
                    V0[r, i] = (lam1 * lam2_mat[i, j] / l ** lam3 * s[i] / s[j]) ** 2
    return B0, V0, s


def corr_lam2(Y, p, weight=0.8):
    """lam2(i,j) = weight * corr(Delta log y_i, Delta log y_j), used for cross shrinkage."""
    C = np.corrcoef(Y[p:].T)
    return weight * C


# --------------------------------------------------------------------------- #
# Gaussian draw from a given precision matrix
# --------------------------------------------------------------------------- #

def _draw_from_precision(rhs, Prec, rng):
    """Sample x ~ N(Prec^{-1} rhs, Prec^{-1}) via Cholesky."""
    L = np.linalg.cholesky(Prec)                         # Prec = L L^T
    mean = np.linalg.solve(L.T, np.linalg.solve(L, rhs))
    z = rng.standard_normal(rhs.shape[0])
    return mean + np.linalg.solve(L.T, z)


# --------------------------------------------------------------------------- #
# Model 1 - Minnesota reduced-form BVAR (2-block Gibbs)
# --------------------------------------------------------------------------- #

def gibbs_minnesota(Y, p=1, lam1=0.2, lam2=None, lam3=1.0, own_mean=0.0,
                    ndraw=3000, burn=1000, thin=1, seed=0, iw_df=None, iw_scale=None):
    """2-block Gibbs for y_t = c + sum_l Pi_l y_{t-l} + e, e~N(0,Sigma)."""
    rng = np.random.default_rng(seed)
    Yt, X = build_regressors(Y, p, intercept=True)
    rows, m = Yt.shape
    k = X.shape[1]
    if lam2 is None:
        lam2 = corr_lam2(Y, p)
    B0, V0, s = minnesota_moments(Y, p, lam1, lam2, lam3, own_mean, intercept=True)
    b0 = B0.reshape(-1, order="F")                       # vec by equation (column-major)
    V0inv = np.diag(1.0 / V0.reshape(-1, order="F"))
    XtX, XtY = X.T @ X, X.T @ Yt

    v0 = iw_df if iw_df else m + 2
    S0 = iw_scale if iw_scale is not None else np.diag(s ** 2)
    B = B0.copy()
    Sig = np.diag(s ** 2)

    kept_B, kept_S = [], []
    for it in range(ndraw + burn):
        Sinv = np.linalg.inv(Sig)
        Prec = V0inv + np.kron(Sinv, XtX)
        rhs = V0inv @ b0 + (XtY @ Sinv).reshape(-1, order="F")
        beta = _draw_from_precision(rhs, Prec, rng)
        B = beta.reshape(k, m, order="F")
        E = Yt - X @ B
        Sig = invwishart.rvs(df=v0 + rows, scale=S0 + E.T @ E, random_state=rng)
        if it >= burn and (it - burn) % thin == 0:
            kept_B.append(B.copy()); kept_S.append(Sig.copy())
    return {"B": np.array(kept_B), "Sigma": np.array(kept_S), "p": p,
            "intercept": True, "kind": "minnesota"}


# --------------------------------------------------------------------------- #
# Model 2 - Villani steady-state (mean-adjusted) BVAR (3-block Gibbs)
# --------------------------------------------------------------------------- #

def gibbs_steady_state(Y, p=1, lam1=0.2, lam2=None, lam3=1.0, own_mean=0.0,
                       psi_prior_mean=None, psi_prior_sd=10.0,
                       ndraw=3000, burn=1000, thin=1, seed=0, iw_df=None, iw_scale=None):
    """3-block Gibbs for y_t = psi + sum_l Pi_l (y_{t-l}-psi) + e.

    psi_prior_mean : (m,) informative steady-state centre (e.g. SSA Alt-2 ultimate
        Delta log m). Default 0 (diffuse). psi_prior_sd : scalar or (m,) prior sd.
    """
    rng = np.random.default_rng(seed)
    T, m = Y.shape
    rows = T - p
    if lam2 is None:
        lam2 = corr_lam2(Y, p)
    B0, V0, s = minnesota_moments(Y, p, lam1, lam2, lam3, own_mean, intercept=False)
    b0 = B0.reshape(-1, order="F")
    V0inv = np.diag(1.0 / V0.reshape(-1, order="F"))
    k = m * p

    psi0 = np.zeros(m) if psi_prior_mean is None else np.asarray(psi_prior_mean, float)
    psd = np.full(m, psi_prior_sd) if np.ndim(psi_prior_sd) == 0 else np.asarray(psi_prior_sd)
    Lam0inv = np.diag(1.0 / psd ** 2)

    v0 = iw_df if iw_df else m + 2
    S0 = iw_scale if iw_scale is not None else np.diag(s ** 2)

    Ylag_full = np.hstack([Y[p - l:T - l] for l in range(1, p + 1)])    # (rows, m*p) original units
    Yt = Y[p:]
    mu = Y.mean(axis=0).copy() if psi_prior_mean is None else psi0.copy()
    Pi = B0.copy()                                        # (k, m)
    Sig = np.diag(s ** 2)

    kept_Pi, kept_mu, kept_S = [], [], []
    for it in range(ndraw + burn):
        Sinv = np.linalg.inv(Sig)
        # --- block 1: Pi | mu, Sigma  (regress demeaned y on demeaned lags) ---
        Z = Y - mu                                       # (T, m) demeaned
        Zt, Zlag = build_regressors(Z, p, intercept=False)
        ZtZ, ZtY = Zlag.T @ Zlag, Zlag.T @ Zt
        Prec = V0inv + np.kron(Sinv, ZtZ)
        rhs = V0inv @ b0 + (ZtY @ Sinv).reshape(-1, order="F")
        Pi = _draw_from_precision(rhs, Prec, rng).reshape(k, m, order="F")

        # --- block 2: mu | Pi, Sigma  (u_t = (I - sum Pi_l) mu + e) ---
        U = Yt - Ylag_full @ Pi                          # (rows, m)
        sumPi = sum(Pi[l * m:(l + 1) * m, :] for l in range(p))   # sum of blocks
        D = np.eye(m) - sumPi.T
        Prec_mu = Lam0inv + rows * (D.T @ Sinv @ D)
        rhs_mu = Lam0inv @ psi0 + D.T @ Sinv @ U.sum(axis=0)
        mu = _draw_from_precision(rhs_mu, Prec_mu, rng)

        # --- block 3: Sigma | Pi, mu ---
        Z = Y - mu
        Zt, Zlag = build_regressors(Z, p, intercept=False)
        E = Zt - Zlag @ Pi
        Sig = invwishart.rvs(df=v0 + rows, scale=S0 + E.T @ E, random_state=rng)

        if it >= burn and (it - burn) % thin == 0:
            kept_Pi.append(Pi.copy()); kept_mu.append(mu.copy()); kept_S.append(Sig.copy())
    return {"Pi": np.array(kept_Pi), "mu": np.array(kept_mu), "Sigma": np.array(kept_S),
            "p": p, "kind": "steady_state"}


# --------------------------------------------------------------------------- #
# forecasting -> draws of future Delta log m
# --------------------------------------------------------------------------- #

def stationary_mask(fit):
    """Boolean mask of posterior draws whose VAR is covariance-stationary
    (all companion-matrix eigenvalues inside the unit circle)."""
    p = fit["p"]
    coefs = fit["B"] if fit["kind"] == "minnesota" else fit["Pi"]
    off = 1 if fit["kind"] == "minnesota" else 0
    m = fit["Sigma"].shape[1]
    mask = np.ones(len(coefs), bool)
    eye = np.eye(m * (p - 1)) if p > 1 else None
    for d, B in enumerate(coefs):
        Pis = [B[off + (l - 1) * m: off + l * m].T for l in range(1, p + 1)]  # Pi_l = block^T
        comp = np.zeros((m * p, m * p))
        comp[:m] = np.hstack(Pis)
        if p > 1:
            comp[m:, :m * (p - 1)] = eye
        mask[d] = np.max(np.abs(np.linalg.eigvals(comp))) < 1.0
    return mask


def filter_stationary(fit):
    """Return a copy of `fit` keeping only covariance-stationary draws."""
    mask = stationary_mask(fit)
    out = dict(fit)
    for key in ("B", "Pi", "mu", "Sigma"):
        if key in fit:
            out[key] = fit[key][mask]
    out["stationary_frac"] = float(mask.mean())
    return out


def forecast(fit, Y, H, seed=1):
    """Simulate H-step-ahead paths of the modelled series (Delta log m) for every
    posterior draw. Returns array (ndraw, H, m)."""
    rng = np.random.default_rng(seed)
    p = fit["p"]
    m = Y.shape[1]
    hist0 = Y[-p:]                                        # last p observations
    if fit["kind"] == "minnesota":
        Bs, Ss = fit["B"], fit["Sigma"]
    else:
        Bs, Ss, mus = fit["Pi"], fit["Sigma"], fit["mu"]
    nd = len(Ss)
    out = np.empty((nd, H, m))
    for d in range(nd):
        Sig = Ss[d]
        chol = np.linalg.cholesky(Sig)
        hist = [hist0[i].copy() for i in range(p)]       # hist[-1] = most recent
        for h in range(H):
            eps = chol @ rng.standard_normal(m)
            if fit["kind"] == "minnesota":
                B = Bs[d]
                x = B[0].copy()                          # intercept
                for l in range(1, p + 1):
                    x += B[1 + (l - 1) * m:1 + l * m].T @ hist[-l]
                y = x + eps
            else:
                Pi, mu = Bs[d], mus[d]
                x = mu.copy()
                for l in range(1, p + 1):
                    x += Pi[(l - 1) * m:l * m].T @ (hist[-l] - mu)
                y = x + eps
            out[d, h] = y
            hist.append(y)
        # keep only last p in hist via slicing not needed (indexing by -l)
    return out