"""
mvr_conjugate.py
----------------
Conjugate posterior sampler for the multivariate regression model

    Y = X B + E,   E_t ~ N(0, Sigma)

Prior:  B | Sigma ~ MN(Bbar, Sigma x A^{-1})
        Sigma ~ IW(nu, V)

rmultireg_np() mirrors R bayesm::rmultireg exactly (column-major B storage).
"""

import numpy as np
import pandas as pd
from scipy.stats import invwishart


def rmultireg_np(Y, X, Bbar, A, nu, V, rng):
    """
    One IID draw from the conjugate Normal-IW posterior.

    Parameters
    ----------
    Y    : (N, M) array of dependent variables
    X    : (N, K) design matrix  (same for all equations)
    Bbar : (K, M) prior mean of B
    A    : (K, K) prior precision on B  (small = diffuse)
    nu   : scalar IW degrees of freedom  (must be > M-1)
    V    : (M, M) IW scale matrix
    rng  : numpy Generator  (e.g. np.random.default_rng(42))

    Returns
    -------
    B     : (K, M) draw from posterior
    Sigma : (M, M) draw from posterior
    """
    N, M = Y.shape
    K    = X.shape[1]

    XtX = X.T @ X
    XtY = X.T @ Y

    # Posterior hyperparameters
    A_tilde = A + XtX
    A_tilde_inv = np.linalg.inv(A_tilde)
    B_tilde = A_tilde_inv @ (A @ Bbar + XtY)

    V_tilde = (V
               + Y.T @ Y
               + Bbar.T @ A @ Bbar
               - B_tilde.T @ A_tilde @ B_tilde)
    nu_tilde = nu + N

    # Draw Sigma ~ IW(nu_tilde, V_tilde)
    Sigma = invwishart.rvs(df=nu_tilde, scale=V_tilde, random_state=rng)

    # Draw B | Sigma ~ MN(B_tilde, Sigma x A_tilde^{-1})
    L_Ainv = np.linalg.cholesky(A_tilde_inv)
    L_Sig  = np.linalg.cholesky(Sigma)
    Z      = rng.standard_normal((K, M))
    B      = B_tilde + L_Ainv @ Z @ L_Sig.T

    return B, Sigma


def rmultireg_np_batch(Y, X, Bbar, A, nu, V, R, seed=42):
    """
    Draw R IID samples from the conjugate posterior.

    Returns
    -------
    B_draws     : (R, K*M)  column-major, matching R as.vector(B)
    Sigma_draws : (R, M*M)  column-major
    """
    N, M = Y.shape
    K    = X.shape[1]
    rng  = np.random.default_rng(seed)

    # Precompute posterior hyperparameters once
    XtX     = X.T @ X
    XtY     = X.T @ Y
    A_tilde = A + XtX
    A_tilde_inv = np.linalg.inv(A_tilde)
    B_tilde = A_tilde_inv @ (A @ Bbar + XtY)
    V_tilde = (V
               + Y.T @ Y
               + Bbar.T @ A @ Bbar
               - B_tilde.T @ A_tilde @ B_tilde)
    nu_tilde = nu + N

    L_Ainv = np.linalg.cholesky(A_tilde_inv)

    B_draws     = np.empty((R, K * M))
    Sigma_draws = np.empty((R, M * M))

    for r in range(R):
        Sigma = invwishart.rvs(df=nu_tilde, scale=V_tilde, random_state=rng)
        L_Sig = np.linalg.cholesky(Sigma)
        Z     = rng.standard_normal((K, M))
        B     = B_tilde + L_Ainv @ Z @ L_Sig.T
        B_draws[r]     = B.ravel(order='F')   # column-major = R's as.vector
        Sigma_draws[r] = Sigma.ravel(order='F')

    return B_draws, Sigma_draws


def posterior_summary(B_draws, K, M, eq_names=None, param_names=None,
                      B_true=None, ols_B=None):
    """
    Build a tidy DataFrame summarising B posterior draws.

    B_draws : (R, K*M) column-major array (output of rmultireg_np_batch)
    """
    if eq_names    is None: eq_names    = [f'y{m+1}' for m in range(M)]
    if param_names is None: param_names = [f'x{k}'   for k in range(K)]

    rows = []
    for m in range(M):
        for k in range(K):
            d = B_draws[:, m * K + k]
            row = dict(
                equation  = eq_names[m],
                regressor = param_names[k],
                bayes_mean= d.mean(),
                bayes_sd  = d.std(),
                CI_lo     = np.quantile(d, 0.025),
                CI_hi     = np.quantile(d, 0.975),
            )
            if B_true is not None:
                row['truth'] = B_true[k, m]
            if ols_B is not None:
                row['ols'] = ols_B[k, m]
            rows.append(row)

    df = pd.DataFrame(rows)
    if B_true is not None:
        df['covers'] = (df['CI_lo'] <= df['truth']) & (df['truth'] <= df['CI_hi'])
    return df