"""
blacklitterman.py -- the Black-Litterman model as Bayesian updating.

Backs the notebooks in  "The Black-Litterman Model".
Follows Meucci, "Risk and Asset Allocation" (Springer 2005), chapter 9-B, and the
original Black & Litterman (1992) / He & Litterman (1999) construction.

THE IDEA IN ONE PARAGRAPH
-------------------------
Plugging historical average returns into a mean-variance optimiser produces crazy,
unstable portfolios (Project 2). Black-Litterman fixes the *inputs* instead of the
optimiser. It starts from a sensible PRIOR for expected returns -- the returns that
would make today's MARKET-CAP portfolio optimal (the "equilibrium" or "implied"
returns) -- and then nudges that prior in the direction of the investor's own
VIEWS, by an amount that reflects how confident the investor is. The result is a
posterior expected-return vector that is close to equilibrium for assets you have
no opinion on, and tilted where you do. It is exactly a Bayesian (normal-normal)
update, which is why it belongs in this arc.

THE MATH
--------
Inputs:
  Sigma  (N,N)  covariance of returns (from data).
  w_mkt  (N,)   market-capitalisation weights.
  delta         market risk-aversion  = E[r_mkt - r_f] / Var[r_mkt].
  tau           small scalar: uncertainty in the equilibrium prior (e.g. 0.05).
  P (K,N),Q(K,) K views:  P @ mu = Q + eps,  eps ~ N(0, Omega).
  Omega (K,K)   view uncertainty (diagonal).

Equilibrium prior on expected returns:
  Pi = delta * Sigma @ w_mkt                 (reverse optimisation)
  mu ~ N(Pi, tau*Sigma)

Posterior (the Black-Litterman "master formula"):
  M     = ( (tau*Sigma)^-1 + P' Omega^-1 P )^-1        posterior covariance of the mean
  mu_BL = M ( (tau*Sigma)^-1 Pi + P' Omega^-1 Q )       posterior mean of expected returns

Allocation (optimal risky portfolio):
  w_BL = (1/delta) Sigma^-1 mu_BL
With NO views, mu_BL = Pi and w_BL = w_mkt exactly: Black-Litterman recovers the
market portfolio, and only tilts it where views are expressed. That self-consistency
is the property naive mean-variance lacks.
"""

import numpy as np


def risk_aversion(mkt_excess_mean, mkt_var):
    """delta = E[r_mkt - r_f] / Var[r_mkt].  The market's price of risk."""
    return mkt_excess_mean / mkt_var


def implied_returns(delta, Sigma, w_mkt):
    """Reverse optimisation: the expected returns that make w_mkt optimal."""
    return delta * Sigma @ np.asarray(w_mkt, float)


def omega_proportional(P, tau, Sigma):
    """Default view-uncertainty (He-Litterman): Omega = diag( P (tau Sigma) P' ).
    Each view is as uncertain as its variance under the prior -- so no free
    confidence parameter is needed."""
    P = np.atleast_2d(P)
    return np.diag(np.diag(P @ (tau * Sigma) @ P.T))


def black_litterman(Pi, Sigma, tau, P, Q, Omega):
    """Return the Black-Litterman posterior.

    dict keys:
      mu_bl      posterior mean of expected returns
      M          posterior covariance of the mean estimate ( (tauS)^-1 + P'O^-1P )^-1
      Sigma_post Sigma + M   (predictive covariance of returns, for allocation)
    """
    P = np.atleast_2d(P); Q = np.atleast_1d(Q).astype(float)
    tauS = tau * Sigma
    inv_tauS = np.linalg.inv(tauS)
    inv_O = np.linalg.inv(np.atleast_2d(Omega))
    M = np.linalg.inv(inv_tauS + P.T @ inv_O @ P)
    mu_bl = M @ (inv_tauS @ Pi + P.T @ inv_O @ Q)
    return dict(mu_bl=mu_bl, M=M, Sigma_post=Sigma + M)


def mv_weights(mu, Sigma, delta):
    """Optimal risky-portfolio weights  w = (1/delta) Sigma^-1 mu  (unnormalised;
    with mu = Pi this returns w_mkt exactly)."""
    return (1.0 / delta) * np.linalg.solve(Sigma, mu)


def mv_weights_budget(mu, Sigma, gamma):
    """Mean-variance optimal weights subject to full investment w'1 = 1."""
    Si = np.linalg.inv(Sigma); one = np.ones(len(mu))
    A = one @ Si @ one; B = one @ Si @ mu
    lam = (B - gamma) / A
    return Si @ (mu - lam * one) / gamma


def view_matrix(names, views):
    """Build (P, Q) from a list of human-readable views.

    Each view is a dict:
      absolute : {"assets": ["XLK"],          "weights": [1],      "q": 0.4}
      relative : {"assets": ["XLK","XLP"],    "weights": [1,-1],   "q": 0.5}
    `names` is the ordered list of asset tickers.
    """
    N = len(names); idx = {n: i for i, n in enumerate(names)}
    P = np.zeros((len(views), N)); Q = np.zeros(len(views))
    for k, v in enumerate(views):
        for a, w in zip(v["assets"], v["weights"]):
            P[k, idx[a]] = w
        Q[k] = v["q"]
    return P, Q
