"""
Simple utils to save and load from disk.
"""
import joblib
import gzip
import pickle
import os
import tempfile
import tarfile
import zipfile
import logging
from urllib.request import urlretrieve
from typing import Any, Iterator, List, Optional, Tuple, Union, cast, IO

import pandas as pd
import numpy as np

import deepchem as dc

logger = logging.getLogger(__name__)


def pad_array(x: np.ndarray,
              shape: Union[Tuple, int],
              fill: float = 0.0,
              both: bool = False) -> np.ndarray:
    """
    Pad an array with a fill value.

    Parameters
    ----------
    x: np.ndarray
        A numpy array.
    shape: Tuple or int
        Desired shape. If int, all dimensions are padded to that size.
    fill: float, optional (default 0.0)
        The padded value.
    both: bool, optional (default False)
        If True, split the padding on both sides of each axis. If False,
        padding is applied to the end of each axis.

    Returns
    -------
    np.ndarray
        A padded numpy array
    """
    x = np.asarray(x)
    if not isinstance(shape, tuple):
        shape = tuple(shape for _ in range(x.ndim))
    pad = []
    for i in range(x.ndim):
        diff = shape[i] - x.shape[i]
        assert diff >= 0
        if both:
            a, b = divmod(diff, 2)
            b += a
            pad.append((a, b))
        else:
            pad.append((0, diff))
    pad = tuple(pad)  # type: ignore
    x = np.pad(x, pad, mode='constant', constant_values=fill)
    return x


def get_data_dir() -> str:
    """Get the DeepChem data directory.

    Returns
    -------
    str
        The default path to store DeepChem data. If you want to
        change this path, please set your own path to `DEEPCHEM_DATA_DIR`
        as an environment variable.
    """
    if 'DEEPCHEM_DATA_DIR' in os.environ:
        return os.environ['DEEPCHEM_DATA_DIR']
    return tempfile.gettempdir()


def download_url(url: str,
                 dest_dir: str = get_data_dir(),
                 name: Optional[str] = None):
    """Download a file to disk.

    Parameters
    ----------
    url: str
        The URL to download from
    dest_dir: str
        The directory to save the file in
    name: str
        The file name to save it as.  If omitted, it will try to extract a file name from the URL
    """
    if name is None:
        name = url
        if '?' in name:
            name = name[:name.find('?')]
        if '/' in name:
            name = name[name.rfind('/') + 1:]
    if not os.path.exists(dest_dir):
        os.makedirs(dest_dir)
    urlretrieve(url, os.path.join(dest_dir, name))


def untargz_file(file: str,
                 dest_dir: str = get_data_dir(),
                 name: Optional[str] = None):
    """Untar and unzip a .tar.gz file to disk.

    Parameters
    ----------
    file: str
        The filepath to decompress
    dest_dir: str
        The directory to save the file in
    name: str
        The file name to save it as.  If omitted, it will use the file name
    """
    if name is None:
        name = file
    tar = tarfile.open(name)
    tar.extractall(path=dest_dir)
    tar.close()


def unzip_file(file: str,
               dest_dir: str = get_data_dir(),
               name: Optional[str] = None):
    """Unzip a .zip file to disk.

    Parameters
    ----------
    file: str
        The filepath to decompress
    dest_dir: str
        The directory to save the file in
    name: str
        The directory name to unzip it to.  If omitted, it will use the file name
    """
    if name is None:
        name = file
    if dest_dir is None:
        dest_dir = os.path.join(get_data_dir, name)
    with zipfile.ZipFile(file, "r") as zip_ref:
        zip_ref.extractall(dest_dir)


class UniversalNamedTemporaryFile:
    """The implementation for cross platform NamedTemporaryFile.

    `tempfile.NamedTemporaryFile` causes a permission error on Windows.
    This implementation avoids the error, please see threads on the stackoverflow [1]_.

    References
    ----------
    .. [1] https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file
    """

    def __init__(self, mode='wb', delete=True):
        self._mode = mode
        self._delete = delete

    def __enter__(self):
        # Generate a random temporary file name
        file_name = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())
        # Ensure the file is created
        open(file_name, "x").close()
        # Open the file in the given mode
        self._tempFile = open(file_name, self._mode)
        return self._tempFile

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._tempFile.close()
        if self._delete:
            os.remove(self._tempFile.name)


