# -*- coding: utf-8 -*-
"""
This module includes implementations of popular spatially constrained 
hierarchical agglomerative clustering (SHAC) methods as described by Carvalho et al. (2009). 
These functions require imaging data as input, where the voxels are organized 
on a regular grid structure. The idea is to start with each voxel in a separate 
cluster and to merge in each step of the algorithms those two neighboring 
clusters, which have the closest distance according to the agglomeration method. 

A SHAC method that is especially designed to cluster variables is the SPARTACUS 
(SPAtial hieRarchical agglomeraTive vAriable ClUStering) method (Tietz et al, 2021).
The SPARTACUS distance between two clusters is the overall loss in explained 
total variance by all clusters first principal components that would be caused, 
if these two clusters are merged (also see Vigneau and Qannari (2003)). The 
SPARTACUS method is also implemented in this module.

Finally, this module includes an implementation of spatial hierarchical ensemble 
clustering (SEC) including the Hellinger method as introduced by Tietz et al. (2021).


References
----------
Carvalho AXY, Albuquerque PHM, de Almeida Junior GZ, Guimaraes RD (2009)
        Spatial hierarchical clustering. Revista Brasileira de Biometria 
        27(3):411-442
Vigneau E, Qannari EM (2003) Clustering of variables around latent components.
        Communications in Statistics-Simulation and Computation 32(4):1131-1150
Tietz et al. (2021) (Publication in progress.)

"""

import warnings
import numpy as np
import scipy
import scipy.cluster.hierarchy as hierarchy
import scipy.spatial.distance as distance
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from scipy.spatial.distance import cdist
import sklearn.metrics as metrics

_LINKAGE_METHODS = ('single', 'complete', 'average', 'centroid', 'median', 'ward')
_DISTANCE_METHODS = ('squared_correlation', 'correlation', 'euclidean', 'spartacus')
_ENSEMBLE_LINKAGE_METHODS = {'single': 0, 'complete': 1, 'average': 2, 'centroid': 3,
                    'hellinger': 4}
_ENSEMBLE_DISTANCE_METHODS = ('jaccard')

###############################################################################
##### SHAC methods including SPARTACUS
###############################################################################

def shac(X, matXYZ, method='ward', metric='spartacus', diag_neighbor = False, 
         standardize = True, print_progress = False):
    """    
    Perform spatial hierarchical agglomerative clustering (SHAC), including
    SPARTACUS (SPAtial hieRarchical agglomeraTive vAriable ClUStering). Note 
    that the objects to be clustered are the variables and not the subjects.
    
    Parameters
    ----------
    X : ndarray, shape(N, V)
        Input matrix including, e.g., grey matter volume values, where N is 
        the number of subjects and V is the number of variables, e.g. voxels.
    matXYZ : ndarray, shape(V, 3)
        Matrix of variable coordinates. Rows include variable locations sorted 
        as in the column of X, i.e. the coordinates stored in the j-th row 
        of matXYZ correspond to the variable stored in the j-th column of X.
    method : str, optional
        The linkage method. One of 'single', 'complete, 'average', 'centroid',
        'median' or 'ward'. Ignored, if metric is 'spartacus'. Default is 'ward'.
    metric : str, optional
        The linkage metric. One of 'euclidean', 'squared_correlation',
        'correlation' or 'spartacus'. If equal to 'spartacus', the SPARTACUS
        method is performed, ignoring the selection of method. 
        Default is 'spartacus'.
    diag_neighbor : bool, optional
        If False, a maximum of six voxels are considered as neighbors of 
        each voxel. If True, a maximum of 26 voxels belong to each voxel's 
        neighborhood. Default is False.
    standardize : bool, optional
        Should the columns of the input matrix be standardized? 
        Default is True.
    print_progress : bool, optional
        If True, a progress message is printed with every ten thousandth 
        iteration of initializing the sparse distance matrix and with every 
        ten thousandth iteration of the main algorithm. Only needed, if input 
        data consists of a very large number of variables (larger than 10000).
        Default is False.
        
    Returns
    -------
    Z : ndarray, shape (V - a, 4)
        Computed linkage matrix, encoding the hierarchical clustering, where 
        'a' is the number of contiguous regions of variables included in the 
        input data. Identically structured as the linkage matrix generated by 
        ?scipy.cluster.hierarchy.linkage.
    
    Examples
    --------
    Random input data with V = 24 spatial variables on a 4x3x2 grid and ten 
    subjects:
    
    >>> import numpy as np    
    >>> V = 24
    >>> X = np.random.normal(size = (10, V))
    >>> matXYZ = np.argwhere(np.zeros((4,3,2)) == 0)
    
    SPARTACUS based partition with four clusters:
        
    >>> Z = shac(X, matXYZ, metric = 'spartacus', standardize = False)
    >>> labels = get_cluster(Z, V, n_init_cluster = 4)
    >>> labels
    array([1, 4, 1, 4, 4, 4, 4, 4, 3, 4, 3, 3, 4, 2, 4, 3, 3, 3, 4, 2, 4, 3, 
           4, 3])
    
    Ward based partition with three clusters:
        
    >>> Z = shac(X, matXYZ, method = "ward", metric = 'euclidean', 
                 standardize = True)
    >>> labels = get_cluster(Z, V, n_init_cluster = 3)
    >>> labels
    array([2, 3, 2, 2, 3, 3, 1, 3, 2, 3, 2, 3, 1, 3, 3, 3, 2, 2, 1, 1, 3, 1,
           3, 1])
    
    References
    ----------
    Carvalho AXY, Albuquerque PHM, de Almeida Junior GZ, Guimaraes RD (2009)
        Spatial hierarchical clustering. Revista Brasileira de Biometria 
        27(3):411-442
    Tietz et al. (2021) (Publication in progress.)
        
    """ 
    if metric not in _DISTANCE_METHODS:
        raise ValueError("Invalid metric: {0}".format(metric))
    if metric != "spartacus":    
        if method not in _LINKAGE_METHODS:
            raise ValueError("Invalid method: {0}".format(method))
        if (method in ["median", "ward", "centroid"]) and (metric != 'euclidean'):
            raise ValueError("Method {0} is only defined, if euclidean pairwise "
                             "metric is used.".format(method))
    if X.ndim == 2:
        X = hierarchy._convert_to_double(np.asarray(X, order='c'))
        if X.shape[0] == X.shape[1] and np.allclose(np.diag(X), 0):
            if np.all(X >= 0) and np.allclose(X, X.T):
                warnings.warn('The symmetric non-negative hollow observation '
                              'matrix looks suspiciously like an uncondensed '
                              'distance matrix.')
        if X.shape[0] > X.shape[1]:
            warnings.warn('For the two dimensional intensity matrix the number '
                          'of colums (i.e. the number of voxels) is smaller than '
                          'the number of rows (i.e. the number of pictures). '
                          'You may want to consider the transpose of X instead.')
    else:
        raise ValueError("'X' must be 2 dimensional.")
    if not np.all(np.isfinite(X)):
        raise ValueError("The condensed distance matrix must contain only "
                         "finite values.")
    # Checks for matXYZ 
    if matXYZ.ndim != 2:
        raise ValueError("Neighbor matrix 'matXYZ' must be 2 dimensional.")
    if matXYZ.shape[1] != 3:
        raise ValueError("Neighbor matrix 'matXYZ' must have three spatial coordinates.")
    if not (np.round(matXYZ)==matXYZ).all():
        raise ValueError("All entries in matXYZ must be integers.")
    # Standardize the columns of the input matrix?
    if standardize:
        X = StandardScaler().fit_transform(X)
        X = X / np.sqrt(X.shape[0] / (X.shape[0] - 1))
    # Perform either SHAC or SPARTACUS
    if metric in ("squared_correlation", "correlation", "euclidean"):
        Z = shac_linkage(X, matXYZ, method, metric, diag_neighbor, print_progress)
    else:
        Z = shac_spartacus(X, matXYZ, diag_neighbor, print_progress)
    # Return linkage matrix
    return Z

