"""
copulas.py -- Copulas and tail dependence for asset returns.

Backs the notebooks in  "Copulas and Tail Dependence".
Follows Meucci, "Risk and Asset Allocation" (Springer 2005), chapter 2-B.

THE IDEA (Sklar's theorem)
--------------------------
Any joint distribution of returns splits into two independent pieces:
  * the MARGINALS  F_1,...,F_N  (each asset's own distribution), and
  * the COPULA  C  (how the assets move TOGETHER), a distribution on the unit
    cube whose own marginals are uniform.
Formally  F(x_1,...,x_N) = C(F_1(x_1),...,F_N(x_N)).  Modelling the two pieces
separately is powerful: we can fit fat-tailed marginals to each asset and choose
the dependence structure independently.

Why it matters for risk: the popular GAUSSIAN copula has *zero tail dependence* --
it says extreme joint moves (everything crashing at once) are vanishingly rare.
Real markets have strong lower-tail dependence, which the Student-t and Clayton
copulas capture. Mispricing that dependence is, famously, part of what blew up
credit models in 2008.

CONVENTIONS
-----------
X : (T, N) returns.   U : (T, N) 'pseudo-observations' in (0,1) = ranks/(T+1).
Bivariate helpers take two columns; elliptical copulas (Gaussian, t) take a
correlation matrix R.
"""

import numpy as np
from scipy.stats import norm, t as student_t, rankdata, kendalltau


# --------------------------------------------------------------------------- #
#  Pseudo-observations and rank-dependence measures                             #
# --------------------------------------------------------------------------- #

def pseudo_obs(X):
    """Map each column to uniform pseudo-observations via its empirical CDF:
    u = rank(x)/(T+1). This strips away the marginals, leaving only the
    dependence structure -- the copula's data."""
    X = np.asarray(X, float); T = X.shape[0]
    return np.column_stack([rankdata(X[:, j]) / (T + 1) for j in range(X.shape[1])])


def kendall_tau(u, v):
    return kendalltau(u, v)[0]


def spearman_rho(u, v):
    return np.corrcoef(rankdata(u), rankdata(v))[0, 1]


# --------------------------------------------------------------------------- #
#  Gaussian copula                                                              #
# --------------------------------------------------------------------------- #

def fit_gaussian_copula(U):
    """Correlation of the normal scores z = Phi^{-1}(u)."""
    Z = norm.ppf(np.clip(U, 1e-6, 1 - 1e-6))
    return np.corrcoef(Z.T)


def sim_gaussian_copula(R, n, rng):
    Z = rng.multivariate_normal(np.zeros(len(R)), R, size=n)
    return norm.cdf(Z)


def gaussian_copula_loglik(U, R):
    """Sum of log copula density.  c(u) = |R|^{-1/2} exp(-1/2 z'(R^{-1}-I)z)."""
    Z = norm.ppf(np.clip(U, 1e-9, 1 - 1e-9))
    Ri = np.linalg.inv(R); N = len(R)
    sign, logdet = np.linalg.slogdet(R)
    quad = np.einsum("ti,ij,tj->t", Z, Ri - np.eye(N), Z)
    return np.sum(-0.5 * logdet - 0.5 * quad)


# --------------------------------------------------------------------------- #
#  Student-t copula (symmetric tail dependence)                                 #
# --------------------------------------------------------------------------- #

def _t_copula_loglik_given(U, R, nu):
    Z = student_t.ppf(np.clip(U, 1e-9, 1 - 1e-9), nu)
    N = len(R); Ri = np.linalg.inv(R)
    sign, logdet = np.linalg.slogdet(R)
    from scipy.special import gammaln
    q = np.einsum("ti,ij,tj->t", Z, Ri, Z)
    # log multivariate-t density minus sum of log univariate-t densities
    log_num = (gammaln((nu + N) / 2) + (N - 1) * gammaln(nu / 2)
               - N * gammaln((nu + 1) / 2) - 0.5 * logdet)
    log_c = log_num - (nu + N) / 2 * np.log1p(q / nu) \
        + (nu + 1) / 2 * np.sum(np.log1p(Z ** 2 / nu), axis=1)
    return np.sum(log_c)