def load_image_files(input_files: List[str]) -> np.ndarray:
    """Loads a set of images from disk.

    Parameters
    ----------
    input_files: List[str]
        List of image filenames.

    Returns
    -------
    np.ndarray
        A numpy array that contains loaded images. The shape is, `(N,...)`.

    Notes
    -----
    This method requires Pillow to be installed.
    The supported file types are PNG and TIF.
    """
    try:
        from PIL import Image
    except ModuleNotFoundError:
        raise ImportError("This function requires Pillow to be installed.")

    images = []
    for input_file in input_files:
        _, extension = os.path.splitext(input_file)
        extension = extension.lower()
        if extension == ".png":
            image = np.array(Image.open(input_file))
            images.append(image)
        elif extension == ".tif":
            im = Image.open(input_file)
            imarray = np.array(im)
            images.append(imarray)
        else:
            raise ValueError("Unsupported image filetype for %s" % input_file)
    return np.array(images)


def load_sdf_files(input_files: List[str],
                   clean_mols: bool = True,
                   tasks: List[str] = [],
                   shard_size: Optional[int] = None) -> Iterator[pd.DataFrame]:
    """Load SDF file into dataframe.

    Parameters
    ----------
    input_files: List[str]
        List of filenames
    clean_mols: bool, default True
        Whether to sanitize molecules.
    tasks: List[str], default []
        Each entry in `tasks` is treated as a property in the SDF file and is
        retrieved with `mol.GetProp(str(task))` where `mol` is the RDKit mol
        loaded from a given SDF entry.
    shard_size: int, default None
        The shard size to yield at one time.

    Returns
    -------
    Iterator[pd.DataFrame]
        Generator which yields the dataframe which is the same shard size.

    Notes
    -----
    This function requires RDKit to be installed.
    """
    try:
        from rdkit import Chem
    except ModuleNotFoundError:
        raise ImportError("This function requires RDKit to be installed.")

    df_rows = []
    for input_file in input_files:
        # Tasks are either in .sdf.csv file or in the .sdf file itself for QM9 dataset
        has_csv = os.path.isfile(input_file + ".csv")
        # Structures are stored in .sdf file
        logger.info("Reading structures from %s." % input_file)

        suppl = Chem.SDMolSupplier(str(input_file),
                                   sanitize=clean_mols,
                                   removeHs=False,
                                   strictParsing=False)
        for ind, mol in enumerate(suppl):
            if mol is None:
                continue
            smiles = Chem.MolToSmiles(mol)
            df_row = [ind, smiles, mol]
            if not has_csv:  # Get task targets from .sdf file
                for task in tasks:
                    df_row.append(mol.GetProp(str(task)))

            conf = mol.GetConformer()
            positions = conf.GetPositions()
            pos_x, pos_y, pos_z = zip(*positions)
            df_row.append(str(pos_x))
            df_row.append(str(pos_y))
            df_row.append(str(pos_z))
            df_rows.append(df_row)

            if shard_size is not None and len(df_rows) == shard_size:
                if has_csv:
                    mol_df = pd.DataFrame(df_rows,
                                          columns=('mol_id', 'smiles', 'mol',
                                                   'pos_x', 'pos_y', 'pos_z'))
                    raw_df = next(
                        load_csv_files([input_file + ".csv"], shard_size=None))
                    yield pd.concat([mol_df, raw_df], axis=1, join='inner')
                else:
                    # Note: Here, the order of columns is based on the order in which the values
                    # are appended to `df_row`. Since pos_x, pos_y, pos_z are appended after appending
                    # tasks above, they occur after `tasks` here.
                    # FIXME Ideally, we should use something like a dictionary here to keep it independent
                    # of column ordering.
                    mol_df = pd.DataFrame(df_rows,
                                          columns=('mol_id', 'smiles', 'mol') +
                                          tuple(tasks) +
                                          ('pos_x', 'pos_y', 'pos_z'))
                    yield mol_df
                # Reset aggregator
                df_rows = []

        # Handle final leftovers for this file
        if len(df_rows) > 0:
            if has_csv:
                mol_df = pd.DataFrame(df_rows,
                                      columns=('mol_id', 'smiles', 'mol',
                                               'pos_x', 'pos_y', 'pos_z'))
                raw_df = next(
                    load_csv_files([input_file + ".csv"], shard_size=None))
                yield pd.concat([mol_df, raw_df], axis=1, join='inner')
            else:
                mol_df = pd.DataFrame(df_rows,
                                      columns=('mol_id', 'smiles', 'mol') +
                                      tuple(tasks) +
                                      ('pos_x', 'pos_y', 'pos_z'))
                yield mol_df
            df_rows = []


def load_csv_files(input_files: List[str],
                   shard_size: Optional[int] = None) -> Iterator[pd.DataFrame]:
    """Load data as pandas dataframe from CSV files.

    Parameters
    ----------
    input_files: List[str]
        List of filenames
    shard_size: int, default None
        The shard size to yield at one time.

    Returns
    -------
    Iterator[pd.DataFrame]
        Generator which yields the dataframe which is the same shard size.
    """
    # First line of user-specified CSV *must* be header.
    shard_num = 1
    for input_file in input_files:
        if shard_size is None:
            yield pd.read_csv(input_file)
        else:
            logger.info("About to start loading CSV from %s" % input_file)
            for df in pd.read_csv(input_file, chunksize=shard_size):
                logger.info("Loading shard %d of size %s." %
                            (shard_num, str(shard_size)))
                df = df.replace(np.nan, str(""), regex=True)
                shard_num += 1
                yield df


