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

##################
# global parameters
DIAGNOSTIC = False
BATCH_SIZE = 1024

###################
# helper functions
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"{desc}{unit}{name}"
                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
        embedder = T5Embedder(train_data, checkpoint_path)
        # embedder = BertEmbedder(train_data, checkpoint_path)
        embedder.make_embedding(batch_size=BATCH_SIZE)
        return embedder.embeddings




# the selection function takes in the full cos_sim_matrix then subsets the
# matrix according to the test_candidates_mask and train_candidates_mask that we
# give it
# it returns the most likely source candidate index and score among the source
# candidate list
# we then map the local idx to the ship-level idx
def selection(cos_sim_matrix, source_mask, target_mask):
    # subset_matrix = cos_sim_matrix[condition_source]
    # except we are subsetting 2D matrix (row, column)
    subset_matrix = cos_sim_matrix[np.ix_(source_mask, target_mask)]
    # we select top-k here
    # Get the indices of the top-k maximum values along axis 1
    top_k = 1
    # returns a potential 2d matrix of which columns have the highest values
    # top_k_indices = np.argsort(subset_matrix, axis=1)[:, -top_k:]  # Get indices of top k values
    # this partial sorts and ensures we care only top_k are correctly sorted
    top_k_indices = np.argpartition(subset_matrix, -top_k, axis=1)[:, -top_k:]
    
    # Get the values of the top 5 maximum scores
    top_k_values = np.take_along_axis(subset_matrix, top_k_indices, axis=1)
    
    # Calculate the average of the top-k scores along axis 1
    y_scores = np.mean(top_k_values, axis=1)
    max_idx = np.argmax(y_scores)
    max_score = y_scores[max_idx]

    # convert boolean to indices
    condition_indices = np.where(source_mask)[0]
    max_idx = condition_indices[max_idx]
    

    return max_idx, max_score



####################
# global level
# obtain the full mdm_list
data_path = '../../data_import/exports/data_mapping_mdm.csv'
full_df = pd.read_csv(data_path, skipinitialspace=True)
full_df['mapping'] = full_df['thing'] + ' ' + full_df['property']
full_mdm_mapping_list = sorted(list((set(full_df['mapping']))))


#####################
# fold level

