import os
import os.path as osp

import numpy as np
from joblib import Parallel, delayed
from tensorflow.keras.utils import get_file
from tqdm import tqdm

from spektral.data import Dataset, Graph
from spektral.utils import label_to_one_hot, sparse
from spektral.utils.io import load_csv, load_sdf

ATOM_TYPES = [1, 6, 7, 8, 9]
BOND_TYPES = [1, 2, 3, 4]


class QM9(Dataset):
    """
    The QM9 chemical data set of small molecules.

    In this dataset, nodes represent atoms and edges represent chemical bonds.
    There are 5 possible atom types (H, C, N, O, F) and 4 bond types (single,
    double, triple, aromatic).

    Node features represent the chemical properties of each atom and include:

    - The atomic number, one-hot encoded;
    - The atom's position in the X, Y, and Z dimensions;
    - The atomic charge;
    - The mass difference from the monoisotope;

    The edge features represent the type of chemical bond between two atoms,
    one-hot encoded.

    Each graph has an 19-dimensional label for regression.

    **Arguments**

    - `amount`: int, load this many molecules instead of the full dataset
    (useful for debugging).
    - `n_jobs`: number of CPU cores to use for reading the data (-1, to use all
    available cores).
    """

    url = "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/gdb9.tar.gz"

    def __init__(self, amount=None, n_jobs=1, **kwargs):
        self.amount = amount
        self.n_jobs = n_jobs
        super().__init__(**kwargs)

    def download(self):
        get_file(
            "qm9.tar.gz",
            self.url,
            extract=True,
            cache_dir=self.path,
            cache_subdir=self.path,
        )
        os.remove(osp.join(self.path, "qm9.tar.gz"))

    def read(self):
        print("Loading QM9 dataset.")
        sdf_file = osp.join(self.path, "gdb9.sdf")
        data = load_sdf(sdf_file, amount=self.amount)  # Internal SDF format

        def read_mol(mol):
            x = np.array([atom_to_feature(atom) for atom in mol["atoms"]])
            a, e = mol_to_adj(mol)
            return x, a, e

        data = Parallel(n_jobs=self.n_jobs)(
            delayed(read_mol)(mol) for mol in tqdm(data, ncols=80)
        )
        x_list, a_list, e_list = list(zip(*data))

        # Load labels
        labels_file = osp.join(self.path, "gdb9.sdf.csv")
        labels = load_csv(labels_file)
        labels = labels.set_index("mol_id").values
        if self.amount is not None:
            labels = labels[: self.amount]

        return [
            Graph(x=x, a=a, e=e, y=y)
            for x, a, e, y in zip(x_list, a_list, e_list, labels)
        ]


def atom_to_feature(atom):
    atomic_num = label_to_one_hot(atom["atomic_num"], ATOM_TYPES)
    coords = atom["coords"]
    charge = atom["charge"]
    iso = atom["iso"]

    return np.concatenate((atomic_num, coords, [charge, iso]), -1)


def mol_to_adj(mol):
    row, col, edge_features = [], [], []
    for bond in mol["bonds"]:
        start, end = bond["start_atom"], bond["end_atom"]
        row += [start, end]
        col += [end, start]
        edge_features += [bond["type"]] * 2

    a, e = sparse.edge_index_to_matrix(
        edge_index=np.array((row, col)).T,
        edge_weight=np.ones_like(row),
        edge_features=label_to_one_hot(edge_features, BOND_TYPES),
    )

    return a, e

[((91, 17, 91, 65), 'spektral.utils.label_to_one_hot', 'label_to_one_hot', ({(91, 34, 91, 52): "atom['atomic_num']", (91, 54, 91, 64): 'ATOM_TYPES'}, {}), "(atom['atomic_num'], ATOM_TYPES)", False, 'from spektral.utils import label_to_one_hot, sparse\n'), ((96, 11, 96, 66), 'numpy.concatenate', 'np.concatenate', ({(96, 26, 96, 61): '(atomic_num, coords, [charge, iso])', (96, 63, 96, 65): '(-1)'}, {}), '((atomic_num, coords, [charge, iso]), -1)', True, 'import numpy as np\n'), ((53, 8, 59, 9), 'tensorflow.keras.utils.get_file', 'get_file', (), '', False, 'from tensorflow.keras.utils import get_file\n'), ((64, 19, 64, 50), 'os.path.join', 'osp.join', ({(64, 28, 64, 37): 'self.path', (64, 39, 64, 49): '"""gdb9.sdf"""'}, {}), "(self.path, 'gdb9.sdf')", True, 'import os.path as osp\n'), ((65, 15, 65, 53), 'spektral.utils.io.load_sdf', 'load_sdf', (), '', False, 'from spektral.utils.io import load_csv, load_sdf\n'), ((78, 22, 78, 57), 'os.path.join', 'osp.join', ({(78, 31, 78, 40): 'self.path', (78, 42, 78, 56): '"""gdb9.sdf.csv"""'}, {}), "(self.path, 'gdb9.sdf.csv')", True, 'import os.path as osp\n'), ((79, 17, 79, 38), 'spektral.utils.io.load_csv', 'load_csv', ({(79, 26, 79, 37): 'labels_file'}, {}), '(labels_file)', False, 'from spektral.utils.io import load_csv, load_sdf\n'), ((60, 18, 60, 51), 'os.path.join', 'osp.join', ({(60, 27, 60, 36): 'self.path', (60, 38, 60, 50): '"""qm9.tar.gz"""'}, {}), "(self.path, 'qm9.tar.gz')", True, 'import os.path as osp\n'), ((72, 15, 72, 43), 'joblib.Parallel', 'Parallel', (), '', False, 'from joblib import Parallel, delayed\n'), ((85, 12, 85, 37), 'spektral.data.Graph', 'Graph', (), '', False, 'from spektral.data import Dataset, Graph\n'), ((109, 20, 109, 37), 'numpy.ones_like', 'np.ones_like', ({(109, 33, 109, 36): 'row'}, {}), '(row)', True, 'import numpy as np\n'), ((110, 22, 110, 65), 'spektral.utils.label_to_one_hot', 'label_to_one_hot', ({(110, 39, 110, 52): 'edge_features', (110, 54, 110, 64): 'BOND_TYPES'}, {}), '(edge_features, BOND_TYPES)', False, 'from spektral.utils import label_to_one_hot, sparse\n'), ((108, 19, 108, 39), 'numpy.array', 'np.array', ({(108, 28, 108, 38): '(row, col)'}, {}), '((row, col))', True, 'import numpy as np\n'), ((73, 12, 73, 29), 'joblib.delayed', 'delayed', ({(73, 20, 73, 28): 'read_mol'}, {}), '(read_mol)', False, 'from joblib import Parallel, delayed\n'), ((73, 46, 73, 66), 'tqdm.tqdm', 'tqdm', (), '', False, 'from tqdm import tqdm\n')]