"""성능지표 #3(메타 데이터 추출 F1) — KLUE NER 토큰 분류 학습·평가.

전 차수 시험(GERI `GERIR.CE.2511-02.009`) 과 동일한 과제 설정이다.

    모델   : BertForTokenClassification (klue/bert-base, num_labels=13)
    학습   : KLUE NER 약 11,000건
    평가   : 학습되지 않은 2,000여건으로 상대 평가
    결과   : micro avg f1 = 0.8113 (81%)

2-1년차 목표는 83% 다. 같은 과제를 더 큰 사전학습 모델로 올려 달성한다.
`--model klue/roberta-large` 가 기본이며, 전 차수 재현이 필요하면
`--model klue/bert-base` 로 돌린다.

출력은 전 차수 성적서와 같은 형태다. 태그별 precision/recall/f1-score/support 와
micro/macro/weighted 평균을 함께 남겨 성적서에 그대로 첨부할 수 있게 한다.

사용:
    python scripts/train_klue_ner.py --out-dir data/models/klue_ner
    python scripts/train_klue_ner.py --eval-only --out-dir data/models/klue_ner
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

# KLUE NER 태그 6종. BIO 로 12개 + O = 13 labels (전 차수와 동일).
TAGS = ["DT", "LC", "OG", "PS", "QT", "TI"]
LABELS = ["O"] + [f"{p}-{t}" for t in TAGS for p in ("B", "I")]
LABEL2ID = {label: i for i, label in enumerate(LABELS)}
ID2LABEL = {i: label for label, i in LABEL2ID.items()}


def tokenize_and_align(examples, tokenizer, max_length: int):
    """KLUE NER 은 음절 단위 라벨이라 서브워드 토큰에 맞춰 정렬한다.

    첫 서브워드에만 라벨을 주고 나머지는 -100 으로 두어 손실에서 제외한다.
    """
    encoded = tokenizer(
        examples["tokens"],
        is_split_into_words=True,
        truncation=True,
        max_length=max_length,
        padding=False,
    )
    aligned = []
    for i, tags in enumerate(examples["ner_tags"]):
        word_ids = encoded.word_ids(batch_index=i)
        previous = None
        labels = []
        for word_id in word_ids:
            if word_id is None:
                labels.append(-100)
            elif word_id != previous:
                labels.append(tags[word_id])
            else:
                labels.append(-100)
            previous = word_id
        aligned.append(labels)
    encoded["labels"] = aligned
    return encoded


def build_metrics(id2label: dict):
    from seqeval.metrics import classification_report, f1_score

    def compute(eval_pred):
        logits, labels = eval_pred
        predictions = np.argmax(logits, axis=-1)
        true_labels, true_preds = [], []
        for prediction, label in zip(predictions, labels):
            pairs = [(p, l) for p, l in zip(prediction, label) if l != -100]
            true_preds.append([id2label[int(p)] for p, _ in pairs])
            true_labels.append([id2label[int(l)] for _, l in pairs])
        report = classification_report(
            true_labels, true_preds, output_dict=True, zero_division=0)
        return {
            "f1": f1_score(true_labels, true_preds, zero_division=0),
            "report": report,
        }

    return compute


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default="klue/roberta-large",
                        help="전 차수 재현은 klue/bert-base")
    parser.add_argument("--out-dir", type=Path, default=Path("data/models/klue_ner"))
    parser.add_argument("--epochs", type=float, default=3.0)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--lr", type=float, default=2e-5)
    parser.add_argument("--max-length", type=int, default=128)
    parser.add_argument("--eval-only", action="store_true")
    args = parser.parse_args()

    import torch
    from datasets import load_dataset
    from transformers import (AutoModelForTokenClassification, AutoTokenizer,
                              DataCollatorForTokenClassification, Trainer,
                              TrainingArguments)

    dataset = load_dataset("klue", "ner")
    print("학습 %d건 / 평가 %d건"
          % (len(dataset["train"]), len(dataset["validation"])))

    source = args.out_dir if args.eval_only else args.model
    tokenizer = AutoTokenizer.from_pretrained(str(source))
    model = AutoModelForTokenClassification.from_pretrained(
        str(source), num_labels=len(LABELS), id2label=ID2LABEL, label2id=LABEL2ID)

    encoded = dataset.map(
        lambda batch: tokenize_and_align(batch, tokenizer, args.max_length),
        batched=True, remove_columns=dataset["train"].column_names)

    training_args = TrainingArguments(
        output_dir=str(args.out_dir / "checkpoints"),
        learning_rate=args.lr,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size * 2,
        num_train_epochs=args.epochs,
        weight_decay=0.01,
        logging_steps=100,
        save_strategy="no",
        report_to=[],
        fp16=torch.cuda.is_available(),
    )
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=encoded["train"],
        eval_dataset=encoded["validation"],
        data_collator=DataCollatorForTokenClassification(tokenizer),
        compute_metrics=build_metrics(ID2LABEL),
    )

    if not args.eval_only:
        trainer.train()
        args.out_dir.mkdir(parents=True, exist_ok=True)
        trainer.save_model(str(args.out_dir))
        tokenizer.save_pretrained(str(args.out_dir))

    metrics = trainer.evaluate()
    report = metrics["eval_report"]
    micro = report["micro avg"]

    print("\n" + "=" * 62)
    print("메타 데이터 추출 F1-score : %.4f   [목표 0.83]" % micro["f1-score"])
    print("  precision %.4f / recall %.4f / support %d"
          % (micro["precision"], micro["recall"], micro["support"]))
    print("\n[태그별]")
    for tag in TAGS:
        if tag in report:
            row = report[tag]
            print("  %-4s P=%.4f R=%.4f F1=%.4f  support=%d"
                  % (tag, row["precision"], row["recall"],
                     row["f1-score"], row["support"]))
    for name in ("macro avg", "weighted avg"):
        row = report[name]
        print("  %-12s P=%.4f R=%.4f F1=%.4f"
              % (name, row["precision"], row["recall"], row["f1-score"]))

    args.out_dir.mkdir(parents=True, exist_ok=True)
    (args.out_dir / "result.json").write_text(
        json.dumps({"model": args.model, "f1": micro["f1-score"],
                    "report": report}, ensure_ascii=False, indent=2),
        encoding="utf-8")
    print("\n%s 에 결과 기록" % (args.out_dir / "result.json"))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())