def run_selection(fold):

    # set the fold
    # import test data
    # data_path = f"../binary_classifier/classification_prediction/exports/result_group_{fold}.csv"
    data_path = f"../similarity_classifier/exports/result_group_{fold}.csv"
    df = pd.read_csv(data_path, skipinitialspace=True)
    predicted_mdm = df['p_mdm'].to_numpy().astype(bool)

    data_path = f"../../train/mapping_t5_complete_desc_unit_name/mapping_prediction/exports/result_group_{fold}.csv"
    df = pd.read_csv(data_path, skipinitialspace=True)
    df['p_mdm'] = predicted_mdm
    df['p_mapping'] = df['p_thing'] + " " + df['p_property']

    # get target data
    data_path = f"../../data_preprocess/exports/dataset/group_{fold}/train_all.csv"
    train_df = pd.read_csv(data_path, skipinitialspace=True)
    train_df['mapping'] = train_df['thing'] + " " + train_df['property']

    # generate your embeddings
    # checkpoint_directory defined at global level
    # checkpoint_directory = "../../train/classification_bert_pattern_desc_unit"
    checkpoint_directory = "../../train/mapping_t5_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]

    # we can generate the train embeddings once and re-use for every ship
    train_embedder = Embedder(input_df=train_df)
    train_embeds = train_embedder.make_embedding(checkpoint_path)

    # generate new embeddings for each ship
    test_embedder = Embedder(input_df=df)
    global_test_embeds = test_embedder.make_embedding(checkpoint_path)


    # create global_answer array
    # the purpose of this array is to track the classification state at the global
    # level
    global_answer = np.zeros(len(df), dtype=bool)

    #############################
    # ship level
    # we have to split into per-ship analysis
    ships_list = sorted(list(set(df['ships_idx'])))

    for ship_idx in tqdm(ships_list):
        # ship_df = df[df['ships_idx'] == ship_idx]
        # required to map local ship_answer array to global_answer array
        # map_local_index_to_global_index = ship_df.index.to_numpy()

        # we want to subset the ship and only p_mdm values
        ship_mask = df['ships_idx'] == ship_idx
        p_mdm_mask = df['p_mdm']
        map_local_index_to_global_index = np.where(ship_mask & p_mdm_mask)[0]
        ship_df = df[ship_mask & p_mdm_mask].reset_index(drop=True)

        # subset the test embeds
        test_embeds = global_test_embeds[map_local_index_to_global_index]

        # generate the cosine sim matrix for the ship level
        cos_sim_matrix = cosine_similarity_chunked(test_embeds, train_embeds, chunk_size=1024).cpu().numpy()

        ##############################
        # selection level
        # The general idea:
        # step 1: keep only pattern generations that belong to mdm list
        # -> this removes totally wrong datasets that mapped to totally wrong things
        # step 2: loop through the mdm list and isolate data in both train and test that
        # belong to the same pattern class
        # -> this is more tricky, because we have non-mdm mapping to correct classes
        # -> so we have to find which candidate is most similar to the training data

        # it is very tricky to keep track of classification across multiple stages so we
        # will use a boolean answer list to map answers back to the global answer list

        # initialize the local answer list
        ship_answer_list = np.ones(len(ship_df), dtype=bool)

        ###########
        # STEP 1A: ensure that the predicted mapping labels are valid
        pattern_match_mask = ship_df['p_mapping'].apply(lambda x: x in full_mdm_mapping_list).to_numpy()
        pattern_match_mask = pattern_match_mask.astype(bool)
        # anything not in the pattern_match_mask are hallucinations
        # this has the same effect as setting any wrong generations as non-mdm
        ship_answer_list[~pattern_match_mask] = False

        # # STEP 1B: subset our de-duplication to use only predicted_mdm labels
        # p_mdm_mask = ship_df['p_mdm']
        # # assign false to any non p_mdm entries
        # ship_answer_list[~p_mdm_mask] = False
        # # modify pattern_match_mask to remove any non p_mdm values
        # pattern_match_mask = pattern_match_mask & p_mdm_mask

        ###########
        # STEP 2
        # we now go through each class found in our generated set

        # we want to identify per-ship mdm classes
        ship_predicted_classes = sorted(set(ship_df['p_mapping'][pattern_match_mask].to_list()))

        # this function performs the selection given a class
        # it takes in the cos_sim_matrix
        # it returns the selection by mutating the answer_list
        # it sets all relevant idxs to False initially, then sets the selected values to True
        def selection_for_class(select_class, cos_sim_matrix, answer_list):

            # create local copy of answer_list
            ship_answer_list = answer_list.copy()
            # sample_df = ship_df[ship_df['p_mapping'] == select_class]


            # we need to set all idx of chosen entries as False in answer_list -> assume wrong by default
            # selected_idx_list = sample_df.index.to_numpy()
            selected_idx_list = np.where(ship_df['p_mapping'] == select_class)[0]

            # basic assumption check

            # generate the masking arrays for both test and train embeddings
            # we select a tuple from each group, and use that as a candidate for selection
            test_candidates_mask = ship_df['p_mapping'] == select_class
            # we make candidates to compare against in the data sharing the same class
            train_candidates_mask = train_df['mapping'] == select_class

            if sum(train_candidates_mask) == 0:
                # it can be the case that the mdm-valid mapping class is not found in training data
                # print("not found in training data", select_class)
                ship_answer_list[selected_idx_list] = False
                return ship_answer_list

            # perform selection
            # max_idx is the id
            max_idx, max_score = selection(cos_sim_matrix, test_candidates_mask, train_candidates_mask)


            # set the duplicate entries to False
            ship_answer_list[selected_idx_list] = False
            # then only set the one unique chosen value as True
            ship_answer_list[max_idx] = True

            return ship_answer_list 

        # we choose one mdm class
        for select_class in ship_predicted_classes:
            # this resulted in big improvement
            if (sum(ship_df['p_mapping'] == select_class)) > 0:
                ship_answer_list = selection_for_class(select_class, cos_sim_matrix, ship_answer_list)

        # we want to write back to global_answer
        # first we convert local indices to global indices
        ship_local_indices = np.where(ship_answer_list)[0]
        ship_global_indices = map_local_index_to_global_index[ship_local_indices]
        global_answer[ship_global_indices] = True


        if DIAGNOSTIC:
            # evaluation at per-ship level
            y_true = ship_df['MDM'].to_list()
            y_pred = ship_answer_list
            tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel()
            print(f"tp: {tp}")
            print(f"tn: {tn}")
            print(f"fp: {fp}")
            print(f"fn: {fn}")

            # 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)

            # Print the results
            print(f'Accuracy: {accuracy:.5f}')
            print(f'F1 Score: {f1:.5f}')
            print(f'Precision: {precision:.5f}')
            print(f'Recall: {recall:.5f}')



    with open("output.txt", "a") as f:
        print(80 * '*', file=f)
        print(f'Statistics for fold {fold}', file=f)

        y_true = df['MDM'].to_list()
        y_pred = global_answer

        tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel()
        print(f"tp: {tp}", file=f)
        print(f"tn: {tn}", file=f)
        print(f"fp: {fp}", file=f)
        print(f"fn: {fn}", file=f)

        # 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)

        # 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(f'Perform selection for fold {fold}')
    run_selection(fold)


# %%