def shac_spartacus(X, matXYZ, diag_neighbor = False, print_progress = False):
    """
    SPARTACUS method. Please use shac() function to perform SPARTACUS.
    
    Parameters
    ----------
    X : ndarray, shape(N, V)
        Input matrix including, e.g., grey matter volume values, where N is 
        the number of subjects and V is the number of variables, e.g. voxels.
    matXYZ : ndarray, shape(V, 3)
        Matrix of variable coordinates. Rows include variable locations sorted 
        as in the column of X, i.e. the coordinates stored in the j-th row 
        of matXYZ correspond to the variable stored in the j-th column of X.
    diag_neighbor : bool, optional
        If False, a maximum of six voxels are considered as neighbors of 
        each voxel. If True, a maximum of 26 voxels belong to each voxel's 
        neighborhood. Default is False.
    print_progress : bool, optional
        If True, a progress message is printed with every ten thousandth 
        iteration of initializing the sparse distance matrix and with every 
        ten thousandth iteration of the main algorithm. Default is False.
    
    Returns
    -------
    Z : ndarray, shape (V - a, 4)
        Computed linkage matrix, where 'a' is the number of contiguous regions
        of variables included in the input data.
        
    """    
    V = matXYZ.shape[0]
    X_cluster = []
    vec_lambda = np.var(X, axis = 0, ddof = 1)
    for i in range(V):
        X_cluster.append(np.array([i]))
    vec_dist = np.array([])
    vec_index1 = np.array([], dtype=np.int)
    vec_index2 = np.array([], dtype=np.int)
    # Initialization of sparse distance matrix calculating only distances 
    # between neighbor-variables
    for i in np.nditer(np.arange(V)):
        if print_progress:
            if (i % 10000) == 0:
                print("Initial spartacus: Iteration " + str(i))
        if diag_neighbor:
            id_neighbor = np.where(np.max(abs(matXYZ[i,] - matXYZ[(i+1):,]),1) <= 1)[0]+i+1
        else:
            id_neighbor = np.where(np.sum(abs(matXYZ[i,] - matXYZ[(i+1):,]), axis = 1) <= 1)[0]+i+1
        vec_index1 = np.append(vec_index1, np.array([i]*len(id_neighbor), dtype=np.int))
        vec_index2 = np.append(vec_index2, id_neighbor)
        if len(id_neighbor) > 0:
            for j in np.nditer(id_neighbor):
                vec_dist = np.append(vec_dist, get_dist_spartacus(X, i, j, 
                                                vec_lambda[i], vec_lambda[j]))
    # Main SPARTACUS algorithm
    Z_arr = np.empty((V - 1, 4))
    size = np.ones(V, dtype=np.int)  # Sizes of clusters.
    for k in range(V - 1):
        if print_progress:
            if (k % 10000) == 0:
                print("Main spartacus: Iteration " + str(k))
        if len(vec_dist) == 0:
            for l in range(k, V - 1):
                Z_arr = np.delete(Z_arr, k, 0)
            print(V - k, "contiguous regions in data set.")
            break
        index = np.argmin(vec_dist)
        current_min = vec_dist[index]
        x = vec_index1[index]
        y = vec_index2[index]
        # get the original numbers of points in clusters x and y
        nx = size[x]
        ny = size[y]
        # Record the new node.
        Z_arr[k, 0] = x
        Z_arr[k, 1] = y
        Z_arr[k, 2] = current_min
        Z_arr[k, 3] = nx + ny
        size[x] = 0  # Cluster x will be dropped.
        size[y] = nx + ny  # Cluster y will be replaced with the new cluster
        # Update X_cluster and then vec_lambda
        X_cluster[y] = np.concatenate((X_cluster[y], X_cluster[x]))
        pca = PCA(n_components=1)
        pca.fit(X[:, X_cluster[y]])
        vec_lambda[y] = pca.explained_variance_
        if k == V - 2:
            break
        # Find neighbors of x and y:
        xy_neighbor = np.concatenate((vec_index2[np.where(vec_index1 == x)[0]], 
                                      vec_index2[np.where(vec_index1 == y)[0]],
                                      vec_index1[np.where(vec_index2 == x)[0]],
                                      vec_index1[np.where(vec_index2 == y)[0]]))    
        # Remove x and y from this list:
        xy_neighbor = xy_neighbor[xy_neighbor != x] 
        xy_neighbor = xy_neighbor[xy_neighbor != y]   
        # Get unique neighbors
        xy_neighbor = np.unique(xy_neighbor)
        # Calculate distance of cluster y to its neighbors
        tmp_vec_index1 = np.array([], dtype = int)
        tmp_vec_index2 = np.array([], dtype = int)
        tmp_vec_dist = np.array([], dtype = int)
        if len(xy_neighbor) > 0:
            for i in np.nditer(xy_neighbor):
                dst = get_dist_spartacus(X, X_cluster[i], X_cluster[y], 
                                      vec_lambda[i], vec_lambda[y])
                if i < y:
                    tmp_vec_index1 = np.append(tmp_vec_index1, i)
                    tmp_vec_index2 = np.append(tmp_vec_index2, y)
                else:
                    tmp_vec_index1 = np.append(tmp_vec_index1, y)
                    tmp_vec_index2 = np.append(tmp_vec_index2, i)
                tmp_vec_dist = np.append(tmp_vec_dist, dst)
        # Remove all distances of x and y to its respective neighbors 
        id1 = np.unique(np.concatenate((np.where(vec_index1 == x)[0],
                                         np.where(vec_index1 == y)[0],
                                         np.where(vec_index2 == x)[0],
                                         np.where(vec_index2 == y)[0])))
        vec_index1 = np.delete(vec_index1, id1)
        vec_index2 = np.delete(vec_index2, id1)
        vec_dist = np.delete(vec_dist, id1)
        # Add new distances
        vec_index1 = np.append(vec_index1, tmp_vec_index1)
        vec_index2 = np.append(vec_index2, tmp_vec_index2)
        vec_dist = np.append(vec_dist, tmp_vec_dist)  
    # Return linkage matrix
    Z = get_label(Z_arr, V)
    return Z

