from tensorflow.python.platform import gfile


import numpy as np

from tensorflow.contrib.learn.python.learn.datasets import base
from tensorflow.python.platform import gfile

DBPEDIA_URL = 'https://github.com/le-scientifique/torchDatasets/blob/master/dbpedia_csv.tar.gz'


def maybe_download_dbpedia(data_dir):
  """Download if DBpedia data is not present."""
  train_path = os.path.join(data_dir, 'dbpedia_csv/train.csv')
  test_path = os.path.join(data_dir, 'dbpedia_csv/test.csv')
  if not (gfile.Exists(train_path) and gfile.Exists(test_path)):
    archive_path = base.maybe_download(
        'dbpedia_csv.tar.gz', data_dir, DBPEDIA_URL)
    tfile = tarfile.open(archive_path, 'r:*')
    tfile.extractall(data_dir)


def load_dbpedia(size='small', test_with_fake_data=False):
  """Get DBpedia datasets from CSV files."""
  if not test_with_fake_data:
    data_dir = os.path.join(os.getenv('TF_EXP_BASE_DIR', ''), 'dbpedia_data')
    maybe_download_dbpedia(data_dir)

