import pandas as pd
import numpy as np
from scipy.sparse import csr_matrix, lil_matrix
import streamlit as st

# multi_project_matching
def calc_matches(filtered_df, project_df, similarity_matrix, top_x):
    st.write(filtered_df.shape)
    st.write(project_df.shape)
    st.write(similarity_matrix.shape)

    # Ensure the matrix is in a suitable format for manipulation
    #if not isinstance(similarity_matrix, csr_matrix):
    #    similarity_matrix = csr_matrix(similarity_matrix)

    # Get indices from dataframes
    filtered_df_indices = filtered_df.index.to_list()
    project_df_indices = project_df.index.to_list()

    # Create mapping dictionaries
    filtered_df_index_map = {i: index for i, index in enumerate(filtered_df_indices)}
    project_df_index_map = {i: index for i, index in enumerate(project_df_indices)}

    # Select submatrix based on indices from both dataframes
    match_matrix = similarity_matrix[filtered_df_indices, :][:, project_df_indices]

    st.write(match_matrix.shape)

    # Get the linear indices of the top 'top_x' values
    # (flattened index to handle the sparse matrix more effectively)
    #linear_indices = np.argsort(match_matrix.data)[-top_x:]
    #if len(linear_indices) < top_x:
    #    top_x = len(linear_indices)

    # Convert flat indices to 2D indices using the shape of the submatrix
    #top_indices = np.unravel_index(linear_indices, match_matrix.shape)

    # Get the corresponding similarity values
    #top_values = match_matrix.data[linear_indices]

    flat_indices = np.argpartition(match_matrix.flatten(), -3)[-3:]

    # Convert flat indices to 2D row and column indices
    row_indices, col_indices = np.unravel_index(flat_indices, match_matrix.shape)

    # Get the values corresponding to the top k indices
    top_values = match_matrix[row_indices, col_indices]

    top_filtered_df_indices = [filtered_df_index_map[i] for i in col_indices]
    top_project_df_indices = [project_df_index_map[i] for i in row_indices]

    st.write(top_filtered_df_indices)

    # Create resulting dataframes with top matches and their similarity scores
    p1_df = filtered_df.loc[top_filtered_df_indices].copy()
    p1_df['similarity'] = top_values

    p2_df = project_df.loc[top_project_df_indices].copy()
    p2_df['similarity'] = top_values
    print("finished calc matches")

    return p1_df, p2_df