def get_dist_spartacus(X, clus_i, clus_j, lambda_i, lambda_j):
    """
    Calculates the SPARTACUS distance between two (adjacent) clusters, i.e.
    calculates eigenvalue of cluster i plus eigenvalue of cluster j minus 
    the eigenvalue of the combined cluster.
    """
    if type(clus_i) != np.ndarray or type(clus_j) != np.ndarray:
        raise ValueError("clus_i and clus_j must be of type numpy array.")
    pca = PCA(n_components=1)
    pca.fit(X[:, np.append(clus_i, clus_j)])
    lambda_ij = pca.explained_variance_
    return lambda_i + lambda_j - lambda_ij

def shac_linkage(X, matXYZ, method, metric, diag_neighbor = False, 
                 print_progress = False):
    """
    Spatial hierarchical agglomerative clustering (SHAC) using popular linkage 
    methods. please use shac() method to perform SHAC.
    
    Parameters
    ----------
     X : ndarray, shape(N, V)
        Input matrix including, e.g., grey matter volume values, where N is 
        the number of subjects and V is the number of variables, e.g. voxels.
    matXYZ : ndarray, shape(V, 3)
        Matrix of variable coordinates. Rows include variable locations sorted 
        as in the column of X, i.e. the coordinates stored in the j-th row 
        of matXYZ correspond to the variable stored in the j-th column of X.
    method : str
        The linkage method. One of 'single', 'complete, 'average', 'centroid',
        'median' or 'ward'.
    metric : str
        The linkage metric. One of 'euclidean', 'squared_correlation' or
        'correlation'.    
    diag_neighbor : bool, optional
        If False, a maximum of six voxels are considered as neighbors of 
        each voxel. If True, a maximum of 26 voxels belong to each voxel's 
        neighborhood. Default is False.
    print_progress : bool, optional
        If True, a progress message is printed with every ten thousandth 
        iteration of initializing the sparse distance matrix and with every 
        ten thousandth iteration of the main algorithm. Default is False.
        
    Returns
    -------
    Z : ndarray, shape (V - a, 4)
        Computed linkage matrix, where 'a' is the number of contiguous regions
        of variables included in the input data.
        
    """
    if (method in ["median", "ward", "centroid"]) and (metric != 'euclidean'):
        raise ValueError("Method {0} is only defined, if euclidean pairwise "
                         "metric is used.".format(method))
    V = matXYZ.shape[0]
    if method in ["ward", "centroid"]:
        X_cluster = X.copy()
    elif method in ["single", "complete", "average"]:
        X_cluster = []
        for i in range(V):
            X_cluster.append([i])
    elif method == "median":
        X_cluster = []
        for i in range(V):
            X_cluster.append(X[:, i])
    else:
        ValueError("Method '{0}' not implemented so far.".format(method))
    # Initialization of sparse distance matrix calculating only distances 
    # between neighbor-variables
    vec_dist = np.array([])
    vec_index1 = np.array([], dtype=np.int)
    vec_index2 = np.array([], dtype=np.int)
    for i in range(V):
        if print_progress:
            if (i % 10000) == 0:
                print("Initial linkage: Iteration " + str(i))
        if diag_neighbor:
            id_neighbor = np.where(np.max(abs(matXYZ[i,] - matXYZ[(i+1):,]),1) <= 1)[0]+i+1
        else:
            id_neighbor = np.where(np.sum(abs(matXYZ[i,] - matXYZ[(i+1):,]), axis = 1) <= 1)[0]+i+1
        vec_index1 = np.append(vec_index1, np.array([i]*len(id_neighbor), dtype=np.int))
        vec_index2 = np.append(vec_index2, id_neighbor)
        vec_dist = np.append(vec_dist, get_dist(X, i, id_neighbor, metric))
    if method == "ward":
        vec_dist = vec_dist**2/2
    # Main SHAC algorithm
    Z_arr = np.empty((V - 1, 4))
    size = np.ones(V, dtype=np.int)  # Sizes of clusters.
    for k in range(V - 1):
        if print_progress:
            if (k % 10000) == 0:
                print("Main linkage: Iteration " + str(k))
        if len(vec_dist) == 0:
            for l in range(k, V - 1):
                Z_arr = np.delete(Z_arr, k, 0)
            print(V - k, "contiguous regions in data set.")
            break
        index = np.argmin(vec_dist)
        current_min = vec_dist[index]
        x = vec_index1[index]
        y = vec_index2[index]
        # get the original numbers of points in clusters x and y
        nx = size[x]
        ny = size[y]
        # Record the new node.
        Z_arr[k, 0] = x
        Z_arr[k, 1] = y
        Z_arr[k, 2] = current_min
        Z_arr[k, 3] = nx + ny
        size[x] = 0  # Cluster x will be dropped.
        size[y] = nx + ny  # Cluster y will be replaced with the new cluster
        # Break, if all variables are in the same cluster
        if k == V - 2:
            break
        # Find neighbors of x and y
        xy_neighbor = np.concatenate((vec_index2[np.where(vec_index1 == x)[0]], 
                                      vec_index2[np.where(vec_index1 == y)[0]],
                                      vec_index1[np.where(vec_index2 == x)[0]],
                                      vec_index1[np.where(vec_index2 == y)[0]]))    
        xy_neighbor = xy_neighbor[xy_neighbor != x]   # Remove x and y from list xy_neighbor
        xy_neighbor = xy_neighbor[xy_neighbor != y]  
        xy_neighbor = np.unique(xy_neighbor)   # Calculate unique neighbors   
        # Calculated distance of new cluster to its neighbors
        if method in ["ward", "centroid"]:
            X_cluster[:,y] = X_cluster[:,x] + X_cluster[:,y]
        elif method == "median":
            X_cluster[y] = 1/2 * (X_cluster[x] + X_cluster[y])
        tmp_vec_index1 = np.array([], dtype = int)
        tmp_vec_index2 = np.array([], dtype = int)
        tmp_vec_dist = np.array([], dtype = int)
        if len(xy_neighbor) > 0:
            for i in np.nditer(xy_neighbor):
                if method in ["ward", "centroid"]:
                    dst = get_cluster_dist_mean(X_cluster, size, i, y, method)
                elif method == "median":
                    dst = get_cluster_dist_median(X_cluster, i, y)
                else: 
                    # Find out, whether i is neighbor to x and/or y
                    if i < x:
                        ix = np.where((vec_index1==i) * (vec_index2 == x))[0]
                        iy = np.where((vec_index1==i) * (vec_index2 == y))[0]
                    elif i > y:
                        ix = np.where((vec_index1==x) * (vec_index2 == i))[0]
                        iy = np.where((vec_index1==y) * (vec_index2 == i))[0]
                    else:
                        ix = np.where((vec_index1==x) * (vec_index2 == i))[0]
                        iy = np.where((vec_index1==i) * (vec_index2 == y))[0]
                    # Calculate distance of i to x and i to y
                    if len(ix) == 1:
                        dist_ix = vec_dist[ix]
                    else:
                        dist_ix = get_cluster_dist_linkage(X, X_cluster, i, x, method, metric)
                    if len(iy) == 1:
                        dist_iy = vec_dist[iy]
                    else:
                        dist_iy = get_cluster_dist_linkage(X, X_cluster, i, y, method, metric)
                    # Use dist_ix and dist_iy to calculate distance of i to the 
                    # newly merged cluster. 
                    # single linkage
                    if method == "single":
                        dst = min(dist_ix, dist_iy)
                    # complete linkage
                    elif method == "complete":
                        dst = max(dist_ix, dist_iy)
                    # average linkage
                    else:
                        N_x = len(X_cluster[x])
                        N_y = len(X_cluster[y])
                        dst = N_x/(N_x+N_y)*dist_ix + N_y/(N_x+N_y)*dist_iy               
                if i < y:
                    tmp_vec_index1 = np.append(tmp_vec_index1, i)
                    tmp_vec_index2 = np.append(tmp_vec_index2, y)
                else:
                    tmp_vec_index1 = np.append(tmp_vec_index1, y)
                    tmp_vec_index2 = np.append(tmp_vec_index2, i)
                tmp_vec_dist = np.append(tmp_vec_dist, dst)
        # Remove all distances of x and y to their neighbors from vec_dist
        id1 = np.unique(np.concatenate((np.where(vec_index1 == x)[0],
                                         np.where(vec_index1 == y)[0],
                                         np.where(vec_index2 == x)[0],
                                         np.where(vec_index2 == y)[0])))
        vec_index1 = np.delete(vec_index1, id1)
        vec_index2 = np.delete(vec_index2, id1)
        vec_dist = np.delete(vec_dist, id1)
        # Add new distances to vec_dist 
        vec_index1 = np.append(vec_index1, tmp_vec_index1)
        vec_index2 = np.append(vec_index2, tmp_vec_index2)
        vec_dist = np.append(vec_dist, tmp_vec_dist)
        # Update X_cluster
        if method not in ["ward", "centroid", "median"]:
            X_cluster[y].extend(X_cluster[x])
    # Return linkage matrix
    Z = get_label(Z_arr, V)
    return Z

