"""
Molecule images.
"""

__author__ = "Steven Kearnes"
__copyright__ = "Copyright 2014, Stanford University"
__license__ = "BSD 3-clause"

import os
import subprocess

from rdkit import Chem

from pande_gas.features import Featurizer
from pande_gas.utils import image_utils


class MolImage(Featurizer):
    """
    Molecule images.

    Parameters
    ----------
    size : int, optional (default 32)
        Size (in any dimension) of generated images.
    flatten : bool, optional (default False)
        Whether to flatten the pixel array. If False, the features for each
        molecule will be a 3D array.
    """
    name = 'image'

    def __init__(self, size=32, flatten=False):
        self.size = size
        if not flatten:
            self.topo_view = True
        self.flatten = flatten

    def _featurize(self, mol):
        """
        Generate a 2D depiction of a molecule.

        Parameters
        ----------
        mol : RDKit Mol
            Molecule.
        """
        smiles = Chem.MolToSmiles(mol, isomericSmiles=True, canonical=True)
        image = self.image_from_smiles(smiles)
        pixels = image_utils.get_pixels(image)
        if self.flatten:
            pixels = pixels.ravel()
        return pixels

    def image_from_smiles(self, smiles):
        """
        Generate a PNG image from a SMILES string using OpenBabel.

        Note that we call OpenBabel using subprocss to avoid entanglements
        with the GNU GPL.

        Parameters
        ----------
        smiles : str
            SMILES string.
        """
        png_args = ['obabel', '-ican', '-opng', '-xd', '-xC',
                    '-xp {}'.format(self.size)]
        p = subprocess.Popen(png_args, stdin=subprocess.PIPE, 
                             stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        png, _ = p.communicate(smiles)
        im = image_utils.load(png)
        return im

@@ -137,37 +137,8 @@
 e%22%0A%0A
-import os%0Aimport subprocess%0A%0A
 from
@@ -143,24 +143,29 @@
 om rdkit
+.Chem
  import 
 Chem%0A%0Afr
@@ -156,20 +156,20 @@
  import 
-Chem
+Draw
 %0A%0Afrom p
@@ -973,121 +973,92 @@
     
-smiles = Chem.MolToSmiles(mol, isomericSmiles=True, canonical=True)%0A        image = self.image_from_smiles(smiles
+dim = (self.size, self.size)%0A        image = Draw.MolToImage(mol, dim, fitImage=True
 )%0A  
@@ -1189,668 +1189,4 @@
 els%0A
-%0A    def image_from_smiles(self, smiles):%0A        %22%22%22%0A        Generate a PNG image from a SMILES string using OpenBabel.%0A%0A        Note that we call OpenBabel using subprocss to avoid entanglements%0A        with the GNU GPL.%0A%0A        Parameters%0A        ----------%0A        smiles : str%0A            SMILES string.%0A        %22%22%22%0A        png_args = %5B'obabel', '-ican', '-opng', '-xd', '-xC',%0A                    '-xp %7B%7D'.format(self.size)%5D%0A        p = subprocess.Popen(png_args, stdin=subprocess.PIPE, %0A                             stdout=subprocess.PIPE, stderr=subprocess.PIPE)%0A        png, _ = p.communicate(smiles)%0A        im = image_utils.load(png)%0A        return im%0A

Use RDKit to generate images