# %%

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

# %%

# we need to create the mdm_list
# import the full mdm-only file
data_path = '../../data_import/exports/data_mapping_mdm.csv'
full_df = pd.read_csv(data_path, skipinitialspace=True)
mdm_list = sorted(list((set(full_df['pattern']))))

# %%
id2label = {}
label2id = {}
for idx, val in enumerate(mdm_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, mdm_list):
    output_list = []
    for _, row in df.iterrows():
        desc = f"<DESC>{row['tag_description']}<DESC>"
        pattern = row['pattern']
        try:
            index = mdm_list.index(pattern)
        except ValueError:
            index = -1
        element = {
            'text' : f"{desc}",
            'label': index,
        }
        output_list.append(element)

    return output_list


def create_split_dataset(fold, mdm_list):
    # train 
    data_path = f"../../data_preprocess/exports/dataset/group_{fold}/train.csv"
    train_df = pd.read_csv(data_path, skipinitialspace=True)

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

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


# %%

# function to perform training for a given fold
# def train(fold):
fold = 1

save_path = f'checkpoint_fold_{fold}'
split_datasets = create_split_dataset(fold, mdm_list)

# prepare tokenizer

model_checkpoint = "distilbert/distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, 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
# 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()
# %% temp
tokenized_datasets['train']['input_ids']

# %%
# 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(mdm_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",
    logging_dir="tensorboard-log",
    logging_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    learning_rate=2e-5,
    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=40,
    bf16=True,
    push_to_hub=False,
    remove_unused_columns=False,
)


trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    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
# for fold in [1,2,3,4,5]:
#     print(fold)
#     train(fold)


# %%