def get_dist(X, i, id_neighbor, metric):
    """
    Calculates distance between variable i and each variable in id_neighbor
    """
    dst = np.array([])
    if len(id_neighbor) > 0:
        for j in np.nditer(id_neighbor):
            if metric == "euclidean":
                dst = np.append(dst, distance.euclidean(X[:,i], X[:,j]))
            elif metric == "squared_correlation":
                dst = np.append(dst, 1 - scipy.stats.pearsonr(X[:,i], X[:,j])[0]**2)
            elif metric == "correlation":
                dst = np.append(dst, 1 - scipy.stats.pearsonr(X[:,i], X[:,j])[0])
            else:
                raise ValueError("Metric '{0}' currently not "
                                     "implemented".format(metric))
    return dst

def get_cluster_dist_mean(X_cluster, size, i, y, method):
    """
    Calculates ward or centroid distance between cluster i and y.
    """
    dst = distance.euclidean(X_cluster[:,i]/size[i], 
                             X_cluster[:,y]/size[y])
    if method == "ward":
        return dst**2/(1/size[i]+1/size[y])
    else:
        return dst

def get_cluster_dist_median(X_cluster, i, y):
    """
    Calculates median euclidean distance between cluster i and y.
    """
    centroid_i = X_cluster[i]
    centroid_y = X_cluster[y]
    return distance.euclidean(centroid_i, centroid_y)
    
