
# %%
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 tqdm import tqdm

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    AutoModelForSeq2SeqLM,
)

##################
# global parameters
##################



class BertEmbedder:
    def __init__(self, input_texts, model_checkpoint):
        # we need to generate the embedding from list of input strings
        self.embeddings = []
        self.inputs = input_texts
        model_checkpoint = model_checkpoint 
        self.tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, return_tensors="pt", clean_up_tokenization_spaces=True)

        model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # device = "cpu"
        self.model = model.to(self.device)
        self.model = self.model.eval()
        # self.model = torch.compile(self.model)


    def make_embedding(self, batch_size=128):
        all_embeddings = self.embeddings
        input_texts = self.inputs

        for i in range(0, len(input_texts), batch_size):
            batch_texts = input_texts[i:i+batch_size]
            # Tokenize the input text
            inputs = self.tokenizer(batch_texts, return_tensors="pt", padding=True, truncation=True, max_length=120)
            input_ids = inputs.input_ids.to(self.device)
            attention_mask = inputs.attention_mask.to(self.device)


            # Pass the input through the encoder and retrieve the embeddings
            with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
                with torch.no_grad():
                    encoder_outputs = self.model(input_ids, attention_mask=attention_mask, output_hidden_states=True)
                    # get last layer
                    embeddings = encoder_outputs.hidden_states[-1]
                    # get cls token embedding
                    cls_embeddings = embeddings[:, 0, :]  # Shape: (batch_size, hidden_size)
                    all_embeddings.append(cls_embeddings)
        
        # remove the batch list and makes a single large tensor, dim=0 increases row-wise
        all_embeddings = torch.cat(all_embeddings, dim=0)

        self.embeddings = all_embeddings

class T5Embedder:
    def __init__(self, input_texts, model_checkpoint):
        # we need to generate the embedding from list of input strings
        self.embeddings = []
        self.inputs = input_texts
        model_checkpoint = model_checkpoint 
        self.tokenizer = AutoTokenizer.from_pretrained("t5-base", return_tensors="pt", clean_up_tokenization_spaces=True)
        # define additional special tokens
        additional_special_tokens = ["<THING_START>", "<THING_END>", "<PROPERTY_START>", "<PROPERTY_END>", "<NAME>", "<DESC>", "<SIG>", "<UNIT>", "<DATA_TYPE>"]
        # add the additional special tokens to the tokenizer
        self.tokenizer.add_special_tokens({"additional_special_tokens": additional_special_tokens})

        model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # device = "cpu"
        model.to(self.device)
        self.model = model.eval()
        self.model = torch.compile(self.model)




    def make_embedding(self, batch_size=128):
        all_embeddings = self.embeddings
        input_texts = self.inputs

        for i in range(0, len(input_texts), batch_size):
            batch_texts = input_texts[i:i+batch_size]
            # Tokenize the input text
            inputs = self.tokenizer(batch_texts, return_tensors="pt", padding=True, truncation=True, max_length=128)
            input_ids = inputs.input_ids.to(self.device)
            attention_mask = inputs.attention_mask.to(self.device)


            # Pass the input through the encoder and retrieve the embeddings

            with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
                with torch.no_grad():
                    encoder_outputs = self.model.encoder(input_ids, attention_mask=attention_mask)
                    embeddings = encoder_outputs.last_hidden_state

            # Compute the mean pooling of the token embeddings
            # mean_embedding = embeddings.mean(dim=1)
            mean_embedding = (embeddings * attention_mask.unsqueeze(-1)).sum(dim=1) / attention_mask.sum(dim=1, keepdim=True)
            all_embeddings.append(mean_embedding)
        
        # remove the batch list and makes a single large tensor, dim=0 increases row-wise
        all_embeddings = torch.cat(all_embeddings, dim=0)

        self.embeddings = all_embeddings


def cosine_similarity_chunked(batch1, batch2, chunk_size=1024):
    device = 'cuda'
    batch1_size = batch1.size(0)
    batch2_size = batch2.size(0)
    batch2.to(device)
    
    # Prepare an empty tensor to store results
    cos_sim = torch.empty(batch1_size, batch2_size, device=device)

    # Process batch1 in chunks
    for i in range(0, batch1_size, chunk_size):
        batch1_chunk = batch1[i:i + chunk_size]  # Get chunk of batch1
        
        batch1_chunk.to(device)
        # Expand batch1 chunk and entire batch2 for comparison
        # batch1_chunk_exp = batch1_chunk.unsqueeze(1)  # Shape: (chunk_size, 1, seq_len)
        # batch2_exp = batch2.unsqueeze(0)  # Shape: (1, batch2_size, seq_len)
        batch2_norms = batch2.norm(dim=1, keepdim=True)

        
        # Compute cosine similarity for the chunk and store it in the final tensor
        # cos_sim[i:i + chunk_size] = F.cosine_similarity(batch1_chunk_exp, batch2_exp, dim=-1)

        # Compute cosine similarity by matrix multiplication and normalizing
        sim_chunk = torch.mm(batch1_chunk, batch2.T) / (batch1_chunk.norm(dim=1, keepdim=True) * batch2_norms.T + 1e-8)
        
        # Store the results in the appropriate part of the final tensor
        cos_sim[i:i + chunk_size] = sim_chunk
    
    return cos_sim




###################
# helper functions
class Embedder():
    input_df: pd.DataFrame
    fold: int
    batch_size: int

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


    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}"
                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=self.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

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

def run_deduplication(
    test_df,
    train_df,
    batch_size=1024,
    threshold=0.9,
    diagnostic=False
    ):

    # TODO: replace this with a list of values to import
    # too wasteful to just import everything
    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']))))

    # set the fold
    # import test data
    df = test_df
    df['p_mapping'] = df['p_thing'] + " " + df['p_property']

    # get target data
    data_path = "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_path = 'models/bert_model'

    # cache embeddings
    file_path = "train_embeds.pt"
    if os.path.exists(file_path):
        # Load the tensor if the file exists
        tensor = torch.load(file_path, weights_only=True)
        print("Loaded tensor")
    else:
        # Create and save the tensor if the file doesn't exist
        print('generate train embeddings')
        train_embedder = Embedder(input_df=train_df, batch_size=batch_size)
        tensor = train_embedder.make_embedding(checkpoint_path)
        torch.save(tensor, file_path, weights_only=True)
        print("Tensor saved to file.")
    
    train_embeds = tensor


    # if we can, we can cache the train embeddings and load directly
    # we can generate the train embeddings once and re-use for every ship

    # generate new embeddings for each ship
    print('generate test embeddings')
    test_embedder = Embedder(input_df=df, batch_size=batch_size)
    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
        map_local_index_to_global_index = np.where(ship_mask)[0]
        ship_df = df[ship_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
            if max_score > threshold:
                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

    # we set all unselected values to None
    df.loc[~global_answer, 'p_thing'] = None
    df.loc[~global_answer, 'p_property'] = None


    if diagnostic:
        print(80 * '*')

        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}")
        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}')


    return df


