"""
McCulloch & Tsay (1993) Bayesian AR(p) with Mean and Variance Shifts.

Reference:
  McCulloch, R.E. & Tsay, R.S. (1993).
  "Bayesian Inference and Prediction for Mean and Variance Shifts
   in Autoregressive Processes."
  Journal of the American Statistical Association, 88(423), 968-978.

Gibbs sampler — four blocks per iteration:
  Block 1 : AR coefficients phi         (conjugate Normal)
  Block 2 : Mean-shift indicators delta + mean levels mu   (single-site)
  Block 3 : Var-shift  indicators gamma + variance levels h2 (single-site)
  Block 4 : Shift probabilities pi_mu, pi_h               (conjugate Beta)

Key modelling assumption (long-segment approximation):
  r_t = y_t - sum_j phi_j y_{t-j}  ~  N(kappa * mu_t, h2_t)
where kappa = 1 - sum phi_j.  This is exact within a segment when all AR
lags also belong to the same mean segment; cross-segment corrections are O(phi^p)
and negligible for stationary AR.
"""

import time
import numpy as np
import pandas as pd
from scipy.special import gammaln
from scipy.stats import invgamma


# ─────────────────────────────────────────────────────────────────────────────
# Residual helpers
# ─────────────────────────────────────────────────────────────────────────────

def _raw_ar_resid(y, phi):
    """r_t = y_t - sum_{j=1}^p phi_j y_{t-j}  (no mean subtraction)."""
    T, p = len(y), len(phi)
    r = np.full(T, np.nan)
    for t in range(p, T):
        r[t] = y[t] - phi @ y[t - p:t][::-1]
    return r


def _full_resid(y, mu, phi):
    """e_t = (y_t - mu_t) - sum_j phi_j (y_{t-j} - mu_{t-j})."""
    T, p = len(y), len(phi)
    z = y - mu
    e = np.full(T, np.nan)
    for t in range(p, T):
        e[t] = z[t] - phi @ z[t - p:t][::-1]
    return e


# ─────────────────────────────────────────────────────────────────────────────
# Block 1 — AR coefficients
# ─────────────────────────────────────────────────────────────────────────────

def _sample_phi(y, mu, h2, p, b0, B0inv):
    """Normal posterior for phi.

    Demeaned series z_t = y_t - mu_t follows AR(p) with zero mean.
    """
    T = len(y)
    z = y - mu
    X  = np.array([z[t - p:t][::-1] for t in range(p, T)])  # (T-p, p)
    zp = z[p:]                                                # (T-p,)
    w  = 1.0 / h2[p:]                                        # (T-p,)

    XtWX   = (X * w[:, None]).T @ X
    XtWz   = (X * w[:, None]).T @ zp
    Bbar   = np.linalg.inv(B0inv + XtWX)
    bbar   = Bbar @ (B0inv @ b0 + XtWz)
    return np.random.multivariate_normal(bbar, Bbar)


# ─────────────────────────────────────────────────────────────────────────────
# Normal segment utilities  (Block 2)
# ─────────────────────────────────────────────────────────────────────────────

def _norm_posterior(r_seg, h2_seg, kappa, c0, C0):
    """Posterior N(c_bar, C_bar) for segment mean mu.

    Prior:       mu ~ N(c0, C0)
    Likelihood:  r_t | mu ~ N(kappa * mu, h2_t)
    """
    prec0   = 1.0 / C0
    prec_lik = kappa ** 2 * np.sum(1.0 / h2_seg)
    C_bar   = 1.0 / (prec0 + prec_lik)
    c_bar   = C_bar * (prec0 * c0 + kappa * np.sum(r_seg / h2_seg))
    return c_bar, C_bar


def _norm_log_mlik(r_seg, h2_seg, kappa, c0, C0):
    """Log marginal likelihood for mean segment (mu integrated out).

    Uses the Gaussian identity:
      log m = log p(r|mu_bar) + log p(mu_bar|prior) - log p(mu_bar|posterior)
    evaluated at the posterior mean mu_bar = c_bar.
    """
    n = len(r_seg)
    if n == 0:
        return 0.0
    c_bar, C_bar = _norm_posterior(r_seg, h2_seg, kappa, c0, C0)
    resid   = r_seg - kappa * c_bar
    ll_lik  = (-0.5 * n * np.log(2 * np.pi)
               - 0.5 * np.sum(np.log(h2_seg))
               - 0.5 * np.sum(resid ** 2 / h2_seg))
    ll_adj  = ((-0.5 * np.log(2 * np.pi * C0) - (c_bar - c0) ** 2 / (2 * C0))
               - (-0.5 * np.log(2 * np.pi * C_bar)))
    return ll_lik + ll_adj


