# AUTOGENERATED! DO NOT EDIT! File to edit: 00_core.ipynb (unless otherwise specified).

__all__ = ['Config', 'URLs', 'download_data', 'file_extract', 'download_file_from_google_drive', 'untar_data']

# Cell
import os
import shutil
import requests
from pathlib import Path
from tqdm.notebook import tqdm
import zipfile, tarfile

# Cell
class Config:
    config_path = Path('~/.aiadv').expanduser()
    def __init__(self):
        self.config_path.mkdir(parents=True, exist_ok=True)

# Cell
class URLs:
    # Datasets
    YELP_REIVEWS      = {'id': '1G42LXv72DrhK4QKJoFhabVL4IU6v2ZvB',
                         'fname': 'yelp_reveiw.csv'}

    MOVIE_LENS_SAMPLE = {'id': '1k2y0qC0E3oHeGA5a427hRgfbW7hnQBgF',
                         'fname': 'movie_lens_sample.zip'}

    ENG_FRA           = {'id': '1dU-cTcPxHlpoFMnWe21jB4n6GRJdLdJO',
                         'fname': 'eng_fra.txt'}

    def path(ds=None):
        fname = ds['fname']
        path = Config().config_path/fname
        return path

    def stem(path):
        if str(path).endswith('gz') or str(path).endswith('zip'):
            parent = path.parent
            return parent/path.stem
        else: return path

# Cell
def download_data(ds, force_download=False):
    "Download `url` to `fname`."
    dest = URLs.path(ds)
    dest.parent.mkdir(parents=True, exist_ok=True)
    if not dest.exists() or force_download:
        download_file_from_google_drive(ds['id'], dest, overwrite=force_download)
    return dest

# Cell
def file_extract(fname):
    "Extract `fname` using `tarfile` or `zipfile"
    fname_str = str(fname)

    if fname_str.endswith('gz'):
        dest = URLs.stem(fname)
        tarfile.open(fname, 'r:gz').extractall(dest)
        os.remove(fname)
        return dest

    elif fname_str.endswith('zip'):
        dest = URLs.stem(fname)
        zipfile.ZipFile(fname).extractall(dest)
        os.remove(fname)
        return dest

    elif fname_str.endswith('csv') or fname_str.endswith('txt'): return fname

    else: raise Exception(f'Unrecognized archive: {fname}')

# Cell
def download_file_from_google_drive(id, dest, overwrite=False):
    "Download `url` to `dest` unless it exists and not `overwrite`"
    if os.path.exists(dest) and not overwrite: return

    print("Trying to fetch {}".format(dest.name))

    def get_confirm_token(response):
        for key, value in response.cookies.items():
            if key.startswith('download_warning'):
                return value
        return None

    def save_response_content(response, destination, pbar=None):
        CHUNK_SIZE = 1024*1024
        show_progress = True
        try: file_size = int(response.headers["Content-Length"])
        except: show_progress = False

        with open(destination, "wb") as f:
            if show_progress: pbar = tqdm(unit="MB", total=int(file_size/CHUNK_SIZE))
            else: pbar = tqdm(unit="MB")

            for chunk in response.iter_content(CHUNK_SIZE):
                if chunk: # filter out keep-alive new chunks
                    pbar.update()
                    f.write(chunk)

    URL = "https://docs.google.com/uc?export=download"
    session = requests.Session()
    response = session.get(URL, params = { 'id' : id }, stream = True)
    token = get_confirm_token(response)

    if token:
        params = { 'id' : id, 'confirm' : token }
        response = session.get(URL, params = params, stream = True)

    save_response_content(response, dest)

# Cell
def untar_data(ds, force_download=False, extract_func=file_extract):
    dest = URLs.path(ds)
    stem = URLs.stem(dest)
    fname = ds['fname']
    if force_download:
        if stem.exists():
            try: os.remove(stem)
            except: shutil.rmtree(stem)
    if not stem.exists():
        download_data(ds)
        path=extract_func(dest)
        return path
    return stem