def get_cluster_dist_linkage(X, X_cluster, i, z, method, metric, iter_max = 10000):
    """
    Calculates either single, complete or average linkage distance between 
    cluster i and cluster z. In order to avoid a memory error, the distance 
    calculation is partitioned, if at least one of the two clusters includes
    more than iter_max = 10000 variables. 
    """
    voxel_i = X_cluster[i]
    voxel_z = X_cluster[z]
    size_i = len(voxel_i)
    size_z = len(voxel_z)
    if (size_i <= iter_max) and (size_z <= iter_max):
        if metric == "euclidean":
            dst = cdist(X[:,voxel_i].T, X[:,voxel_z].T)
        elif metric == "squared_correlation":
            dst = 1 - (1 - cdist(X[:,voxel_i].T, X[:,voxel_z].T, "correlation"))**2
        elif metric == "correlation":
            dst = cdist(X[:,voxel_i].T, X[:,voxel_z].T, "correlation")
    else:
        n_groups_i = np.ceil(size_i/iter_max).astype(int)
        n_groups_z = np.ceil(size_z/iter_max).astype(int)
        dst = np.array([])
        for j in np.nditer(np.arange(n_groups_i)):
            tmp_vi = voxel_i[(j * iter_max):(np.min(((j + 1) * iter_max, size_i)))]
            for l in np.nditer(np.arange(n_groups_z)):
                tmp_vz = voxel_z[(l * iter_max):(np.min(((l + 1) * iter_max, size_z)))]
                if metric == "euclidean":
                    dst_mat = cdist(X[:,tmp_vi].T, X[:,tmp_vz].T)
                elif metric == "squared_correlation":
                    dst_mat = 1 - (1 - cdist(X[:,tmp_vi].T, X[:,tmp_vz].T, "correlation"))**2 
                elif metric == "correlation":
                    dst_mat = cdist(X[:,tmp_vi].T, X[:,tmp_vz].T, "correlation") 
                if method == "single":
                    dst = np.append(dst, np.min(dst_mat))
                elif method == "complete":
                    dst = np.append(dst, np.max(dst_mat))
                else:
                    dst = np.append(dst, np.sum(dst_mat))     
    if method == "single":
        return np.min(dst)
    elif method == "complete":
        return np.max(dst)
    else:
        return 1/size_i*1/size_z*np.sum(dst)
    
def get_label(Z, V):
    """
    Transforms Z_arr into the desired form. Also works if data set is 
    not contiguous.
    """
    xlim = Z.shape[0]
    dict_labels = dict(zip(range(V), range(V)))
    for i in range(xlim):
        Z1 = int(Z[i,0])
        Z2 = int(Z[i,1])
        Z[i,0] = min(dict_labels[Z1], dict_labels[Z2])
        Z[i,1] = max(dict_labels[Z1], dict_labels[Z2])
        dict_labels[Z2] = int(V + i) 
    return Z 

def get_cluster(Z, V, n_init_cluster):
    """
    Returns vector that indicates the cluster membership of each variable. In 
    total n_init_cluster cluster are generated. Note, that the function even 
    works, if the input data, based on which the linkage matrix is calculated,
    is structured into multiple spatially separated regions of variables.
    
    Parameters
    ----------
    Z : ndarray, shape (V - a, 4)
        Linkage matrix, where 'a' is the number of contiguous regions
        of variables included in the input data.
    V : int
        Number of variables of the input data set, based on which the 
        linkage matrix is genereated.
    n_init_cluster : int
        Number of clusters.

    Returns
    -------
    init_cluster : ndarray, shape (V,)
        Array of integers indicating the cluster membership of each variable, 
        where each variable is assigned to one of in total n_init_cluster 
        clusters.
    """
    
    if n_init_cluster < V - Z.shape[0]:
        n_init_cluster = V - Z.shape[0]
        warnings.warn("n_init_cluster is set to {0} since this is the " 
              "minimum number of contiguous regions in the "
              "data set.".format(V - Z.shape[0]))
    cl = {}
    for i in range(V - n_init_cluster):
        if Z[i, 0] < V and Z[i, 1] < V:
            cl[V + i] = (int(Z[i, 0]), int(Z[i, 1]))
        elif Z[i, 0] < V and not Z[i, 1] < V:
            cl[V + i] = (int(Z[i, 0]), *cl[Z[i, 1]])
            del cl[Z[i, 1]]
        elif not Z[i, 0] < V and Z[i, 1] < V:
            cl[V + i] = (*cl[Z[i, 0]], int(Z[i, 1]))
            del cl[Z[i, 0]]
        else:
            cl[V + i] = (*cl[Z[i, 0]], *cl[Z[i, 1]])
            del cl[Z[i, 0]]
            del cl[Z[i, 1]]
    counter = 1
    init_cluster = np.zeros(V, dtype = int)
    for key in cl:
        init_cluster[list(cl[key])] = counter 
        counter += 1
    del cl
    if counter <= n_init_cluster:
        init_cluster[init_cluster == 0] = range(counter, n_init_cluster + 1)
    return init_cluster

