"""
Bayesian Markov-switching AR(p) — from-scratch Gibbs sampler (FFBS).

Model (switching intercept, constant AR coefficients and variance):
    y_t = c_{S_t} + sum_{j=1}^p phi_j y_{t-j} + e_t,   e_t ~ N(0, sigma^2)
    S_t in {0,...,K-1} a first-order Markov chain with transition matrix P.

Sampler (fully conjugate Gibbs), one sweep:
  1. States S_{p:T-1}  -- Forward-Filter Backward-Sample (Chib 1996)
  2. phi | S, c, sigma2 -- Normal-Normal (GLS on lagged levels)
  3. c_k | S, phi, sigma2 -- Normal-Normal (per regime)
  4. sigma2 | .          -- Inverse-Gamma
  5. P rows | S          -- Dirichlet (conjugate to transition counts)
Identification: regimes relabelled each sweep so intercepts are ascending
(c_0 < c_1 < ...), which removes label switching.

No external MCMC packages -- NumPy only. Companion to mts_gibbs.py.
"""
import numpy as np


def simulate_msar(T, c, phi, sigma, P, seed=1):
    """Simulate a switching-intercept AR(p). Returns (y, S_true)."""
    rng = np.random.default_rng(seed)
    c = np.asarray(c, float); phi = np.asarray(phi, float)
    K = len(c); p = len(phi)
    S = np.zeros(T, int)
    for t in range(1, T):
        S[t] = rng.choice(K, p=P[S[t - 1]])
    y = np.zeros(T)
    for t in range(T):
        ar = sum(phi[j] * y[t - 1 - j] for j in range(p) if t - 1 - j >= 0)
        y[t] = c[S[t]] + ar + rng.normal(0, sigma)
    return y, S


def _stationary(P):
    """Stationary distribution of transition matrix P (rows = from-state)."""
    K = P.shape[0]
    A = np.vstack([P.T - np.eye(K), np.ones(K)])
    b = np.append(np.zeros(K), 1.0)
    pi, *_ = np.linalg.lstsq(A, b, rcond=None)
    pi = np.clip(pi, 1e-12, None)
    return pi / pi.sum()


def msar_gibbs(y, p=2, K=2, R=4000, burn=1000, seed=0,
               m0=0.0, v0=25.0, Vphi=0.25, a0=2.0, b0=1.0,
               alpha_diag=8.0, alpha_off=2.0, nprint=0):
    """Gibbs sampler for the switching-intercept MS-AR(p). Returns dict of draws."""
    rng = np.random.default_rng(seed)
    y = np.asarray(y, float); T = len(y)
    idx = np.arange(p, T); n = len(idx)
    X = np.column_stack([y[idx - j - 1] for j in range(p)])   # (n, p) lagged levels
    yt = y[idx]

    # --- inits ---
    phi = np.zeros(p)
    sigma2 = float(np.var(yt))
    c = np.quantile(yt, np.linspace(0.3, 0.7, K))
    P = np.full((K, K), 0.1); np.fill_diagonal(P, 0.9); P /= P.sum(1, keepdims=True)
    alpha = np.full((K, K), alpha_off); np.fill_diagonal(alpha, alpha_diag)

    keep = R - burn
    Sd = np.zeros((keep, n), np.int8); cD = np.zeros((keep, K))
    phiD = np.zeros((keep, p)); sigD = np.zeros(keep); Pd = np.zeros((keep, K, K))

    for g in range(R):
        # 1. FFBS -----------------------------------------------------------
        mu_tk = c[None, :] + (X @ phi)[:, None]               # (n, K)
        loglik = -0.5 * np.log(2 * np.pi * sigma2) - 0.5 * (yt[:, None] - mu_tk) ** 2 / sigma2
        filt = np.zeros((n, K))
        a_ = np.log(_stationary(P)) + loglik[0]; a_ -= a_.max()
        w = np.exp(a_); filt[0] = w / w.sum()
        for t in range(1, n):
            pred = filt[t - 1] @ P
            a_ = np.log(pred + 1e-300) + loglik[t]; a_ -= a_.max()
            w = np.exp(a_); filt[t] = w / w.sum()
        S = np.zeros(n, int)
        S[-1] = rng.choice(K, p=filt[-1])
        for t in range(n - 2, -1, -1):
            w = filt[t] * P[:, S[t + 1]]; w /= w.sum()
            S[t] = rng.choice(K, p=w)

        # 2. phi | S, c, sigma2  (Normal-Normal) ----------------------------
        z = yt - c[S]
        Bcov = np.linalg.inv(np.eye(p) / Vphi + X.T @ X / sigma2)
        phi = rng.multivariate_normal(Bcov @ (X.T @ z / sigma2), Bcov)

        # 3. c_k | S, phi, sigma2 ------------------------------------------
        u = yt - X @ phi
        for k in range(K):
            mk = (S == k); nk = int(mk.sum())
            prec = 1.0 / v0 + nk / sigma2
            mean = (m0 / v0 + u[mk].sum() / sigma2) / prec
            c[k] = rng.normal(mean, np.sqrt(1.0 / prec))

        # 4. sigma2 | .  (Inverse-Gamma) -----------------------------------
        e = yt - c[S] - X @ phi
        sigma2 = 1.0 / rng.gamma(a0 + n / 2.0, 1.0 / (b0 + 0.5 * np.sum(e ** 2)))

        # 5. P rows | S  (Dirichlet) ---------------------------------------
        Nc = np.zeros((K, K))
        np.add.at(Nc, (S[:-1], S[1:]), 1.0)
        for i in range(K):
            P[i] = rng.dirichlet(alpha[i] + Nc[i])

        # identification: relabel so intercepts ascending ------------------
        order = np.argsort(c)
        if not np.all(order == np.arange(K)):
            c = c[order]; P = P[np.ix_(order, order)]
            S = np.argsort(order)[S]

        if g >= burn:
            j = g - burn
            Sd[j] = S; cD[j] = c; phiD[j] = phi; sigD[j] = sigma2; Pd[j] = P
        if nprint and (g + 1) % nprint == 0:
            print(f"  [msar] {g+1}/{R}", flush=True)

    return dict(S=Sd, c=cD, phi=phiD, sigma2=sigD, P=Pd, idx=idx, p=p, K=K)


def regime_probs(out):
    """Smoothed P(S_t = k | y), shape (n, K), by averaging sampled state paths."""
    S = out['S']; K = out['K']
    return np.stack([(S == k).mean(0) for k in range(K)], axis=1)


def summary(out):
    """Posterior means of the structural parameters."""
    return dict(c=out['c'].mean(0), phi=out['phi'].mean(0),
                sigma=np.sqrt(out['sigma2'].mean()),
                P=out['P'].mean(0))