""" Pandas extension for chemical analysis """
from __future__ import absolute_import
import pandas as pd

import oddt

pd.set_option("display.max_colwidth", 999999)


def _mol_reader(fmt='sdf',
                filepath_or_buffer=None,
                usecols=None,
                molecule_column='mol',
                molecule_name='mol_name',
                smiles_column=None,
                skip_bad_mols=False,
                **kwargs):
    out = []
    for mol in oddt.toolkit.readfile(fmt, filepath_or_buffer):
        if skip_bad_mols and mol is None:
            continue  # add warning with number of skipped molecules
        if usecols:
            mol_data = dict((k, v) for k, v in mol.data.items() if k in usecols)
        else:
            mol_data = dict(mol.data)
        if molecule_column:
            mol_data[molecule_column] = mol
        if molecule_name:
            mol_data[molecule_name] = mol.title
        if smiles_column:
            mol_data[smiles_column] = mol.write('smi').split()[0]
        out.append(mol_data)
    return ChemDataFrame(out, **kwargs)


def _mol_chunk_reader(fmt='sdf',
                      filepath_or_buffer=None,
                      usecols=None,
                      molecule_column='mol',
                      molecule_name='mol_name',
                      smiles_column=None,
                      skip_bad_mols=False,
                      chunksize=1000,
                      **kwargs):
    chunk = []
    for n, mol in enumerate(oddt.toolkit.readfile(fmt, filepath_or_buffer)):
        if skip_bad_mols and mol is None:
            continue  # add warning with number of skipped molecules
        if usecols:
            mol_data = dict((k, v) for k, v in mol.data.items() if k in usecols)
        else:
            mol_data = dict(mol.data)
        if molecule_column:
            mol_data[molecule_column] = mol
        if molecule_name:
            mol_data[molecule_name] = mol.title
        if smiles_column:
            mol_data[smiles_column] = mol.write('smi').split()[0]
        chunk.append(mol_data)
        if (n + 1) % chunksize == 0:
            yield ChemDataFrame(chunk, **kwargs)
            chunk = []
    if chunk:
        yield ChemDataFrame(chunk, **kwargs)


def read_sdf(filepath_or_buffer=None,
             usecols=None,
             molecule_column='mol',
             molecule_name='mol_name',
             smiles_column=None,
             skip_bad_mols=False,
             chunksize=1000,
             **kwargs):
    if chunksize:
        return _mol_chunk_reader(fmt='sdf',
                                 filepath_or_buffer=filepath_or_buffer,
                                 usecols=usecols,
                                 molecule_column=molecule_column,
                                 molecule_name=molecule_name,
                                 smiles_column=smiles_column,
                                 skip_bad_mols=skip_bad_mols,
                                 chunksize=chunksize,
                                 **kwargs)
    else:
        return _mol_reader(fmt='sdf',
                           filepath_or_buffer=filepath_or_buffer,
                           usecols=usecols,
                           molecule_column=molecule_column,
                           molecule_name=molecule_name,
                           smiles_column=smiles_column,
                           skip_bad_mols=skip_bad_mols,
                           **kwargs)


class ChemSeries(pd.Series):
    @property
    def _constructor(self):
        return ChemSeries


class ChemDataFrame(pd.DataFrame):
    """

    Note:
    Thanks to: http://blog.snapdragon.cc/2015/05/05/subclass-pandas-dataframe-to-save-custom-attributes/
    """
    def to_sdf(self,
               filepath_or_buffer=None,
               update_properties=True,
               molecule_column='mol',
               columns=None):
        out = oddt.toolkit.Outputfile('sdf', filepath_or_buffer, overwrite=True)
        for ix, row in self.iterrows():
            mol = row[molecule_column].clone
            if update_properties:
                new_data = row.to_dict()
                del new_data[molecule_column]
                mol.data.update(new_data)
            if columns:
                for k in mol.data.keys():
                    if k not in columns:
                        del mol.data[k]
            out.write(mol)
        out.close()

    def to_html(self, *args, **kwargs):
        kwargs['escape'] = False
        return super(ChemDataFrame, self).to_html(*args, **kwargs)

    @property
    def _constructor(self):
        """ Force new class to be usead as sconstructor when slicing """
        return ChemDataFrame

    _constructor_sliced = ChemSeries

@@ -433,1006 +433,24 @@
     
-**kwargs):%0A    out = %5B%5D%0A    for mol in oddt.toolkit.readfile(fmt, filepath_or_buffer):%0A        if skip_bad_mols and mol is None:%0A            continue  # add warning with number of skipped molecules%0A        if usecols:%0A            mol_data = dict((k, v) for k, v in mol.data.items() if k in usecols)%0A        else:%0A            mol_data = dict(mol.data)%0A        if molecule_column:%0A            mol_data%5Bmolecule_column%5D = mol%0A        if molecule_name:%0A            mol_data%5Bmolecule_name%5D = mol.title%0A        if smiles_column:%0A            mol_data%5Bsmiles_column%5D = mol.write('smi').split()%5B0%5D%0A        out.append(mol_data)%0A    return ChemDataFrame(out, **kwargs)%0A%0A%0Adef _mol_chunk_reader(fmt='sdf',%0A                      filepath_or_buffer=None,%0A                      usecols=None,%0A                      molecule_column='mol',%0A                      molecule_name='mol_name',%0A                      smiles_column=None,%0A                      skip_bad_mols=False,%0A                      chunksize=1000,%0A      
+chunksize=None,%0A
     
@@ -1108,16 +1108,30 @@
      if 
+chunksize and 
 (n + 1) 
@@ -1515,12 +1515,12 @@
 ize=
-1000
+None
 ,%0A  
@@ -1549,47 +1549,21 @@
     
-if chunksize:%0A        return _mol_chunk
+result = _mol
 _rea
@@ -1573,32 +1573,24 @@
 (fmt='sdf',%0A
-        
             
@@ -1662,32 +1662,24 @@
             
-        
 usecols=usec
@@ -1704,32 +1704,24 @@
             
-  
-      
 molecule_col
@@ -1762,32 +1762,24 @@
             
-        
 molecule_nam
@@ -1816,32 +1816,24 @@
             
-  
-      
 smiles_colum
@@ -1845,32 +1845,24 @@
 les_column,%0A
-        
             
@@ -1924,32 +1924,24 @@
             
- 
-       
 chunksize=ch
@@ -1949,24 +1949,16 @@
 nksize,%0A
-        
         
@@ -1984,27 +1984,35 @@
 kwargs)%0A    
-els
+if chunksiz
 e:%0A        r
@@ -2021,404 +2021,51 @@
 urn 
-_mol_reader(fmt='sdf',%0A                           filepath_or_buffer=filepath_or_buffer,%0A                           usecols=usecols,%0A                           molecule_column=molecule_column,%0A                           molecule_name=molecule_name,%0A                           smiles_column=smiles_column,%0A                           skip_bad_mols=skip_bad_mols,%0A                           **kwargs
+result%0A    else:%0A        return next(result
 )%0A%0A%0A

Unify chunked reading