# Base Dependencies

# -----------------

import numpy as np

import re

import time

from copy import deepcopy

from functools import partial

from os.path import join

from pathlib import Path

from typing import Optional, Dict

# Package Dependencies

# --------------------

from .base import BaseTrainer

from .config import PLExperimentConfig, BaalExperimentConfig

from .utils import get_baal_query_strategy, tokenize, tokenize_pairs

# Local Dependencies

# ------------------

from extensions.baal import my_active_huggingface_dataset, MyActiveLearningLoop

from extensions.transformers import WeightedLossTrainer

from ml_models.bert import ClinicalBERT, ClinicalBERTTokenizer, ClinicalBERTConfig

# 3rd-Party Dependencies

# ----------------------

import neptune

from baal.transformers_trainer_wrapper import BaalTransformersTrainer

from baal.bayesian.dropout import patch_module

from torch.utils.data import Dataset

from transformers import (

    EarlyStoppingCallback,

    EvalPrediction,

    IntervalStrategy,

    TrainingArguments,

)

# Constants

# ---------

from constants import BaalQueryStrategy

from config import NEPTUNE_API_TOKEN, NEPTUNE_PROJECT

class BertTrainer(BaseTrainer):

    """Trainer for the BERT method"""

    def __init__(

        self,

        dataset: str,

        train_dataset: Dataset,

        test_dataset: Dataset,

        pairs: bool = False,

        relation_type: Optional[str] = None,

    ):

        """

        dataset (str): name of the dataset, e.g., "n2c2".

        train_dataset (Dataset): train split of the dataset.

        test_dataset (Dataset): test split of the dataset.

        relation_type (str, optional): relation type.

        Raises:

            ValueError: if the name dataset provided is not supported

        """

        super().__init__(dataset, train_dataset, test_dataset, relation_type)

        self.pairs = pairs

        # tokenizer

        self.tokenizer = ClinicalBERTTokenizer()

        # tokenize datasets

        if not pairs:

            self.train_dataset = tokenize(self.tokenizer, self.train_dataset)

            self.test_dataset = tokenize(self.tokenizer, self.test_dataset)

        else:

            self.train_dataset = tokenize_pairs(self.tokenizer, self.train_dataset)

            self.test_dataset = tokenize_pairs(self.tokenizer, self.test_dataset)

    @property

    def method_name(self) -> str:

        if self.pairs:

            name = "bert-pairs"

        else:

            name = "bert"

        return name

    @property

    def method_name_pretty(self) -> str:

        if self.pairs:

            name = "Paired Clinical BERT"

        else:

            name = "Clinical BERT"

        return name

    def _init_model(self, patch: bool = False) -> ClinicalBERT:

        config = ClinicalBERTConfig

        config.num_labels = self.num_classes

        model = ClinicalBERT(config=ClinicalBERTConfig)

        if patch:

            model = patch_module(model)

        return model

    def compute_metrics_transformer(self, eval_preds: EvalPrediction) -> Dict:

        """Computes metrics from a Transformer's prediction.

        Args:

            eval_preds (EvalPrediction): transformer's prediction

        Returns:

            Dict: precision, recall and F1-score

        """

        logits, labels = eval_preds

        predictions = np.argmax(logits, axis=-1)

        return self.compute_metrics(y_true=labels, y_pred=predictions)

    def train_passive_learning(

        self, config: PLExperimentConfig, verbose: bool = True, logging: bool = True

    ):

        """Trains the BiLSTM model using passive learning and early stopping

        Args:

            config (PLExperimentConfig): cofiguration

            verbose (bool): determines if information is printed during training. Daults to True.

            logging (bool): log the test metrics on Neptune. Defaults to True.

        """

        if logging:

            run = neptune.init_run(project=NEPTUNE_PROJECT, api_token=NEPTUNE_API_TOKEN)

        # setup

        train_val_split = self.train_dataset.train_test_split(

            test_size=config.val_size, stratify_by_column="label"

        )

        train_set = train_val_split["train"]

        val_set = train_val_split["test"]

        test_set = self.test_dataset

        model = self._init_model()

        training_args = TrainingArguments(

            output_dir=self.pl_checkpoint_path,  # output directory

            optim="adamw_torch",  # optimizer

            weight_decay=0.01,  # strength of weight decay

            learning_rate=5e-5,  # learning rate

            evaluation_strategy=IntervalStrategy.EPOCH,

            save_strategy=IntervalStrategy.EPOCH,

            num_train_epochs=config.max_epoch,

            per_device_train_batch_size=config.batch_size,

            per_device_eval_batch_size=config.batch_size,  # batch size for evaluation

            log_level="warning",  # logging level

            logging_dir=".logs/n2c2/bert/",  # directory for storing logs

            report_to="none",

            metric_for_best_model="f1",

            load_best_model_at_end=True,

        )

        trainer = WeightedLossTrainer(

            model=model,

            args=training_args,

            seed=config.seed,

            train_dataset=train_set,

            eval_dataset=val_set,

            tokenizer=self.tokenizer,

            compute_metrics=self.compute_metrics_transformer,

            callbacks=[

                EarlyStoppingCallback(early_stopping_patience=config.es_patience)

            ],

        )

        labels = train_set["label"].numpy()

        trainer.class_weights = self.compute_class_weights(labels)

        # print info

        if verbose:

            self.print_info_passive_learning()

        # train model

        trainer.train()

        eval_loss_values = trainer.eval_loss

        train_loss_values = trainer.training_loss

        # evaluate model on test set

        test_metrics = trainer.evaluate(test_set)

        if verbose:

            self.print_test_metrics(test_metrics)

        # log to Neptune

        if logging:

            run["method"] = self.method_name

            run["dataset"] = self.dataset

            run["relation"] = self.relation_type

            run["strategy"] = "passive learning"

            run["config"] = config.__dict__

            run["epoch"] = len(eval_loss_values)

            for loss in train_loss_values:

                run["loss/train"].append(loss)

            for loss in eval_loss_values:

                run["loss/val"].append(loss)

            for key, value in test_metrics.items():

                key2 = re.sub(r"eval_", "", key)

                run["test/" + key2] = value

            run.stop()

        return model

    def train_active_learning(

        self,

        query_strategy: BaalQueryStrategy,

        config: BaalExperimentConfig,

        verbose: bool = True,

        save_models: bool = False,

        logging: bool = True,

    ):

        """Trains the BiLSTM model using active learning

        Args:

            query_strategy (str): name of the query strategy to be used in the experiment.

            config (BaalExperimentConfig): experiment configuration.

            verbose (bool): determines if information is printed during trainig or not. Defaults to True.s

            logging (bool): log the test metrics on Neptune. Defaults to True.

        """

        if logging:

            run = neptune.init_run(project=NEPTUNE_PROJECT, api_token=NEPTUNE_API_TOKEN)

            run["model"] = self.method_name

            run["dataset"] = self.dataset

            run["relation"] = self.relation_type

            run["strategy"] = query_strategy.value

            run["bayesian"] = config.all_bayesian or (

                query_strategy == BaalQueryStrategy.BATCH_BALD

            )

            run["params"] = config.__dict__

        # setup quering 

        INIT_QUERY_SIZE = self.compute_init_q_size(config)

        QUERY_SIZE = self.compute_q_size(config)

        AL_STEPS = self.compute_al_steps(config)

        f_query_strategy = get_baal_query_strategy(

            name=query_strategy.value,

            shuffle_prop=config.shuffle_prop,

            query_size=QUERY_SIZE,

        )

        # setup model

        PATCH = config.all_bayesian or (query_strategy == BaalQueryStrategy.BATCH_BALD)

        if not PATCH:

            config.iterations = 1     

        # setup active set

        active_set = my_active_huggingface_dataset(self.train_dataset)

        active_set.can_label = False

        active_set.label_randomly(INIT_QUERY_SIZE)

        # print info

        if verbose:

            self.print_info_active_learning(

                q_strategy=query_strategy.value,

                pool_size=self.n_instances,

                init_q_size=INIT_QUERY_SIZE,

                q_size=QUERY_SIZE,

            )

        training_args = TrainingArguments(

            output_dir=self.al_checkpoint_path,

            optim="adamw_torch",  # optimizer

            weight_decay=0.01,  # strength of weight decay

            learning_rate=5e-5,  # learning rate

            num_train_epochs=config.max_epoch,

            per_device_train_batch_size=config.batch_size,

            per_device_eval_batch_size=config.batch_size,  # batch size for evaluation

            log_level="warning",  # logging level

            logging_dir=".logs/n2c2/bert/",  # directory for storing logs

            report_to="none",

        )

        # create the trainer through Baal Wrapper

        baal_trainer = BaalTransformersTrainer(

            model_init=partial(self._init_model, PATCH),

            seed=config.seed,

            args=training_args,

            train_dataset=active_set,

            tokenizer=None,

            compute_metrics=self.compute_metrics_transformer,

        )

        # create Active Learning loop

        active_loop = MyActiveLearningLoop(

            dataset=active_set,

            get_probabilities=baal_trainer.predict_on_dataset,

            heuristic=f_query_strategy,

            query_size=QUERY_SIZE,

            iterations=config.iterations,

            max_sample=config.max_sample,

        )

        init_weights = deepcopy(baal_trainer.model.state_dict())

        # Active Learning loop

        for step in range(AL_STEPS):

            init_step_time = time.time()

            # reset the model to the initial state

            baal_trainer.model.load_state_dict(init_weights)

            # train model on current active set

            init_train_time = time.time()

            baal_trainer.train()

            train_time = time.time() - init_train_time

            if save_models:

                # save model

                path = Path(join(self.al_checkpoint_path, "model_{}.ck".format(step)))

                baal_trainer.model.save_pretrained(path)

            # evaluate model on test set

            metrics = baal_trainer.evaluate(self.test_dataset)

            metrics["dataset_size"] = active_set.n_labelled

            # print step metrics

            if verbose:

                self.print_al_iteration_metrics(step + 1, metrics)

            # query new instances

            init_query_time = time.time()

            should_continue = active_loop.step()

            query_time = time.time() - init_query_time

            step_time = time.time() - init_step_time

            if logging:

                run["times/step_time"].append(step_time)

                run["times/train_time"].append(train_time)

                run["times/query_time"].append(query_time)

                run["annotation/instance_ann"].append(

                    active_set.n_labelled / self.n_instances

                )

                run["annotation/token_ann"].append(

                    active_set.n_labelled_tokens / self.n_tokens

                )

                run["annotation/char_ann"].append(

                    active_set.n_labelled_chars / self.n_characters

                )

                for key, value in metrics.items():

                    f_key = key.replace("test_", "test/").replace("train_", "train/")

                    run[f_key].append(value)

            if not should_continue:

                break

            # We reset the model weights to relearn from the new train set.

            baal_trainer.load_state_dict(init_weights)

            baal_trainer.lr_scheduler = None

        # log to Neptune

        if logging:

            for step_acc in active_loop.step_acc:

                run["train/step_acc"].append(step_acc)

            for step_score in active_loop.step_score:

                run["train/step_score"].append(step_score)

            run.stop()