#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Oct 20 12:04:30 2025

@author: samd
"""

import json
import torch
import numpy as np
from pathlib import Path
from pprint import pprint
from typing import (
    List,
    Dict,
    Tuple,
)
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    TrainingArguments,
    Trainer,
    DataCollatorForTokenClassification,
    pipeline,
)
from datasets import Dataset, DatasetDict
from seqeval.metrics import classification_report, f1_score, precision_score, recall_score
from seqeval.scheme import IOB2


def load_json_data(file_path: str) -> List[Dict]:
    with open(file_path, "r", encoding="utf-8") as f:
        return json.load(f)


def convert_to_iob(sample: Dict) -> Tuple[List[str], List[str]]:
    tokens = sample["text"].split()
    labels = ["O"] * len(tokens)

    for entity in sample["entities"]:
        start, end = entity["start"], entity["end"]
        entity_type = entity["type"]
        labels[start] = f"B-{entity_type}"
        for i in range(start + 1, end + 1):
            if i < len(labels):
                labels[i] = f"I-{entity_type}"
    return tokens, labels


def build_label_mappings(datasets: List[List[Dict]]) -> Tuple[Dict[str, int], Dict[int, str]]:
    all_labels = set(["O"])
    for dataset in datasets:
        for sample in dataset:
            _, labels = convert_to_iob(sample)
            all_labels.update(labels)

    sorted_labels = sorted(list(all_labels))
    label2id = {label: i for i, label in enumerate(sorted_labels)}
    id2label = {i: label for label, i in label2id.items()}
    print(f"Labels ({len(label2id)}): {sorted_labels}")
    return label2id, id2label


def convert_dataset(data: List[Dict], label2id: Dict[str, int]) -> Dataset:
    tokens_list, labels_list = [], []
    for sample in data:
        tokens, labels = convert_to_iob(sample)
        tokens_list.append(tokens)
        labels_list.append([label2id[label] for label in labels])
    return Dataset.from_dict(
        {
            "tokens": tokens_list,
            "ner_tags": labels_list,
        }
    )

def tokenize_and_align_labels(examples, tokenizer):
    """Tokenize and align labels for token classification.

    Args:
        examples: Dataset examples with "tokens" and "ner_tags" fields.

    Returns:
        tokenized_inputs: Tokenized inputs with aligned labels.
    """
    tokenized_inputs = tokenizer(
        examples["tokens"],
        truncation=True,
        is_split_into_words=True,
        max_length=512,
    )
    labels = []
    for i, label in enumerate(examples["ner_tags"]):
        word_ids = tokenized_inputs.word_ids(batch_index=i)
        label_ids = []
        previous_word_idx = None
        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

def tvt2dict(data_sets, label2id = None):
    if not label2id:
        label2id, _ = build_label_mappings(data_sets)
        
    dataset_dict = DatasetDict(
    {
        "train": convert_dataset(data_sets[0], label2id),
        "validation": convert_dataset(data_sets[1], label2id),
        "test": convert_dataset(data_sets[2], label2id),
    }
    )
    return dataset_dict