## Copyleft 2021-2022, Alex Markham, see https://medil.causal.dev/license.html
# last updated 19-09-2022
# Tested with versions:
# python: 3.9.5
# numpy: 1.20.3
# scipy: 1.6.3
import numpy as np
from numpy import linalg as LA
from scipy.spatial.distance import pdist, squareform
from scipy.stats import chi2


def dep_contrib_kernel(X, alpha=None):
    num_samps, num_feats = X.shape
    thresh = np.eye(num_feats)
    if alpha is not None:
        thresh[thresh == 0] = (
            chi2(1).ppf(1 - alpha) / num_samps
        )  # critical value corresponding to alpha
        thresh[thresh == 1] = 0
    Z = np.zeros((num_feats, num_samps, num_samps))
    for j in range(num_feats):
        D = squareform(pdist(X[:, j].reshape(-1, 1), "cityblock"))
        # doubly center and standardized:
        Z[j] = ((D - D.mean(0) - D.mean(1).reshape(-1, 1)) / D.mean()) + 1

    F = Z.reshape(num_feats * num_samps, num_samps)
    left = np.tensordot(Z, thresh, axes=([0], [0]))
    left_right = np.tensordot(left, Z, axes=([2, 1], [0, 1]))
    gamma = (F.T @ F) ** 2 - 2 * (left_right) + LA.norm(thresh)  # helper kernel

    diag = np.diag(gamma)
    kappa = gamma / np.sqrt(np.outer(diag, diag))  # cosine similarity
    kappa[kappa > 1] = 1  # correct numerical errors
    return kappa