"""
Bayesian SUR (Seemingly Unrelated Regressions) -- Gibbs sampler.
Replicates bayesm::rsurGibbs (Rossi, Allenby & McCulloch 2005, Ch. 3).

Model
-----
  y_j = X_j @ beta_j + eps_j,    j = 1,...,M
  eps_t ~ N_M(0, Omega)           (cross-equation correlation per time period)

Stacked: y = X_block @ beta + eps,  X_block = diag(X_1,...,X_M)

Priors (independent)
--------------------
  beta  ~ N(betabar, A^{-1})
  Omega ~ IW(nu, V)

2-block Gibbs
-------------
  (i)  beta  | Omega, Y : N(Q^{-1} b, Q^{-1})
       Q_{ij} = A_{ij} + omega^{ij} X_i' X_j   (K x K block matrix)
       b_i    = (A betabar)_i + sum_j omega^{ij} X_i' y_j
  (ii) Omega | beta,  Y : IW(nu + N, V + E'E)
       E_tj = y_{jt} - x_{jt}' beta_j
"""

import numpy as np
from scipy.stats import invwishart
from scipy.linalg import cholesky, solve_triangular


def _chol_draw(Q, b, rng):
    """Draw x ~ N(Q^{-1} b, Q^{-1}) via Cholesky of the precision matrix."""
    L   = cholesky(Q, lower=True)
    v   = solve_triangular(L,   b, lower=True)
    mu  = solve_triangular(L.T, v, lower=False)
    return mu + solve_triangular(L.T, rng.standard_normal(len(b)), lower=False)


def rsur_gibbs(regdata, Prior=None, Mcmc=None, seed=42):
    """
    2-block Gibbs sampler for Bayesian SUR.

    Parameters
    ----------
    regdata : list of dicts  {'y': (N,), 'X': (N, k_j)}
              M elements, one per equation.  k_j may differ across equations.
    Prior   : dict with optional keys
                betabar : (K,)     prior mean,         default 0
                A       : (K, K)   prior precision,    default 0.01*I
                nu      : int      IW df for Omega,    default M+3
                V       : (M, M)   IW scale for Omega, default (M+2)*I
    Mcmc    : dict with optional keys
                R      : int   total draws,          default 10000
                burn   : int   burn-in,              default 0.2*R
                keep   : int   thinning,             default 1
                nprint : int   print frequency,      default R//10

    Returns
    -------
    dict with
      betadraw  : (R_kept, K)    posterior draws, equations stacked [eq1 | eq2 | ...]
      Omegadraw : (R_kept, M*M)  posterior draws of vec(Omega), column-major (matches R)
    """
    if Prior is None:
        Prior = {}
    if Mcmc is None:
        Mcmc = {'R': 10000}

    rng = np.random.default_rng(seed)

    M      = len(regdata)
    ks     = [d['X'].shape[1] for d in regdata]
    K      = sum(ks)
    N      = len(regdata[0]['y'])
    starts = np.concatenate([[0], np.cumsum(ks)]).astype(int)

    betabar = np.asarray(Prior.get('betabar', np.zeros(K)),         float)
    A       = np.asarray(Prior.get('A',       0.01 * np.eye(K)),    float)
    nu      = int(Prior.get('nu',  M + 3))
    V       = np.asarray(Prior.get('V',  (M + 2.0) * np.eye(M)),    float)

    R      = int(Mcmc['R'])
    burn   = int(Mcmc.get('burn',   0.2 * R))
    keep   = int(Mcmc.get('keep',   1))
    nprint = int(Mcmc.get('nprint', max(R // 10, 1)))
    R_kept = (R - burn) // keep

    # Precompute sufficient statistics once
    ys   = [d['y'] for d in regdata]
    Xs   = [d['X'] for d in regdata]
    S    = [[Xs[i].T @ Xs[j] for j in range(M)] for i in range(M)]   # S[i][j] = X_i'X_j
    XtY  = [[Xs[i].T @ ys[j] for j in range(M)] for i in range(M)]   # XtY[i][j] = X_i'y_j

    A_betabar = A @ betabar

    # Storage
    betadraw  = np.zeros((R_kept, K))
    Omegadraw = np.zeros((R_kept, M * M))

    # Initialise
    beta  = np.zeros(K)
    Omega = V / max(nu - M - 1, 1)

    i_kept = 0
    for r in range(R):
        Omega_inv = np.linalg.inv(Omega)

        # Block 1: beta | Omega, Y
        Q = A.copy()
        b = A_betabar.copy()
        for i in range(M):
            si = starts[i]; ei = starts[i + 1]
            for j in range(M):
                sj = starts[j]; ej = starts[j + 1]
                Q[si:ei, sj:ej] += Omega_inv[i, j] * S[i][j]
                b[si:ei]        += Omega_inv[i, j] * XtY[i][j]
        beta = _chol_draw(Q, b, rng)

        # Block 2: Omega | beta, Y
        E     = np.column_stack([ys[j] - Xs[j] @ beta[starts[j]:starts[j + 1]]
                                 for j in range(M)])              # N x M
        Omega = invwishart.rvs(df=nu + N, scale=V + E.T @ E, random_state=rng)

        if nprint and (r + 1) % nprint == 0:
            print(f'  Iter {r+1:6d}/{R}')

        if r >= burn and (r - burn) % keep == 0:
            betadraw[i_kept]  = beta
            Omegadraw[i_kept] = Omega.ravel(order='F')   # column-major, matches R vec()
            i_kept += 1

    return {'betadraw': betadraw, 'Omegadraw': Omegadraw}


def simulate_sur(M, k, N, true_beta, true_Omega, seed=0):
    """
    Simulate data from a SUR model with known parameters.

    Parameters
    ----------
    M          : number of equations
    k          : regressors per equation (int, same for all) or list of ints
    N          : number of observations
    true_beta  : (K,) true coefficient vector, equations stacked
    true_Omega : (M, M) true error covariance
    seed       : random seed

    Returns
    -------
    regdata : list of M dicts {'y': (N,), 'X': (N, k_j)}
    """
    rng = np.random.default_rng(seed)
    ks  = [k] * M if np.isscalar(k) else list(k)
    K   = sum(ks)
    starts = np.concatenate([[0], np.cumsum(ks)]).astype(int)

    L = np.linalg.cholesky(true_Omega)
    E = rng.standard_normal((N, M)) @ L.T          # N x M correlated errors

    regdata = []
    for j in range(M):
        X_j = np.column_stack([np.ones(N),
                                rng.standard_normal((N, ks[j] - 1))])
        y_j = X_j @ true_beta[starts[j]:starts[j + 1]] + E[:, j]
        regdata.append({'y': y_j, 'X': X_j})

    return regdata


def posterior_summary(draws, param_names=None, true_vals=None):
    """Posterior mean, SD, 95% CrI for each column of a (R, p) draws array."""
    import pandas as pd
    p = draws.shape[1]
    if param_names is None:
        param_names = [f'p{j}' for j in range(p)]
    rows = []
    for j, name in enumerate(param_names):
        d      = draws[:, j]
        lo, hi = np.percentile(d, [2.5, 97.5])
        row    = {'param': name, 'mean': d.mean(), 'sd': d.std(),
                  'q025': lo, 'q975': hi}
        if true_vals is not None:
            row['true']    = float(true_vals[j])
            row['covered'] = bool(lo <= true_vals[j] <= hi)
        rows.append(row)
    return pd.DataFrame(rows).set_index('param')