"""
niw.py -- Normal-Inverse-Wishart Bayesian estimation, the predictive distribution,
a from-scratch MCMC sampler, the Markowitz frontier, and the "estimation-risk"
(deception-of-the-sample-frontier) experiment.

Backs the notebooks in
    Bayesian Estimation & Estimation Risk

Follows Meucci, "Risk and Asset Allocation" (Springer 2005), chapters 7 (Bayesian
estimation), 8 (evaluating estimation risk) and 9-A (Bayesian allocation).

THE MODEL
---------
Returns are i.i.d. normal:      x_t | mu, Sigma ~ N(mu, Sigma),  t = 1..T.
Conjugate Normal-Inverse-Wishart prior on the parameters:

    Sigma          ~ InverseWishart(nu0, nu0 * Sigma0)      (belief about covariance)
    mu | Sigma     ~ Normal(mu0, Sigma / T0)                (belief about the mean)

Interpretation of the prior strengths:
    T0  = "how many observations' worth" of confidence we place in the prior mean mu0.
    nu0 = the analogous confidence in the prior covariance Sigma0.

POSTERIOR (conjugate, closed form)
----------------------------------
With sample mean mu_hat and MLE covariance S_hat = (1/T) sum (x-mu_hat)(x-mu_hat)':

    T1  = T0 + T
    mu1 = (T0*mu0 + T*mu_hat) / T1                                  <- shrinkage of the mean!
    nu1 = nu0 + T
    Sigma1 = ( nu0*Sigma0 + T*S_hat
               + (T*T0/T1) * (mu0-mu_hat)(mu0-mu_hat)' ) / nu1      <- shrinkage of covariance

    Sigma      | data ~ InverseWishart(nu1, nu1*Sigma1)
    mu | Sigma , data ~ Normal(mu1, Sigma/T1)

POSTERIOR PREDICTIVE (the distribution of a *future* return, parameters integrated out)
--------------------------------------------------------------------------------------
    E[x_new | data]   = mu1
    Cov[x_new | data] = (1 + 1/T1) * nu1*Sigma1/(nu1 - N - 1)

The predictive covariance is LARGER than the plug-in covariance: it adds the
uncertainty about the parameters themselves. Optimising against the predictive
moments is what "Bayesian allocation" means, and it is automatically more
conservative than plugging in point estimates.
"""

import numpy as np
from scipy.stats import invwishart


# --------------------------------------------------------------------------- #
#  Moments and the conjugate NIW posterior                                      #
# --------------------------------------------------------------------------- #

def sample_moments(X, mle=True):
    X = np.asarray(X, float); T = X.shape[0]
    mu = X.mean(0); Xc = X - mu
    return mu, Xc.T @ Xc / (T if mle else T - 1)


def niw_posterior(X, mu0, T0, Sigma0, nu0):
    """Return the posterior hyper-parameters dict(mu1, T1, Sigma1, nu1)."""
    X = np.asarray(X, float); T, N = X.shape
    mu_hat, S_hat = sample_moments(X, mle=True)
    T1 = T0 + T
    mu1 = (T0 * mu0 + T * mu_hat) / T1
    nu1 = nu0 + T
    d = (mu0 - mu_hat).reshape(-1, 1)
    Sigma1 = (nu0 * Sigma0 + T * S_hat + (T * T0 / T1) * (d @ d.T)) / nu1
    return dict(mu1=mu1, T1=T1, Sigma1=Sigma1, nu1=nu1)


def niw_rand(post, size, rng):
    """Draw `size` samples of (mu, Sigma) from a NIW distribution given its
    hyper-parameters. Returns (mus [size,N], Sigmas [size,N,N])."""
    mu1, T1, Sigma1, nu1 = post["mu1"], post["T1"], post["Sigma1"], post["nu1"]
    N = len(mu1)
    Sigs = invwishart.rvs(df=nu1, scale=nu1 * Sigma1, size=size, random_state=rng)
    Sigs = Sigs.reshape(size, N, N)
    mus = np.empty((size, N))
    for i in range(size):
        mus[i] = rng.multivariate_normal(mu1, Sigs[i] / T1)
    return mus, Sigs


