"""
Contains an abstract base class that supports data transformations.
"""
from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals
import os
import numpy as np
import warnings
from functools import partial
from deepchem.utils.save import save_to_disk
from deepchem.utils.save import load_from_disk
from deepchem.utils import pad_array
import shutil
from deepchem.data import DiskDataset, NumpyDataset
def undo_transforms(y, transformers):
  """Undoes all transformations applied."""
  # Note that transformers have to be undone in reversed order
  for transformer in reversed(transformers):
    if transformer.transform_y:
      y = transformer.untransform(y)
  return y
def undo_grad_transforms(grad, tasks, transformers):
  for transformer in reversed(transformers):
    if transformer.transform_y:
      grad = transformer.untransform_grad(grad, tasks)
  return grad
def get_grad_statistics(dataset):
  """Computes and returns statistics of a dataset
  This function assumes that the first task of a dataset holds the energy for
  an input system, and that the remaining tasks holds the gradient for the
  system.
  """
  if len(dataset) == 0:
    return None, None, None, None
  y = dataset.y
  energy = y[:,0]
  grad = y[:,1:]
  for i in range(energy.size):
    grad[i] *= energy[i]
  ydely_means = np.sum(grad, axis=0)/len(energy)
  return grad, ydely_means
class Transformer(object):
  """
  Abstract base class for different ML models.
  """
  # Hack to allow for easy unpickling:
  # http://stefaanlippens.net/pickleproblem
  __module__ = os.path.splitext(os.path.basename(__file__))[0]
  def __init__(self, transform_X=False, transform_y=False, transform_w=False,
               dataset=None):
    """Initializes transformation based on dataset statistics."""
    self.dataset = dataset
    self.transform_X = transform_X
    self.transform_y = transform_y
    self.transform_w = transform_w
    # One, but not both, transform_X or tranform_y is true
    assert transform_X or transform_y or transform_w
    # Use fact that bools add as ints in python
    assert (transform_X + transform_y + transform_w) == 1 
  def transform_array(self, X, y, w):
    """Transform the data in a set of (X, y, w) arrays."""
    raise NotImplementedError(
      "Each Transformer is responsible for its own transform_array method.")
  def untransform(self, z):
    """Reverses stored transformation on provided data."""
    raise NotImplementedError(
      "Each Transformer is responsible for its own untransfomr method.")
  def transform(self, dataset, parallel=False):
    """
    Transforms all internally stored data.
    Adds X-transform, y-transform columns to metadata.
    """
    return dataset.transform(lambda X, y, w: self.transform_array(X, y, w))
  def transform_on_array(self, X, y, w):
    """
    Transforms numpy arrays X, y, and w
    """
    X, y, w = self.transform_array(X, y, w)    
    return X, y, w
class NormalizationTransformer(Transformer):
  def __init__(self, transform_X=False, transform_y=False, transform_w=False,
               dataset=None, transform_gradients=False):
    """Initialize normalization transformation."""
    super(NormalizationTransformer, self).__init__(
        transform_X=transform_X, transform_y=transform_y,
        transform_w=transform_w, dataset=dataset)
    if transform_X:
      X_means, X_stds = dataset.get_statistics(X_stats=True, y_stats=False)
      self.X_means = X_means 
      self.X_stds = X_stds
    elif transform_y:
      y_means, y_stds = dataset.get_statistics(X_stats=False, y_stats=True)
      self.y_means = y_means 
      # Control for pathological case with no variance.
      y_stds[y_stds == 0] = 1.
      self.y_stds = y_stds
    self.transform_gradients = transform_gradients
    if self.transform_gradients:
      true_grad, ydely_means = get_grad_statistics(dataset)
      self.grad = np.reshape(true_grad, (true_grad.shape[0],-1,3))
      self.ydely_means = ydely_means
  def transform(self, dataset, parallel=False):
    return super(NormalizationTransformer, self).transform(
        dataset, parallel=parallel)
  def transform_array(self, X, y, w):
    """Transform the data in a set of (X, y, w) arrays."""
    if self.transform_X:
      X = np.nan_to_num((X - self.X_means) / self.X_stds)
    if self.transform_y:
      y = np.nan_to_num((y - self.y_means) / self.y_stds)
    return (X, y, w)
  def untransform(self, z):
    """
    Undo transformation on provided data.
    """
    if self.transform_X:
      return z * self.X_stds + self.X_means
    elif self.transform_y:
      return z * self.y_stds + self.y_means
  def untransform_grad(self, grad, tasks):
    """
    Undo transformation on gradient.
    """
    if self.transform_y:
      grad_means = self.y_means[1:]
      energy_var = self.y_stds[0]        
      grad_var = 1/energy_var*(self.ydely_means-self.y_means[0]*self.y_means[1:])
      energy = tasks[:,0]
      transformed_grad = []
      for i in range(energy.size):
        Etf = energy[i]
        grad_Etf = grad[i].flatten()
        grad_E = Etf*grad_var+energy_var*grad_Etf+grad_means
        grad_E = np.reshape(grad_E, (-1,3))
        transformed_grad.append(grad_E)   
      transformed_grad = np.asarray(transformed_grad)
      return transformed_grad
