import numpy as np

u_values = [ 1, 2, 8, 9 ]
v_values = [ 2, 4, 6, 8 ]
u_weights = [ 0.375, 0.125, 0.125, 0.375 ]
v_weights = [ 0.25, 0.25, 0.25, 0.25 ]


def wasserstein(u, v, fu, fv):
    n, m = len(u), len(v)
    su = sorted(enumerate(u), key=lambda x:x[1])
    sv = sorted(enumerate(v), key=lambda x:x[1])
    res = 0
    i, j = 0, 0
    tot = []
    u_arr, v_arr = [], []
    fsu, fsv = [0], [0]
    cum_fsu, cum_fsv = [], []
    deltas = []
    while i < n or j < m:
        if i == n:
            tot.append(sv[j][1])
            j += 1
        elif j == m:
            tot.append(su[i][1])
            i += 1            
        else:
            if su[i][1] < sv[j][1]:
                tot.append(su[i][1])
                i += 1
            else:
                tot.append(sv[j][1])
                j += 1
        if len(tot) > 1:
            deltas.append(tot[-1] - tot[-2])
    i, k = 0, 0
    while k < n + m:
        if i < n and su[i][1] <= tot[k]:            
            i += 1
        k += 1
        u_arr.append(i)

    j, k = 0, 0
    while k < n + m:
        if j < m and sv[j][1] <= tot[k]:            
            j += 1
        k += 1
        v_arr.append(j)
    
    cum_u = [0]
    for i in range(n):
        cum_u.append(cum_u[-1] + fu[su[i][0]])
    cum_v = [0]
    for j in range(m):
        cum_v.append(cum_v[-1] + fv[sv[j][0]])
    res = sum(abs(cum_u[u_arr[i]] - cum_v[v_arr[i]]) * deltas[i] for i in range(n + m - 1))
    print("res :", res)
    return res


wasserstein(u_values, v_values, u_weights, v_weights)
# def wasserstein_distance(u_values, v_values, u_weights, v_weights):
#     u_sorter = np.argsort(u_values)
#     v_sorter = np.argsort(v_values)
#     print(u_sorter, v_sorter)
#     all_values = np.concatenate((u_values, v_values))
#     all_values.sort(kind='mergesort')
#     print("all :", all_values)
#     # Compute the differences between pairs of successive values of u and v.
#     deltas = np.diff(all_values)
#     print("deltas :", deltas)
#     # Get the respective positions of the values of u and v among the values of
#     # both distributions.
#     u_cdf_indices = u_values[u_sorter].searchsorted(all_values[:-1], 'right')
#     v_cdf_indices = v_values[v_sorter].searchsorted(all_values[:-1], 'right')
#     print("u indices", u_values[u_sorter], u_cdf_indices)
#     print("v indices", v_values[v_sorter], v_cdf_indices)
#     # Calculate the CDFs of u and v using their weights, if specified.

#     u_sorted_cumweights = np.concatenate(([0], np.cumsum(u_weights[u_sorter])))    
#     u_cdf = u_sorted_cumweights[u_cdf_indices] 
#     print(u_sorted_cumweights, u_cdf)

#     v_sorted_cumweights = np.concatenate(([0], np.cumsum(v_weights[v_sorter])))
#     v_cdf = v_sorted_cumweights[v_cdf_indices] 
#     print(v_sorted_cumweights, v_cdf)

#     return np.sum(np.multiply(np.abs(u_cdf - v_cdf), deltas))

# res = wasserstein_distance(*map(np.array, (u_values, v_values, u_weights, v_weights)))
# print(res)