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

__all__ = ['list_datasets', 'datasetsdf', 'list_dataset_files', 'read_download_info', 'download_all_datasets',
           'download_dataset_file', 'downloaddf', 'read_dataset_file', 'download_urls_file', 'read_urls_file', 'urlsdf',
           'get_rows_from_datasetdf', 'show_first_rows', 'get_textblocks_from_dataset', 'show_first_textblocks',
           'get_text_from_rowindex', 'get_url_from_rowindex', 'find_textblocks_with_chars', 'getContextAroundWord']

# Cell
import numpy as np
import pandas as pd
pd.options.display.max_rows = 100
pd.options.display.max_columns = 50

# Cell
from .core import *

# Cell
def list_datasets():
    return pd.read_csv(config.libdata /"datasets"/"datasets.csv",sep=';')

datasetsdf = list_datasets()

# Cell
def list_dataset_files():
    return datasetsdf["DatasetFile"].unique().tolist()

# Cell
def read_download_info():
    return pd.read_csv(config.libdata /"datasets"/"download_info.csv",sep=';')

downloaddf = read_download_info()

def download_all_datasets():
    for datasetfile in datasetsdf['DatasetFile'].unique():
        download_dataset_file(datasetfile)

import math

def download_dataset_file(datasetfile):
    row = downloaddf[downloaddf["DatasetFile"]==datasetfile].iloc[0]
    url = row["DownloadUrl"]
    size = row["FileSize"]
    print(f"Downloading dataset file : {datasetfile} ({math.floor(size/1024/1024)} MB)")
    resfile = config.datasets / (datasetfile+".dataset.feather")
    if not resfile.exists():
        download_url(url, config.datasets / (datasetfile+".dataset.zip"), size)

# Cell
def read_dataset_file(datasetfile):
    datasetdffile = config.datasets / (datasetfile+".dataset.feather")
    if(datasetdffile.exists()):
        datasetdf = pd.read_feather(datasetdffile)
        print(f"Loaded dataframe for dataset {datasetfile} : {len(datasetdf)} text blocks")
        return datasetdf
    else:
        raise Exception(f"No dataframe for dataset {datasetfile}")

# Cell
def download_urls_file():
    row = downloaddf[downloaddf["DatasetFile"]=="urls"].iloc[0]
    url = row["DownloadUrl"]
    size = row["FileSize"]
    print(f"Downloading datasets urls file ({math.floor(size/1024/1024)} MB)")
    download_url(url, config.datasets / "datasets.urls.zip", size)
    print("Done")

urlsdf = None

def read_urls_file():
    global urlsdf
    if urlsdf is None:
        download_urls_file()
        urlsdf = pd.read_feather(config.datasets / "datasets.urls.feather")
    print(f"Loaded datasets urls : {len(urlsdf)} urls")
    return urlsdf

# Cell
def get_rows_from_datasetdf(datasetdf, minwords=5, maxwords=None, lang=None):
    condition = datasetdf.index>=0
    if isinstance(minwords,int):
        condition = condition & (datasetdf["Words"] >= minwords)
    if isinstance(maxwords,int):
        condition = condition & (datasetdf["Words"] <= maxwords)
    if isinstance(lang,str):
        condition = condition & (datasetdf["Lang"] == lang)
    return datasetdf[condition].iterrows()

def show_first_rows(rowsiterator, count=5, skip=0):
    i=0
    for rowidx,row in rowsiterator:
        i = i+1
        if(i <= skip): continue
        text = row["Text"]
        print(f"{rowidx} - {text}")
        if(i >= skip+count): break

# Cell
def get_textblocks_from_dataset(dataset, minwords=5, maxwords=None, lang=None):
    for datasetfile in datasetsdf.loc[datasetsdf["Dataset"]==dataset,"DatasetFile"].unique():
        datasetdf = read_dataset_file(datasetfile)
        for rowidx,row in get_rows_from_datasetdf(datasetdf,minwords,maxwords,lang):
            yield row["Text"]

def show_first_textblocks(textiterator, count=5, skip=0):
    i=0
    for text in textiterator:
        i = i+1
        if(i <= skip): continue
        print(f"{i} - {text}")
        if(i >= (skip+count)): break

# Cell
def get_text_from_rowindex(datasetdf, rowidx):
    return datasetdf.iloc[rowidx]["Text"]

def get_url_from_rowindex(datasetdf, rowidx):
    global urlsdf
    row = datasetdf.iloc[rowidx]
    website = row["Website"]
    docid = row["DocId"]
    if urlsdf is None:
        urlsdf = read_urls_file()
    url = urlsdf.loc[(urlsdf["Website"] == website) & (urlsdf["DocId"] == docid),"DocUrl"].values[0]
    return url

# Cell
def find_textblocks_with_chars(datasetdf, chars, count=100, ctxsize=20, wrap=False):
    if isinstance(chars,int): chars = chr(chars)
    textsWithWord = datasetdf[datasetdf["Text"].str.contains(chars,regex=False)]
    if count>0 and (textsWithWord["Text"].count()>count):
        textsWithWord = textsWithWord.sample(count)
    return textsWithWord["Text"].apply(lambda t: getContextAroundWord(t,chars,ctxsize,wrap))

def getContextAroundWord(text,word,ctxsize=20,wrap=False):
    idx = text.index(word)
    start = max(idx-ctxsize,0)
    padbefore = -min(idx-ctxsize,0)
    end = min(idx+len(word)+ctxsize,len(text))
    passage = text[start:end+1].replace(chr(160)," ")
    if(wrap):
        passage = passage.replace(word,"["+word+"]")
    padafter = (2*ctxsize+len(word)+1)-len(passage)-padbefore
    if(wrap):
        padafter = padafter+2
    paddedpassage = (" "*padbefore) + passage + (" "*padafter)
    return paddedpassage