import numpy as np
from scipy.optimize import linear_sum_assignment
from concurrent.futures import ProcessPoolExecutor
from HistogramBuilder import HistogramBuilder
from scipy.spatial import Delaunay
from HistogramComparsion_functions import sinkhorn_solver, get_EMD
from env_info import num_workers
from Parser_functions import Kabsch_Umeyama
from scipy.spatial.distance import cdist
import math
import random
np.random.seed(0)
random.seed(0)

def find_atoms(vertices, coords):
    # Obtain atoms within individual hull
    hull = Delaunay(vertices)
    ins = hull.find_simplex(coords) >= 0
    return ins

def bin_biomolecule(hex_bins, center, coords, features=None):
    # Iterate through all bins in each sphere of histogram to find atoms
    bin_dict = {}
    total_atom_count = 0
    number_of_bins = int(len(hex_bins) / 6)  # Total number of bins in all spheres

    # Ignore center atom
    retained_coords = (coords != center).all(axis=1)
    coords_without_center = coords[retained_coords]  # old coords = coords[retained_coords]
    # do I really need the atoms?
    # atoms_temp = coords[retained_coords]

    # Retreive lists of attributes used to calculate the EMD between histograms
    w = np.zeros([number_of_bins]) # Weights, or fraction of atoms within each bin
    p = np.zeros((number_of_bins, 3))
    c = np.zeros([number_of_bins]) # Partial charges within bins

    for i in range(number_of_bins):
        # Obtain a list of boolean variables with true corresponding to indices of found atom
        k = i * 6
        vertices = hex_bins[k : k + 6]
        ins = find_atoms(vertices, coords_without_center)
        idx = np.where(ins)

        # Retrieve binned atom count and average vector to center
        binned_atoms = coords_without_center[idx]
        atom_count = len(binned_atoms)
        total_atom_count += atom_count
        binned_vectors = coords_without_center[idx] - center

        # Append atoms, vectors, and potentially features to dictionary / lists
        bin_key = f"bin{i}"

        if atom_count > 0:
            w[i] = atom_count
            p[i] = np.mean(binned_vectors, axis=0)

        if features is not None:
            retained_features = np.array(features)[retained_coords]
            binned_features = retained_features[idx]
            bin_dict[bin_key] = (binned_atoms, binned_vectors, binned_features)
            if len(binned_features) > 0:
                c[i] = np.mean(binned_features)
        else:
            bin_dict[bin_key] = (binned_atoms, binned_vectors)

    if total_atom_count != 0:
        w = w / total_atom_count

    # Get metrics to represent histogram bins in comparison methods
    bin_metrics = (w, p, c, total_atom_count)
    return bin_dict, bin_metrics
    
def get_histogram(coords, center=[0, 0, 0], features = None,  lrf=None, 
                  number_of_spheres = None, max_radius = None, number_of_bins = None):
    histogram_builder = HistogramBuilder.get_histogram_builder(center, number_of_spheres, max_radius+1, number_of_bins)
    angles = histogram_builder.get_angles()
    points = histogram_builder.get_points(angles)
    hex_bins, poly3d = histogram_builder.get_bins(points, lrf)
    bin_dict, bin_metrics = bin_biomolecule(
                hex_bins, center, coords, features
            )

    return hex_bins, poly3d, bin_dict, bin_metrics

def compare_site_pair(args):
    weights_A, bin_metrics_A, weights_B, bin_metrics_B, weights, solver, ratio = args
    if solver == 'Sinkhorn':
        cost = sinkhorn_solver(weights_A, bin_metrics_A, weights_B, bin_metrics_B, weights)
    else:
        cost = get_EMD(weights_A, bin_metrics_A, weights_B, bin_metrics_B, weights)
    return cost * ratio