# ─────────────────────────────────────────────────────────────────────────────
# Inverse-Gamma segment utilities  (Block 3)
# ─────────────────────────────────────────────────────────────────────────────

def _ig_posterior(e_seg, alpha0, beta0):
    """Posterior IG(alpha1, beta1) for segment variance h2.

    Prior:       h2 ~ IG(alpha0, beta0)  [shape, scale]
    Likelihood:  e_t | h2 ~ N(0, h2)
    """
    n = len(e_seg)
    return alpha0 + n / 2.0, beta0 + np.sum(e_seg ** 2) / 2.0


def _ig_log_mlik(e_seg, alpha0, beta0):
    """Log marginal likelihood for variance segment (h2 integrated out)."""
    n = len(e_seg)
    if n == 0:
        return 0.0
    alpha1, beta1 = _ig_posterior(e_seg, alpha0, beta0)
    return (gammaln(alpha1) - gammaln(alpha0)
            + alpha0 * np.log(beta0) - alpha1 * np.log(beta1)
            - 0.5 * n * np.log(2 * np.pi))


# ─────────────────────────────────────────────────────────────────────────────
# Segment-boundary finding (shared)
# ─────────────────────────────────────────────────────────────────────────────

def _prev_next(indicators, t, p, T):
    """Return (prev_break, next_break) given indicator array."""
    prev = p
    for s in range(t - 1, p - 1, -1):
        if indicators[s]:
            prev = s
            break
    nxt = T
    for s in range(t + 1, T):
        if indicators[s]:
            nxt = s
            break
    return prev, nxt


# ─────────────────────────────────────────────────────────────────────────────
# Block 2 — Mean shifts
# ─────────────────────────────────────────────────────────────────────────────

def _sample_mean_levels(y, phi, h2, pi_mu, c0, C0, p, delta_prev=None):
    """Single-site Gibbs for mean-shift indicators (delta) and levels (mu).

    Starts from delta_prev (previous iteration's state) so that each delta[t]
    is updated conditional on all OTHER indicators — i.e. correct single-site Gibbs.
    Resetting to zeros each call would give wrong conditionals for early t.

    Returns
    -------
    delta : (T,) int   -- 1 if a new mean segment starts at t
    mu    : (T,) float -- piecewise-constant mean level
    """
    T     = len(y)
    kappa = 1.0 - phi.sum()
    r     = _raw_ar_resid(y, phi)

    if delta_prev is not None:
        delta = delta_prev.copy()
    else:
        delta = np.zeros(T, dtype=int)
    delta[p] = 1  # first active time always opens a segment

    log_po = np.log(pi_mu + 1e-15) - np.log(1 - pi_mu + 1e-15)

    for t in range(p + 1, T):
        prev, nxt = _prev_next(delta, t, p, T)
        ll0 = _norm_log_mlik(r[prev:nxt], h2[prev:nxt], kappa, c0, C0)
        ll1 = (_norm_log_mlik(r[prev:t], h2[prev:t],  kappa, c0, C0)
               + _norm_log_mlik(r[t:nxt], h2[t:nxt],   kappa, c0, C0))
        p1       = 1.0 / (1.0 + np.exp(-np.clip(log_po + ll1 - ll0, -30, 30)))
        delta[t] = int(np.random.rand() < p1)

    # Sample mean for each segment
    mu = np.empty(T)
    t  = p
    while t < T:
        nxt = T
        for s in range(t + 1, T):
            if delta[s]:
                nxt = s
                break
        c_bar, C_bar = _norm_posterior(r[t:nxt], h2[t:nxt], kappa, c0, C0)
        mu[t:nxt]    = np.random.normal(c_bar, np.sqrt(C_bar))
        t = nxt
    mu[:p] = mu[p]
    return delta, mu


# ─────────────────────────────────────────────────────────────────────────────
# Block 3 — Variance shifts
# ─────────────────────────────────────────────────────────────────────────────

