"""
qm7 dataset loader.
"""
import os
import deepchem as dc
from deepchem.molnet.load_function.molnet_loader import TransformerGenerator, _MolnetLoader
from deepchem.data import Dataset
from typing import List, Optional, Tuple, Union

QM7_MAT_UTL = "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/qm7.mat"
QM7_CSV_URL = "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/qm7.csv"
QM7B_MAT_URL = "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/qm7b.mat"
GDB7_URL = "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/gdb7.tar.gz"
QM7_TASKS = ["u0_atom"]


class _QM7Loader(_MolnetLoader):

  def create_dataset(self) -> Dataset:
    dataset_file = os.path.join(self.data_dir, "gdb7.sdf")
    if not os.path.exists(dataset_file):
      dc.utils.data_utils.download_url(url=GDB7_URL, dest_dir=self.data_dir)
      dc.utils.data_utils.untargz_file(
          os.path.join(self.data_dir, "gdb7.tar.gz"), self.data_dir)
    loader = dc.data.SDFLoader(tasks=self.tasks, featurizer=self.featurizer)
    return loader.create_dataset(dataset_file, shard_size=8192)


def load_qm7(
    featurizer: Union[dc.feat.Featurizer, str] = dc.feat.CoulombMatrix(23),
    splitter: Union[dc.splits.Splitter, str, None] = 'random',
    transformers: List[Union[TransformerGenerator, str]] = ['normalization'],
    reload: bool = True,
    data_dir: Optional[str] = None,
    save_dir: Optional[str] = None,
    **kwargs
) -> Tuple[List[str], Tuple[Dataset, ...], List[dc.trans.Transformer]]:
  """Load QM7 dataset

  QM7 is a subset of GDB-13 (a database of nearly 1 billion
  stable and synthetically accessible organic molecules)
  containing up to 7 heavy atoms C, N, O, and S. The 3D
  Cartesian coordinates of the most stable conformations and
  their atomization energies were determined using ab-initio
  density functional theory (PBE0/tier2 basis set). This dataset
  also provided Coulomb matrices as calculated in [Rupp et al.
  PRL, 2012]:

  Stratified splitting is recommended for this dataset.

  The data file (.mat format, we recommend using `scipy.io.loadmat`
  for python users to load this original data) contains five arrays:

  - "X" - (7165 x 23 x 23), Coulomb matrices
  - "T" - (7165), atomization energies (unit: kcal/mol)
  - "P" - (5 x 1433), cross-validation splits as used in [Montavon et al.
    NIPS, 2012]
  - "Z" - (7165 x 23), atomic charges
  - "R" - (7165 x 23 x 3), cartesian coordinate (unit: Bohr) of each atom in
    the molecules

  Parameters
  ----------
  featurizer: Featurizer or str
    the featurizer to use for processing the data.  Alternatively you can pass
    one of the names from dc.molnet.featurizers as a shortcut.
  splitter: Splitter or str
    the splitter to use for splitting the data into training, validation, and
    test sets.  Alternatively you can pass one of the names from
    dc.molnet.splitters as a shortcut.  If this is None, all the data
    will be included in a single dataset.
  transformers: list of TransformerGenerators or strings
    the Transformers to apply to the data.  Each one is specified by a
    TransformerGenerator or, as a shortcut, one of the names from
    dc.molnet.transformers.
  reload: bool
    if True, the first call for a particular featurizer and splitter will cache
    the datasets to disk, and subsequent calls will reload the cached datasets.
  data_dir: str
    a directory to save the raw data in
  save_dir: str
    a directory to save the dataset in

  References
  ----------
  .. [1] Rupp, Matthias, et al. "Fast and accurate modeling of molecular
     atomization energies with machine learning." Physical review letters
     108.5 (2012): 058301.
  .. [2] Montavon, Grégoire, et al. "Learning invariant representations of
     molecules for atomization energy prediction." Advances in Neural
     Information Proccessing Systems. 2012.
  """
  loader = _QM7Loader(featurizer, splitter, transformers, QM7_TASKS, data_dir,
                      save_dir, **kwargs)
  return loader.load_dataset('qm7', reload)
