# %%

# from datasets import load_from_disk
import os

os.environ['NCCL_P2P_DISABLE'] = '1'
os.environ['NCCL_IB_DISABLE'] = '1'
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    DataCollatorWithPadding,
    Trainer,
    EarlyStoppingCallback,
    TrainingArguments
)
import evaluate
import numpy as np
import pandas as pd
# import matplotlib.pyplot as plt
from datasets import Dataset, DatasetDict



torch.set_float32_matmul_precision('high')

# %%

# import training file
data_path = '../../data_import/train.csv'
train_df = pd.read_csv(data_path, skipinitialspace=True)
# rather than use pattern, we use the real thing and property
entity_ids = train_df['entity_id'].to_list()
target_id_list = sorted(list(set(entity_ids)))


# %%
id2label = {}
label2id = {}
for idx, val in enumerate(target_id_list):
    id2label[idx] = val
    label2id[val] = idx

# %%

# outputs a list of dictionaries
# processes dataframe into lists of dictionaries
# each element maps input to output
# input: tag_description
# output: class label
def process_df_to_dict(df):
    output_list = []
    for _, row in df.iterrows():
        desc = row['mention']
        index = row['entity_id']
        element = {
            'text' : f"{desc}",
            'label': label2id[index], # ensure labels starts from 0
        }
        output_list.append(element)

    return output_list


def create_dataset():
    # train 
    data_path = '../../data_import/train.csv'
    train_df = pd.read_csv(data_path, skipinitialspace=True)


    combined_data = DatasetDict({
        'train': Dataset.from_list(process_df_to_dict(train_df)),
    })
    return combined_data


# %%

def train():

    save_path = f'checkpoint'
    split_datasets = create_dataset()

    # prepare tokenizer

    # model_checkpoint = "distilbert/distilbert-base-uncased"
    model_checkpoint = 'google-bert/bert-base-cased'
    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, return_tensors="pt", clean_up_tokenization_spaces=True)
    # Define additional special tokens
    # additional_special_tokens = ["<DESC>"]
    # Add the additional special tokens to the tokenizer
    # tokenizer.add_special_tokens({"additional_special_tokens": additional_special_tokens})

    max_length = 120

    # given a dataset entry, run it through the tokenizer
    def preprocess_function(example):
        input = example['text']
        # text_target sets the corresponding label to inputs
        # there is no need to create a separate 'labels'
        model_inputs = tokenizer(
            input,
            max_length=max_length,
            truncation=True,
            padding=True
        )
        return model_inputs

    # map maps function to each "row" in the dataset
    # aka the data in the immediate nesting
    tokenized_datasets = split_datasets.map(
        preprocess_function,
        batched=True,
        num_proc=8,
        remove_columns="text",
    )

    # %% temp
    # tokenized_datasets['train'].rename_columns()

    # %%
    # create data collator

    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    # %%
    # compute metrics
    metric = evaluate.load("accuracy")


    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        preds = np.argmax(preds, axis=1)
        return metric.compute(predictions=preds, references=labels)

    # %%
    # create id2label and label2id


    # %%
    model = AutoModelForSequenceClassification.from_pretrained(
        model_checkpoint,
        num_labels=len(target_id_list),
        id2label=id2label,
        label2id=label2id)
    # important! after extending tokens vocab
    model.resize_token_embeddings(len(tokenizer))

    # model = torch.compile(model, backend="inductor", dynamic=True)


    # %%
    # Trainer

    training_args = TrainingArguments(
        output_dir=f"{save_path}",
        # eval_strategy="epoch",
        eval_strategy="no",
        logging_dir="tensorboard-log",
        logging_strategy="epoch",
        # save_strategy="epoch",
        load_best_model_at_end=False,
        learning_rate=1e-4,
        per_device_train_batch_size=64,
        per_device_eval_batch_size=64,
        auto_find_batch_size=False,
        ddp_find_unused_parameters=False,
        weight_decay=0.01,
        save_total_limit=1,
        num_train_epochs=250,
        bf16=True,
        push_to_hub=False,
        remove_unused_columns=False,
    )


    trainer = Trainer(
        model,
        training_args,
        train_dataset=tokenized_datasets["train"],
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
        # callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
    )

    # uncomment to load training from checkpoint
    # checkpoint_path = 'default_40_1/checkpoint-5600'
    # trainer.train(resume_from_checkpoint=checkpoint_path)

    trainer.train()

# execute training
train()


# %%