def fit_t_copula(U, nu_grid=None):
    """Fit correlation R (from Kendall's tau, robust for elliptical copulas) and
    degrees of freedom nu by profile maximum likelihood over a grid."""
    N = U.shape[1]
    R = np.eye(N)
    for i in range(N):
        for j in range(i + 1, N):
            tau = kendalltau(U[:, i], U[:, j])[0]
            R[i, j] = R[j, i] = np.sin(np.pi * tau / 2)      # tau -> linear corr
    if nu_grid is None:
        nu_grid = np.array([3, 4, 5, 6, 8, 10, 15, 20, 30, 50, 100])
    lls = [_t_copula_loglik_given(U, R, nu) for nu in nu_grid]
    nu_hat = nu_grid[int(np.argmax(lls))]
    return R, float(nu_hat)


def sim_t_copula(R, nu, n, rng):
    Z = rng.multivariate_normal(np.zeros(len(R)), R, size=n)
    W = rng.chisquare(nu, size=n) / nu
    Tvar = Z / np.sqrt(W)[:, None]
    return student_t.cdf(Tvar, nu)


# --------------------------------------------------------------------------- #
#  Archimedean copulas (bivariate): Clayton, Gumbel, Frank                       #
# --------------------------------------------------------------------------- #

def clayton_theta(tau):  return 2 * tau / (1 - tau)          # tau = theta/(theta+2)
def gumbel_theta(tau):   return 1 / (1 - tau)                # tau = 1 - 1/theta


def frank_theta(tau):
    """Solve tau = 1 - 4/theta (1 - D1(theta)) for theta (Debye function D1)."""
    from scipy.optimize import brentq
    from scipy.integrate import quad
    def D1(a):
        return quad(lambda t: t / (np.exp(t) - 1), 0, a)[0] / a
    if abs(tau) < 1e-4:
        return 1e-4
    f = lambda th: (1 - 4 / th * (1 - D1(th))) - tau
    lo, hi = (1e-4, 60) if tau > 0 else (-60, -1e-4)
    return brentq(f, lo, hi)


def sim_clayton(theta, n, rng):
    """Conditional-inversion sampling (closed form). Lower-tail dependent."""
    u = rng.uniform(size=n); w = rng.uniform(size=n)
    v = ((w ** (-theta / (1 + theta)) - 1) * u ** (-theta) + 1) ** (-1 / theta)
    return np.column_stack([u, v])


def sim_frank(theta, n, rng):
    u = rng.uniform(size=n); w = rng.uniform(size=n)
    v = -1 / theta * np.log1p(w * (1 - np.exp(-theta)) /
                              (w * (np.exp(-theta * u) - 1) - np.exp(-theta * u)))
    return np.column_stack([u, v])


def _rpstable(alpha, size, rng):
    """Positive alpha-stable subordinator (Laplace transform exp(-t^alpha)),
    alpha in (0,1], via Chambers-Mallows-Stuck."""
    U = rng.uniform(0, np.pi, size=size)
    W = rng.exponential(1.0, size=size)
    return (np.sin(alpha * U) / np.sin(U) ** (1 / alpha)) * \
           (np.sin((1 - alpha) * U) / W) ** ((1 - alpha) / alpha)


def sim_gumbel(theta, n, rng):
    """Marshall-Olkin frailty sampling. Upper-tail dependent."""
    alpha = 1 / theta
    M = _rpstable(alpha, n, rng)
    E = rng.exponential(1.0, size=(n, 2))
    return np.exp(-(E / M[:, None]) ** (1 / theta))


# --------------------------------------------------------------------------- #
#  Tail dependence coefficients                                                 #
# --------------------------------------------------------------------------- #

def tail_dep_t(rho, nu):
    """Lower = upper tail dependence of the bivariate Student-t copula."""
    x = -np.sqrt((nu + 1) * (1 - rho) / (1 + rho))
    return float(2 * student_t.cdf(x, nu + 1))


def tail_dep_clayton(theta):  return dict(lower=2 ** (-1 / theta), upper=0.0)
def tail_dep_gumbel(theta):   return dict(lower=0.0, upper=2 - 2 ** (1 / theta))
def tail_dep_gaussian():      return dict(lower=0.0, upper=0.0)


def empirical_tail_dep(u, v, q=0.05, lower=True):
    """Nonparametric estimate of tail dependence: P(V<q | U<q) (lower) or
    P(V>1-q | U>1-q) (upper), from pseudo-observations."""
    u = np.asarray(u); v = np.asarray(v)
    if lower:
        cond = u < q
        return float(np.mean(v[cond] < q)) if cond.sum() else np.nan
    cond = u > 1 - q
    return float(np.mean(v[cond] > 1 - q)) if cond.sum() else np.nan


# --------------------------------------------------------------------------- #
#  Model comparison                                                             #
# --------------------------------------------------------------------------- #

def aic(loglik, k):
    return 2 * k - 2 * loglik