"""
robust.py -- Robust Bayesian mean-variance allocation.

Backs the notebooks in  "Robust Bayesian Allocation".
Follows Meucci, "Risk and Asset Allocation" (Springer 2005), chapter 9-C/D
(robust and robust-Bayesian allocation).

THE IDEA
--------
Projects 1-3 shrank the *inputs* to a mean-variance optimiser. But even a shrunk /
Bayesian estimate is only a point, surrounded by a cloud of uncertainty (the
posterior). A standard optimiser trusts that point completely and happily loads up
on whatever direction looks best -- precisely the directions where the estimate is
least reliable. ROBUST allocation refuses to be fooled: it optimises the WORST case
over an uncertainty region around the estimate.

For the expected-return vector, take the ellipsoid
    U = { mu : (mu - mu_hat)' Theta^-1 (mu - mu_hat) <= q^2 }
where mu_hat is the (Bayesian) estimate, Theta its estimation-error covariance
(for the NIW posterior, Theta = Sigma1 / T1), and q sets the size of the region.
The worst-case expected return of a portfolio w over U has a closed form:
    min_{mu in U} w' mu = w' mu_hat - q * sqrt(w' Theta w).
So the robust mean-variance problem is
    max_w   w' mu_hat - q * sqrt(w' Theta w) - (gamma/2) w' Sigma w   s.t.  w'1 = 1.
The extra term  -q*sqrt(w' Theta w)  is a penalty on positions whose expected return
is uncertain. It is a SECOND-ORDER CONE program (Meucci solves it with SeDuMi; we use
scipy's SLSQP). Two limits:
    q -> 0        : ordinary (Bayesian) mean-variance -- trust the estimate fully.
    q -> infinity : the penalty dominates the mean; the allocation collapses toward the
                    minimum-variance portfolio -- ignore the means entirely.
Between them sits a family of increasingly cautious, increasingly stable portfolios.
"""

import numpy as np
from scipy.optimize import minimize


# --------------------------------------------------------------------------- #
#  Moments and the NIW posterior (compact, self-contained copy)                 #
# --------------------------------------------------------------------------- #

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):
    X = np.asarray(X, float); T, N = X.shape
    mu_hat, S_hat = sample_moments(X)
    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 posterior_moments(post):
    """Posterior mean of mu, posterior mean of Sigma, and the estimation-error
    covariance Theta = Sigma1/T1 (the size of the location ellipsoid)."""
    mu1, T1, Sigma1, nu1 = post["mu1"], post["T1"], post["Sigma1"], post["nu1"]
    N = len(mu1)
    Sigma_mean = nu1 * Sigma1 / (nu1 - N - 1)
    Theta = Sigma1 / T1
    return mu1, Sigma_mean, Theta


# --------------------------------------------------------------------------- #
#  Plain and robust mean-variance optimisers                                    #
# --------------------------------------------------------------------------- #

def _budget_constraint():
    return {"type": "eq", "fun": lambda w: np.sum(w) - 1.0}


def mv_weights(mu, Sigma, gamma, long_only=False):
    """Ordinary mean-variance: max w'mu - gamma/2 w'Sigma w  s.t. w'1 = 1."""
    N = len(mu)
    obj = lambda w: -(w @ mu - 0.5 * gamma * (w @ Sigma @ w))
    bounds = [(0, None)] * N if long_only else None
    res = minimize(obj, np.ones(N) / N, method="SLSQP",
                   constraints=[_budget_constraint()], bounds=bounds,
                   options={"maxiter": 500, "ftol": 1e-10})
    return res.x


def robust_mv(mu, Sigma, Theta, q, gamma, long_only=False):
    """Robust mean-variance with an ellipsoidal uncertainty set on mu.

        max_w  w'mu - q*sqrt(w'Theta w) - gamma/2 w'Sigma w   s.t.  w'1 = 1.

    q is the radius of the location ellipsoid (aversion to estimation error).
    """
    N = len(mu)
    def neg_obj(w):
        pen = q * np.sqrt(max(w @ Theta @ w, 1e-18))
        return -(w @ mu - pen - 0.5 * gamma * (w @ Sigma @ w))
    bounds = [(0, None)] * N if long_only else None
    res = minimize(neg_obj, np.ones(N) / N, method="SLSQP",
                   constraints=[_budget_constraint()], bounds=bounds,
                   options={"maxiter": 1000, "ftol": 1e-11})
    return res.x


def min_var_weights(Sigma, long_only=False):
    N = len(Sigma)
    if not long_only:
        Si = np.linalg.inv(Sigma); one = np.ones(N)
        return Si @ one / (one @ Si @ one)
    obj = lambda w: w @ Sigma @ w
    res = minimize(obj, np.ones(N) / N, method="SLSQP",
                   constraints=[_budget_constraint()], bounds=[(0, None)] * N,
                   options={"maxiter": 500, "ftol": 1e-12})
    return res.x


# --------------------------------------------------------------------------- #
#  Robust frontier: sweep the estimation-risk aversion q                        #
# --------------------------------------------------------------------------- #

def robust_path(mu, Sigma, Theta, gamma, qs, long_only=False):
    """Return the weights for a grid of robustness radii q (0 = Bayesian MV,
    large q -> minimum variance)."""
    return np.array([robust_mv(mu, Sigma, Theta, q, gamma, long_only) for q in qs])


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

def make_true_market(N, rho=0.5, vol_lo=1.0, vol_hi=4.0, mu_scale=0.6, seed=0):
    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