class AtomicNormalizationTransformer(Transformer):
  """
  TODO(rbharath): Needs more discussion of what a gradient is semantically.
  It's evident that not every Dataset has meaningful gradient information, so
  this transformer can't be applied to all data. Should there be a subclass of
  Dataset named GradientDataset perhaps?
  """
  def __init__(self, transform_X=False, transform_y=False, transform_w=False,
               dataset=None):
    """Initialize normalization transformation."""
    super(AtomicNormalizationTransformer, self).__init__(
        transform_X=transform_X, transform_y=transform_y,
        transform_w=transform_w, dataset=dataset)
    X_means, X_stds, y_means, y_stds = dataset.get_statistics()
    self.X_means = X_means 
    self.X_stds = X_stds
    self.y_means = y_means 
    # Control for pathological case with no variance.
    y_stds[y_stds == 0] = 1.
    self.y_stds = y_stds
    true_grad, ydely_means = get_grad_statistics(dataset)
    self.grad = np.reshape(true_grad, (true_grad.shape[0],-1,3))
    self.ydely_means = ydely_means
  def transform(self, dataset, parallel=False):
    return super(AtomicNormalizationTransformer, self).transform(
        dataset, parallel=parallel)
    
  def transform_row(self, i, df, data_dir):
    """
    Normalizes the data (X, y, w, ...) in a single row).
    """
    row = df.iloc[i]
    if self.transform_X:
      X = load_from_disk(
          os.path.join(data_dir, row['X-transformed']))
      X = np.nan_to_num((X - self.X_means) / self.X_stds)
      save_to_disk(X, os.path.join(data_dir, row['X-transformed']))
    if self.transform_y:
      y = load_from_disk(os.path.join(data_dir, row['y-transformed']))
      # transform tasks as normal
      y = np.nan_to_num((y - self.y_means) / self.y_stds)
      # add 2nd order correction term to gradients
      grad_var = 1/self.y_stds[0]*(self.ydely_means-self.y_means[0]*self.y_means[1:])
      for i in range(y.shape[0]):
        y[i,1:] = y[i,1:] - grad_var*y[i,0]/self.y_stds[0]
      save_to_disk(y, os.path.join(data_dir, row['y-transformed']))
  def transform_array(self, X, y, w):
    """Transform the data in a set of (X, y, w) arrays."""
    if self.transform_X:
      X = np.nan_to_num((X - self.X_means) / self.X_stds)
    if self.transform_y:
      # transform tasks as normal
      y = np.nan_to_num((y - self.y_means) / self.y_stds)
      # add 2nd order correction term to gradients
      grad_var = 1/self.y_stds[0]*(self.ydely_means-self.y_means[0]*self.y_means[1:])
      for i in range(y.shape[0]):
        y[i,1:] = y[i,1:] - grad_var*y[i,0]/self.y_stds[0]
    return (X, y, w)
  def untransform(self, z):
    """
    Undo transformation on provided data.
    """
    if self.transform_X:
      return z * self.X_stds + self.X_means
    elif self.transform_y:
      # untransform grad
      grad_var = 1/self.y_stds[0]*(self.ydely_means-self.y_means[0]*self.y_means[1:])
      for i in range(z.shape[0]):
        z[i,1:] = z[i,0]*grad_var + self.y_stds[0]*z[i,1:] + self.y_means[1:] 
      # untransform energy
      z[:,0] = z[:,0] * self.y_stds[0] + self.y_means[0]
      return z
  def untransform_grad(self, grad, tasks):
    """
    Undo transformation on gradient.
    """
    if self.transform_y:
      grad_means = self.y_means[1:]
      energy_var = self.y_stds[0]        
      grad_var = 1/energy_var*(self.ydely_means-self.y_means[0]*self.y_means[1:])
      energy = tasks[:,0]
      transformed_grad = []


Answer:       for i in range(energy.size):