# %%
import os
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
import warnings
warnings.filterwarnings("ignore")
DEBUG=False

USE_SIMPLIFIED_LABELS = True # <<< CHANGE THIS TO False TO USE ORIGINAL BIO LABELS

# %%
# !pip install -q seqeval

# %%
import gc
import json
import copy
import time
import pickle
import numpy as np
import pandas as pd
from itertools import chain
from transformers import AutoTokenizer
from transformers import AutoModelForTokenClassification
from transformers import get_cosine_schedule_with_warmup
from tqdm import tqdm

from seqeval.metrics import recall_score, precision_score
from seqeval.metrics import f1_score

import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
from torch.nn import CrossEntropyLoss
from torch.optim import AdamW

import matplotlib.pyplot as plt
from pylab import rcParams

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# %% [markdown]
# # Config

# %%
exp_name_base = 'deberta_v3_large_w10_ado1_wd0_bs2_lr1e5_5epo' # Changed ado4 to ado1
exp_name = f"{exp_name_base}_{'simplified_labels' if USE_SIMPLIFIED_LABELS else 'original_bio_labels'}"


backbone = 'microsoft/deberta-v3-large'

TRAINING_MAX_LENGTH = 4096
batch_size = 2
lr = 1e-5
ds_rate = 0.3
attention_probs_dropout_prob = 0.4
hidden_dropout_prob = 0.1
wd = 0.
pos_weight = 10
n_epochs = 5

use_amp = True

# %% [markdown]
# # Data

# %%
from sklearn.model_selection import train_test_split

data = json.load(open('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/pii-data-detection-3rd-place-solution/datasets/train.json'))
for sample in data:
    if len(set(sample['labels'])) > 1:
        sample['is_pos'] = 1
    else:
        sample['is_pos'] = 0
print("original datapoints: ", len(data))


moredata = json.load(open('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/pii-data-detection-3rd-place-solution/datasets/moredata_dataset_fixed.json'))
print("moredata datapoints: ", len(moredata))


# %%
from sklearn.model_selection import GroupKFold
import random
random.seed(0)
kf = GroupKFold(n_splits=5)
folds = []
for i, (_, valid_index) in enumerate(kf.split(np.array(data), groups=[x['document'] for x in data])):
    folds.append([data[i]['document'] for i in valid_index])

# %%
# Define simplified labels
pii_types = [
    'EMAIL', 'ID_NUM', 'NAME_STUDENT', 'PHONE_NUM', 
    'STREET_ADDRESS', 'URL_PERSONAL', 'USERNAME'
]
all_labels_simplified = sorted(pii_types + ['O'])
label2id_simplified = {l: i for i, l in enumerate(all_labels_simplified)}
id2label_simplified = {v: k for k, v in label2id_simplified.items()}

# Original target BIO labels
original_target_bio = sorted([
    'B-EMAIL', 'B-ID_NUM', 'B-NAME_STUDENT', 'B-PHONE_NUM', 
    'B-STREET_ADDRESS', 'B-URL_PERSONAL', 'B-USERNAME', 'I-ID_NUM', 
    'I-NAME_STUDENT', 'I-PHONE_NUM', 'I-STREET_ADDRESS', 'I-URL_PERSONAL',
    'O' 
])

if 'O' not in original_target_bio:
    original_target_bio_with_o = sorted(list(set(original_target_bio + ['O'])))
else:
    original_target_bio_with_o = sorted(list(set(original_target_bio)))

label2id_original_bio = {l: i for i, l in enumerate(original_target_bio_with_o)}
id2label_original_bio = {v: k for k, v in label2id_original_bio.items()}


if USE_SIMPLIFIED_LABELS:
    all_labels = all_labels_simplified
    label2id = label2id_simplified
    id2label = id2label_simplified
    print(f"Using SIMPLIFIED labels ({len(all_labels)}): {id2label}")
else:
    all_labels = original_target_bio_with_o
    label2id = label2id_original_bio
    id2label = id2label_original_bio
    print(f"Using ORIGINAL BIO labels ({len(all_labels)}): {id2label}")


# %% [markdown]
# # Dataset