###############################################################################
##### Spatial ensemble clustering methods
###############################################################################

def spatial_ensemble_clustering(X, matXYZ, method='hellinger', metric='jaccard',
                     diag_neighbor = False, print_progress = False):
    """    
    Perform spatial hierarchical ensemble clustering. Also includes the 
    Hellinger method as introduced by Tietz et al. (2021).
    
    Parameters
    ----------
    X : ndarray, shape(K, V)
        Input cluster ensemble matrix, where K is the number of different base 
        partitions and V is the number of variables, e.g., voxels. All entries 
        must be positive integers.
    matXYZ : ndarray, shape(V, 3)
        Matrix of variable coordinates. Rows include variable locations sorted 
        as in the column of X, i.e. the coordinates stored in the j-th row 
        of matXYZ correspond to the variable stored in the j-th column of X.
    method : str, optional
        The linkage method. One of 'single', 'complete', 'average', 
        'centroid' or 'hellinger'. Default is 'hellinger'.
    metric : str, optional
        The linkage metric. Only 'jaccard' distance is implemented so far. See 
        ?scipy.spatial.distance.cdist for further details on the 'jaccard' 
        distance. Ignored, if method is equal to 'hellinger'. 
        Default is 'jaccard'.
    diag_neighbor : bool, optional
        If False, a maximum of six voxels are considered as neighbors for 
        each voxel. If True, a maximum of 26 voxels belong to each voxels 
        neighborhood. Default is False.
    print_progress : bool, optional
        If True, a progress message is printed with every ten thousandth 
        iteration of initializing the sparse ensemble matrix and with every 
        ten thousandth iteration of the main algorithm. Only needed, if input 
        data consists of a very large number of variables (larger than 10000). 
        Default is False.
        
    Returns
    -------
    Z : ndarray, shape (V - a, 4)
        Computed linkage matrix, encoding the hierarchical clustering, where 
        'a' is the number of contiguous regions of variables included in the 
        input data. Identically structured as the linkage matrix generated by 
        ?scipy.cluster.hierarchy.linkage.
    
    Examples
    --------
    Input cluster ensemble with V = 8 spatial variables and on a 2x2x2 grid and 
    six subjects:
        
    >>> import numpy as np    
    >>> X = np.array([[1,1,2,2,3,3,4,4],
                      [1,1,2,2,3,3,4,4],
                      [1,1,2,2,3,3,4,4],
                      [1,1,2,2,5,5,6,6],
                      [1,1,1,2,3,3,3,4],
                      [1,1,1,2,3,3,3,4]])
    >>> matXYZ = np.argwhere(np.zeros((2,2,2)) == 0)
    
    Hellinger based partition with four clusters:
        
    >>> Z = spatial_ensemble_clustering(X, matXYZ, method='hellinger')
    >>> labels = get_cluster(Z, V = 8, n_init_cluster = 4)
    >>> labels
    array([1, 1, 3, 3, 2, 2, 4, 4])
    
    Average linkage based partition with two clusters:
        
    >>> Z = spatial_ensemble_clustering(X, matXYZ, method='average')
    >>> labels = get_cluster(Z, V = 8, n_init_cluster = 2)
    >>> labels
    array([1, 1, 1, 1, 2, 2, 2, 2])
    
    References
    ----------
    Tietz et al. (2021) (Publication in progress)
    
    """   
    # Checks for 'method' and 'metric'     
    if method not in _ENSEMBLE_LINKAGE_METHODS:
        raise ValueError("Invalid method: {0}".format(method))
    if metric not in _ENSEMBLE_DISTANCE_METHODS:
        raise ValueError("Invalid metric: {0}".format(metric))
    # Checks for ensemble matrix
    if X.ndim != 2:
        raise ValueError("X must be 2 dimensional.")
    if X.shape[0] > X.shape[1]:
        warnings.warn('For the two dimensional ensemble matrix the number '
                      'of colums (i.e. the number of variables) is smaller than '
                      'the number of rows (i.e. the number of base partitions). '
                      'You may want to consider the transpose of X instead.')
    if not np.all(np.isfinite(X)):
        raise ValueError("The ensemble matrix must contain only "
                         "finite values.")
    # Checks for matXYZ 
    if matXYZ.ndim != 2:
        raise ValueError("Neighbor matrix matXYZ must be 2 dimensional.")
    if matXYZ.shape[1] != 3:
        raise ValueError("Neighbor matrix matXYZ must have three spatial "
                         "coordinates.")
    if not (np.round(matXYZ)==matXYZ).all():
        raise ValueError("All entries in matXYZ must be integers.")
    # Calculate dendrogram
    Z = ensemble_linkage(X, matXYZ, method, metric, diag_neighbor, print_progress)
    # Return dendrogram
    return Z



