# %%
import pandas as pd
from utils import Retriever, cosine_similarity_chunked
import os
import glob
import numpy as np

# %%
data_path = f'../data_preprocess/exports/preprocessed_data.csv'
df_pre = pd.read_csv(data_path, skipinitialspace=True)

# %%
# this should be >0 if we are using abbreviations processed data
desc_list = df_pre['tag_description'].to_list()

# %%
[ elem for elem in desc_list if isinstance(elem, float)]
##########################################
# %%
fold = 1
data_path = f'../train/mapping_pattern/mapping_prediction/exports/result_group_{fold}.csv'
df = pd.read_csv(data_path, skipinitialspace=True)

# %%
# subset to mdm
df = df[df['MDM']]

thing_condition = df['p_thing'] == df['thing_pattern']
error_thing_df = df[~thing_condition][['tag_description', 'thing_pattern','p_thing']]

property_condition = df['p_property'] == df['property_pattern']
error_property_df = df[~property_condition][['tag_description', 'property_pattern','p_property']]

correct_df = df[thing_condition & property_condition][['tag_description', 'property_pattern', 'p_property']]

test_df = df

# %%
# thing_df.to_html('thing_errors.html')
# property_df.to_html('property_errors.html')

##########################################
# what we need now is understand why the model is making these mispredictions
# import train data and test data
# %%
class Embedder():
    input_df: pd.DataFrame
    fold: int

    def __init__(self, input_df):
        self.input_df = input_df


    def make_embedding(self, checkpoint_path):

        def generate_input_list(df):
            input_list = []
            for _, row in df.iterrows():
                # name = f"<NAME>{row['tag_name']}<NAME>"
                desc = f"<DESC>{row['tag_description']}<DESC>"
                # element = f"{name}{desc}"
                element = f"{desc}"
                input_list.append(element)
            return input_list

        # prepare reference embed
        train_data = list(generate_input_list(self.input_df))
        # Define the directory and the pattern
        retriever_train = Retriever(train_data, checkpoint_path)
        retriever_train.make_mean_embedding(batch_size=64)
        return retriever_train.embeddings.to('cpu')



# %%
data_path = f"../data_preprocess/exports/dataset/group_{fold}/train.csv"
train_df = pd.read_csv(data_path, skipinitialspace=True)

checkpoint_directory = "../train/mapping_pattern"
directory = os.path.join(checkpoint_directory, f'checkpoint_fold_{fold}')
# Use glob to find matching paths
# path is usually checkpoint_fold_1/checkpoint-<step number>
# we are guaranteed to save only 1 checkpoint from training
pattern = 'checkpoint-*'
checkpoint_path = glob.glob(os.path.join(directory, pattern))[0]

train_embedder = Embedder(input_df=train_df)
train_embeds = train_embedder.make_embedding(checkpoint_path)

test_embedder = Embedder(input_df=test_df)
test_embeds = test_embedder.make_embedding(checkpoint_path)



# %%
# test embeds are inputs since we are looking back at train data
cos_sim_matrix = cosine_similarity_chunked(test_embeds, train_embeds, chunk_size=8).cpu().numpy()

# %%
# the following function takes in a full cos_sim_matrix
# condition_source: boolean selectors of the source embedding
# condition_target: boolean selectors of the target embedding
def find_closest(cos_sim_matrix, condition_source, condition_target):
    # subset_matrix = cos_sim_matrix[condition_source]
    # except we are subsetting 2D matrix (row, column)
    subset_matrix = cos_sim_matrix[np.ix_(condition_source, condition_target)]
    # we select top k here
    # Get the indices of the top 5 maximum values along axis 1
    top_k = 5
    top_k_indices = np.argsort(subset_matrix, axis=1)[:, -top_k:]  # Get indices of top k values
    # note that top_k_indices is a nested list because of the 2d nature of the matrix
    # the result is flipped
    top_k_indices[0] = top_k_indices[0][::-1]
    
    # Get the values of the top 5 maximum scores
    top_k_values = np.take_along_axis(subset_matrix, top_k_indices, axis=1)
    

    return top_k_indices, top_k_values

# %%
error_thing_df.index

####################################################
# special find-back code
# %%
select_idx = 2884
condition_source = test_df['tag_description'] == test_df[test_df.index == select_idx]['tag_description'].tolist()[0]
condition_target = np.ones(train_embeds.shape[0], dtype=bool)

top_k_indices, top_k_values = find_closest(
    cos_sim_matrix=cos_sim_matrix,
    condition_source=condition_source,
    condition_target=condition_target)

print(top_k_indices)
print(top_k_values)
# %%
train_df.iloc[top_k_indices[0]]
# %%
test_df[test_df.index == select_idx]
####################################################

# %%
# make function to compute similarity of closest retrieved result
def compute_similarity(select_idx):
    condition_source = test_df['tag_description'] == test_df[test_df.index == select_idx]['tag_description'].tolist()[0]
    condition_target = np.ones(train_embeds.shape[0], dtype=bool)
    top_k_indices, top_k_values = find_closest(
        cos_sim_matrix=cos_sim_matrix,
        condition_source=condition_source,
        condition_target=condition_target)

    return np.mean(top_k_values[0])

# %%
def print_summary(similarity_scores):
    # Convert list to numpy array for additional stats
    np_array = np.array(similarity_scores)

    # Get stats
    mean_value = np.mean(np_array)
    percentiles = np.percentile(np_array, [25, 50, 75])  # 25th, 50th, and 75th percentiles

    # Display numpy results
    print("Mean:", mean_value)
    print("25th, 50th, 75th Percentiles:", percentiles)


# %%
##########################################
# Analyze the degree of similarity differences between correct and incorrect results

# %%
# compute similarity scores for all values in error_thing_df
similarity_thing_scores = []
for idx in error_thing_df.index:
    similarity_thing_scores.append(compute_similarity(idx))
print_summary(similarity_thing_scores)


# %%
similarity_property_scores = []
for idx in error_property_df.index:
    similarity_property_scores.append(compute_similarity(idx))
print_summary(similarity_property_scores)

# %%
similarity_correct_scores = []
for idx in correct_df.index:
    similarity_correct_scores.append(compute_similarity(idx))
print_summary(similarity_correct_scores)



# %%
import matplotlib.pyplot as plt

# Sample data
list1 = similarity_thing_scores
list2 = similarity_property_scores
list3 = similarity_correct_scores

# Plot histograms
bins = 50
plt.hist(list1, bins=bins, alpha=0.5, label='List 1', density=True)
plt.hist(list2, bins=bins, alpha=0.5, label='List 2', density=True)
plt.hist(list3, bins=bins, alpha=0.5, label='List 3', density=True)

# Labels and legend
plt.xlabel('Value')
plt.ylabel('Frequency')
plt.legend(loc='upper right')
plt.title('Histograms of Three Lists')

# Show plot
plt.show()

###########################################
# %%
# why do similarities of 97% still map correctly?
score_array = np.array(similarity_correct_scores)
# %%
sum(score_array < 0.95)
# %%
correct_df[score_array < 0.95]['tag_description'].index.to_list()
# %%