def load_json_files(input_files: List[str],
                    shard_size: Optional[int] = None) -> Iterator[pd.DataFrame]:
    """Load data as pandas dataframe.

    Parameters
    ----------
    input_files: List[str]
        List of json filenames.
    shard_size: int, default None
        Chunksize for reading json files.

    Returns
    -------
    Iterator[pd.DataFrame]
        Generator which yields the dataframe which is the same shard size.

    Notes
    -----
    To load shards from a json file into a Pandas dataframe, the file
    must be originally saved with ``df.to_json('filename.json', orient='records', lines=True)``
    """
    shard_num = 1
    for input_file in input_files:
        if shard_size is None:
            yield pd.read_json(input_file, orient='records', lines=True)
        else:
            logger.info("About to start loading json from %s." % input_file)
            for df in pd.read_json(input_file,
                                   orient='records',
                                   chunksize=shard_size,
                                   lines=True):
                logger.info("Loading shard %d of size %s." %
                            (shard_num, str(shard_size)))
                df = df.replace(np.nan, str(""), regex=True)
                shard_num += 1
                yield df


def load_pickle_file(input_file: str) -> Any:
    """Load from single, possibly gzipped, pickle file.

    Parameters
    ----------
    input_file: str
        The filename of pickle file. This function can load from
        gzipped pickle file like `XXXX.pkl.gz`.

    Returns
    -------
    Any
        The object which is loaded from the pickle file.
    """
    if ".gz" in input_file:
        with gzip.open(input_file, "rb") as unzipped_file:
            return pickle.load(cast(IO[bytes], unzipped_file))
    else:
        with open(input_file, "rb") as opened_file:
            return pickle.load(opened_file)


def load_pickle_files(input_files: List[str]) -> Iterator[Any]:
    """Load dataset from pickle files.

    Parameters
    ----------
    input_files: List[str]
        The list of filenames of pickle file. This function can load from
        gzipped pickle file like `XXXX.pkl.gz`.

    Returns
    -------
    Iterator[Any]
        Generator which yields the objects which is loaded from each pickle file.
    """
    for input_file in input_files:
        yield load_pickle_file(input_file)


def load_data(input_files: List[str],
              shard_size: Optional[int] = None) -> Iterator[Any]:
    """Loads data from files.

    Parameters
    ----------
    input_files: List[str]
        List of filenames.
    shard_size: int, default None
        Size of shard to yield

    Returns
    -------
    Iterator[Any]
        Iterator which iterates over provided files.

    Notes
    -----
    The supported file types are SDF, CSV and Pickle.
    """
    if len(input_files) == 0:
        raise ValueError("The length of `filenames` must be more than 1.")

    file_type = _get_file_type(input_files[0])
    if file_type == "sdf":
        if shard_size is not None:
            logger.info("Ignoring shard_size for sdf input.")
        for value in load_sdf_files(input_files):
            yield value
    elif file_type == "csv":
        for value in load_csv_files(input_files, shard_size):
            yield value
    elif file_type == "pickle":
        if shard_size is not None:
            logger.info("Ignoring shard_size for pickle input.")
        for value in load_pickle_files(input_files):
            yield value


def _get_file_type(input_file: str) -> str:
    """Get type of input file. Must be csv/pkl/sdf/joblib file."""
    filename, file_extension = os.path.splitext(input_file)
    # If gzipped, need to compute extension again
    if file_extension == ".gz":
        filename, file_extension = os.path.splitext(filename)
    if file_extension == ".csv":
        return "csv"
    elif file_extension == ".pkl":
        return "pickle"
    elif file_extension == ".joblib":
        return "joblib"
    elif file_extension == ".sdf":
        return "sdf"
    else:
        raise ValueError("Unrecognized extension %s" % file_extension)


def save_to_disk(dataset: Any, filename: str, compress: int = 3):
    """Save a dataset to file.

    Parameters
    ----------
    dataset: str
        A data saved
    filename: str
        Path to save data.
    compress: int, default 3
        The compress option when dumping joblib file.
    """
    if filename.endswith('.joblib'):
        joblib.dump(dataset, filename, compress=compress)
    elif filename.endswith('.npy'):
        np.save(filename, dataset)
    else:
        raise ValueError("Filename with unsupported extension: %s" % filename)