def ensemble_linkage(X, matXYZ, method, metric, diag_neighbor = False, 
                     print_progress = False):
    """
    Performs spatial hierarchical clustering. Please use the 
    spatial_ensemble_clustering() function instead of this function.
    
    Parameters
    ----------
    X : ndarray, shape(B, V)
        Input cluster ensemble matrix, where K is the number of different base 
        partitions and V is the number of variables, e.g., voxels. All entries 
        must be positive integers.
    matXYZ : ndarray, shape(V, 3)
        Matrix of voxel coordinates. Rows include voxel locations sorted 
        as in the column of X.
    method : int
        The linkage method. One of 'single', 'complete', 'average', 
        'centroid' or 'hellinger'.
    metric : str
        The linkage metric. Only 'jaccard' distance is implemented so far. See 
        ?scipy.spatial.distance.cdist for further details on the 'jaccard' 
        distance. Ignored, if method is equal to 'hellinger'.   
    diag_neighbor : bool, optional
        If False, a maximum of six voxels are considered as neighbors for 
        each voxel. If True, a maximum of 26 voxels belong to each voxels 
        neighborhood.
    print_progress : bool, optional
        If True, a progress message is printed with every ten thousandth 
        iteration of initializing the sparse ensemble matrix and with every 
        ten thousandth iteration of the main algorithm. Default is False.
        
    Returns
    -------
    Z : ndarray, shape (V - 1, 4)
        Computed linkage matrix.
        
    """
    # Make entries of ensemble matrix to be 1,2,...,K.
    X = check_ensemble_matrix(X)
    # Number of voxels
    V = matXYZ.shape[0]
    # Initialize clusters, where in the beginning each voxel is its own cluster
    X_cluster = []
    for i in range(V):
        X_cluster.append([i])
    # Initialization of sparse ensemble matrix calculating only distances 
    # between neighbor-variables
    vec_dist = np.array([])
    vec_index1 = np.array([], dtype=np.int)
    vec_index2 = np.array([], dtype=np.int)
    for i in range(V):
        if print_progress:
            if (i % 10000) == 0:
                print("Initial ensemble linkage: Iteration " + str(i))
        if diag_neighbor:
            id_neighbor = np.where(np.max(abs(matXYZ[i,] - matXYZ[(i+1):,]),1) <= 1)[0]+i+1
        else:
            id_neighbor = np.where(np.sum(abs(matXYZ[i,] - matXYZ[(i+1):,]), axis = 1) <= 1)[0]+i+1
        vec_index1 = np.append(vec_index1, np.array([i]*len(id_neighbor), dtype=np.int))
        vec_index2 = np.append(vec_index2, id_neighbor)
        vec_dist = np.append(vec_dist, get_dist_ensemble(X, i, id_neighbor, metric))
    # Main ensemble clustering
    Z_arr = np.empty((V - 1, 4))
    size = np.ones(V, dtype=np.int)  # Sizes of clusters.
    for k in range(V - 1):
        if print_progress:
            if (k % 10000) == 0:
                print("Main ensemble linkage: Iteration " + str(k))
        if len(vec_dist) == 0:
            for l in range(k, V - 1):
                Z_arr = np.delete(Z_arr, k, 0)
            print(V - k, "contiguous regions in data set.")
            break
        index = np.argmin(vec_dist)
        current_min = vec_dist[index]
        x = vec_index1[index]
        y = vec_index2[index]
        # get the original numbers of points in clusters x and y
        nx = size[x]
        ny = size[y]
        # Record the new node.
        Z_arr[k, 0] = x
        Z_arr[k, 1] = y
        Z_arr[k, 2] = current_min
        Z_arr[k, 3] = nx + ny
        size[x] = 0  # Cluster x will be dropped.
        size[y] = nx + ny  # Cluster y will be replaced with the new cluster
        #  Break, if all variables are in the same cluster
        if k == V - 2:
            break
        # Find neighbors of x and y:
        xy_neighbor = np.concatenate((vec_index2[np.where(vec_index1 == x)[0]], 
                                      vec_index2[np.where(vec_index1 == y)[0]],
                                      vec_index1[np.where(vec_index2 == x)[0]],
                                      vec_index1[np.where(vec_index2 == y)[0]]))    
        # Remove x and y from that neighbor list:
        xy_neighbor = xy_neighbor[xy_neighbor != x] 
        xy_neighbor = xy_neighbor[xy_neighbor != y]       
        # Determine unique neighbors and consider if they are neighbors to both 
        # x and y or just to one of them.
        xy_neighbor = np.unique(xy_neighbor)        
        # For centroid or hellinger method update X_cluster already
        if method in ["centroid", "hellinger"]:
            X_cluster[y].extend(X_cluster[x])
        # Determine distance of new cluster to its neighbors
        tmp_vec_index1 = np.array([], dtype = int)
        tmp_vec_index2 = np.array([], dtype = int)
        tmp_vec_dist = np.array([], dtype = int)
        if len(xy_neighbor) > 0:
            for i in np.nditer(xy_neighbor):
                if method in ["centroid", "hellinger"]:
                    K = np.max(X)
                    dst = get_new_dist_ensemble_centroid(X[:, X_cluster[i]], X[:, X_cluster[y]], method, K)
                else:
                    # Find out, whether i is neighbor to x and/or y
                    if i < x:
                        ix = np.where((vec_index1==i) * (vec_index2 == x))[0]
                        iy = np.where((vec_index1==i) * (vec_index2 == y))[0]
                    elif i > y:
                        ix = np.where((vec_index1==x) * (vec_index2 == i))[0]
                        iy = np.where((vec_index1==y) * (vec_index2 == i))[0]
                    else:
                        ix = np.where((vec_index1==x) * (vec_index2 == i))[0]
                        iy = np.where((vec_index1==i) * (vec_index2 == y))[0]
                    # Calculate distance of i to x and i to y
                    if len(ix) == 1:
                        dist_ix = vec_dist[ix]
                    else:
                        dist_ix = get_new_dist_ensemble(X, X_cluster, i, x, method, metric)
                    if len(iy) == 1:
                        dist_iy = vec_dist[iy]
                    else:
                        dist_iy = get_new_dist_ensemble(X, X_cluster, i, y, method, metric)
                    # Use dist_ix and dist_iy to calculate distance of i to the 
                    # newly merged cluster. 
                    # single linkage
                    if method == "single":
                        dst = min(dist_ix, dist_iy)
                    # complete linkage
                    elif method == "complete":
                        dst = max(dist_ix, dist_iy)
                    # average linkage
                    else:
                        N_x = len(X_cluster[x])
                        N_y = len(X_cluster[y])
                        dst = N_x/(N_x+N_y)*dist_ix + N_y/(N_x+N_y)*dist_iy               
                if i < y:
                    tmp_vec_index1 = np.append(tmp_vec_index1, i)
                    tmp_vec_index2 = np.append(tmp_vec_index2, y)
                else:
                    tmp_vec_index1 = np.append(tmp_vec_index1, y)
                    tmp_vec_index2 = np.append(tmp_vec_index2, i)
                tmp_vec_dist = np.append(tmp_vec_dist, dst)
        # Remove all distances of x and y to their neighbors from vec_dist
        id1 = np.unique(np.concatenate((np.where(vec_index1 == x)[0],
                                         np.where(vec_index1 == y)[0],
                                         np.where(vec_index2 == x)[0],
                                         np.where(vec_index2 == y)[0])))
        vec_index1 = np.delete(vec_index1, id1)
        vec_index2 = np.delete(vec_index2, id1)
        vec_dist = np.delete(vec_dist, id1)
        # Add new distances to vec_dist 
        vec_index1 = np.append(vec_index1, tmp_vec_index1)
        vec_index2 = np.append(vec_index2, tmp_vec_index2)
        vec_dist = np.append(vec_dist, tmp_vec_dist)
        # Update X_cluster
        if method not in ["centroid", "hellinger"]:
            X_cluster[y].extend(X_cluster[x])
    # Return linkage matrix    
    Z = get_label(Z_arr, V)
    return Z

