# Base Dependencies

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

import numpy as np

import time

from copy import deepcopy

from functools import partial

from tqdm import tqdm

from typing import Dict, Optional

from pathlib import Path

from os.path import join

# Package Dependencies

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

from .base import BaseTrainer

from .config import PLExperimentConfig, BaalExperimentConfig

from .early_stopping import EarlyStopping

from .utils import get_baal_query_strategy

# Local Dependencies

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

from extensions.baal import (

    MyModelWrapperBilstm,

    MyActiveLearningDatasetBilstm,

    MyActiveLearningLoop,

)

from extensions.torchmetrics import (

    DetectionF1Score,

    DetectionPrecision,

    DetectionRecall,

)

from ml_models.bilstm import (

    HasanModel,

    EmbeddingConfig,

    LSTMConfig,

    RDEmbeddingConfig,

)

from re_datasets.bilstm_utils import pad_and_sort_batch, custom_collate

from vocabulary import Vocabulary, read_list_from_file

# 3rd-Party Dependencies

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

import neptune

import torch

from baal.bayesian.dropout import patch_module

from datasets import Dataset

from torch.optim import Adam

from torch.nn import CrossEntropyLoss, Module

from torch.utils.data import DataLoader

from torch.utils.data.sampler import BatchSampler, RandomSampler

from torchmetrics import Accuracy

from torchmetrics.classification import F1Score, Precision, Recall

# Constants

# ---------

from constants import (

    N2C2_VOCAB_PATH,

    DDI_VOCAB_PATH,

    N2C2_IOB_TAGS,

    DDI_IOB_TAGS,

    N2C2_RD_MAX,

    DDI_RD_MAX,

    RD_EMB_DIM,

    IOB_EMB_DIM,

    BIOWV_EMB_DIM,

    POS_EMB_DIM,

    DEP_EMB_DIM,

    BIOWORD2VEC_PATH,

    U_POS_TAGS,

    DEP_TAGS,

    BaalQueryStrategy,

)

from config import NEPTUNE_API_TOKEN, NEPTUNE_PROJECT