# %%
def tokenize(example, tokenizer, current_label2id, max_length, use_simplified): # Pass current_label2id and use_simplified flag

    # rebuild text from tokens
    text = []
    char_labels = [] 

    for t, l, ws in zip(
        example["tokens"], example["provided_labels"], example["trailing_whitespace"]
    ):
        text.append(t)
        current_label = l
        if use_simplified and l != "O":
            current_label = l[2:] # Strip "B-" or "I-"
        
        # Ensure current_label is in the target label set (either simplified or original BIO)
        # If a label from data isn't in the defined 'all_labels', map it to 'O'.
        if current_label not in current_label2id:
            current_label = "O"
            
        char_labels.extend([current_label] * len(t))

        if ws:
            text.append(" ")
            char_labels.append("O") # Whitespace is always "O"

    tokenized = tokenizer("".join(text), return_offsets_mapping=True, truncation=False)

    if len(tokenized['input_ids']) > max_length:
        # "first N/2 + last N/2" truncation strategy
        first_half_len = max_length // 2
        second_half_len = max_length - first_half_len 

        keys_to_truncate = ['input_ids', 'attention_mask', 'offset_mapping']
        for key in keys_to_truncate:
            if key in tokenized:
                original_list = tokenized[key]
                # takes the first 'first_half_len' elements and the last 'second_half_len' elements
                tokenized[key] = original_list[:first_half_len] + original_list[-second_half_len:]
            

    char_labels_arr = np.array(char_labels) 

    text_str = "".join(text)
    token_labels = []


    current_offset_mapping = tokenized["offset_mapping"] 

    for start_idx, end_idx in current_offset_mapping:
        # Handle special tokens like [CLS], [SEP] which have (0,0) offset_mapping

        if start_idx == end_idx:
            token_labels.append(current_label2id["O"]) # Assign 'O' to special/empty tokens
            continue

        # Ensure indices are within the bounds of char_labels_arr and form a valid span
        # end_idx is exclusive for slicing, so it can be == len(char_labels_arr)
        if not (start_idx < end_idx and start_idx < len(char_labels_arr) and end_idx <= len(char_labels_arr)):
            token_labels.append(current_label2id["O"]) # Assign 'O' if token is outside char_labels range or invalid span
            continue
        
        char_span_labels = char_labels_arr[start_idx:end_idx]
        

        if not char_span_labels.size: 
            token_labels.append(current_label2id["O"])
            continue

        unique_labels, counts = np.unique(char_span_labels, return_counts=True)
        
        pii_labels_present_in_token = []
        pii_counts_present_in_token = []

        for lbl, count in zip(unique_labels, counts): # lbl is already simplified or original based on earlier logic
            if lbl != "O": 
                pii_labels_present_in_token.append(lbl)
                pii_counts_present_in_token.append(count)
        
        if pii_labels_present_in_token:
            # If PII labels are present, choose the most frequent PII label
            most_frequent_pii_idx = np.argmax(pii_counts_present_in_token)
            most_frequent_pii_label = pii_labels_present_in_token[most_frequent_pii_idx]
            token_labels.append(current_label2id[most_frequent_pii_label])
        else:
            # If only \'O\' is present or no PII labels found in the span
            token_labels.append(current_label2id["O"])

    length = len(tokenized["input_ids"]) # Use dictionary-style access

    return {**tokenized, "labels": token_labels, "length": length}

# %%
tokenizer = AutoTokenizer.from_pretrained(backbone)

class PIIDataset(Dataset):
    def __init__(self, tokenizer, data, current_label2id, max_length, use_simplified): # Pass current_label2id and use_simplified
        self.tokenizer = tokenizer
        self.data = copy.deepcopy(data)
        self.current_label2id = current_label2id # Store the active label2id
        self.max_length = max_length
        self.use_simplified = use_simplified

        for i in tqdm(range(len(self.data))):
            if self.data[i].get('input_ids') is None:
                self.data[i]['provided_labels'] = self.data[i]['labels'] 
                self.data[i] = tokenize(
                    self.data[i],
                    self.tokenizer,
                    self.current_label2id, # Pass the active label2id
                    self.max_length,
                    self.use_simplified # Pass the flag
                )

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


def collate_fn(batch):
    input_ids = [torch.tensor(x['input_ids']) for x in batch]
    attention_masks = [torch.tensor(x['attention_mask']) for x in batch]
    labels = [torch.tensor(x['labels']) for x in batch]

    input_ids_padded = pad_sequence(input_ids, batch_first=True, padding_value=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0)
    attention_masks_padded = pad_sequence(attention_masks, batch_first=True, padding_value=0)
    labels_padded = pad_sequence(labels, batch_first=True, padding_value=-100)

    return {
        'input_ids': input_ids_padded,
        'attention_mask': attention_masks_padded,
        'labels': labels_padded,
    }

# Test with simplified labels