def predictive_moments(post):
    """Analytic posterior-predictive mean and covariance of a future return."""
    mu1, T1, Sigma1, nu1 = post["mu1"], post["T1"], post["Sigma1"], post["nu1"]
    N = len(mu1)
    E_Sigma = nu1 * Sigma1 / (nu1 - N - 1)
    pred_cov = (1.0 + 1.0 / T1) * E_Sigma
    return mu1.copy(), pred_cov


def post_mean_cov(post):
    """Posterior mean of mu and of Sigma (point estimates)."""
    mu1, Sigma1, nu1 = post["mu1"], post["Sigma1"], post["nu1"]
    N = len(mu1)
    return mu1.copy(), nu1 * Sigma1 / (nu1 - N - 1)


# --------------------------------------------------------------------------- #
#  From-scratch MCMC (random-walk Metropolis) to validate the analytic posterior#
# --------------------------------------------------------------------------- #

def _logpost(mu, L, X, mu0, T0, Sigma0, nu0):
    """Un-normalised log NIW posterior at (mu, Sigma=L L').  L lower-triangular
    with positive diagonal (Cholesky), so Sigma is guaranteed positive-definite."""
    N = len(mu)
    diagL = np.diag(L)
    if np.any(diagL <= 0):
        return -np.inf
    Sigma = L @ L.T
    # log|Sigma| and Sigma^{-1} from the Cholesky factor
    logdet = 2.0 * np.sum(np.log(diagL))
    Linv = np.linalg.inv(L)
    Sinv = Linv.T @ Linv
    T = X.shape[0]
    Xc = X - mu
    # Gaussian log-likelihood
    ll = -0.5 * T * logdet - 0.5 * np.sum((Xc @ Sinv) * Xc)
    # prior mu | Sigma ~ N(mu0, Sigma/T0)
    dm = (mu - mu0)
    lp_mu = -0.5 * logdet - 0.5 * T0 * (dm @ Sinv @ dm) + 0.5 * N * np.log(T0)
    # prior Sigma ~ IW(nu0, nu0*Sigma0):  logpdf up to const  = -(nu0+N+1)/2 log|Sigma| - 1/2 tr(nu0 Sigma0 Sigma^{-1})
    lp_S = -0.5 * (nu0 + N + 1) * logdet - 0.5 * np.trace(nu0 * Sigma0 @ Sinv)
    return ll + lp_mu + lp_S


def metropolis_niw(X, mu0, T0, Sigma0, nu0, n_draws=20000, burn=4000,
                   step_mu=None, step_L=None, rng=None, thin=1):
    """Random-walk Metropolis over (mu, chol(Sigma)). Returns dict with arrays
    'mu' [n,N] and 'Sigma' [n,N,N] plus the acceptance rate. A generic sampler
    whose output should match the analytic NIW posterior -- the point of the
    validation."""
    rng = np.random.default_rng() if rng is None else rng
    X = np.asarray(X, float); T, N = X.shape
    mu_hat, S_hat = sample_moments(X)
    mu = mu_hat.copy()
    L = np.linalg.cholesky(S_hat + 1e-6 * np.eye(N))
    tril = np.tril_indices(N)
    if step_mu is None:
        step_mu = 0.5 * np.sqrt(np.diag(S_hat) / T)
    if step_L is None:
        step_L = 0.15 * np.mean(np.sqrt(np.diag(S_hat)))
    lp = _logpost(mu, L, X, mu0, T0, Sigma0, nu0)
    keep_mu, keep_S = [], []
    acc = 0
    total = burn + n_draws
    for it in range(total):
        # propose mu
        mu_p = mu + step_mu * rng.standard_normal(N)
        lp_p = _logpost(mu_p, L, X, mu0, T0, Sigma0, nu0)
        if np.log(rng.random()) < lp_p - lp:
            mu, lp = mu_p, lp_p; acc += 1
        # propose L (perturb the free lower-triangular entries)
        L_p = L.copy()
        L_p[tril] += step_L * rng.standard_normal(len(tril[0]))
        lp_p = _logpost(mu, L_p, X, mu0, T0, Sigma0, nu0)
        if np.log(rng.random()) < lp_p - lp:
            L, lp = L_p, lp_p; acc += 1
        if it >= burn and (it - burn) % thin == 0:
            keep_mu.append(mu.copy())
            keep_S.append(L @ L.T)
    return dict(mu=np.array(keep_mu), Sigma=np.array(keep_S),
                accept=acc / (2 * total))