def obtain_cost_matrix(weights_A, bin_metrics_A, total_atoms_A,  weights_B, bin_metrics_B, total_atoms_B, weights, solver='Sinkhorn', num_workers=12):
    args_list = []
    for i in range(len(weights_A)):
        for j in range(len(weights_B)):
            atoms_A = total_atoms_A[i]
            atoms_B = total_atoms_B[j]
            try:
                ratio = atoms_A/atoms_B if atoms_A > atoms_B else atoms_B/atoms_A
            except:
                ratio = 1
            args_list.append((weights_A[i], bin_metrics_A[i], weights_B[j], bin_metrics_B[j], weights, solver, ratio))

    # Obtain cost matrix by comparing distributions within all histograms
    # Use ProcessPoolExecutor for parallel execution
    with ProcessPoolExecutor(max_workers=num_workers) as executor:
        results = list(executor.map(compare_site_pair, args_list))
        cost_matrix = np.array(results).reshape(len(weights_A), len(weights_B))
        
    return cost_matrix

def compare_sites(coords_A: np.ndarray, coords_B: np.ndarray, lrf_A: np.ndarray, lrf_B: np.ndarray, max_radius=20,
                  features_A=None, features_B=None, weights=[1, 40]):
    # Get weights and bin metrics (in this case, the average vector) for histograms centered around each atom in their respective site
    weights_A, bin_metrics_A, weights_B, bin_metrics_B = [], [], [], []
    total_atoms_A, total_atoms_B = [], []
    for ai in range(len(coords_A)): # Each row entry
        _, _, _, bin_metrics = get_histogram(coords_A, coords_A[ai], features_A, lrf_A[ai], max_radius = max_radius)
        weights_A.append(bin_metrics[0])
        total_atoms_A.append(bin_metrics[3])
        if features_A is not None:
            bin_metrics_A.append(np.hstack((bin_metrics[1], bin_metrics[2].reshape(-1, 1))))
        else:
            bin_metrics_A.append(bin_metrics[1])

    for bj in range(len(coords_B)): # Each column entry
        _, _, _, bin_metrics = get_histogram(coords_B, coords_B[bj], features_B, lrf_B[bj], max_radius = max_radius)
        weights_B.append(bin_metrics[0])
        total_atoms_B.append(bin_metrics[3])
        if features_B is not None:
            bin_metrics_B.append(np.hstack((bin_metrics[1], bin_metrics[2].reshape(-1, 1))))
        else:
            bin_metrics_B.append(bin_metrics[1])
    # Generate all pairwise combinations of histogram distributions
    results = {}
        
    
    cost_matrix = obtain_cost_matrix(weights_A, bin_metrics_A, total_atoms_A, weights_B, bin_metrics_B, total_atoms_B, weights, "Sinkhorn", num_workers)
    # Collapse cost matrix using Hungarian algorithm
    min_row_idx, min_col_idx = linear_sum_assignment(cost_matrix)
    cost = cost_matrix[min_row_idx, min_col_idx].sum()
    # Obtain corresponding atoms between the two sites
    correspondance_A, correspondance_B = coords_A[min_row_idx], coords_B[min_col_idx]
    R, t = Kabsch_Umeyama(correspondance_A, correspondance_B)
    # correspondance_B_aligned = np.array([t + R @ b for b in correspondance_B])
    correspondance_B_aligned = correspondance_B @ R.T + t
    new_RMSD = np.sqrt(np.mean((correspondance_A - correspondance_B_aligned)**2))
    results['cost'] = cost
    results['row_idx'] = min_row_idx
    results['col_idx'] = min_col_idx
    results['un_col_idx'] = np.setdiff1d(np.arange(len(coords_B)), min_col_idx)
    results['RMSD'] = new_RMSD
    results['R'] = R
    results['t'] = t
    results['length'] = len(min_row_idx)
    results['RMSD_normalized'] = new_RMSD/len(min_row_idx)
    results['cost_normalized'] = cost/len(min_row_idx)
    
    each_distane = correspondance_A - correspondance_B_aligned
    distances = np.linalg.norm(each_distane, axis=1)
    results['distances'] = np.sort(distances)
    sorted_indices = np.argsort(distances)
    num_total = len(distances)
    results['corr_10'] = sorted_indices[:math.ceil(0.1 * num_total)]
    results['corr_25'] = sorted_indices[:math.ceil(0.25 * num_total)]
    results['corr_50'] = sorted_indices[:math.ceil(0.5 * num_total)]
    
    return results
