import json
import logging
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
from datasets import load_dataset, DatasetDict
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    TrainingArguments,
    Trainer,
    DataCollatorForTokenClassification,
    HfArgumentParser,
    EarlyStoppingCallback
)

from utils import compute_metrics, label_tokenized, chunk_dataset, LABEL2ID, ID2LABEL

logger = logging.getLogger(__name__)


@dataclass
class ModelArguments:
    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )

    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
    )

    early_stopping_patience: int = field(
        default=1,
    )

    early_stopping_threshold: int = field(
        default=1e-3,
    )


@dataclass
class DataTrainingArguments:
    train_dataset: str = field(
        default="bigcode/pseudo-labeled-python-data-pii-detection-filtered",
        metadata={"help": "The train dataset"}
    )

    dev_dataset: str = field(
        default="bigcode/pii-for-code-v2",
        metadata={"help": "The validation dataset"}
    )

    max_seq_length: int = field(
        default=512,
        metadata={
            "help": (
                "The maximum input sequence length after tokenization. Sequences longer "
                "than this will be chunked into pieces of this length."
            )
        },
    )


def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))

    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    model = AutoModelForTokenClassification.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        num_labels=len(ID2LABEL)
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        add_prefix_space=True)

    model.config.id2label = {str(i): label for i, label in enumerate(ID2LABEL)}
    model.config.label2id = LABEL2ID

    tokenizer.model_max_length = data_args.max_seq_length

    train_dataset = load_dataset(data_args.train_dataset, use_auth_token=True)['train']
    dev_dataset = load_dataset(data_args.dev_dataset, use_auth_token=True)['train']

    def tokenize_and_label(entry, tokenizer=tokenizer):
        inputs = tokenizer.encode_plus(entry['content'], return_offsets_mapping=True, add_special_tokens=False)
        entry.update(inputs)
        return label_tokenized(entry)

    dev_dataset = dev_dataset.map(lambda x: dict(pii=json.loads(x['pii'])))
    dev_dataset = dev_dataset.map(tokenize_and_label)

    train_dataset = train_dataset.map(lambda x: dict(pii=json.loads(x['pii'])))
    train_dataset = train_dataset.map(tokenize_and_label, num_proc=8)
    train_dataset = chunk_dataset(train_dataset, tokenizer)

    ner_dataset = DatasetDict(
        train=train_dataset,
        validation=chunk_dataset(dev_dataset, tokenizer),
    )

    trainer = Trainer(
        model,
        training_args,
        train_dataset=ner_dataset["train"],
        eval_dataset=ner_dataset["validation"],
        data_collator=DataCollatorForTokenClassification(tokenizer),
        tokenizer=tokenizer,
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=model_args.early_stopping_patience,
                                         early_stopping_threshold=model_args.early_stopping_threshold)]
    )

    trainer.train()