import os
import numpy as np
import transformers
from transformers import (
    AutoTokenizer, 
    AutoModelForTokenClassification, 
    TrainingArguments, 
    Trainer,
    DataCollatorForTokenClassification
)
from datasets import load_dataset
import evaluate

# 1. Константы и параметры
MODEL_CHECKPOINT = "Eraly-ml/KazBERT"
TASK = "ner"
BATCH_SIZE = 32
LEARNING_RATE = 5e-5
EPOCHS = 8
SAVE_PATH = "./kazbert_ner_finetuned"

# 2. Загрузка данных и метрик
datasets = load_dataset("issai/kaznerd", trust_remote_code=True)
metric = evaluate.load("seqeval")
label_list = datasets["train"].features[f"{TASK}_tags"].feature.names

# 3. Токенизация с выравниванием меток
tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT, trust_remote_code=True)

def tokenize_and_align_labels(examples):
    tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
    labels = []
    for i, label in enumerate(examples[f"{TASK}_tags"]):
        word_ids = tokenized_inputs.word_ids(batch_index=i)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:
            if word_idx is None:
                label_ids.append(-100) # Игнорируем спецтокены
            elif word_idx != previous_word_idx:
                label_ids.append(label[word_idx]) # Первая часть слова
            else:
                label_ids.append(-100) # Остальные части слова (игнорируем для оценки)
            previous_word_idx = word_idx
        labels.append(label_ids)
    tokenized_inputs["labels"] = labels
    return tokenized_inputs

tokenized_datasets = datasets.map(tokenize_and_align_labels, batched=True)

# 4. Модель с маппингом меток
id2label = {i: label for i, label in enumerate(label_list)}
label2id = {label: i for i, label in enumerate(label_list)}

model = AutoModelForTokenClassification.from_pretrained(
    MODEL_CHECKPOINT, 
    num_labels=len(label_list),
    id2label=id2label,
    label2id=label2id,
    trust_remote_code=True
)

# 5. Функция для вычисления метрик (Precision, Recall, F1)
def compute_metrics(p):
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)
    true_predictions = [[label_list[p] for (p, l) in zip(prediction, label) if l != -100]
                        for prediction, label in zip(predictions, labels)]
    true_labels = [[label_list[l] for (p, l) in zip(prediction, label) if l != -100]
                   for prediction, label in zip(predictions, labels)]
    results = metric.compute(predictions=true_predictions, references=true_labels, scheme="IOB2")
    return {
        "precision": results["overall_precision"],
        "recall": results["overall_recall"],
        "f1": results["overall_f1"],
        "accuracy": results["overall_accuracy"],
    }

# 6. Настройка и запуск Trainer
args = TrainingArguments(
    output_dir=SAVE_PATH,
    eval_strategy="epoch",
    learning_rate=LEARNING_RATE,
    per_device_train_batch_size=BATCH_SIZE,
    per_device_eval_batch_size=BATCH_SIZE,
    num_train_epochs=EPOCHS,
    weight_decay=0.01,
    save_strategy="no",
    report_to=[]
)

trainer = Trainer(
    model,
    args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=DataCollatorForTokenClassification(tokenizer),
    tokenizer=tokenizer,
    compute_metrics=compute_metrics
)

trainer.train()

# 7. Финальное сохранение
model.save_pretrained(SAVE_PATH)
tokenizer.save_pretrained(SAVE_PATH)