import pandas as pd
import numpy as np
from typing import List
from tqdm import tqdm
from utils import Retriever, cosine_similarity_chunked

class Selector():
    input_df: pd.DataFrame
    reference_df: pd.DataFrame
    ships_list: List[int]
    fold: int

    def __init__(self, input_df, reference_df, fold):
        self.ships_list = sorted(list(set(input_df['ships_idx'])))
        self.input_df = input_df
        self.reference_df = reference_df
        self.fold = fold


    def run_selection(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

        # given a dataframe, return a single idx of the entry has the highest match with
        # the embedding
        def selection(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 = 1
            top_k_indices = np.argsort(subset_matrix, axis=1)[:, -top_k:]  # Get indices of top k values
            
            # 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 5 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 (1,2,3)
            condition_indices = np.where(condition_source)[0]
            max_idx = condition_indices[max_idx]

            return max_idx, max_score


        # prepare reference embed
        train_data = list(generate_input_list(self.reference_df))
        # Define the directory and the pattern
        retriever_train = Retriever(train_data, checkpoint_path)
        retriever_train.make_mean_embedding(batch_size=64)
        train_embed = retriever_train.embeddings

        # take the inputs for df_sub
        test_data = list(generate_input_list(self.input_df))
        retriever_test = Retriever(test_data, checkpoint_path)
        retriever_test.make_mean_embedding(batch_size=64)
        test_embed = retriever_test.embeddings



        # precision_list = []
        # recall_list = []
        tp_accumulate = 0
        tn_accumulate = 0
        fp_accumulate = 0
        fn_accumulate = 0
        THRESHOLD = 0.9
        for ship_idx in self.ships_list:
            print(ship_idx)
            # we select a ship and select only data exhibiting MDM pattern in the predictions
            ship_mask = (self.input_df['ships_idx'] == ship_idx) & (self.input_df['p_MDM'])
            df_ship = self.input_df[ship_mask].reset_index(drop=True)
            # we then try to make a dataframe for each thing_property attribute
            df_ship['thing_property'] = df_ship['p_thing'] + " " + df_ship['p_property']
            unique_patterns = list(set(df_ship['thing_property']))
            condition_list = []
            for pattern in unique_patterns:
                # we obtain the boolean mask to subset the source and target entries
                condition_source = (df_ship['thing_property'] == pattern)
                condition_target = (self.reference_df['thing_property'] == pattern)
                item = {'condition_source': condition_source,
                        'condition_target': condition_target}
                condition_list.append(item)

            # subset part of self.input_df that belongs to the ship 
            test_embed_subset = test_embed[ship_mask]
            cos_sim_matrix = cosine_similarity_chunked(test_embed_subset, train_embed, chunk_size=8).cpu().numpy()


            # for each sub_df, we have to select the best candidate
            # we will do this by finding which desc input has the highest similarity with train data
            all_idx_list = []
            selected_idx_list = []
            similarity_score = []
            for item in tqdm(condition_list):
                condition_source = item['condition_source']
                condition_target = item['condition_target']
                # if there is no equivalent data in target, we skip
                if sum(condition_target) == 0:
                    pass
                # if there is equivalent data in target, we perform selection among source
                # by top-k highest similarity with targets
                else:
                    # idx is with respect
                    max_idx, max_score = selection(
                        cos_sim_matrix, condition_source, condition_target 
                    )
                    all_idx_list.append(max_idx)
                    similarity_score.append(max_score)
                    # implement thresholding
                    if max_score > THRESHOLD:
                        selected_idx_list.append(max_idx)

            # let us tag the df_ship with the respective 'selected' and 'ood' tags
            df_ship['selected'] = False
            df_ship.loc[all_idx_list, 'selected'] = True
            df_ship['ood'] = 0.0
            df_ship.loc[all_idx_list, 'ood'] = similarity_score

            # we now split the dataframe by p_mdm
            # explanation:
            # we first separated our ship into p_mdm and non p_mdm
            # we only select final in-mdm prediction from p_mdm subset
            # anything that is not selected and from non-p_mdm is predicted not in mdm

            # get our final prediction
            df_subset_predicted_true = df_ship.loc[selected_idx_list]
            # take the set difference between df_ship's index and the given list
            inverse_list = df_ship.index.difference(selected_idx_list).to_list()
            df_subset_predicted_false = df_ship.loc[inverse_list]

            not_p_mdm_mask = (self.input_df['ships_idx'] == ship_idx) & (~self.input_df['p_MDM'])
            # this is the part we don't care
            df_not_p_mdm = self.input_df[not_p_mdm_mask].reset_index(drop=True)

            # concat
            df_false = pd.concat([df_subset_predicted_false, df_not_p_mdm], axis=0)
            assert(len(df_false) + len(df_subset_predicted_true) == sum(self.input_df['ships_idx'] == ship_idx))

            # we want to return a df with the final prediction
            # a bit dirty, but we re-use the fields
            df_false['p_MDM'] = False
            df_subset_predicted_true['p_MDM'] = True


            # save ship for analysis later
            # df_return = pd.concat([df_false, df_subset_predicted_true], axis=0)
            # df_return.to_parquet(f'exports/fold_{self.fold}/ship_{ship_idx}.parquet')




            # true positive -> predicted in mdm, actual in mdm
            # we get all the final predictions that are also found in MDM
            true_positive = sum(df_subset_predicted_true['MDM'])
            # true negative -> predicted not in mdm, and not found in MDM
            # we negate the condition to get those that are not found in MDM
            true_negative = sum(~df_false['MDM'])
            # false positive -> predicted in mdm, not found in mdm
            false_positive = sum(~df_subset_predicted_true['MDM'])
            # false negative -> predicted not in mdm, found in mdm
            false_negative = sum(df_false['MDM'])


            tp_accumulate = tp_accumulate + true_positive
            tn_accumulate = tn_accumulate + true_negative
            fp_accumulate = fp_accumulate + false_positive
            fn_accumulate = fn_accumulate + false_negative

             

        total_sum = (tp_accumulate + tn_accumulate + fp_accumulate + fn_accumulate)
        # ensure that all entries are accounted for
        assert(total_sum == len(self.input_df))
        return tp_accumulate, tn_accumulate, fp_accumulate, fn_accumulate

