import numpy as np
from scipy.sparse import csr_matrix

"""
Function to calculate the multi project matching results

The Multi-Project Matching Feature uncovers synergy opportunities among various development banks and organizations by facilitating the search for similar projects 
within a selected filter setting (filtered_df) and all projects (project_df).
"""

def calc_multi_matches(filtered_df, project_df, similarity_matrix, top_x, identical_country=False):
    """
    filtered_df: df with applied filters
    project_df: df with all projects
    similarity_matrix: np sparse matrix with all similarities between projects
    top_x: top x project which should be displayed
    identical_country: boolean flag to filter matches where country is identical
    """

    # convert npz sparse matrix into csr matrix
    if not isinstance(similarity_matrix, csr_matrix):
        similarity_matrix = csr_matrix(similarity_matrix)

    # extract indices of the projects
    filtered_indices = filtered_df.index.to_list()
    project_indices = project_df.index.to_list()

    # size down the matrix to only projects within the filter and convert to dense matrix and flatten it
    match_matrix = similarity_matrix[project_indices, :][:, filtered_indices] # row / column
    dense_match_matrix = match_matrix.toarray()
    flat_matrix = dense_match_matrix.flatten()
    
    # get the indices of the top X values in the flattened matrix
    top_indices = np.argsort(flat_matrix)[-top_x:]

    # Convert flat indices back to 2D indices
    top_2d_indices = np.unravel_index(top_indices, dense_match_matrix.shape)
    
    # Extract the corresponding values
    top_values = flat_matrix[top_indices]

    # Prepare the result with row and column indices from original dataframes
    org_rows = []
    org_cols = []
    for value, row, col in zip(top_values, top_2d_indices[0], top_2d_indices[1]):
        original_row_index = project_indices[row]
        original_col_index = filtered_indices[col]
        org_rows.append(original_row_index)
        org_cols.append(original_col_index)
    
    # create two result dataframes

    """
    p1_df: first results of match
    p2_df: matching result

    matches are displayed through the indices of p1 and p2 dfs

    match1 p1_df.iloc[0] & p2_df.iloc[0]
    match2 p1_df.iloc[1] & p2_df.iloc[1]
    """
    p1_df = filtered_df.loc[org_cols].copy()
    p1_df['similarity'] = top_values
    # filter out rows with similarity score less than 50
    p1_df = p1_df[p1_df['similarity'] > 0.50]

    p2_df = project_df.loc[org_rows].copy()
    p2_df['similarity'] = top_values
    p2_df = p2_df[p2_df['similarity'] > 0.50]

    if identical_country:
        # Reset indices before comparison
        p1_df = p1_df.reset_index(drop=True)
        p2_df = p2_df.reset_index(drop=True)
        # Filter to only include matches with identical countries
        identical_country_mask = p1_df['country'] == p2_df['country']
        p1_df = p1_df[identical_country_mask]
        p2_df = p2_df[identical_country_mask]

    # return both results df with matching projects
    return p1_df, p2_df