# --------------------------------------------------------------------------- #
#  Markowitz mean-variance frontier (budget-constrained, shorts allowed)        #
# --------------------------------------------------------------------------- #

def _abc(mu, Sigma):
    Si = np.linalg.inv(Sigma); one = np.ones(len(mu))
    A = one @ Si @ one; B = one @ Si @ mu; C = mu @ Si @ mu
    return Si, one, A, B, C, A * C - B * B


def frontier_weights(mu, Sigma, m):
    """Minimum-variance weights achieving expected return m, fully invested
    (w'1 = 1), short sales allowed (closed-form Markowitz solution)."""
    Si, one, A, B, C, D = _abc(mu, Sigma)
    return Si @ (one * (C - B * m) + mu * (A * m - B)) / D


def min_var_weights(Sigma):
    Si = np.linalg.inv(Sigma); one = np.ones(len(Sigma))
    return Si @ one / (one @ Si @ one)


def frontier(mu, Sigma, n=40, lo=None, hi=None):
    """Return (means, vols, W) sweeping target returns across the frontier."""
    Si, one, A, B, C, D = _abc(mu, Sigma)
    r_mv = B / A                                   # min-variance return
    spread = np.sqrt(max(C / A - r_mv ** 2, 1e-12))
    lo = r_mv - 0.2 * spread if lo is None else lo
    hi = r_mv + 3.0 * spread if hi is None else hi
    ms = np.linspace(lo, hi, n)
    W = np.array([frontier_weights(mu, Sigma, m) for m in ms])
    vols = np.sqrt(np.einsum("ij,jk,ik->i", W, Sigma, W))
    return ms, vols, W


def certainty_equivalent(w, mu, Sigma, gamma):
    """Quadratic (mean-variance) utility / certainty-equivalent return."""
    return w @ mu - 0.5 * gamma * (w @ Sigma @ w)


def optimal_mv(mu, Sigma, gamma):
    """Weights maximising w'mu - gamma/2 w'Sigma w subject to w'1 = 1."""
    Si, one, A, B, C, D = _abc(mu, Sigma)
    lam = (B - gamma) / A                           # Lagrange multiplier for budget
    return Si @ (mu - lam * one) / gamma


# --------------------------------------------------------------------------- #
#  Estimation risk: the "deception" of the sample efficient frontier            #
# --------------------------------------------------------------------------- #