def _sample_var_levels(y, phi, mu, pi_h, alpha0, beta0, p, gamma_prev=None):
    """Single-site Gibbs for variance-shift indicators (gamma) and levels (h2).

    Starts from gamma_prev (previous iteration's state) — see _sample_mean_levels.

    Returns
    -------
    gamma : (T,) int   -- 1 if a new variance segment starts at t
    h2    : (T,) float -- piecewise-constant innovation variance
    """
    T = len(y)
    e = _full_resid(y, mu, phi)

    if gamma_prev is not None:
        gamma = gamma_prev.copy()
    else:
        gamma = np.zeros(T, dtype=int)
    gamma[p] = 1

    log_po = np.log(pi_h + 1e-15) - np.log(1 - pi_h + 1e-15)

    for t in range(p + 1, T):
        prev, nxt = _prev_next(gamma, t, p, T)
        ll0 = _ig_log_mlik(e[prev:nxt], alpha0, beta0)
        ll1 = (_ig_log_mlik(e[prev:t], alpha0, beta0)
               + _ig_log_mlik(e[t:nxt], alpha0, beta0))
        p1       = 1.0 / (1.0 + np.exp(-np.clip(log_po + ll1 - ll0, -30, 30)))
        gamma[t] = int(np.random.rand() < p1)

    h2 = np.empty(T)
    t  = p
    while t < T:
        nxt = T
        for s in range(t + 1, T):
            if gamma[s]:
                nxt = s
                break
        alpha1, beta1 = _ig_posterior(e[t:nxt], alpha0, beta0)
        h2[t:nxt]     = invgamma.rvs(alpha1, scale=beta1)
        t = nxt
    h2[:p] = h2[p]
    return gamma, h2


# ─────────────────────────────────────────────────────────────────────────────
# Block 4 — Shift probabilities
# ─────────────────────────────────────────────────────────────────────────────

def _sample_pi(indicators, p, a, b):
    """Beta posterior for shift probability (excludes initial forced break at p)."""
    active = indicators[p + 1:]
    n1 = int(active.sum())
    n0 = len(active) - n1
    return np.random.beta(a + n1, b + n0)


# ─────────────────────────────────────────────────────────────────────────────
# Main sampler
# ─────────────────────────────────────────────────────────────────────────────

def mts_gibbs(y, p=2, R=10000, burn=2000, seed=None,
              b0=None, B0=None,
              pi_mu_prior=(1, 1),
              mu_shift_prior=None,
              pi_h_prior=(1, 1),
              h2_shift_prior=None,
              nprint=1000):
    """Gibbs sampler for McCulloch & Tsay (1993) AR(p) with mean and variance shifts.

    Parameters
    ----------
    y               : (T,) array-like
    p               : AR order
    R               : total MCMC iterations
    burn            : burn-in (discarded)
    b0, B0          : prior mean / covariance for phi  (default: 0, 10*I)
    pi_mu_prior     : (a, b) Beta prior for mean-shift probability
    mu_shift_prior  : (c0, C0) Normal prior for new means; default (mean(y), 10*var(y))
    pi_h_prior      : (a, b) Beta prior for variance-shift probability
    h2_shift_prior  : (alpha0, beta0) IG prior for new variances; default (3, OLS_resid_var)
    nprint          : print frequency (0 = silent)

    Returns
    -------
    dict with posterior arrays (only kept iterations):
      phi   (R-burn, p)   AR coefficients
      mu    (R-burn, T)   mean levels
      h2    (R-burn, T)   innovation variances
      delta (R-burn, T)   mean-shift indicators
      gamma (R-burn, T)   variance-shift indicators
      pi_mu (R-burn,)     mean-shift probability
      pi_h  (R-burn,)     variance-shift probability
    """
    if seed is not None:
        np.random.seed(seed)

    y = np.asarray(y, dtype=float)
    T = len(y)

    if b0 is None:
        b0 = np.zeros(p)
    if B0 is None:
        B0 = 10.0 * np.eye(p)
    B0inv = np.linalg.inv(B0)

    # OLS AR(p) fit for initialization and prior defaults.
    # Using OLS residual variance avoids inflation from structural breaks.
    X_ols = np.array([y[t - p:t][::-1] for t in range(p, T)])
    phi_ols = np.linalg.lstsq(X_ols, y[p:], rcond=None)[0]
    e_ols   = y[p:] - X_ols @ phi_ols
    h2_ols  = max(float(np.var(e_ols)), 1e-6)
    y_var   = float(np.var(y)) if np.var(y) > 0 else 1.0

    if mu_shift_prior is None:
        mu_shift_prior = (float(np.mean(y)), 10.0 * y_var)
    if h2_shift_prior is None:
        # Use OLS residual variance, not y_var, to avoid prior inflation from breaks
        h2_shift_prior = (3.0, h2_ols)

    c0, C0       = float(mu_shift_prior[0]),  float(mu_shift_prior[1])
    alpha0, beta0 = float(h2_shift_prior[0]), float(h2_shift_prior[1])
    a_mu, b_mu   = float(pi_mu_prior[0]), float(pi_mu_prior[1])
    a_h,  b_h    = float(pi_h_prior[0]),  float(pi_h_prior[1])

    # Initialise from OLS — much better than zeros
    phi   = phi_ols.copy()
    mu    = np.full(T, np.mean(y))
    h2    = np.full(T, h2_ols)
    pi_mu = 0.05
    pi_h  = 0.05
    # Indicator state — carried across iterations for correct single-site Gibbs
    delta = None
    gamma = None

    kept = R - burn
    out  = dict(
        phi   = np.empty((kept, p)),
        mu    = np.empty((kept, T)),
        h2    = np.empty((kept, T)),
        delta = np.empty((kept, T), dtype=int),
        gamma = np.empty((kept, T), dtype=int),
        pi_mu = np.empty(kept),
        pi_h  = np.empty(kept),
    )

    t0 = time.time()
    for i in range(R):
        phi          = _sample_phi(y, mu, h2, p, b0, B0inv)
        delta, mu    = _sample_mean_levels(y, phi, h2, pi_mu, c0, C0, p, delta)
        gamma, h2    = _sample_var_levels(y, phi, mu, pi_h, alpha0, beta0, p, gamma)
        pi_mu        = _sample_pi(delta, p, a_mu, b_mu)
        pi_h         = _sample_pi(gamma, p, a_h, b_h)

        if i >= burn:
            k = i - burn
            out['phi'][k]   = phi
            out['mu'][k]    = mu
            out['h2'][k]    = h2
            out['delta'][k] = delta
            out['gamma'][k] = gamma
            out['pi_mu'][k] = pi_mu
            out['pi_h'][k]  = pi_h

        if nprint > 0 and (i + 1) % nprint == 0:
            print(f'  [mts]  iter {i+1:>6}/{R}  ({time.time()-t0:.1f}s)')

    elapsed = time.time() - t0
    print(f'[mts] Done in {elapsed:.1f}s  |  kept: {kept}')
    return out


