#!/usr/bin/env python

# -*- coding: utf-8 -*-

"""

Experiments on the BERT model and the different datasets (i.e. n2c2, DDI)

"""

# Base Dependencies

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

from copy import deepcopy

from pathlib import Path

from os.path import join as pjoin

# Package Dependencies

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

from .common import final_repetition

# Local Dependencies

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

from training.config import PLExperimentConfig, BaalExperimentConfig

from training.bert import BertTrainer

from utils import set_seed

# 3rd-Party Dependencies

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

from datasets import load_from_disk

# Constants

# ----------

from constants import (

    DDI_HF_TEST_PATH,

    DDI_HF_TRAIN_PATH,

    N2C2_HF_TRAIN_PATH,

    N2C2_HF_TEST_PATH,

    N2C2_REL_TYPES,

    EXP_RANDOM_SEEDS,

    BaalQueryStrategy

)

MODEL_NAME = "bert"

def bert_passive_learning_n2c2(init_repetition: int = 0, n_repetitions: int = 5, pairs: bool = False, logging: bool = True):

    config = PLExperimentConfig(

        max_epoch=25, batch_size=32, val_size=0.2, es_patience=3

    )

    for repetition in range(init_repetition, final_repetition(init_repetition, n_repetitions)):

        # set random seed

        random_seed: int = EXP_RANDOM_SEEDS[repetition]

        set_seed(random_seed)

        config.seed = random_seed

        for rel_type in N2C2_REL_TYPES:

            # load datasets

            train_dataset = load_from_disk(

                str(Path(pjoin(N2C2_HF_TRAIN_PATH, MODEL_NAME, rel_type)))

            )

            test_dataset = load_from_disk(

                str(Path(pjoin(N2C2_HF_TEST_PATH, MODEL_NAME, rel_type)))

            )

            # create trainer

            trainer = BertTrainer(

                dataset="n2c2",

                train_dataset=train_dataset,

                test_dataset=test_dataset,

                pairs=pairs,

                relation_type=rel_type,

            )

            # train passive learning

            trainer.train_passive_learning(config=config, logging=logging)

def bert_active_learning_n2c2(init_repetition: int = 0, n_repetitions: int = 5, pairs: bool = False, logging: bool = True):

    config = BaalExperimentConfig(max_epoch=10, batch_size=32)

    for repetition in range(init_repetition, final_repetition(init_repetition, n_repetitions)):

        # set random seed

        random_seed: int = EXP_RANDOM_SEEDS[repetition]

        set_seed(random_seed)

        config.seed = random_seed

        for rel_type in N2C2_REL_TYPES:

            # load datasets

            train_dataset = load_from_disk(

                str(Path(pjoin(N2C2_HF_TRAIN_PATH, MODEL_NAME, rel_type)))

            )

            test_dataset = load_from_disk(

                str(Path(pjoin(N2C2_HF_TEST_PATH, MODEL_NAME, rel_type)))

            )

            # create trainer

            trainer = BertTrainer(

                dataset="n2c2",

                train_dataset=train_dataset,

                test_dataset=test_dataset,

                pairs=pairs,

                relation_type=rel_type,

            )

            for query_strategy in BaalQueryStrategy:

                exp_config = deepcopy(config)

                trainer.train_active_learning(query_strategy, exp_config, logging=logging)

def bert_passive_learning_ddi(init_repetition: int = 0, n_repetitions: int = 5, pairs: bool = False, logging: bool = True):

    config = PLExperimentConfig(

        max_epoch=25, batch_size=32, val_size=0.2, es_patience=3

    )

    for repetition in range(init_repetition, final_repetition(init_repetition, n_repetitions)):

        # set random seed

        random_seed: int = EXP_RANDOM_SEEDS[repetition]

        set_seed(random_seed)

        config.seed = random_seed

        # load datasets

        train_dataset = load_from_disk(str(Path(pjoin(DDI_HF_TRAIN_PATH, MODEL_NAME))))

        test_dataset = load_from_disk(str(Path(pjoin(DDI_HF_TEST_PATH, MODEL_NAME))))

        # create trainer

        trainer = BertTrainer(

            dataset="ddi",

            train_dataset=train_dataset,

            test_dataset=test_dataset,

            pairs=pairs,

        )

        # train passive learning

        trainer.train_passive_learning(config=config, logging=logging)

def bert_active_learning_ddi(init_repetition: int = 0, n_repetitions: int = 5, pairs: bool = False, logging: bool = True):

    config = BaalExperimentConfig(max_epoch=15, batch_size=32,)

    for repetition in range(init_repetition, final_repetition(init_repetition, n_repetitions)):

        # set random seed

        random_seed: int = EXP_RANDOM_SEEDS[repetition]

        set_seed(random_seed)

        config.seed = random_seed

        # load datasets

        train_dataset = load_from_disk(str(Path(pjoin(DDI_HF_TRAIN_PATH, MODEL_NAME))))

        test_dataset = load_from_disk(str(Path(pjoin(DDI_HF_TEST_PATH, MODEL_NAME))))

        # create trainer

        trainer = BertTrainer(

            dataset="ddi",

            train_dataset=train_dataset,

            test_dataset=test_dataset,

            pairs=pairs,

        )

        for query_strategy in BaalQueryStrategy:

            exp_config = deepcopy(config)

            trainer.train_active_learning(query_strategy, exp_config, logging=logging)