# %%
import pandas as pd
from utils import Retriever, cosine_similarity_chunked
import os
import glob
import numpy as np
from tqdm import tqdm
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, confusion_matrix

##################################################
# helper functions


# 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 k maximum values along axis 1
    top_k = 3
    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




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():
                desc = f"<DESC>{row['tag_description']}<DESC>"
                unit = f"<UNIT>{row['unit']}<UNIT>"
                name = f"<NAME>{row['tag_name']}<NAME"
                element = f"{name}{desc}{unit}"
                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_embedding(batch_size=64)
        return retriever_train.embeddings.to('cpu')



def run_similarity_classifier(fold):
    data_path = f'../../train/mapping_t5_complete_desc_unit_name/mapping_prediction/exports/result_group_{fold}.csv'
    test_df = pd.read_csv(data_path, skipinitialspace=True)



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

    checkpoint_directory = "../../train/classification_bert_complete_desc_unit_name"
    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)

    def compute_top_k(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_values = find_closest(
            cos_sim_matrix=cos_sim_matrix,
            condition_source=condition_source,
            condition_target=condition_target)

        return top_k_values[0][0]



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


    sim_list = []
    for select_idx in tqdm(test_df.index):
        top_sim_value = compute_top_k(select_idx)
        sim_list.append(top_sim_value)

    # analysis 1: using threshold to perform find-back prediction success
    # threshold_values = np.linspace(0.85, 1.00, 21) # test 20 values, 21 to get nice round numbers
    # best_threshold = 0
    # best_f1 = 0
    # for threshold in threshold_values:
    #     predict_list = [ elem > threshold for elem in sim_list ]

    #     y_true = test_df['MDM'].to_list()
    #     y_pred = predict_list

    #     # Compute metrics
    #     accuracy = accuracy_score(y_true, y_pred)
    #     f1 = f1_score(y_true, y_pred)
    #     precision = precision_score(y_true, y_pred)
    #     recall = recall_score(y_true, y_pred)

    #     if f1 > best_f1:
    #         best_threshold = threshold
    #         best_f1 = f1
    
    # just manually set best_threshold
    best_threshold = 0.90

    # compute metrics again with best threshold
    predict_list = [ elem > best_threshold for elem in sim_list ]

    # save
    pred_labels = np.array(predict_list, dtype=bool)

    # append the mdm prediction to the test_df for analysis later
    df_out = pd.DataFrame({
        'p_mdm': pred_labels, 
    })
    df_out.to_csv(f"exports/result_group_{fold}.csv", index=False)


    y_true = test_df['MDM'].to_list()
    y_pred = predict_list
    # Compute metrics
    accuracy = accuracy_score(y_true, y_pred)
    f1 = f1_score(y_true, y_pred)
    precision = precision_score(y_true, y_pred)
    recall = recall_score(y_true, y_pred)



    with open("output.txt", "a") as f:

        print(f'Fold: {fold}', file=f)
        print(f'Best threshold: {best_threshold}', file=f)
        # Print the results
        print(f'Accuracy: {accuracy:.5f}', file=f)
        print(f'F1 Score: {f1:.5f}', file=f)
        print(f'Precision: {precision:.5f}', file=f)
        print(f'Recall: {recall:.5f}', file=f)



# %%
# reset file before writing to it
with open("output.txt", "w") as f:
    print('', file=f)

for fold in [1,2,3,4,5]:
    print(fold)
    run_similarity_classifier(fold)