"""

Script that performs Enumeration based augmentation for chemical SMILES

Implementation borrow from: https://github.com/EBjerrum/SMILES-enumeration

"""

from __future__ import print_function

from __future__ import division

from __future__ import unicode_literals

import os

import shutil

import numpy as np

import deepchem as dc

from deepchem.molnet import load_muv

from sklearn.ensemble import RandomForestClassifier

import pandas as pd

# Experimental Class for Smiles Enumeration, Iterator and SmilesIterator adapted from Keras 2.6.0

from rdkit import Chem

import threading

np.random.seed(123)

class Iterator(object):

    """Abstract base class for data iterators.

    # Arguments

        n: Integer, total number of samples in the dataset to loop over.

        batch_size: Integer, size of a batch.

        shuffle: Boolean, whether to shuffle the data between epochs.

        seed: Random seeding for data shuffling.

    """

    def __init__(self, n, batch_size, shuffle, seed):

        self.n = n

        self.batch_size = batch_size

        self.shuffle = shuffle

        self.batch_index = 0

        self.total_batches_seen = 0

        self.lock = threading.Lock()

        self.index_generator = self._flow_index(n, batch_size, shuffle, seed)

        if n < batch_size:

            raise ValueError(

                "Input data length is shorter than batch_size\nAdjust batch_size"

            )

    def reset(self):

        self.batch_index = 0

    def _flow_index(self, n, batch_size=32, shuffle=False, seed=None):

        # Ensure self.batch_index is 0.

        self.reset()

        while 1:

            if seed is not None:

                np.random.seed(seed + self.total_batches_seen)

            if self.batch_index == 0:

                index_array = np.arange(n)

                if shuffle:

                    index_array = np.random.permutation(n)

            current_index = (self.batch_index * batch_size) % n

            if n > current_index + batch_size:

                current_batch_size = batch_size

                self.batch_index += 1

            else:

                current_batch_size = n - current_index

                self.batch_index = 0

            self.total_batches_seen += 1

            yield (

                index_array[current_index : current_index + current_batch_size],

                current_index,

                current_batch_size,

            )

    def __iter__(self):

        # Needed if we want to do something like:

        # for x, y in data_gen.flow(...):

        return self

    def __next__(self, *args, **kwargs):

        return self.next(*args, **kwargs)

class SmilesIterator(Iterator):

    """Iterator yielding data from a SMILES array.

    # Arguments

        x: Numpy array of SMILES input data.

        y: Numpy array of targets data.

        smiles_data_generator: Instance of `SmilesEnumerator`

            to use for random SMILES generation.

        batch_size: Integer, size of a batch.

        shuffle: Boolean, whether to shuffle the data between epochs.

        seed: Random seed for data shuffling.

        dtype: dtype to use for returned batch. Set to keras.backend.floatx if using Keras

    """

    def __init__(

        self,

        x,

        y,

        smiles_data_generator,

        batch_size=32,

        shuffle=False,

        seed=None,

        dtype=np.float32,

    ):

        if y is not None and len(x) != len(y):

            raise ValueError(

                "X (images tensor) and y (labels) "

                "should have the same length. "

                "Found: X.shape = %s, y.shape = %s"

                % (np.asarray(x).shape, np.asarray(y).shape)

            )

        self.x = np.asarray(x)

        if y is not None:

            self.y = np.asarray(y)

        else:

            self.y = None

        self.smiles_data_generator = smiles_data_generator

        self.dtype = dtype

        super(SmilesIterator, self).__init__(x.shape[0], batch_size, shuffle, seed)

    def next(self):

        """For python 2.x.

        # Returns

            The next batch.

        """

        # Keeps under lock only the mechanism which advances

        # the indexing of each batch.

        with self.lock:

            index_array, current_index, current_batch_size = next(self.index_generator)

        # The transformation of images is not under thread lock

        # so it can be done in parallel

        batch_x = np.zeros(

            tuple(

                [current_batch_size]

                + [self.smiles_data_generator.pad, self.smiles_data_generator._charlen]

            ),

            dtype=self.dtype,

        )

        for i, j in enumerate(index_array):

            smiles = self.x[j : j + 1]

            x = self.smiles_data_generator.transform(smiles)

            batch_x[i] = x

        if self.y is None:

            return batch_x

        batch_y = self.y[index_array]

        return batch_x, batch_y

