import numpy as np
from scipy.optimize import linprog, linear_sum_assignment
from scipy.stats import wasserstein_distance
from scipy.spatial.distance import cdist

def euclidean_dist(coord_1, coord_2):
    # Compute euclidean distance between atomic coordinates
    euc_dist = np.linalg.norm(coord_1 - coord_2)
    return euc_dist

def absolute_value(feat_1, feat_2):
    # Compute absolute value between atomic features
    abs_val = np.abs(feat_1 - feat_2)
    return abs_val

def bin_dist(bin_metric_A, bin_metric_B):
    # Get distances between bins based on bin metrics (either average vector or feature(s))
    distances = []

    # Bin distance for vector metrics
    if bin_metric_A.ndim > 1 and bin_metric_B.ndim > 1:
        for vec_A in bin_metric_A:
            for vec_B in bin_metric_B:
                distances.append(euclidean_dist(vec_A, vec_B))

    # Bin distance for feature metrics (e.g., charge)
    else:
        for feat_A in bin_metric_A:
            for feat_B in bin_metric_B:
                distances.append(absolute_value(feat_A, feat_B))
    return distances


def inequality_coeffs(weights_A, weights_B):
    # Write coefficients for constraints in form acceptable by Scipy.optimize linear programming function
    # Note, non-negative numbers are assumed in linear programming function
    # Write coefficients for constraint from equation 5
    A_ub5_list = []
    b_ub5 = np.array(weights_A).reshape(1, len(weights_A))
    increment = len(weights_A)

    for i in range(0, len(weights_A) ** 2, increment):
        A_ub5 = np.zeros(len(weights_A) ** 2)
        A_ub5[i : i + increment] = [1] * increment
        A_ub5_list.append(A_ub5)

    A_ub5s = np.vstack(A_ub5_list)

    # Write coefficients for constraint from equation 6
    A_ub6s = np.empty([len(weights_B), len(weights_B) ** 2])
    b_ub6 = np.array(weights_B).reshape(1, len(weights_B))
    increment = len(weights_B)

    for j in range(len(weights_B)):
        A_ub6 = np.zeros(len(weights_B) ** 2)
        offset = j

        for k in range(len(A_ub6)):
            if k % increment == 0:
                A_ub6[k + offset] = 1

        A_ub6s[j] = A_ub6

    # Concatenate inequality coefficient arrays
    A_ub = np.vstack([A_ub5s, A_ub6s])
    b_ub = np.hstack([np.array(b_ub5), np.array(b_ub6)]).T
    return A_ub, b_ub


def equality_coeffs(distances, weights_A, weights_B):
    # Get equality constraints
    A_eq = np.ones(len(distances)).reshape(1, len(distances))
    sum_weights_A = np.sum(weights_A)
    sum_weights_A = np.sum(weights_A)
    sum_weights_B = np.sum(weights_B)
    b_eq = np.array([np.min([sum_weights_A, sum_weights_B])])
    return A_eq, b_eq

def get_EMD(weights_A, bin_metrics_A, weights_B, bin_metrics_B, weights):
    # Compare two histograms with EMD
    # distances = bin_dist(bin_metrics_A, bin_metrics_B)
    if bin_metrics_A.shape[-1] <= 3:
        distances = cdist(bin_metrics_A, bin_metrics_B, 'euclidean')
    else:
        A3 = bin_metrics_A[:, :3]
        B3 = bin_metrics_B[:, :3]
        Ac = bin_metrics_A[:, 3:]
        Bc = bin_metrics_B[:, 3:]
        D_3 = cdist(A3, B3, 'euclidean')
        D_c = cdist(Ac, Bc, 'euclidean') # This is the same as abs different
        assert len(weights) == 2
        distances = weights[0]*D_3 + weights[1]*D_c
    # Get coefficients for constraints
    A_ub, b_ub = inequality_coeffs(weights_A, weights_B)
    A_eq, b_eq = equality_coeffs(distances, weights_A, weights_B)

    # Coefficients for objective function
    c = np.array(distances).reshape(1, len(distances))

    # Minimize objective function
    result = linprog(c, A_ub=A_ub, b_ub=b_ub, A_eq=A_eq, b_eq=b_eq, method="highs")
    f = result.x

    # Compute EMD (equation 8)
    num = np.dot(distances, f)
    denom = np.sum(f)

    if denom != 0:
        EMD = num / denom

    else:
        EMD = 0
    return EMD


def sinkhorn_solver(
    weights_A,
    bin_metrics_A,
    weights_B,
    bin_metrics_B,
    weights,
    epsilon=1,
    max_iterations=1000,
):
    # Approximate EMD with sinkhorn algorithm
    # Compute ground distance between bins
    if bin_metrics_A.shape[-1] <= 3:
        D = cdist(bin_metrics_A, bin_metrics_B, 'euclidean')
    else:
        A3 = bin_metrics_A[:, :3]
        B3 = bin_metrics_B[:, :3]
        Ac = bin_metrics_A[:, 3:]
        Bc = bin_metrics_B[:, 3:]
        D_3 = cdist(A3, B3, 'euclidean')
        D_c = cdist(Ac, Bc, 'euclidean') # This is the same as abs different
        assert len(weights) == 2
        D = weights[0]*D_3 + weights[1]*D_c

    # Gibbs Kernel
    K = np.exp(-D / epsilon)

    # Initialize weight arrays to solve for Lagrangian multipliers (u, v)
    a = np.array(weights_A).reshape(-1, 1)
    b = np.array(weights_B).reshape(-1, 1)
    a_flat = a.flatten()
    b_flat = b.flatten()

    u = np.ones_like(a_flat)
    v = np.ones_like(b_flat)

    # Iterate until Lagrangian multiplier converge
    iterations = 0
    while iterations < max_iterations:
        # Calculate F
        u_prev, v_prev = u.copy(), v.copy()  # For convergence check
        u = a_flat / (K @ v)
        v = b_flat / (K.T @ u)
        iterations += 1
        # Break if negligible change in the convergence or max iterations achieved
        if np.allclose(u, u_prev, rtol=1e-6, atol=1e-6) and np.allclose(
            v, v_prev, rtol=1e-6, atol=1e-6
        ):
            break

    # Calculate the minimal cost of transport
    F = np.diag(u) @ K @ np.diag(v)
    return np.sum(F * D)