def load_from_disk(filename: str) -> Any:
    """Load a dataset from file.

    Parameters
    ----------
    filename: str
        A filename you want to load data.

    Returns
    -------
    Any
        A loaded object from file.
    """
    name = filename
    if os.path.splitext(name)[1] == ".gz":
        name = os.path.splitext(name)[0]
    extension = os.path.splitext(name)[1]
    if extension == ".pkl":
        return load_pickle_file(filename)
    elif extension == ".joblib":
        return joblib.load(filename)
    elif extension == ".csv":
        # First line of user-specified CSV *must* be header.
        df = pd.read_csv(filename, header=0)
        df = df.replace(np.nan, str(""), regex=True)
        return df
    elif extension == ".npy":
        return np.load(filename, allow_pickle=True)
    else:
        raise ValueError("Unrecognized filetype for %s" % filename)


def load_dataset_from_disk(
    save_dir: str
) -> Tuple[bool, Optional[Tuple["dc.data.DiskDataset", "dc.data.DiskDataset",
                                "dc.data.DiskDataset"]],
           List["dc.trans.Transformer"]]:
    """Loads MoleculeNet train/valid/test/transformers from disk.

    Expects that data was saved using `save_dataset_to_disk` below. Expects the
    following directory structure for `save_dir`:
    save_dir/
      |
      ---> train_dir/
      |
      ---> valid_dir/
      |
      ---> test_dir/
      |
      ---> transformers.pkl

    Parameters
    ----------
    save_dir: str
        Directory name to load datasets.

    Returns
    -------
    loaded: bool
        Whether the load succeeded
    all_dataset: Tuple[DiskDataset, DiskDataset, DiskDataset]
        The train, valid, test datasets
    transformers: Transformer
        The transformers used for this dataset

    See Also
    --------
    save_dataset_to_disk
    """

    train_dir = os.path.join(save_dir, "train_dir")
    valid_dir = os.path.join(save_dir, "valid_dir")
    test_dir = os.path.join(save_dir, "test_dir")
    if not os.path.exists(train_dir) or not os.path.exists(
            valid_dir) or not os.path.exists(test_dir):
        return False, None, list()
    loaded = True
    train = dc.data.DiskDataset(train_dir)
    valid = dc.data.DiskDataset(valid_dir)
    test = dc.data.DiskDataset(test_dir)
    train.memory_cache_size = 40 * (1 << 20)  # 40 MB
    all_dataset = (train, valid, test)
    transformers = load_transformers(save_dir)
    return loaded, all_dataset, transformers


def save_dataset_to_disk(save_dir: str, train: "dc.data.DiskDataset",
                         valid: "dc.data.DiskDataset",
                         test: "dc.data.DiskDataset",
                         transformers: List["dc.trans.Transformer"]):
    """Utility used by MoleculeNet to save train/valid/test datasets.

    This utility function saves a train/valid/test split of a dataset along
    with transformers in the same directory. The saved datasets will take the
    following structure:
    save_dir/
      |
      ---> train_dir/
      |
      ---> valid_dir/
      |
      ---> test_dir/
      |
      ---> transformers.pkl

    Parameters
    ----------
    save_dir: str
        Directory name to save datasets to.
    train: DiskDataset
        Training dataset to save.
    valid: DiskDataset
        Validation dataset to save.
    test: DiskDataset
        Test dataset to save.
    transformers: List[Transformer]
        List of transformers to save to disk.

    See Also
    --------
    load_dataset_from_disk
    """
    train_dir = os.path.join(save_dir, "train_dir")
    valid_dir = os.path.join(save_dir, "valid_dir")
    test_dir = os.path.join(save_dir, "test_dir")
    train.move(train_dir)
    valid.move(valid_dir)
    test.move(test_dir)
    save_transformers(save_dir, transformers)


def load_transformers(save_dir: str) -> List["dc.trans.Transformer"]:
    """Load the transformers for a MoleculeNet dataset from disk."""
    with open(os.path.join(save_dir, "transformers.pkl"), 'rb') as f:
        return pickle.load(f)


def save_transformers(save_dir: str,
                      transformers: List["dc.trans.Transformer"]):
    """Save the transformers for a MoleculeNet dataset to disk."""
    with open(os.path.join(save_dir, "transformers.pkl"), 'wb') as f:
        pickle.dump(transformers, f)


def remove_missing_entries(dataset):
    """Remove missing entries.

    Some of the datasets have missing entries that sneak in as zero'd out
    feature vectors. Get rid of them.
    """
    for i, (X, y, w, ids) in enumerate(dataset.itershards()):
        available_rows = X.any(axis=1)
        logger.info("Shard %d has %d missing entries." %
                    (i, np.count_nonzero(~available_rows)))
        X = X[available_rows]
        y = y[available_rows]
        w = w[available_rows]
        ids = ids[available_rows]
        dataset.set_shard(i, X, y, w, ids)
