"""
qm9 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

GDB9_URL = "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/gdb9.tar.gz"
QM9_CSV_URL = "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/qm9.csv"
QM9_TASKS = [
    "mu", "alpha", "homo", "lumo", "gap", "r2", "zpve", "cv", "u0", "u298",
    "h298", "g298"
]


class _QM9Loader(_MolnetLoader):

  def create_dataset(self) -> Dataset:
    dataset_file = os.path.join(self.data_dir, "gdb9.sdf")
    if not os.path.exists(dataset_file):
      dc.utils.data_utils.download_url(url=GDB9_URL, dest_dir=self.data_dir)
      dc.utils.data_utils.untargz_file(
          os.path.join(self.data_dir, "gdb9.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_qm9(
    featurizer: Union[dc.feat.Featurizer, str] = dc.feat.CoulombMatrix(29),
    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 QM9 dataset

  QM9 is a comprehensive dataset that provides geometric, energetic,
  electronic and thermodynamic properties for a subset of GDB-17 database,
  comprising 134 thousand stable organic molecules with up to 9 heavy atoms.
  All molecules are modeled using density functional theory
  (B3LYP/6-31G(2df,p) based DFT).

  Random splitting is recommended for this dataset.

  The source data contain:

  - qm9.sdf: molecular structures
  - qm9.sdf.csv: tables for molecular properties

    - "mol_id" - Molecule ID (gdb9 index) mapping to the .sdf file
    - "A" - Rotational constant (unit: GHz)
    - "B" - Rotational constant (unit: GHz)
    - "C" - Rotational constant (unit: GHz)
    - "mu" - Dipole moment (unit: D)
    - "alpha" - Isotropic polarizability (unit: Bohr^3)
    - "homo" - Highest occupied molecular orbital energy (unit: Hartree)
    - "lumo" - Lowest unoccupied molecular orbital energy (unit: Hartree)
    - "gap" - Gap between HOMO and LUMO (unit: Hartree)
    - "r2" - Electronic spatial extent (unit: Bohr^2)
    - "zpve" - Zero point vibrational energy (unit: Hartree)
    - "u0" - Internal energy at 0K (unit: Hartree)
    - "u298" - Internal energy at 298.15K (unit: Hartree)
    - "h298" - Enthalpy at 298.15K (unit: Hartree)
    - "g298" - Free energy at 298.15K (unit: Hartree)
    - "cv" - Heat capavity at 298.15K (unit: cal/(mol*K))
    - "u0_atom" - Atomization energy at 0K (unit: kcal/mol)
    - "u298_atom" - Atomization energy at 298.15K (unit: kcal/mol)
    - "h298_atom" - Atomization enthalpy at 298.15K (unit: kcal/mol)
    - "g298_atom" - Atomization free energy at 298.15K (unit: kcal/mol)

  "u0_atom" ~ "g298_atom" (used in MoleculeNet) are calculated from the
  differences between "u0" ~ "g298" and sum of reference energies of all
  atoms in the molecules, as given in
  https://figshare.com/articles/Atomref%3A_Reference_thermochemical_energies_of_H%2C_C%2C_N%2C_O%2C_F_atoms./1057643

  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] Blum, Lorenz C., and Jean-Louis Reymond. "970 million druglike small
     molecules for virtual screening in the chemical universe database GDB-13."
     Journal of the American Chemical Society 131.25 (2009): 8732-8733.
  .. [2] Ramakrishnan, Raghunathan, et al. "Quantum chemistry structures and
     properties of 134 kilo molecules." Scientific data 1 (2014): 140022.
  """
  loader = _QM9Loader(featurizer, splitter, transformers, QM9_TASKS, data_dir,
                      save_dir, **kwargs)
  return loader.load_dataset('qm9', reload)
