# Source: https://github.com/EBjerrum/molvecgen
import threading

import numpy as np
from tensorflow.keras.utils import Sequence


class Iterator(object):
    """Abstract base class for data iterators.

    :parameter n: Integer, total number of samples in the dataset to loop over.
    :parameter batch_size: Integer, size of a batch.
    :parameter shuffle: Boolean, whether to shuffle the data between epochs.
    :parameter 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 SmilesGenerator(Iterator):
    """Iterator yielding data from a SMILES array.

    :parameter x: Numpy array of SMILES input data.
    :parameter y: Numpy array of targets data.
    :parameter vectorizer: Instance of molecular vectorizer
    :parameter batch_size: Integer, size of a batch.
    :parameter shuffle: Boolean, whether to shuffle the data between epochs.
    :parameter seed: Random seed for data shuffling.
    :parameter dtype: dtype to use for returned batch. Set to keras.backend.floatx if using Keras
    """

    def __init__(
        self,
        x,
        y,
        vectorizer,
        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.vectorizer = vectorizer
        self.dtype = dtype
        # print(type(self))
        # print(type(SmilesGenerator))
        super(SmilesGenerator, self).__init__(len(x), batch_size, shuffle, seed)

    def next(self):
        """For python 2.x.
        returns the next batch. The X is directly the vectorized format and y is as supplied.
        """
        # 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] + list(self.vectorizer.dims)), dtype=self.dtype
        )
        for i, j in enumerate(index_array):
            smiles = self.x[j : j + 1]
            x = self.vectorizer.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 HetSmilesGenerator2(SmilesGenerator):
    """Hetero (maybe) generator class, for use to train the autoencoder.
    
    smilesvectorizer creates the input for the encoder
        Can be left_padded
    smilesvectorizer_2 creates the teacher input for the decoder + output.
        Must be right_padded. Output for decoder left shifted 1 pos, so no startchar.
    """

    def __init__(
        self,
        x,
        y,
        smilesvectorizer,
        smilesvectorizer_2,
        batch_size=32,
        shuffle=False,
        seed=None,
        dtype=np.float32,
    ):
        super(HetSmilesGenerator, self).__init__(
            x,
            y,
            smilesvectorizer,
            batch_size=batch_size,
            shuffle=shuffle,
            seed=seed,
            dtype=dtype,
        )
        self.smilesvectorizer = smilesvectorizer
        self.smilesvectorizer_2 = smilesvectorizer_2

        self.enc_dims = list(self.smilesvectorizer.dims)
        # Subtract one from the output dims to prepare for the left shifting of output
        self.dec_dims = list(self.smilesvectorizer.dims)
        self.dec_dims[0] = self.dec_dims[0] - 1

    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)

        # Prepare output arrays
        batch_1D = np.zeros(
            tuple([current_batch_size] + self.enc_dims), dtype=self.dtype
        )
        batch_1D_i = np.zeros(
            tuple([current_batch_size] + self.dec_dims), dtype=self.dtype
        )
        batch_1D_o = np.zeros(
            tuple([current_batch_size] + self.dec_dims), dtype=self.dtype
        )

        # TODO Maybe vectorize this, transform already has a for loop
        for i, j in enumerate(index_array):
            mol = self.x[j : j + 1]

            chem1d_enc = self.smilesvectorizer.transform(mol)
            chem1d_dec = self.smilesvectorizer_2.transform(mol)

            batch_1D[i] = chem1d_enc
            batch_1D_i[i] = chem1d_dec[:, 0:-1, :]  # Including start_char
            batch_1D_o[i] = chem1d_dec[:, 1:, :]  # No start_char

        return [batch_1D, batch_1D_i], batch_1D_o


class SmilesGenerator2(SmilesGenerator):
    """Generator class, for use to train the unbiased SMILES RNN.
    
    smilesvectorizer creates the input for the encoder (not used BUT kept for compatibility)
        Can be left_padded. 
    smilesvectorizer_2 creates the teacher input for the decoder + output.
        Must be right_padded. Output for decoder left shifted 1 pos, so no startchar.
    """

    def __init__(
        self,
        x,
        y,
        smilesvectorizer,
        smilesvectorizer_2,
        batch_size=32,
        shuffle=False,
        seed=None,
        dtype=np.float32,
    ):
        super(SmilesGenerator2, self).__init__(
            x,
            y,
            smilesvectorizer,
            batch_size=batch_size,
            shuffle=shuffle,
            seed=seed,
            dtype=dtype,
        )
        self.smilesvectorizer = smilesvectorizer
        self.smilesvectorizer_2 = smilesvectorizer_2

        self.enc_dims = list(self.smilesvectorizer.dims)
        # Subtract one from the output dims to prepare for the left shifting of output
        self.dec_dims = list(self.smilesvectorizer.dims)
        self.dec_dims[0] = self.dec_dims[0] - 1

    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)

        # Prepare output arrays
        batch_1D = np.zeros(
            tuple([current_batch_size] + self.enc_dims), dtype=self.dtype
        )
        batch_1D_i = np.zeros(
            tuple([current_batch_size] + self.dec_dims), dtype=self.dtype
        )
        batch_1D_o = np.zeros(
            tuple([current_batch_size] + self.dec_dims), dtype=self.dtype
        )

        # TODO Maybe vectorize this, transform already has a for loop
        for i, j in enumerate(index_array):
            mol = self.x[j : j + 1]

            chem1d_enc = self.smilesvectorizer.transform(mol)
            chem1d_dec = self.smilesvectorizer_2.transform(mol)

            batch_1D[i] = chem1d_enc
            batch_1D_i[i] = chem1d_dec[:, 0:-1, :]  # Including start_char
            batch_1D_o[i] = chem1d_dec[:, 1:, :]  # No start_char

        return [batch_1D_i], batch_1D_o


class CodeGenerator2(SmilesGenerator):
    """Code generator class to train a DDC.
    :parameter x: Numpy array of encoded input data.
    :parameter y: Numpy array of SMILES output data.
    :parameter vectorizer: Instance of molecular vectorizer
    :parameter batch_size: Integer, size of a batch.
    :parameter shuffle: Boolean, whether to shuffle the data between epochs.
    :parameter seed: Random seed for data shuffling.
    :parameter dtype: dtype to use for returned batch. Set to keras.backend.floatx if using Keras
    """

    def __init__(
        self,
        x,
        y,
        smilesvectorizer,
        smilesvectorizer_2,
        batch_size=32,
        shuffle=False,
        seed=None,
        dtype=np.float32,
    ):
        super(CodeGenerator, self).__init__(
            x,
            y,
            smilesvectorizer,
            batch_size=batch_size,
            shuffle=shuffle,
            seed=seed,
            dtype=dtype,
        )
        self.smilesvectorizer = smilesvectorizer
        self.smilesvectorizer_2 = smilesvectorizer_2

        self.input_dims = [self.x.shape[1]]
        # Subtract one from the output dims to prepare for the left shifting of output
        self.dec_dims = list(self.smilesvectorizer.dims)
        self.dec_dims[0] = self.dec_dims[0] - 1

    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)

        # Prepare output arrays
        batch_1D = np.zeros(
            tuple([current_batch_size] + self.input_dims), dtype=self.dtype
        )
        batch_1D_i = np.zeros(
            tuple([current_batch_size] + self.dec_dims), dtype=self.dtype
        )
        batch_1D_o = np.zeros(
            tuple([current_batch_size] + self.dec_dims), dtype=self.dtype
        )

        # TODO Maybe vectorize this, transform already has a for loop
        for i, j in enumerate(index_array):

            mol = self.y[j : j + 1]
            chem1d_dec = self.smilesvectorizer_2.transform(mol)

            batch_1D[i] = self.x[j : j + 1]
            batch_1D_i[i] = chem1d_dec[:, 0:-1, :]  # Including start_char
            batch_1D_o[i] = chem1d_dec[:, 1:, :]  # No start_char

        return [batch_1D, batch_1D_i], batch_1D_o


class HetSmilesGenerator(Sequence):
    def __init__(
        self,
        x,
        y,  # TODO: this is unused, remove from everywhere
        smilesvectorizer,
        smilesvectorizer_2,
        batch_size=32,
        shuffle=False,
        seed=None,
        dtype=np.float32,
    ):
        super().__init__()
        self.smilesvectorizer = smilesvectorizer
        self.smilesvectorizer_2 = smilesvectorizer_2
        self.x = x
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.dtype = dtype

        # Assume that all binmols are to be used
        # Otherwise use list_IDs as an argument to pass an arbitrary selection of IDs to use
        self.list_IDs = np.arange(len(self.x))
        self.num_batches = int(np.floor(len(self.x) / self.batch_size))

        self.input_dims = list(self.smilesvectorizer.dims)
        self.dec_dims = list(self.smilesvectorizer.dims)
        self.dec_dims[0] = self.dec_dims[0] - 1

        self.__shuffle()

    def __len__(self):
        return self.num_batches

    def __getitem__(self, batch_id):
        indices_of_batch = self.indices[
            batch_id * self.batch_size : (batch_id + 1) * self.batch_size
        ]
        list_IDs_temp = [self.list_IDs[i] for i in indices_of_batch]
        return self.__get_batch(list_IDs_temp)

    def __get_batch(self, list_IDs_temp):
        batch_1D = np.zeros(
            tuple([self.batch_size] + self.input_dims), dtype=self.dtype
        )
        batch_1D_i = np.zeros(
            tuple([self.batch_size] + self.dec_dims), dtype=self.dtype
        )
        batch_1D_o = np.zeros(
            tuple([self.batch_size] + self.dec_dims), dtype=self.dtype
        )

        for i, ID in enumerate(list_IDs_temp):
            binmol = self.x[ID : ID + 1]
            chem1d_enc = self.smilesvectorizer.transform(binmol)
            chem1d_dec = self.smilesvectorizer_2.transform(binmol)
            batch_1D[i] = chem1d_enc
            batch_1D_i[i] = chem1d_dec[:, 0:-1, :]  # Including start_char
            batch_1D_o[i] = chem1d_dec[:, 1:, :]  # No start_char

        return [batch_1D, batch_1D_i], batch_1D_o

    def __shuffle(self):
        self.indices = np.arange(len(self.list_IDs))
        if self.shuffle == True:
            np.random.shuffle(self.indices)

    def on_epoch_end(self):
        self.__shuffle()


class CodeGenerator(Sequence):
    def __init__(
        self,
        x,
        y,
        smilesvectorizer,
        smilesvectorizer_2,
        batch_size=32,
        shuffle=False,
        seed=None,
        dtype=np.float32,
    ):
        super().__init__()
        self.smilesvectorizer = smilesvectorizer
        self.smilesvectorizer_2 = smilesvectorizer_2
        self.x = x
        self.y = y
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.dtype = dtype

        # Assume that all binmols are to be used
        # Otherwise use list_IDs as an argument to pass an arbitrary selection of IDs to use
        self.list_IDs = np.arange(len(y))
        self.num_batches = int(np.floor(len(self.y) / self.batch_size))
        self.input_dims = [self.x.shape[1]]
        self.dec_dims = list(self.smilesvectorizer.dims)
        self.dec_dims[0] = self.dec_dims[0] - 1

        self.__shuffle()

    def __len__(self):
        return self.num_batches

    def __getitem__(self, batch_id):
        indices_of_batch = self.indices[
            batch_id * self.batch_size : (batch_id + 1) * self.batch_size
        ]
        list_IDs_temp = [self.list_IDs[i] for i in indices_of_batch]
        return self.__get_batch(list_IDs_temp)

    def __get_batch(self, list_IDs_temp):
        batch_1D = np.zeros(
            tuple([self.batch_size] + self.input_dims), dtype=self.dtype
        )
        batch_1D_i = np.zeros(
            tuple([self.batch_size] + self.dec_dims), dtype=self.dtype
        )
        batch_1D_o = np.zeros(
            tuple([self.batch_size] + self.dec_dims), dtype=self.dtype
        )

        for i, ID in enumerate(list_IDs_temp):
            binmol = self.y[ID : ID + 1]
            chem1d_dec = self.smilesvectorizer_2.transform(binmol)

            batch_1D[i] = self.x[ID : ID + 1]
            batch_1D_i[i] = chem1d_dec[:, 0:-1, :]  # Including start_char
            batch_1D_o[i] = chem1d_dec[:, 1:, :]  # No start_char
        return [batch_1D, batch_1D_i], batch_1D_o

    def __shuffle(self):
        self.indices = np.arange(len(self.list_IDs))
        if self.shuffle == True:
            np.random.shuffle(self.indices)

    def on_epoch_end(self):
        self.__shuffle()
