# Base Dependencies

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

import functools

import numpy as np

import operator

import os

import random

import re

from glob import glob

from os.path import join as pjoin

from pathlib import Path

from typing import List, Any, Union

# Local Dependencies

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

from constants import N2C2_PATH, DDI_PATH, N2C2_ANNONYM_PATTERNS, DDI_ALL_TYPES

# 3rd-Party Dependencies

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

import torch

from torch import nn

from transformers import set_seed as transformers_set_seed

def set_seed(seed: int) -> None:

    """Sets the random seed for modules torch, numpy and random.

    Args:

        seed (int): random seed

    """

    transformers_set_seed(seed)

    torch.manual_seed(seed)

    np.random.seed(seed)

    random.seed(seed)

    if torch.cuda.is_available() and torch.cuda.device_count() > 0:

        torch.backends.cudnn.deterministic = True

        torch.cuda.manual_seed_all(seed)

def flatten(array: List[List[Any]]) -> List[Any]:

    """

    Flattens a nested 2D list. faster even with a very long array than

    [item for subarray in array for item in subarray] or newarray.extend().

    Args:

        array (List[List[Any]]): a nested list

    Returns:

        List[Any]: flattened list

    """

    return functools.reduce(operator.iconcat, array, [])

def write_list_to_file(output_path: Path, array: List[Any]) -> None:

    """

    Writes list of str to file in `output_path`.

    Args:

        output_path (Path): output file path

        array (List[Any]): list of strings

    """

    with output_path.open("w", encoding="utf-8") as opened_file:

        for entry in array:

            opened_file.write(f"{entry}\n")

def read_list_from_file(input_path: Path) -> List[str]:

    """

    Reads list of str from file in `input_path`.

    Args:

        input_path (Path): input file path

    Returns:

        List[str]: list of strings

    """

    if input_path is None:

        return []

    tokens = []

    for line in input_path.read_text(encoding="utf-8").splitlines():

        tokens.append(line.rstrip("\n"))

    return tokens

def make_dir(dirpath: str):

    """Creates a directory if it doesn't exist"""

    if not os.path.exists(dirpath):

        os.makedirs(dirpath)

def freeze_params(module: nn.Module) -> None:

    """

    Freezes the parameters of this module,

    i.e. do not update them during training

    Args:

        module (nn.Module): freeze parameters of this module

    """

    for _, p in module.named_parameters():

        p.requires_grad = False

def ddi_binary_relation(rel_type: Union[str, int]) -> int:

    """Converts a DDI's relation type into binary

    Args:

        rel_type (str): relation type

    Returns:

        int: 0 if the relation type is `"NO-REL"`, `"0"` or `0`,

        1 if the relation type is an string in `["EFFECT", "MECHANISM", "ADVISE", "INT"]` or is an integer `> 0`.

    """

    rt = rel_type

    if isinstance(rt, str):

        if rt in DDI_ALL_TYPES:

            rt = DDI_ALL_TYPES.index(rt)

        else:

            rt = int(rt)

    if rt == 0:

        return 0

    else:

        return 1

def doc_id_n2c2(filepath: str) -> str:

    """Extracts the document id of a n2c2 filepath"""

    return re.findall(r"\d{2,}", filepath)[-1]

def doc_id_ddi(filepath: str) -> str:

    """Extracts the document id of a ddi filepath"""

    file_name = filepath.split()[-1]

    return file_name[:-4]

def clean_text_ddi(text: str) -> str:

    """Cleans text of a text fragment from a ddi document

    Args:

        text (str): text fragment

    Returns:

        str: cleaned text fragment

    """

    # remove more than one space

    text = re.sub(r"[\s]+", " ", text)

    # include space after ;

    text = re.sub(r";", "; ", text)

    return text

def clean_text_n2c2(text: str) -> str:

    """Cleans text of a text fragment from a n2c2 document

    Args:

        text (str): text fragment

    Returns:

        str: cleaned text fragment

    """

    # remove head and tail spaces

    # text = text.strip()

    # remove newlines

    text = re.sub(r"\n", " ", text)

    # substitute annonymizations by their type

    for repl, pattern in N2C2_ANNONYM_PATTERNS.items():

        text = re.sub(pattern, repl, text)

    # remove not matching annonymizations

    text = re.sub(r"\[\*\*[^\*]+\*\*\]", "", text)

    # remove more than one space

    text = re.sub(r"[\s]+", " ", text)

    # replace two points by one

    text = re.sub(r"\.\.", ".", text)

    return text

def files_n2c2():

    """Loads the filepaths of the n2c2 dataset splits"""

    splits = {}

    for split in ["train", "test"]:

        files = glob(pjoin(N2C2_PATH, split, "*.txt"))

        splits[split] = list(map(lambda file: file[:-4], files))

    return splits

def files_ddi():

    """Loads the filepaths of the DDI corpus splits"""

    splits = {}

    for split in ["train", "test"]:

        if split == "train":

            splits["train"] = glob(pjoin(DDI_PATH, split, "DrugBank", "*.xml")) + glob(

                pjoin(DDI_PATH, split, "MedLine", "*.xml")

            )

        else:

            splits["test"] = glob(

                pjoin(DDI_PATH, split, "re", "DrugBank", "*.xml")

            ) + glob(pjoin(DDI_PATH, split, "re", "MedLine", "*.xml"))

    return splits