# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/07_data.ipynb (unless otherwise specified).

__all__ = ['probabilistic_hash_item', 'probabilistic_hash_tensor', 'BotoUploader']


# Cell
import boto3
from collections import deque
from functools import partial
import mmh3
from multiprocessing import Pool
import numpy as np
import os
import pandas as pd
import torch
from torch.utils.data import Dataset
from tqdm.auto import tqdm
import warnings

from htools import auto_repr, save, load


# Cell
def probabilistic_hash_item(x, n_buckets, mode=int, n_hashes=3):
    """Slightly hacky way to probabilistically hash an integer by
    first converting it to a string.

    Parameters
    ----------
    x: int
        The integer or string to hash.
    n_buckets: int
        The number of buckets that items will be mapped to. Typically
        this would occur outside the hashing function, but since
        the intended use case is so narrow here it makes sense to me
        to include it here.
    mode: type
        The type of input you want to hash. This is user-provided to prevent
        accidents where we pass in a different item than intended and hash
        the wrong thing. One of (int, str). When using this inside a
        BloomEmbedding layer, this must be `int` because there are no
        string tensors. When used inside a dataset or as a one-time
        pre-processing step, you can choose either as long as you
        pass in the appropriate inputs.
    n_hashes: int
        The number of times to hash x, each time with a different seed.

    Returns
    -------
    list[int]: A list of integers with length `n_hashes`, where each integer
        is in [0, n_buckets).
    """
    # Check type to ensure we don't accidentally hash Tensor(5) instead of 5.
    assert isinstance(x, mode), f'Input `x` must have type {mode}.'
    return [mmh3.hash(str(x), i, signed=False) % n_buckets
            for i in range(n_hashes)]


# Cell
def probabilistic_hash_tensor(x_r2, n_buckets, n_hashes=3, pad_idx=0):
    """Hash a rank 2 LongTensor.

    Parameters
    ----------
    x_r2: torch.LongTensor
        Rank 2 tensor of integers. Shape: (bs, seq_len)
    n_buckets: int
        Number of buckets to hash items into (i.e. the number of
        rows in the embedding matrix). Typically a moderately large
        prime number, like 251 or 997.
    n_hashes: int
        Number of hashes to take for each input index. This determines
        the number of rows of the embedding matrix that will be summed
        to get the representation for each word. Typically 2-5.
    pad_idx: int or None
        If you want to pad sequences with vectors of zeros, pass in an
        integer (same as the `padding_idx` argument to nn.Embedding).
        If None, no padding index will be used. The sequences must be
        padded before passing them into this function.

    Returns
    -------
    torch.LongTensor: Tensor of indices where each row corresponds
        to one of the input indices. Shape: (bs, seq_len, n_hashes)
    """
    return torch.tensor(
        [[probabilistic_hash_item(x.item(), n_buckets, int, n_hashes)
          if x != pad_idx else [pad_idx]*n_hashes for x in row]
         for row in x_r2]
    )


# Cell
class BotoUploader:
    """Uploads files to S3. Built as a public alternative to Accio. Note to
    self: the interfaces are not identical so be careful to know which you're
    using.
    """

    def __init__(self, bucket, verbose=True):
        """
        Parameters
        ----------
        bucket: str
            Name of s3 bucket to upload to. For a single project, I generally
            stick to a single bucket so we can usually keep this fixed. We can
            always change the attribute later if necessary.
        verbose: bool
            If True, print message when downloading each file.
        """
        self.s3 = boto3.resource('s3')
        self.bucket = bucket
        self.verbose = verbose

    def upload_file(self, path, s3_dir='', retain_tree=True):
        """Upload a single local file. By default, its path in S3 will be the
        same as its local path. Usually, this means your S3 bucket will have a
        single directory called "data" which corresponds exactly to your local
        "data" directory.

        Parameters
        ----------
        path: str or Path
            Local path to the file to upload. This must be relative to the
            project root directory (e.g. "data/models/v1/model.gz").
        s3_dir: str
            If provided, this will pre prepended to each path: {s3_dir}/{path}.
            Otherwise, s3 paths will be the same as local paths.
        retain_tree: bool
            If True, the local file structure will be retained. Otherwise,
            only the base name is kept. All four combinations of retain_tree
            (True/False) and s3_dir (empty/non-empty) are supported.
        """
        path = str(path)
        s3_path = self._convert_local_path(path, s3_dir, retain_tree)
        if self.verbose: print(f'Uploading {path} -> {s3_path}.')
        self.s3.meta.client.upload_file(path, self.bucket, s3_path)

    def upload_files(self, paths, s3_dir='', retain_tree=True):
        """Upload multiple files. Currently does not support parallelized
        uploads: need to figure out a workaround because self.s3 cannot be
        pickled.

        Parameters
        ----------
        paths: Iterable[str or Path]
            Sequence of file paths to upload.
        s3_dir: str
            If provided, this will pre prepended to each path: {s3_dir}/{path}.
            Otherwise, s3 paths will be the same as local paths.
        retain_tree: bool
            If True, the local file structure will be retained. Otherwise,
            only the base name is kept. All four combinations of retain_tree
            (True/False) and s3_dir (empty/non-empty) are supported.
        """
        for p in paths:
            self.upload_file(p, s3_dir=s3_dir, retain_tree=retain_tree)

    def upload_folder(self, dirname, s3_dir, retain_tree=True, recurse=True,
                      keep_fn=None):
        """Upload all files in a directory.

        Parameters
        ----------
        dirname: str or Path
            Directory to upload.
        s3_dir: str
            If provided, this will pre prepended to each path: {s3_dir}/{path}.
            Otherwise, s3 paths will be the same as local paths.
        retain_tree: bool
            If True, the local file structure will be retained. Otherwise,
            only the base name is kept. When uploading recursively, we require
            that the file tree is retained.
        recurse: bool
            If True, upload all files in subdirectories as well.
        keep_fn: None or callable
            If provided, this should be a function that accepts a filename as
            input and returns a boolean specifying whether to include it in the
            upload or not. For example:

            lambda x: os.path.splitext(x)[-1] != '.pkl'

            keeps all files except those with an '.pkl' extension.
        """
        if recurse and not retain_tree:
            raise ValueError('retain_tree must be True when uploading '
                             'recursively.')
        pat = os.path.join(str(dirname), '**' if recurse else '*')
        # glob's recursive option only has an effect when using '**'.
        paths = (o for o in glob(pat, recursive=True) if os.path.isfile(o))
        if keep_fn: paths = filter(keep_fn, paths)
        self.upload_files(paths, s3_dir, retain_tree)

    def _convert_local_path(self, path, s3_dir='', retain_tree=True):
        """Convert local path to s3 path. See public methods for parameter
        documentation.
        """
        path = path if retain_tree else os.path.basename(path)
        return os.path.join(s3_dir, path)