class SmilesEnumerator(object):

    """SMILES Enumerator, vectorizer and devectorizer

    #Arguments

        charset: string containing the characters for the vectorization

        can also be generated via the .fit() method

        pad: Length of the vectorization

        leftpad: Add spaces to the left of the SMILES

        isomericSmiles: Generate SMILES containing information about stereogenic centers

        enum: Enumerate the SMILES during transform

        canonical: use canonical SMILES during transform (overrides enum)

    """

    def __init__(

        self,

        charset="@C)(=cOn1S2/H[N]\\",

        pad=120,

        leftpad=True,

        isomericSmiles=True,

        enum=True,

        canonical=False,

    ):

        self._charset = None

        self.charset = charset

        self.pad = pad

        self.leftpad = leftpad

        self.isomericSmiles = isomericSmiles

        self.enumerate = enum

        self.canonical = canonical

    @property

    def charset(self):

        return self._charset

    @charset.setter

    def charset(self, charset):

        self._charset = charset

        self._charlen = len(charset)

        self._char_to_int = dict((c, i) for i, c in enumerate(charset))

        self._int_to_char = dict((i, c) for i, c in enumerate(charset))

    def fit(self, smiles, extra_chars=[], extra_pad=5):

        """Performs extraction of the charset and length of a SMILES datasets and sets self.pad and self.charset

        #Arguments

            smiles: Numpy array or Pandas series containing smiles as strings

            extra_chars: List of extra chars to add to the charset (e.g. "\\\\" when "/" is present)

            extra_pad: Extra padding to add before or after the SMILES vectorization

        """

        charset = set("".join(list(smiles)))

        self.charset = "".join(charset.union(set(extra_chars)))

        self.pad = max([len(smile) for smile in smiles]) + extra_pad

    def randomize_smiles(self, smiles):

        """Perform a randomization of a SMILES string

        must be RDKit sanitizable"""

        m = Chem.MolFromSmiles(smiles)

        ans = list(range(m.GetNumAtoms()))

        np.random.shuffle(ans)

        nm = Chem.RenumberAtoms(m, ans)

        return Chem.MolToSmiles(

            nm, canonical=self.canonical, isomericSmiles=self.isomericSmiles

        )

    def transform(self, smiles):

        """Perform an enumeration (randomization) and vectorization of a Numpy array of smiles strings

        #Arguments

            smiles: Numpy array or Pandas series containing smiles as strings

        """

        one_hot = np.zeros((smiles.shape[0], self.pad, self._charlen), dtype=np.int8)

        errors = 0

        if self.leftpad:

            for i, ss in enumerate(smiles):

                if self.enumerate:

                    ss = self.randomize_smiles(ss)

                l = len(ss)

                diff = self.pad - l

                for j, c in enumerate(ss):

                    try:

                        one_hot[i, j + diff, self._char_to_int[c]] = 1

                    except:

                        errors += 1

                        break

            # print(f"errors: {errors}")

            return one_hot

        else:

            for i, ss in enumerate(smiles):

                if self.enumerate:

                    ss = self.randomize_smiles(ss)

                for j, c in enumerate(ss):

                    try:

                        one_hot[i, j, self._char_to_int[c]] = 1

                    except:

                        errors += 1

                        break

            # print(f"errors: {errors}")

            return one_hot

    def reverse_transform(self, vect):

        """Performs a conversion of a vectorized SMILES to a smiles strings

        charset must be the same as used for vectorization.

        #Arguments

            vect: Numpy array of vectorized SMILES.

        """

        smiles = []

        for v in vect:

            # mask v

            v = v[v.sum(axis=1) == 1]

            # Find one hot encoded index with argmax, translate to char and join to string

            smile = "".join(self._int_to_char[i] for i in v.argmax(axis=1))

            smiles.append(smile)

        return np.array(smiles)

    def enumerate_smiles(

        self,

        data_reader,

        smiles_col,

        replication_count=2,

        random_pairs=False,

        rand_proba=0.0,

    ):

        """

        Performs enumeration augmentation on the canonical molecular SMILES

        Args:

            dataset (_type_): dataframe containing molecular SMILSS

            smiles_col (_type_): column corresponding to molecular SMILES

            replication_count (int, optional): Number of enumerations for each CHEMICAL SMILE. Defaults to 2.

        """

        smiles = np.repeat(data_reader.dataset[smiles_col].values, replication_count)

        self.fit(smiles, extra_chars=["\\"])

        v = self.transform(smiles)

        transformed = self.reverse_transform(v)

        # print(len(v), len(original_smiles), len(transformed))

        is_enumerated = [1] * len(smiles)

        if random_pairs:

            assert len(smiles) == len(

                transformed

            ), "The length of augmented SMILES must equal original SMILES"

            for idx, _ in enumerate(smiles):

                if round(np.random.uniform(), 1) > rand_proba:

                    continue

                else:

                    transformed[idx] = np.random.choice(smiles)

                    is_enumerated[idx] = 0

        return transformed, list(smiles), is_enumerated

    def enumerate_smiles_df(

        self,

        data_reader,

        smiles_col,

        replication_count=2,

        random_pairs=False,

        rand_proba=0.0,

    ):

        """

        Performs enumeration augmentation on the canonical molecular SMILES

        Args:

            dataset (_type_): dataframe containing molecular SMILSS

            smiles_col (_type_): column corresponding to molecular SMILES

            replication_count (int, optional): Number of enumerations for each CHEMICAL SMILE. Defaults to 2.

        """

        smiles = np.repeat(data_reader[smiles_col].values, replication_count)

        self.fit(smiles, extra_chars=["\\"])

        v = self.transform(smiles)

        transformed = self.reverse_transform(v)

        # print(len(v), len(original_smiles), len(transformed))

        is_enumerated = [1] * len(smiles)

        if random_pairs:

            assert len(smiles) == len(

                transformed

            ), "The length of augmented SMILES must equal original SMILES"

            for idx, _ in enumerate(smiles):

                if round(np.random.uniform(), 1) > rand_proba:

                    continue

                else:

                    transformed[idx] = np.random.choice(smiles)

                    is_enumerated[idx] = 0

        return transformed, is_enumerated

    def smiles_enumeration(self, input_smiles, replication_count=100, n_augment=0):

        """

        Performs enumeration augmentation on the canonical molecular SMILES

        Args:

            dataset (_type_): dataframe containing molecular SMILSS

            smiles_col (_type_): column corresponding to molecular SMILES

            replication_count (int, optional): Number of enumerations for each CHEMICAL SMILE. Defaults to 2.

        """

        enumerations = []

        try:

            smiles = np.repeat([input_smiles], replication_count)

            self.fit(smiles, extra_chars=["\\"])

            v = self.transform(smiles)

            transformed = self.reverse_transform(v)

            for _, enumerated_smiles in enumerate(transformed):

                if len(enumerated_smiles) >= len(input_smiles):

                    enumerations.append(enumerated_smiles)

                if len(enumerations) >= n_augment:

                    break

        except:

            pass

        return enumerations