#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Aug 17 16:54:12 2019

@author: wanxiang.shen@u.nus.edu

@usecase: calculate varies distances
"""
from aggmap.utils.distances import named_distances
from aggmap.utils.multiproc import MultiProcessUnorderedBarRun

import numpy as np
from tqdm import tqdm


def _yield_combinations(N):
    for i1 in range(N):
        for i2 in range(i1):
            yield (i1,i2)
            
def _calculate(i1, i2):
    x1 = data[:, i1]
    x2 = data[:, i2]
    ## dropna
    X = np.stack([x1,x2], axis=1)
    X = X[~np.isnan(X).any(axis=1)]
    x1 = X[:, 0]
    x2 = X[:, 1]
    
#     x1 = np.nan_to_num(x1)
#     x2 = np.nan_to_num(x2)
    if x1.any():
        dist = func(x1, x2)
    else:
        dist = np.nan
    return (i1, i2, dist)

def _fuc(x):
    i1, i2  = x
    return _calculate(i1, i2)


def pairwise_distance(npydata, n_cpus=8, method='correlation'):
    """
    parameters
    ---------------
    method: {'euclidean', 'manhattan', 'canberra', 'chebyshev', 
             'cosine', 'braycurtis', 'correlation',
             'jaccard', 'rogerstanimoto', 'hamming', 'dice', 'kulsinski', 'sokal_sneath'}
    npydata: np.array or np.memmap, Note that the default we will calcuate the vector's distances instead of sample's distances, if you wish to calculate distances between samples, you can pass data.T instead of data

    Usage
    --------------
    >>> import numpy as np
    >>> data = np.random.random_sample(size=(10000,10)
    >>> dist_matrix = pairwise_distance(data)
    >>> dist_matrix.shape
    >>> (10,10)  
    """    
    global data, func
    
    func = named_distances.get(method)
    data = npydata
    N = data.shape[1]
    lst = list(_yield_combinations(N))
    res = MultiProcessUnorderedBarRun(_fuc, lst, n_cpus=n_cpus)
    dist_matrix = np.zeros(shape = (N,N))
    for x,y,v in tqdm(res,ascii=True):
        dist_matrix[x,y] = v
        dist_matrix[y,x] = v
    return dist_matrix

    
    
if __name__ == '__main__':
    
    import numpy as np
    import pandas as pd
    from umap import UMAP
    import matplotlib.pyplot as plt
    
    X = np.random.random_sample(size=(1000000,40))
    distmatrix = pairwise_distance(X, n_cpus=6)
    embedding = UMAP(metric = 'precomputed',random_state=10)
    df = pd.DataFrame(embedding.fit_transform(distmatrix))
    ax = plt.plot(df[0],df[1], 'bo')