# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/utils/utils.common_utils.ipynb (unless otherwise specified).

__all__ = ['wget_download', 'download_url', 'extract_tar', 'extract_zip', 'extract_bz2', 'extract_gz',
           'print_result_as_table', 'list_files', 'seed_everything', 'map_column', 'get_context', 'pad_arr', 'pad_list',
           'mask_list', 'mask_last_elements_list', 'masked_accuracy', 'masked_ce']

# Cell
import sys
import os
import ssl
import os.path as osp
from six.moves import urllib
import errno
import tarfile
import zipfile
import bz2
import gzip
import pandas as pd
from tabulate import tabulate

# Internal Cell
def makedirs(path):
    try:
        os.makedirs(osp.expanduser(osp.normpath(path)))
    except OSError as e:
        if e.errno != errno.EEXIST and osp.isdir(path):
            raise e

# Cell
def wget_download(url, savepath):
    import wget
    wget.download(url, str(savepath))

# Cell
def download_url(url: str, folder: str, log: bool = True):
    r"""Downloads the content of an URL to a specific folder.
    Args:
        url (string): The url.
        folder (string): The folder.
        log (bool, optional): If :obj:`False`, will not print anything to the
            console. (default: :obj:`True`)
    """

    filename = url.rpartition('/')[2]
    filename = filename if filename[0] == '?' else filename.split('?')[0]
    path = osp.join(folder, filename)

    if osp.exists(path):  # pragma: no cover
        if log:
            print(f'Using existing file {filename}', file=sys.stderr)
        return path

    if log:
        print(f'Downloading {url}', file=sys.stderr)

    makedirs(folder)

    context = ssl._create_unverified_context()
    data = urllib.request.urlopen(url, context=context)

    with open(path, 'wb') as f:
        f.write(data.read())

    return path

# Internal Cell
def maybe_log(path, log=True):
    if log:
        print(f'Extracting {path}', file=sys.stderr)

# Cell
def extract_tar(path: str, folder: str, mode: str = 'r:gz', log: bool = True):
    r"""Extracts a tar archive to a specific folder.
    Args:
        path (string): The path to the tar archive.
        folder (string): The folder.
        mode (string, optional): The compression mode. (default: :obj:`"r:gz"`)
        log (bool, optional): If :obj:`False`, will not print anything to the
            console. (default: :obj:`True`)
    """
    maybe_log(path, log)
    with tarfile.open(path, mode) as f:
        f.extractall(folder)

# Cell
def extract_zip(path: str, folder: str, log: bool = True):
    r"""Extracts a zip archive to a specific folder.
    Args:
        path (string): The path to the tar archive.
        folder (string): The folder.
        log (bool, optional): If :obj:`False`, will not print anything to the
            console. (default: :obj:`True`)
    """
    maybe_log(path, log)
    with zipfile.ZipFile(path, 'r') as f:
        f.extractall(folder)

# Cell
def extract_bz2(path: str, folder: str, log: bool = True):
    r"""Extracts a bz2 archive to a specific folder.
    Args:
        path (string): The path to the tar archive.
        folder (string): The folder.
        log (bool, optional): If :obj:`False`, will not print anything to the
            console. (default: :obj:`True`)
    """
    maybe_log(path, log)
    path = osp.abspath(path)
    with bz2.open(path, 'r') as r:
        with open(osp.join(folder, '.'.join(path.split('.')[:-1])), 'wb') as w:
            w.write(r.read())

# Cell
def extract_gz(path: str, folder: str, log: bool = True):
    r"""Extracts a gz archive to a specific folder.
    Args:
        path (string): The path to the tar archive.
        folder (string): The folder.
        log (bool, optional): If :obj:`False`, will not print anything to the
            console. (default: :obj:`True`)
    """
    maybe_log(path, log)
    path = osp.abspath(path)
    with gzip.open(path, 'r') as r:
        with open(osp.join(folder, '.'.join(path.split('.')[:-1])), 'wb') as w:
            w.write(r.read())

# Cell
def print_result_as_table(results, tag=None):
    """Print results as a table."""
    eval_infos = set()
    for result in results:
        eval_infos.update(result.keys())
    eval_infos = list(eval_infos)
    print("-" * 80)
    if tag is not None:
        print(tag)
    for result in results:
        for eval_info in eval_infos:
            if eval_info not in result:
                result[eval_info] = "--"
    df = pd.DataFrame(results)
    df = df.set_index("model")
    df = df.T
    print(tabulate(df, headers=df.columns, tablefmt="psql"))
    print("-" * 80)

# Cell
def list_files(startpath):
    """
    Util function to print the nested structure of a directory
    """
    for root, dirs, files in os.walk(startpath):
        level = root.replace(startpath, "").count(os.sep)
        indent = " " * 4 * (level)
        print("{}{}/".format(indent, os.path.basename(root)))
        subindent = " " * 4 * (level + 1)
        for f in files:
            print("{}{}".format(subindent, f))

# Cell
import sys
import os
import numpy as np
import torch
import random

# Cell
def seed_everything(seed=40):
    """sets the random seed to establish deterministic behaviors

    Args:
        seed (int): the random seed integer
    """
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

# Cell
import random
import numpy as np
import pandas as pd

import torch
from torch.nn import functional as F

# Cell
def map_column(df: pd.DataFrame, col_name: str):
    """Maps column values to integers.
    """
    values = sorted(list(df[col_name].unique()))
    mapping = {k: i + 2 for i, k in enumerate(values)}
    inverse_mapping = {v: k for k, v in mapping.items()}
    df[col_name + "_mapped"] = df[col_name].map(mapping)
    return df, mapping, inverse_mapping

def get_context(df: pd.DataFrame, split: str, context_size: int = 120, val_context_size: int = 5, seed: int = 42):
    """Create a training / validation samples.
    """
    random.seed(seed)
    if split == "train":
        end_index = random.randint(10, df.shape[0] - val_context_size)
    elif split in ["val", "test"]:
        end_index = df.shape[0]
    else:
        raise ValueError
    start_index = max(0, end_index - context_size)
    context = df[start_index:end_index]
    return context

## Padding

def pad_arr(arr: np.ndarray, expected_size: int = 30):
    """Pad top of array when there is not enough history.
    """
    arr = np.pad(arr, [(expected_size - arr.shape[0], 0), (0, 0)], mode="edge")
    return arr

def pad_list(list_integers, history_size: int, pad_val: int = 0, mode="left"):
    """Pad list from left or right
    """
    if len(list_integers) < history_size:
        if mode == "left":
            list_integers = [pad_val] * (history_size - len(list_integers)) + list_integers
        else:
            list_integers = list_integers + [pad_val] * (history_size - len(list_integers))
    return list_integers

# Masking

def mask_list(l1, p=0.8, mask=1, seed=42):
    random.seed(seed)
    l1 = [a if random.random() < p else mask for a in l1]
    return l1

def mask_last_elements_list(l1, val_context_size: int = 5, seed=42):
    l1 = l1[:-val_context_size] + mask_list(l1[-val_context_size:], p=0.5, seed=seed)
    return l1

def masked_accuracy(y_pred: torch.Tensor, y_true: torch.Tensor, mask: torch.Tensor):
    _, predicted = torch.max(y_pred, 1)
    y_true = torch.masked_select(y_true, mask)
    predicted = torch.masked_select(predicted, mask)
    acc = (y_true == predicted).double().mean()
    return acc

def masked_ce(y_pred, y_true, mask):
    loss = F.cross_entropy(y_pred, y_true, reduction="none")
    loss = loss * mask
    return loss.sum() / (mask.sum() + 1e-8)