ds = PIIDataset(tokenizer, data[:20], label2id, TRAINING_MAX_LENGTH, USE_SIMPLIFIED_LABELS)
print(f"Example tokenized input IDs length: {len(ds[9]['input_ids'])}")
print(f"Example tokenized labels: {ds[9]['labels']}")


# %% [markdown]
# # Model

# %%
model = AutoModelForTokenClassification.from_pretrained(
        backbone,
        attention_probs_dropout_prob=attention_probs_dropout_prob,
        hidden_dropout_prob=hidden_dropout_prob,
        num_labels=len(all_labels), # Use length of active labels
        id2label=id2label,         # Use active id2label
        label2id=label2id,         # Use active label2id
        ignore_mismatched_sizes=True      
    )

# %% [markdown]
# # Loss

# %%

weights = torch.ones(len(all_labels)) * pos_weight 
if 'O' in label2id:
    weights[label2id['O']] = 1
else:
    print(f"Warning: 'O' label not found in active label2id for weight assignment: {label2id}. Assuming last index for 'O'.")
    weights[-1] = 1


CE = nn.CrossEntropyLoss(weight=weights.to(device), ignore_index=-100) 
def criterion(logits, labels):


    return CE(logits.view(-1, len(all_labels)), labels.view(-1))

# %% [markdown]
# # Metric

# %%
def compute_metrics(p, current_all_labels_list, th=0.5):
    predictions, labels = p 
    
    o_label_str = "O" 
    o_label_idx = -1

    for i, lbl_str in enumerate(current_all_labels_list): 
        if lbl_str == o_label_str:
            o_label_idx = i
            break
    if o_label_idx == -1:
        print(f"Warning: Could not find 'O' label in {current_all_labels_list} for metric thresholding. Using last index.")
        o_label_idx = len(current_all_labels_list) - 1


    true_predictions = []
    true_labels = []
    for tmp, batch_l in zip(predictions, labels): 
        batch_p_probs = torch.softmax(torch.from_numpy(tmp), dim=-1).numpy() 
        
        o_column_probs_slice = batch_p_probs[:, :, o_label_idx] 

        condition_mask = o_column_probs_slice < th 

        o_column_probs_slice[condition_mask] = -9999 

        batch_p_final_ids = np.argmax(batch_p_probs, axis=-1) 

        true_predictions.extend([current_all_labels_list[p_id] for (p_id, l_id) in zip(batch_p_final_ids.reshape(-1), batch_l.reshape(-1)) if l_id != -100])
        true_labels.extend([current_all_labels_list[l_id] for l_id in batch_l.reshape(-1) if l_id != -100])
    
    recall = recall_score([true_labels], [true_predictions], average='micro') # Use micro for overall
    precision = precision_score([true_labels], [true_predictions], average='micro')

    beta = 5
    f5_score = (1 + beta**2) * precision * recall / ((beta**2 * precision) + recall + 1e-100)
    if precision + recall == 0: # Handle case where P and R are both zero
        f5_score = 0.0

    results = {
        'recall': recall,
        'precision': precision,
        'f1': f5_score
    }
    return results


# %% [markdown]
# # Train & Valid

# %%
def train_one_epoch(model, loader, optimizer, scheduler, scaler=None):
    model.train()
    running_loss = []
    bar = tqdm(loader, total=len(loader))
    for batch in bar:
        optimizer.zero_grad()
        
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        if scaler is not None:
            with torch.cuda.amp.autocast():
                outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
                loss = criterion(outputs['logits'], labels) 
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
            loss = criterion(outputs['logits'], labels) 
            loss.backward()
            optimizer.step()
        
        scheduler.step() 
        running_loss.append(loss.item())
        bar.set_description(f'train loss: {np.mean(running_loss[-20:]):.4f}')

    return np.mean(running_loss)


def valid_one_epoch(model, loader):
    model.eval()
    running_loss = []
    bar = tqdm(loader, total=len(loader))
    predictions_logits = [] 
    true_labels_list = [] 

    with torch.no_grad():
        for batch in bar:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels_batch = batch['labels'].to(device) 

            outputs = model(input_ids, attention_mask=attention_mask)
            loss = criterion(outputs.logits, labels_batch) 
            running_loss.append(loss.item())
            
            predictions_logits.append(outputs.logits.cpu().numpy())
            true_labels_list.append(labels_batch.cpu().numpy()) 
            
            bar.set_description(f'valid loss: {np.mean(running_loss[-20:]):.4f}')
    
    # Pass the active all_labels (list of strings) to compute_metrics
    metrics_results = compute_metrics((predictions_logits, true_labels_list), all_labels)
    return np.mean(running_loss), metrics_results