class BilstmTrainer(BaseTrainer):

    """Trainer for BiLSTM method."""

    def __init__(

        self,

        dataset: str,

        train_dataset: Dataset,

        test_dataset: Dataset,

        relation_type: Optional[str] = None,

    ):

        """

        Args:

            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)

        # vocabulary

        self.vocab = self._init_vocab()

        # transform datasets

        self.transform = partial(

            pad_and_sort_batch, padding_idx=self.vocab.pad_index, rd_max=self.RD_MAX

        )

    @property

    def method_name(self) -> str:

        return "bilstm"

    @property

    def method_name_pretty(self) -> str:

        return "BiLSTM"

    @property

    def task(self) -> str:

        if self.dataset == "n2c2":

            task = "binary"

        else:

            task = "multiclass"

        return task

    @property

    def model_class(self) -> str:

        return HasanModel

    @property

    def RD_MAX(self) -> str:

        if self.dataset == "n2c2":

            rd_max = N2C2_RD_MAX

        else:

            rd_max = DDI_RD_MAX

        return rd_max

    @property

    def IOB_TAGS(self) -> str:

        if self.dataset == "n2c2":

            iob_tags = N2C2_IOB_TAGS

        else:

            iob_tags = DDI_IOB_TAGS

        return iob_tags

    def _init_optimizer(self, model: Module):

        return Adam(model.parameters(), lr=0.0001)

    def _init_vocab(self):

        """Loads the vocabulary of the dataset"""

        if self.dataset == "n2c2":

            vocab_path = N2C2_VOCAB_PATH

        else:

            vocab_path = DDI_VOCAB_PATH

        return Vocabulary(read_list_from_file(vocab_path))

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

        """Builds the BiLSTM model setting the right configuration for the chosen dataset"""

        # word embedding configuration

        biowv_config = EmbeddingConfig(

            embedding_dim=BIOWV_EMB_DIM,

            vocab_size=len(self.vocab),

            emb_path=BIOWORD2VEC_PATH,

            freeze=True,

            padding_idx=self.vocab.pad_index,

        )

        # relative-distance embedding configuration

        rd_config = RDEmbeddingConfig(

            input_dim=self.RD_MAX, embedding_dim=RD_EMB_DIM, freeze=False

        )

        # IOB embedding configuration

        iob_config = EmbeddingConfig(

            embedding_dim=IOB_EMB_DIM, vocab_size=(len(self.IOB_TAGS) + 1), freeze=False

        )

        # Part-of-Speach tag embedding configuration

        pos_config = EmbeddingConfig(

            embedding_dim=POS_EMB_DIM, vocab_size=(len(U_POS_TAGS) + 1), freeze=False

        )

        dep_config = EmbeddingConfig(

            embedding_dim=DEP_EMB_DIM, vocab_size=(len(DEP_TAGS) + 1), freeze=False

        )

        # BiLSTM configuration

        lstm_config = LSTMConfig(

            emb_size=(

                BIOWV_EMB_DIM + 2 * RD_EMB_DIM + POS_EMB_DIM + DEP_EMB_DIM + IOB_EMB_DIM

            )

        )

        model = self.model_class(

            vocab=self.vocab,

            lstm_config=lstm_config,

            bioword2vec_config=biowv_config,

            rd_config=rd_config,

            pos_config=pos_config,

            dep_config=dep_config,

            iob_config=iob_config,

            num_classes=self.num_classes,

        )

        if patch:

            model = patch_module(model)

        return model

    def _reset_trainer(self):

        self.train_dataset.reset_format()

        self.test_dataset.reset_format()

    def create_dataloader(self, dataset: Dataset, batch_size: int = 6) -> DataLoader:

        """Creates a dataloader from a dataset with the adequate configuration

        Args:

            dataset (Dataset): dataset to load

        Returns:

            DataLoader: dataloader for the given dataset

        """

        dataset.set_transform(self.transform)

        # create dataloader

        sampler = BatchSampler(

            RandomSampler(dataset), batch_size=batch_size, drop_last=False

        )

        dataloader = DataLoader(dataset, sampler=sampler, collate_fn=custom_collate)

        return dataloader

    def eval_model(

        self,

        model: Module,

        dataloader: DataLoader,

        criterion: Module,

    ) -> Dict[str, float]:

        """Evaluates the current model on the dev or test set

        Args:

            model (Module): model to use for evaluation.

            dataloader (DataLoader): dataloader of evaluation dataset

        Returns:

            Dict: metrics including loss (`loss`), precision (`p`), recall (`r`) and F1-score (`f1`)

        """

        y_true = np.array([], dtype=np.int8)

        y_pred = np.array([], dtype=np.int8)

        val_loss = 0.0

        with torch.no_grad():

            for inputs, labels in dataloader:

                # send (inputs, labels) to device

                labels = labels.to(self.device)

                for key, value in inputs.items():

                    inputs[key] = value.to(self.device)

                # calculate outputs

                outputs = model(inputs)

                loss = criterion(outputs, labels)

                val_loss += len(inputs) * loss.item()

                # calculate predictions

                _, predicted = torch.max(outputs.data, 1)

                # store labels and predictions

                y_true = np.append(y_true, labels.cpu().detach().numpy())

                y_pred = np.append(y_pred, predicted.cpu().detach().numpy())

        metrics = self.compute_metrics(y_true, y_pred)

        metrics["loss"] = val_loss / len(dataloader)

        return metrics

    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.

        """

        self._reset_trainer()

        # setup

        train_val_split = self.train_dataset.train_test_split(

            test_size=config.val_size, stratify_by_column="label"

        )

        labels = np.array(train_val_split["train"]["label"])

        train_dataloader = self.create_dataloader(

            train_val_split["train"], batch_size=config.batch_size

        )

        val_dataloader = self.create_dataloader(

            train_val_split["test"], batch_size=config.batch_size

        )

        test_dataloader = self.create_dataloader(

            self.test_dataset, batch_size=config.batch_size

        )

        if logging:

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

        model = self._init_model()

        model = model.to(self.device)

        criterion = CrossEntropyLoss(weight=self.compute_class_weights(labels))

        optimizer = self._init_optimizer(model)

        # print info

        if verbose:

            self.print_info_passive_learning()

        # early stopper

        ES = EarlyStopping(

            patience=config.es_patience,

            verbose=True,

            path=Path(join(self.pl_checkpoint_path, "best_model.pt")),

        )

        # training loop

        for epoch in range(config.max_epoch):

            running_loss = 0.0

            for i, (inputs, labels) in tqdm(enumerate(train_dataloader, 0)):

                # get the inputs; data is a list of [inputs, labels]

                labels = labels.to(self.device)

                for key, value in inputs.items():

                    inputs[key] = value.to(self.device)

                # zero the parameter gradients

                optimizer.zero_grad()

                # forward + backward + optimize

                outputs = model(inputs)

                loss = criterion(outputs, labels)

                loss.backward()

                optimizer.step()

                # print statistics

                running_loss += loss.item()

            # evaluate model on validation set

            val_metrics = self.eval_model(model, val_dataloader, criterion)

            train_loss = running_loss / len(train_dataloader)

            val_loss = val_metrics["loss"]

            running_loss = 0.0

            if logging:

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

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

                for key, value in val_metrics.items():

                    if key != "loss":

                        run[f"val/{key}"].append(value)

            if verbose:

                self.print_val_metrics(epoch + 1, val_metrics)

            # check early stopping

            ES(val_loss, model)

            if ES.early_stop:

                break

        # load best model

        model.load_state_dict(

            torch.load(Path(join(self.pl_checkpoint_path, "best_model.pt")))

        )

        # evaluate model on test dataset

        test_metrics = self.eval_model(model, test_dataloader, criterion)

        if verbose:

            self.print_test_metrics(test_metrics)

        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["epochs"] = epoch

            for key, value in test_metrics.items():

                run["test/" + key] = value

            run.stop()

        return model

    def set_al_metrics(self, baal_model: MyModelWrapperBilstm):

        """

        Configures the metrics that are to be computed during the active learning experiment

        Args:

            baal_model (MyModelWrapperBilstm): model wrapper

        """

        # accuracy

        baal_model.add_metric(

            name="acc",

            initializer=lambda: Accuracy(task=self.task, average="micro").to(

                self.device

            ),

        )

        if self.dataset == "n2c2":

            f1 = F1Score(num_classes=self.num_classes, ignore_index=0).to(self.device)

            p = Precision(num_classes=self.num_classes, ignore_index=0).to(self.device)

            r = Recall(num_classes=self.num_classes, ignore_index=0).to(self.device)

            baal_model.add_metric(name="f1", initializer=lambda: f1)

            baal_model.add_metric(name="p", initializer=lambda: p)

            baal_model.add_metric(name="r", initializer=lambda: r)

        else:  # self.dataset == "ddi":

            # detection + classification metrics

            cla_f1_micro = F1Score(

                num_classes=self.num_classes, average="micro", ignore_index=0

            ).to(self.device)

            cla_p_micro = Precision(

                num_classes=self.num_classes, average="micro", ignore_index=0

            ).to(self.device)

            cla_r_micro = Recall(

                num_classes=self.num_classes, average="micro", ignore_index=0

            ).to(self.device)

            cla_f1_macro = F1Score(

                num_classes=self.num_classes, average="macro", ignore_index=0

            ).to(self.device)

            cla_p_macro = Precision(

                num_classes=self.num_classes, average="macro", ignore_index=0

            ).to(self.device)

            cla_r_macro = Recall(

                num_classes=self.num_classes, average="macro", ignore_index=0

            ).to(self.device)

            baal_model.add_metric(name="micro_f1", initializer=lambda: cla_f1_micro)

            baal_model.add_metric(name="micro_p", initializer=lambda: cla_p_micro)

            baal_model.add_metric(name="micro_r", initializer=lambda: cla_r_micro)

            baal_model.add_metric(name="macro_f1", initializer=lambda: cla_f1_macro)

            baal_model.add_metric(name="macro_p", initializer=lambda: cla_p_macro)

            baal_model.add_metric(name="macro_r", initializer=lambda: cla_r_macro)

            # detection metrics

            detect_f1 = DetectionF1Score().to(self.device)

            detect_p = DetectionPrecision().to(self.device)

            detect_r = DetectionRecall().to(self.device)

            baal_model.add_metric(name="detect_f1", initializer=lambda: detect_f1)

            baal_model.add_metric(name="detect_p", initializer=lambda: detect_p)

            baal_model.add_metric(name="detect_r", initializer=lambda: detect_r)

            # per class metrics

            per_class_f1 = F1Score(num_classes=self.num_classes, average="none").to(

                self.device

            )

            per_class_p = Precision(num_classes=self.num_classes, average="none").to(

                self.device

            )

            per_class_r = Recall(num_classes=self.num_classes, average="none").to(

                self.device

            )

            baal_model.add_metric(name="class_f1", initializer=lambda: per_class_f1)

            baal_model.add_metric(name="class_p", initializer=lambda: per_class_p)

            baal_model.add_metric(name="class_r", initializer=lambda: per_class_r)

        return baal_model

    def train_active_learning(

        self,

        query_strategy: BaalQueryStrategy,

        config: BaalExperimentConfig,

        verbose: bool = True,

        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.

        """

        self._reset_trainer()

        if logging:

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

        # setup querying

        INIT_QUERY_SIZE = self.compute_init_q_size(config)

        QUERY_SIZE = self.compute_q_size(config)

        AL_STEPS = 2 # 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,

        )   

        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,

            )

        # setup active set

        self.train_dataset.set_transform(self.transform)

        self.test_dataset.set_transform(self.transform)

        active_set = MyActiveLearningDatasetBilstm(self.train_dataset)

        active_set.can_label = False

        active_set.label_randomly(INIT_QUERY_SIZE)

        # setup model

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

        if not PATCH: 

            config.iterations = 1

        model = self._init_model(PATCH)

        model = model.to(self.device)

        criterion = CrossEntropyLoss(self.compute_class_weights(active_set.labels))

        optimizer = self._init_optimizer(model)

        baal_model = MyModelWrapperBilstm(

            model,

            criterion,

            replicate_in_memory=False,

            min_train_passes=config.min_train_passes,

        )

        baal_model = self.set_al_metrics(baal_model)

        # active loop

        active_loop = MyActiveLearningLoop(

            dataset=active_set,

            get_probabilities=baal_model.predict_on_dataset,

            heuristic=f_query_strategy,

            query_size=QUERY_SIZE,

            batch_size=config.batch_size,

            iterations=config.iterations,

            use_cuda=self.use_cuda,

            verbose=False,

            workers=2,

            collate_fn=custom_collate,

        )

        # We will reset the weights at each active learning step so we make a copy.

        init_weights = deepcopy(baal_model.state_dict())

        if logging:

            run["model"] = self.method_name

            run["dataset"] = self.dataset

            run["relation"] = self.relation_type

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

                query_strategy == BaalQueryStrategy.BATCH_BALD

            )

            run["strategy"] = query_strategy.value

            run["config"] = config.__dict__

            run["annotation/intance_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

            )

        step_acc = []

        # Active learning loop

        for step in tqdm(range(AL_STEPS)):

            init_step_time = time.time()

            # Load the initial weights.

            baal_model.load_state_dict(init_weights)

            # Train the model on the currently labelled dataset.

            init_train_time = time.time()

            _ = baal_model.train_on_dataset(

                dataset=active_set,

                optimizer=optimizer,

                batch_size=config.batch_size,

                use_cuda=self.use_cuda,

                epoch=config.max_epoch,

                collate_fn=custom_collate,

            )

            train_time = time.time() - init_train_time

            # test the model on the test set.

            baal_model.test_on_dataset(

                dataset=self.test_dataset,

                batch_size=config.batch_size,

                use_cuda=self.use_cuda,

                average_predictions=config.iterations,

                collate_fn=custom_collate,

            )

            if verbose:

                self.print_al_iteration_metrics(step + 1, baal_model.get_metrics())

            # query new instances to be labelled

            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/intance_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

                )

            if not should_continue:

                break

            # adjust class weights

            baal_model.criterion = CrossEntropyLoss(

                self.compute_class_weights(active_set.labels)

            )

        # end of active learning loop

        if logging:

            for metrics in baal_model.active_learning_metrics.values():

                for key, value in metrics.items():

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

                    if "class" in key:

                        for i, class_value in enumerate(value):

                            run[f_key + "_" + str(i)].append(class_value)

                    else:

                        run[f_key].append(value)

            run["train/step_acc"].extend(active_loop.step_acc)

            run["train/step_score"].extend(active_loop.step_score)

            run.stop()