def check_ensemble_matrix(X):
    """
    Makes entries of each base partition in ensemble matrix to be 1,2,...,K_b, 
    where b = 1,...,B.
    """
    
    for i in range(X.shape[0]):
        vec_uni_entries = np.unique(X[i,:])
        max_uni_entry = np.max(vec_uni_entries)
        if np.array_equal(vec_uni_entries, np.arange(1,max_uni_entry+1)):
            continue
        elif np.array_equal(vec_uni_entries, np.arange(max_uni_entry)):
            X[i,:] = X[i,:] + 1
        else:
            tmp_row_X = X[i,:].copy()
            counter = 1
            for uni_entry in vec_uni_entries:
                X[i, tmp_row_X==uni_entry] = counter
                counter += 1
    return X


def get_dist_ensemble(X, i, id_neighbor, metric):
    """
    Calculates ensemble distance between voxel i and all voxels in id_neighbor
    """
    dst = np.array([])
    if len(id_neighbor) > 0:
        for j in np.nditer(id_neighbor):
            if metric == "jaccard":
                dst = np.append(dst, 1-metrics.accuracy_score(X[:,i], X[:,j]))
            else:
                raise ValueError("Metric '{0}' currently not "
                                     "implemented".format(metric))
    return dst

def get_new_dist_ensemble_centroid(X_i, X_y, method, K):
    """
    Calculates centroids based or Hellinger based distance between 
    cluster i and cluster y.
    """
    if method == "centroid":
        axis = 1
        # Centroid of cluster i
        u, indices = np.unique(X_i, return_inverse=True)
        centroid_i = u[np.argmax(np.apply_along_axis(np.bincount, axis, indices.reshape(X_i.shape),
                                        None, np.max(indices) + 1), axis=axis)]
        # Centroid of cluster y
        u, indices = np.unique(X_y, return_inverse=True)
        centroid_y = u[np.argmax(np.apply_along_axis(np.bincount, axis, indices.reshape(X_y.shape),
                                        None, np.max(indices) + 1), axis=axis)]
        # Return Jaccard distance between centroids
        return 1-metrics.accuracy_score(centroid_i, centroid_y)
    else:
        # Determine probability matrices with dimensions HxK
        D_i = np.apply_along_axis(get_discrete_distribution, 1, X_i, K = K)
        D_y = np.apply_along_axis(get_discrete_distribution, 1, X_y, K = K)
        # Return mean Hellinger distance
        return np.mean(np.sqrt(np.sum((np.sqrt(D_i)-np.sqrt(D_y))**2, 1))/np.sqrt(2))
        
    
def get_discrete_distribution(X_row, K):
    """
    Determine the sample distribution of the base partition stored in X_row, 
    where X_row is a vector with positive integer values that are smaller than 
    or equal to K. 
    """
    # Initiate distribution vector, which includes sample probabilities
    vec_p = np.zeros(K)
    unique, counts = np.unique(X_row, return_counts=True)
    vec_p[unique-1] = counts/X_row.shape[0]
    return vec_p
    

def get_new_dist_ensemble(X, X_cluster, i, z, method, metric, iter_max = 4):
    """
    Calculates linkage based ensemble distance between cluster i and cluster z. 
    """
    voxel_i = X_cluster[i]
    voxel_z = X_cluster[z]
    
    size_i = len(voxel_i)
    size_z = len(voxel_z)
    
    if (size_i <= iter_max) and (size_z <= iter_max):
        if metric == "jaccard":
            dst = cdist(X[:,voxel_i].T, X[:,voxel_z].T, "jaccard")
        else:
            raise ValueError("Metric '{0}' currently not "
                                 "implemented".format(metric))
    else:
        n_groups_i = np.ceil(size_i/iter_max).astype(int)
        n_groups_z = np.ceil(size_z/iter_max).astype(int)
        dst = np.array([])
        for j in np.nditer(np.arange(n_groups_i)):
            tmp_vi = voxel_i[(j * iter_max):(np.min(((j + 1) * iter_max, size_i)))]
            for l in np.nditer(np.arange(n_groups_z)):
                tmp_vz = voxel_z[(l * iter_max):(np.min(((l + 1) * iter_max, size_z)))]
                if metric == "jaccard":
                    dst_mat = cdist(X[:,tmp_vi].T, X[:,tmp_vz].T, "jaccard")
                else:
                    raise ValueError("Metric '{0}' currently not "
                                         "implemented".format(metric))
                if method == "single":
                    dst = np.append(dst, np.min(dst_mat))
                elif method == "complete":
                    dst = np.append(dst, np.max(dst_mat))
                else:
                    dst = np.append(dst, np.sum(dst_mat))
                
    if method == "single":
        return np.min(dst)
    elif method == "complete":
        return np.max(dst)
    else:
        return 1/size_i*1/size_z*np.sum(dst)