# %% [markdown]
# # Scheduler

# %%
rcParams['figure.figsize'] = 20, 2

def get_scheduler(optimizer, num_training_steps, warmup_steps_ratio=0.1):
    return get_cosine_schedule_with_warmup(
        optimizer,
        int(num_training_steps * warmup_steps_ratio),
        num_training_steps
    )

opt = AdamW(model.parameters(), lr=lr)
sch = get_scheduler(opt, 100, 0.1)

lrs = []
for epoch in range(1, 100+1):
    sch.step()
    lrs.append(opt.param_groups[0]["lr"])

plt.plot(range(len(lrs)), lrs)

# %% [markdown]
# # Run

# %%
def run(fold,model_to_train): 

    log_file = os.path.join('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/single_model/model/logs', f'{exp_name}.txt') 
    
    model_save_dir_best = os.path.join('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/single_model/model', f'{exp_name}_fold{fold}_best_hf') 
    model_save_dir_last = os.path.join('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/single_model/model', f'{exp_name}_fold{fold}_last_hf') 
    
    
    os.makedirs(model_save_dir_best, exist_ok=True)
    os.makedirs(model_save_dir_last, exist_ok=True)
    
    optimizer_etc_file_last = os.path.join(model_save_dir_last, 'optimizer_etc_states.pth')


    valid_doc_ids = folds[fold]
    train_data = [x for x in data    if (x['document'] not in valid_doc_ids) and ((x['is_pos'] == 1) or (x['is_pos'] == 0 and random.random() < ds_rate))] + moredata
    
    valid_data = [x for x in data if x['document'] in valid_doc_ids]

    dataset_train = PIIDataset(tokenizer, train_data, label2id, TRAINING_MAX_LENGTH, USE_SIMPLIFIED_LABELS)
    dataset_valid = PIIDataset(tokenizer, valid_data, label2id, TRAINING_MAX_LENGTH, USE_SIMPLIFIED_LABELS)
    print(len(dataset_train), len(dataset_valid))

    loader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=2, collate_fn=collate_fn)
    loader_valid = DataLoader(dataset_valid, batch_size=batch_size, shuffle=False, num_workers=2, collate_fn=collate_fn)

    current_model = model_to_train.to(device)

    optimizer = AdamW(current_model.parameters(), lr=lr, weight_decay=wd)
    scaler = torch.cuda.amp.GradScaler() if use_amp else None
    from_epoch = 0
    metric_best = 0

    scheduler = get_scheduler(optimizer, n_epochs*len(loader_train), 0.1)

    for epoch in range(1, n_epochs+1):
        print(time.ctime(), 'Epoch:', epoch)

        train_loss = train_one_epoch(current_model, loader_train, optimizer, scheduler, scaler)
        valid_loss, result = valid_one_epoch(current_model, loader_valid)
        metric = result["f1"] # F5 score on simplified labels

        content = time.ctime() + ' ' + \
                f'Fold {fold}, Epoch {epoch}, lr: {optimizer.param_groups[0]["lr"]:.7f}, ' + \
                f'train loss: {train_loss:.5f}, valid loss: {valid_loss:.5f}, ' + \
                f'recall: {result["recall"]:.4f} precision: {result["precision"]:.4f} metric (F5 on simplified): {metric:.4f}.'
        print(content)
        with open(log_file, 'a') as appender:
            appender.write(content + '\n')

        if metric > metric_best:
            print(f'metric_best ({metric_best:.6f} --> {metric:.6f}). Saving model to {model_save_dir_best} ...')
            if not DEBUG:
                current_model.save_pretrained(model_save_dir_best, safe_serialization=True)
                tokenizer.save_pretrained(model_save_dir_best)
            metric_best = metric

        if not DEBUG:
            current_model.save_pretrained(model_save_dir_last, safe_serialization=True)
            tokenizer.save_pretrained(model_save_dir_last)
            torch.save(
                {
                    'epoch': epoch,
                    'optimizer_state_dict': optimizer.state_dict(),
                    'scheduler_state_dict': scheduler.state_dict(),
                    'scaler_state_dict': scaler.state_dict() if scaler else None,
                    'metric_best': metric_best,
                },
                optimizer_etc_file_last
            )
    
    del current_model
    torch.cuda.empty_cache()
    gc.collect()

# %%
run(0, model)

# %%