# ─────────────────────────────────────────────────────────────────────────────
# Data simulation
# ─────────────────────────────────────────────────────────────────────────────

def simulate_ar_shifts(T=200, p=2, phi=None, mean_shifts=None, var_shifts=None, seed=None):
    """Generate AR(p) with piecewise-constant mean and variance.

    Parameters
    ----------
    T            : series length
    p            : AR order
    phi          : (p,) AR coefficients; default [0.5, 0.3]
    mean_shifts  : list of (t_start, mu_value); default [(0, 0.0)]
    var_shifts   : list of (t_start, sigma2_value); default [(0, 1.0)]

    Returns
    -------
    y        : (T,) simulated series
    true_mu  : (T,) true mean level
    true_h2  : (T,) true innovation variance
    """
    if seed is not None:
        np.random.seed(seed)
    if phi is None:
        phi = np.array([0.5, 0.3])
    if mean_shifts is None:
        mean_shifts = [(0, 0.0)]
    if var_shifts is None:
        var_shifts  = [(0, 1.0)]

    true_mu = np.zeros(T)
    true_h2 = np.ones(T)
    for t_start, val in sorted(mean_shifts):
        true_mu[t_start:] = val
    for t_start, val in sorted(var_shifts):
        true_h2[t_start:] = val

    y = np.zeros(T)
    y[:p] = true_mu[:p]  # initialise at true mean
    for t in range(p, T):
        z_lags = np.array([y[t - j - 1] - true_mu[t - j - 1] for j in range(p)])
        y[t]   = true_mu[t] + phi @ z_lags + np.random.normal(0, np.sqrt(true_h2[t]))
    return y, true_mu, true_h2


# ─────────────────────────────────────────────────────────────────────────────
# Post-processing
# ─────────────────────────────────────────────────────────────────────────────

def posterior_summary(samples, names=None):
    """DataFrame of mean / sd / q025 / q975 for 1-D or 2-D samples."""
    if samples.ndim == 1:
        samples = samples[:, None]
    df = pd.DataFrame({
        'mean': samples.mean(0),
        'sd'  : samples.std(0),
        'q025': np.percentile(samples, 2.5, axis=0),
        'q975': np.percentile(samples, 97.5, axis=0),
    }, index=names if names is not None else range(samples.shape[1]))
    return df.round(4)


def shift_probs(out, key='delta'):
    """Posterior mean of shift indicators over time (T,)."""
    return out[key].mean(axis=0)


def credible_band(samples, q=(2.5, 97.5)):
    """Row-wise credible band (lo, hi) for (R, T) posterior array."""
    lo = np.percentile(samples, q[0], axis=0)
    hi = np.percentile(samples, q[1], axis=0)
    return lo, hi