def deception_experiment(mu_true, Sigma_true, T, prior, n_rep=400, n_pts=25, rng=None):
    """For each of n_rep simulated samples of size T from the TRUE market:
       * compute the SAMPLE efficient frontier and the BAYESIAN (predictive) one;
       * for a common grid of target returns, evaluate the TRUE risk/return of the
         weights each method chooses.
    Returns averaged curves:
       claimed_sample : (vol, ret) the manager THINKS the sample frontier delivers
       true_sample    : (vol, ret) it ACTUALLY delivers (evaluated at the truth)
       true_bayes     : (vol, ret) the Bayesian-predictive portfolios actually deliver
       true_frontier  : (vol, ret) the unattainable oracle frontier
    `prior` = dict(mu0, T0, Sigma0, nu0).
    """
    rng = np.random.default_rng() if rng is None else rng
    N = len(mu_true)
    ms, tvols, _ = frontier(mu_true, Sigma_true, n=n_pts)     # common target grid + oracle
    claimed = np.zeros((n_pts, 2)); tru_s = np.zeros((n_pts, 2)); tru_b = np.zeros((n_pts, 2))
    good = 0
    for _ in range(n_rep):
        X = rng.multivariate_normal(mu_true, Sigma_true, size=T)
        mu_h, S_h = sample_moments(X)
        post = niw_posterior(X, prior["mu0"], prior["T0"], prior["Sigma0"], prior["nu0"])
        mu_p, S_p = predictive_moments(post)
        try:
            for i, m in enumerate(ms):
                ws = frontier_weights(mu_h, S_h, m)          # sample-optimal weights
                wb = frontier_weights(mu_p, S_p, m)          # bayes-optimal weights
                claimed[i] += [np.sqrt(ws @ S_h @ ws), ws @ mu_h]
                tru_s[i]   += [np.sqrt(ws @ Sigma_true @ ws), ws @ mu_true]
                tru_b[i]   += [np.sqrt(wb @ Sigma_true @ wb), wb @ mu_true]
            good += 1
        except np.linalg.LinAlgError:
            continue
    claimed /= good; tru_s /= good; tru_b /= good
    return dict(claimed_sample=claimed, true_sample=tru_s, true_bayes=tru_b,
                true_frontier=np.column_stack([tvols, ms]), n_used=good)


def opportunity_cost_experiment(mu_true, Sigma_true, T, prior, gamma,
                                n_rep=1000, rng=None):
    """Certainty-equivalent OPPORTUNITY COST of estimation error.

    The unbeatable benchmark is the investor who knows the truth and holds
    w_true = argmax w'mu_true - gamma/2 w'Sigma_true w, earning CE_true.

    A real investor must estimate. We compare two estimators of the inputs:
      * SAMPLE   : plug in (mu_hat, S_hat).
      * BAYESIAN : plug in the posterior-predictive (mu_pred, S_pred).
    Each chooses weights by the same optimiser, then we score those weights at
    the TRUTH. The opportunity cost is  CE_true - CE_realised  (>= 0; lower is
    better). Returns (oc_sample, oc_bayes, ce_true) as arrays over replications.
    """
    rng = np.random.default_rng() if rng is None else rng
    N = len(mu_true)
    w_star = optimal_mv(mu_true, Sigma_true, gamma)
    ce_true = certainty_equivalent(w_star, mu_true, Sigma_true, gamma)
    oc_s, oc_b = [], []
    for _ in range(n_rep):
        X = rng.multivariate_normal(mu_true, Sigma_true, size=T)
        mu_h, S_h = sample_moments(X)
        post = niw_posterior(X, prior["mu0"], prior["T0"], prior["Sigma0"], prior["nu0"])
        mu_p, S_p = predictive_moments(post)
        w_s = optimal_mv(mu_h, S_h, gamma)
        w_b = optimal_mv(mu_p, S_p, gamma)
        oc_s.append(ce_true - certainty_equivalent(w_s, mu_true, Sigma_true, gamma))
        oc_b.append(ce_true - certainty_equivalent(w_b, mu_true, Sigma_true, gamma))
    return np.array(oc_s), np.array(oc_b), ce_true


# --------------------------------------------------------------------------- #
#  Synthetic market                                                             #
# --------------------------------------------------------------------------- #

def make_true_market(N, rho=0.5, vol_lo=1.0, vol_hi=4.0, mu_scale=0.6, seed=0):
    """Known market with equicorrelation rho, ramped vols, mean proportional to
    Sigma*1 (units chosen for weekly-percent returns)."""
    rng = np.random.default_rng(seed)
    vols = np.linspace(vol_lo, vol_hi, N)
    C = (1 - rho) * np.eye(N) + rho * np.ones((N, N))
    Sigma = np.outer(vols, vols) * C
    mu = mu_scale * Sigma @ np.ones(N) / N
    return mu, Sigma, rng